@@ -168,7 +168,7 @@ module Streamly.Internal.Data.Stream.StreamD
168168 , wordsBy
169169 , splitSuffixBy'
170170
171− , splitOn
171+ , splitOnSeq
172172 , splitSuffixOn
173173
174174 , splitInnerBy
@@ -1786,32 +1786,32 @@ data SplitOptions = SplitOptions
17861786 }
17871787-}
17881788
1789−{-# ANN type SplitOnState Fuse #-}
1790−data SplitOnState fs s a b w rb rh ck =
1791− GO_START
1792− | GO_EMPTY_PAT s
1793− | GO_EMPTY_PAT_WITH s a
1794− | GO_SINGLE_PAT_NEXT !fs s a
1795− | GO_SINGLE_PAT_WITH !fs s a a
1796− | GO_SHORT_PAT_ACCUM Int s !w
1797− | GO_SHORT_PAT_NEXT !fs s !w
1798− | GO_SHORT_PAT_NEXT_WITH !fs s a !w
1799− | GO_SHORT_PAT_DRAIN Int !fs !w
1800− | GO_KARP_RABIN_ACCUM Int s rb !rh
1801− | GO_KARP_RABIN_NEXT !fs s rb !rh !ck
1802− | GO_KARP_RABIN_DRAIN Int !fs rb !rh
1803− | GO_YIELD b (SplitOnState fs s a b w rb rh ck)
1804− | GO_DONE
1805−
1806−{-# INLINE_NORMAL splitOn #-}
1807−splitOn
1789+{-# ANN type SplitOnSeqState Fuse #-}
1790+data SplitOnSeqState rb rh ck w fs s b x =
1791+ SplitOnSeqInit
1792+ | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)
1793+ | SplitOnSeqDone
1794+
1795+ | SplitOnSeqEmpty s
1796+ | SplitOnSeqSingle !fs s x
1797+
1798+ | SplitOnSeqWordInit Int s !w
1799+ | SplitOnSeqWordLoop !fs s !w
1800+ | SplitOnSeqWordDone Int !fs !w
1801+
1802+ | SplitOnSeqKRInit Int s rb !rh
1803+ | SplitOnSeqKRLoop fs s rb !rh !ck
1804+ | SplitOnSeqKRDone Int !fs rb !rh
1805+
1806+{-# INLINE_NORMAL splitOnSeq #-}
1807+splitOnSeq
18081808 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
18091809 => Array a
18101810 -> Fold m a b
18111811 -> Stream m a
18121812 -> Stream m b
1813−splitOn patArr (Fold fstep initial done) (Stream step state) =
1814− Stream stepOuter GO_START
1813+splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
1814+ Stream stepOuter SplitOnSeqInit
18151815
18161816 where
18171817
@@ -1819,26 +1819,19 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =
18191819 maxIndex = patLen - 1
18201820 elemBits = sizeOf (undefined :: a) * 8
18211821
1822− -- XXX I initially put this in the state but I'd like to keep the state as
1823− -- simple as possible to allow more fusion (if that's how it
1824− -- works). Removing this from the state did not change anything. I'll
1825− -- introduce them back once I figure out what's currently wrong. Ideally, I
1826− -- want them to be associated with the state.
1827− mask :: Word
1828− mask = (1 `shiftL` (elemBits * patLen)) - 1
1829−
1830− -- XXX I initially put this in the state but I'd like to keep the state as
1831− -- simple as possible to allow more fusion. Removing this from the state did
1832− -- not change anything. I'll introduce them back once I figure out what's
1833− -- currently wrong.
1834− pat :: Word
1835− -- XXX You dont need .&. here?
1836− pat = mask .&. A.foldl' addToWord 0 patArr
1837−
1838− -- Used only if the pattern is short
1822+ -- For word pattern case
1823+ wordMask :: Word
1824+ wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1825+
1826+ elemMask :: Word
1827+ elemMask = (1 `shiftL` elemBits) - 1
1828+
1829+ wordPat :: Word
1830+ wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1831+
18391832 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
18401833
1841− -- Used in Rabin-Karp
1834+ -- For Rabin-Karp search
18421835 k = 2891336453 :: Word32
18431836 coeff = k ^ patLen
18441837
@@ -1850,136 +1843,125 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =
18501843 -- XXX shall we use a random starting hash or 1 instead of 0?
18511844 patHash = A.foldl' addCksum 0 patArr
18521845
1846+ skip = return . Skip
1847+
18531848 {-# INLINE_LATE stepOuter #-}
1854− stepOuter _ GO_START =
1849+ stepOuter _ SplitOnSeqInit =
18551850 if patLen == 0
1856− then return $ Skip $ GO_EMPTY_PAT state
1851+ then return $ Skip $ SplitOnSeqEmpty state
18571852 else if patLen == 1
18581853 then do
18591854 acc <- initial
1860− pat_ <- liftIO $ A.unsafeIndexIO patArr 0
1861− return $ Skip $ GO_SINGLE_PAT_NEXT acc state pat_
1855+ let !pat = A.unsafeIndex patArr 0
1856+ return $ Skip $ SplitOnSeqSingle acc state pat
18621857 else if sizeOf (undefined :: a) * patLen
18631858 <= sizeOf (undefined :: Word)
1864− then return $ Skip $ GO_SHORT_PAT_ACCUM 0 state (0 :: Word)
1859+ then return $ Skip $ SplitOnSeqWordInit 0 state (0 :: Word)
18651860 else do
18661861 (rb, rhead) <- liftIO $ RB.new patLen
1867− return $ Skip $ GO_KARP_RABIN_ACCUM 0 state rb rhead
1868− stepOuter _ (GO_YIELD x ns) = return $ Yield x ns
1869− stepOuter _ GO_DONE = return Stop
1862+ skip $ SplitOnSeqKRInit 0 state rb rhead
1863+
1864+ stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next
18701865
18711866 ---------------------------
18721867 -- Empty pattern
18731868 ---------------------------
18741869
1875− stepOuter _ (GO_EMPTY_PAT_WITH s x) = do
1876− ini <- initial
1877− res <- fstep ini x
1878− r <- done res
1879− return $ Skip $ GO_YIELD r (GO_EMPTY_PAT s)
1880− stepOuter gst (GO_EMPTY_PAT st) = do
1870+ stepOuter gst (SplitOnSeqEmpty st) = do
18811871 res <- step (adaptState gst) st
1882− return
1883− $ case res of
1884− Yield x s -> Skip $ GO_EMPTY_PAT_WITH s x
1885− Skip s -> Skip $ GO_EMPTY_PAT s
1886− Stop -> Stop
1872+ case res of
1873+ Yield x s -> do
1874+ acc <- initial
1875+ acc1 <- fstep acc x
1876+ r <- done acc1
1877+ skip $ SplitOnSeqYield r (SplitOnSeqEmpty s)
1878+ Skip s -> return $ Skip (SplitOnSeqEmpty s)
1879+ Stop -> return Stop
1880+
1881+ -----------------
1882+ -- Done
1883+ -----------------
1884+
1885+ stepOuter _ SplitOnSeqDone = return Stop
18871886
18881887 -----------------
18891888 -- Single Pattern
18901889 -----------------
18911890
1892− -- XXX These functions cab be designed in a better way? Get rid of go and
1893− -- have more info in the state but you would be relying on GHC for
1894− -- simplification.
1895− -- XXX We can probably move pat_ outside? Or move the things outside, inside
1896− -- the state to make it consistent.
1897− stepOuter _ (GO_SINGLE_PAT_WITH fs s pat_ x) =
1898− if pat_ == x
1899− then do
1900− r <- done fs
1901− ini <- initial
1902− return $ Skip $ GO_YIELD r (GO_SINGLE_PAT_NEXT ini s pat_)
1903− else do
1904− res <- fstep fs x
1905− return $ Skip $ (GO_SINGLE_PAT_NEXT res s pat_)
1906− stepOuter gst (GO_SINGLE_PAT_NEXT fs st pat_) = do
1891+ stepOuter gst (SplitOnSeqSingle fs st pat) = do
19071892 res <- step (adaptState gst) st
19081893 case res of
1909− Yield x s -> return $ Skip $ GO_SINGLE_PAT_WITH fs s pat_ x
1910− Skip s -> return $ Skip $ GO_SINGLE_PAT_NEXT fs s pat_
1894+ Yield x s -> do
1895+ if pat == x
1896+ then do
1897+ r <- done fs
1898+ fs1 <- initial
1899+ return $ Skip $ SplitOnSeqYield r (SplitOnSeqSingle fs1 s pat)
1900+ else do
1901+ fs1 <- fstep fs x
1902+ return $ Skip $ (SplitOnSeqSingle fs1 s pat)
1903+ Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat
19111904 Stop -> do
19121905 r <- done fs
1913− return $ Skip $ GO_YIELD r GO_DONE
1906+ return $ Skip $ SplitOnSeqYield r SplitOnSeqDone
19141907
19151908 ---------------------------
19161909 -- Short Pattern - Shift Or
19171910 ---------------------------
19181911
1919− -- XXX mask can be omitted in most of the places?
1920− stepOuter gst (GO_SHORT_PAT_ACCUM idx st wrd) = do
1912+ stepOuter _ (SplitOnSeqWordDone 0 fs _) = do
1913+ r <- done fs
1914+ skip $ SplitOnSeqYield r SplitOnSeqDone
1915+ stepOuter _ (SplitOnSeqWordDone n fs wrd) = do
1916+ let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1917+ fs1 <- fstep fs (toEnum $ fromIntegral old)
1918+ skip $ SplitOnSeqWordDone (n - 1) fs1 wrd
1919+
1920+ stepOuter gst (SplitOnSeqWordInit idx st wrd) = do
19211921 res <- step (adaptState gst) st
19221922 case res of
19231923 Yield x s -> do
19241924 let wrd1 = addToWord wrd x
1925− if idx == maxIndex
1926− then if wrd1 .&. mask == pat
1927− then do
1928− r <- initial >>= done
1929− return
1930− $ Skip
1931− $ GO_YIELD r $ GO_SHORT_PAT_ACCUM 0 s (0 :: Word)
1932− else do
1933− ini <- initial
1934− return $ Skip $ GO_SHORT_PAT_NEXT ini s wrd1
1935− else return $ Skip $ GO_SHORT_PAT_ACCUM (idx + 1) s wrd1
1936− Skip s -> return $ Skip $ GO_SHORT_PAT_ACCUM idx s wrd
1925+ if idx /= maxIndex
1926+ then skip $ SplitOnSeqWordInit (idx + 1) s wrd1
1927+ else do
1928+ fs <- initial
1929+ if wrd1 .&. wordMask == wordPat
1930+ then do
1931+ r <- done fs
1932+ let next = SplitOnSeqWordInit 0 s (0 :: Word)
1933+ skip $ SplitOnSeqYield r next
1934+ else skip $ SplitOnSeqWordLoop fs s wrd1
1935+ Skip s -> skip $ SplitOnSeqWordInit idx s wrd
19371936 Stop -> do
1938− ini <- initial
1937+ fs <- initial
19391938 if idx /= 0
1940− then return $ Skip $ GO_SHORT_PAT_DRAIN idx ini wrd
1939+ then skip $ SplitOnSeqWordDone idx fs wrd
19411940 else do
1942− r <- done ini
1943− return $ Skip $ GO_YIELD r GO_DONE
1944− stepOuter _ (GO_SHORT_PAT_NEXT_WITH fs s x wrd) = do
1945− let wrd1 = mask .&. addToWord wrd x
1946− old = wrd `shiftR` (elemBits * (patLen - 1))
1947− fres <- fstep fs (toEnum $ fromIntegral old)
1948− -- XXX Removing the then branch increases the
1949− -- performance by 150%
1950− if wrd1 .&. mask == pat
1951− then do
1952− r <- done fres
1953− return
1954− $ Skip $ GO_YIELD r $ GO_SHORT_PAT_ACCUM 0 s (0 :: Word)
1955− else return $ Skip $ GO_SHORT_PAT_NEXT fres s wrd1
1956− -- If the fold terminates then the behaviour is same as if
1957− -- the pattern is matched. We should not ignore the
1958− -- currently stored pattern though.
1959− stepOuter gst (GO_SHORT_PAT_NEXT fs st wrd) = do
1941+ r <- done fs
1942+ skip $ SplitOnSeqYield r SplitOnSeqDone
1943+
1944+ stepOuter gst (SplitOnSeqWordLoop fs st wrd) = do
19601945 res <- step (adaptState gst) st
19611946 case res of
1962− Yield x s -> return $ Skip $ GO_SHORT_PAT_NEXT_WITH fs s x wrd
1963− Skip s -> return $ Skip $ GO_SHORT_PAT_NEXT fs s wrd
1964− Stop -> return $ Skip $ GO_SHORT_PAT_DRAIN patLen fs wrd
1965− -- XXX Check if this is correct? Consider n == 1, fromIntegral (mask
1966− -- .&. word) is considering more than one element. We should change the mask
1967− -- to ignore the previous elements.
1968− stepOuter _ (GO_SHORT_PAT_DRAIN 0 fs _) = do
1969− r <- done fs
1970− return $ Skip $ GO_YIELD r GO_DONE
1971− stepOuter _ (GO_SHORT_PAT_DRAIN n fs wrd) = do
1972− let old = wrd `shiftR` (elemBits * (n - 1))
1973− mask1 = (1 `shiftL` (elemBits * (n - 1))) - 1 :: Word
1974− wrd1 = mask1 .&. wrd
1975− fres <- fstep fs (toEnum $ fromIntegral old)
1976− return $ Skip $ GO_SHORT_PAT_DRAIN (n - 1) fres wrd1
1947+ Yield x s -> do
1948+ let wrd1 = wordMask .&. addToWord wrd x
1949+ old = wrd `shiftR` (elemBits * (patLen - 1))
1950+ fs1 <- fstep fs (toEnum $ fromIntegral old)
1951+ if wrd1 .&. wordMask == wordPat
1952+ then do
1953+ r <- done fs1
1954+ let next = SplitOnSeqWordInit 0 s (0 :: Word)
1955+ skip $ SplitOnSeqYield r next
1956+ else skip $ SplitOnSeqWordLoop fs1 s wrd1
1957+ Skip s -> skip $ SplitOnSeqWordLoop fs s wrd
1958+ Stop -> skip $ SplitOnSeqWordDone patLen fs wrd
19771959
19781960 -------------------------------
19791961 -- General Pattern - Karp Rabin
19801962 -------------------------------
19811963
1982− stepOuter gst (GO_KARP_RABIN_ACCUM idx st rb rh) = do
1964+ stepOuter gst (SplitOnSeqKRInit idx st rb rh) = do
19831965 res <- step (adaptState gst) st
19841966 case res of
19851967 Yield x s -> do
@@ -1991,22 +1973,30 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =
19911973 if ringHash == patHash
19921974 then do
19931975 r <- initial >>= done
1994− return
1995− $ Skip $ GO_YIELD r $ GO_KARP_RABIN_ACCUM 0 s rb rh1
1976+ let next = SplitOnSeqKRInit 0 s rb rh1
1977+ -- XXX We need a direct yield here otherwise GHC is not
1978+ -- able to fuse the code. This GHC issue needs to be
1979+ -- investigated. We should know why.
1980+ return $ Yield r next
19961981 else do
1997− ini <- initial
1998− return $ Skip $ GO_KARP_RABIN_NEXT ini s rb rh1 ringHash
1999− else return $ Skip $ GO_KARP_RABIN_ACCUM (idx + 1) s rb rh1
2000− Skip s -> return $ Skip $ GO_KARP_RABIN_ACCUM idx s rb rh
1982+ fs <- initial
1983+ skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash
1984+ else skip $ SplitOnSeqKRInit (idx + 1) s rb rh1
1985+ Skip s -> skip $ SplitOnSeqKRInit idx s rb rh
20011986 Stop -> do
2002− ini <- initial
1987+ fs <- initial
20031988 let rh1 = RB.moveBy (0 - idx) rb (rh :: Ptr a)
20041989 if idx /= 0
2005− then return $ Skip $ GO_KARP_RABIN_DRAIN idx ini rb rh1
1990+ then skip $ SplitOnSeqKRDone idx fs rb rh1
20061991 else do
2007− r <- done ini
2008− return $ Skip $ GO_YIELD r GO_DONE
2009− stepOuter gst (GO_KARP_RABIN_NEXT fs0 st0 rb rh0 cksum0) =
1992+ r <- done fs
1993+ skip $ SplitOnSeqYield r SplitOnSeqDone
1994+
1995+ -- XXX The recursive "go" is more efficient than the state based recursion
1996+ -- code commented out below. Perhaps its more efficient because of
1997+ -- factoring out "rb" outside the loop.
1998+ --
1999+ stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =
20102000 go SPEC fs0 st0 rh0 cksum0
20112001
20122002 where
@@ -2017,25 +2007,50 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =
20172007 Yield x s -> do
20182008 old <- liftIO $ peek rh
20192009 let cksum1 = deltaCksum cksum old x
2020− fres <- fstep fs old
2010+ fs1 <- fstep fs old
20212011 if (cksum1 == patHash)
20222012 then do
2023− r <- done fres
2024− return $ Skip $ GO_YIELD r $ GO_KARP_RABIN_ACCUM 0 s rb rh
2013+ r <- done fs1
2014+ skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
20252015 else do
20262016 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2027− go SPEC fres s rh1 cksum1
2017+ go SPEC fs1 s rh1 cksum1
20282018 Skip s -> go SPEC fs s rh cksum
2029− Stop -> return $ Skip $ GO_KARP_RABIN_DRAIN patLen fs rb rh
2030− stepOuter _ (GO_KARP_RABIN_DRAIN 0 fs _ _) = do
2019+ Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2020+
2021+ -- XXX The following code is 5 times slower compared to the recursive loop
2022+ -- based code above. Need to investigate why. One possibility is that the
2023+ -- go loop above does not thread around the ring buffer (rb). This code may
2024+ -- be causing the state to bloat and getting allocated on each iteration.
2025+ -- We can check the cmm/asm code to confirm. If so a good GHC solution to
2026+ -- such problem is needed. One way to avoid this could be to use unboxed
2027+ -- mutable state?
2028+ {-
2029+ stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do
2030+ res <- step (adaptState gst) st
2031+ case res of
2032+ Yield x s -> do -- skip $ SplitOnSeqKRYield fs s rb rh cksum x
2033+ old <- liftIO $ peek rh
2034+ let cksum1 = deltaCksum cksum old x
2035+ fs1 <- fstep fs old
2036+ if (cksum1 == patHash)
2037+ then do
2038+ r <- done fs1
2039+ skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
2040+ else do
2041+ rh1 <- liftIO (RB.unsafeInsert rb rh x)
2042+ skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1
2043+ Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum
2044+ Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2045+ -}
2046+ stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do
20312047 r <- done fs
2032− return $ Skip $ GO_YIELD r GO_DONE
2033− stepOuter _ (GO_KARP_RABIN_DRAIN n fs rb rh) = do
2048+ skip $ SplitOnSeqYield r SplitOnSeqDone
2049+ stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do
20342050 old <- liftIO $ peek rh
20352051 let rh1 = RB.advance rb rh
2036− fres <- fstep fs old
2037− return $ Skip $ GO_KARP_RABIN_DRAIN (n - 1) fres rb rh1
2038− stepOuter _ _ = undefined
2052+ fs1 <- fstep fs old
2053+ skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1
20392054
20402055data SplitOnState1 s a =
20412056 GO_START1