2ca6d2a · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor splitOn and make it fuse

2ca6d2aHarendra Kumar6 years ago

src/Streamly/Internal/Data/Stream/IsStream.hs+1−1

compare with master

unifiedsplit

@@ -4152,7 +4152,7 @@ splitWithSuffix predicate f m =

4152splitOnSeq

4153 :: (IsStream t, MonadIO m, Storable a, Enum a, Eq a)

4154 => Array a -> Fold m a b -> t m a -> t m b

4155−splitOnSeq patt f m = D.fromStreamD $ D.splitOn patt f (D.toStreamD m)

4155+splitOnSeq patt f m = D.fromStreamD $ D.splitOnSeq patt f (D.toStreamD m)

4156

4157{-

4158-- This can be implemented easily using Rabin Karp

src/Streamly/Internal/Data/Stream/StreamD.hs+164−149

@@ -168,7 +168,7 @@ module Streamly.Internal.Data.Stream.StreamD

168 , wordsBy

169 , splitSuffixBy'

170

171− , splitOn

171+ , splitOnSeq

172 , splitSuffixOn

173

174 , splitInnerBy

@@ -1786,32 +1786,32 @@ data SplitOptions = SplitOptions

1786 }

1787-}

1788

1789−{-# ANN type SplitOnState Fuse #-}

1789+{-# ANN type SplitOnSeqState Fuse #-}

1790−data SplitOnState fs s a b w rb rh ck =

1790+data SplitOnSeqState rb rh ck w fs s b x =

1791− GO_START

1791+ SplitOnSeqInit

1792− | GO_EMPTY_PAT s

1792+ | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)

1793− | GO_EMPTY_PAT_WITH s a

1793+ | SplitOnSeqDone

1794− | GO_SINGLE_PAT_NEXT !fs s a

1794+

1795− | GO_SINGLE_PAT_WITH !fs s a a

1795+ | SplitOnSeqEmpty s

1796− | GO_SHORT_PAT_ACCUM Int s !w

1796+ | SplitOnSeqSingle !fs s x

1797− | GO_SHORT_PAT_NEXT !fs s !w

1797+

1798− | GO_SHORT_PAT_NEXT_WITH !fs s a !w

1798+ | SplitOnSeqWordInit Int s !w

1799− | GO_SHORT_PAT_DRAIN Int !fs !w

1799+ | SplitOnSeqWordLoop !fs s !w

1800− | GO_KARP_RABIN_ACCUM Int s rb !rh

1800+ | SplitOnSeqWordDone Int !fs !w

1801− | GO_KARP_RABIN_NEXT !fs s rb !rh !ck

1801+

1802− | GO_KARP_RABIN_DRAIN Int !fs rb !rh

1802+ | SplitOnSeqKRInit Int s rb !rh

1803− | GO_YIELD b (SplitOnState fs s a b w rb rh ck)

1803+ | SplitOnSeqKRLoop fs s rb !rh !ck

1804− | GO_DONE

1804+ | SplitOnSeqKRDone Int !fs rb !rh

1805−

1805+

1806−{-# INLINE_NORMAL splitOn #-}

1806+{-# INLINE_NORMAL splitOnSeq #-}

1807−splitOn

1807+splitOnSeq

1808 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)

1809 => Array a

1810 -> Fold m a b

1811 -> Stream m a

1812 -> Stream m b

1813−splitOn patArr (Fold fstep initial done) (Stream step state) =

1813+splitOnSeq patArr (Fold fstep initial done) (Stream step state) =

1814− Stream stepOuter GO_START

1814+ Stream stepOuter SplitOnSeqInit

1815

1816 where

1817

@@ -1819,26 +1819,19 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =

1819 maxIndex = patLen - 1

1820 elemBits = sizeOf (undefined :: a) * 8

1821

1822− -- XXX I initially put this in the state but I'd like to keep the state as

1822+ -- For word pattern case

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

1823+ wordMask :: Word

1828− mask = (1 `shiftL` (elemBits * patLen)) - 1

1824+ wordMask = (1 `shiftL` (elemBits * patLen)) - 1

1829−

1825+

1830− -- XXX I initially put this in the state but I'd like to keep the state as

1826+ elemMask :: Word

1831− -- simple as possible to allow more fusion. Removing this from the state did

1827+ elemMask = (1 `shiftL` elemBits) - 1

1832− -- not change anything. I'll introduce them back once I figure out what's

1828+

1833− -- currently wrong.

1834− pat :: Word

1829+ wordPat :: Word

1835− -- XXX You dont need .&. here?

1836− pat = mask .&. A.foldl' addToWord 0 patArr

1830+ wordPat = wordMask .&. A.foldl' addToWord 0 patArr

1837−

1831+

1838− -- Used only if the pattern is short

1839 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)

1832 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)

1840

1833

1841− -- Used in Rabin-Karp

1834+ -- For Rabin-Karp search

1842 k = 2891336453 :: Word32

1835 k = 2891336453 :: Word32

1843 coeff = k ^ patLen

1836 coeff = k ^ patLen

1844

1837

@@ -1850,136 +1843,125 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =

1850 -- XXX shall we use a random starting hash or 1 instead of 0?

1843 -- XXX shall we use a random starting hash or 1 instead of 0?

1851 patHash = A.foldl' addCksum 0 patArr

1844 patHash = A.foldl' addCksum 0 patArr

1852

1845

1846+ skip = return . Skip

1847+

1853 {-# INLINE_LATE stepOuter #-}

1848 {-# INLINE_LATE stepOuter #-}

1854− stepOuter _ GO_START =

1849+ stepOuter _ SplitOnSeqInit =

1855 if patLen == 0

1850 if patLen == 0

1856− then return $ Skip $ GO_EMPTY_PAT state

1851+ then return $ Skip $ SplitOnSeqEmpty state

1857 else if patLen == 1

1852 else if patLen == 1

1858 then do

1853 then do

1859 acc <- initial

1854 acc <- initial

1860− pat_ <- liftIO $ A.unsafeIndexIO patArr 0

1855+ let !pat = A.unsafeIndex patArr 0

1861− return $ Skip $ GO_SINGLE_PAT_NEXT acc state pat_

1856+ return $ Skip $ SplitOnSeqSingle acc state pat

1862 else if sizeOf (undefined :: a) * patLen

1857 else if sizeOf (undefined :: a) * patLen

1863 <= sizeOf (undefined :: Word)

1858 <= sizeOf (undefined :: Word)

1864− then return $ Skip $ GO_SHORT_PAT_ACCUM 0 state (0 :: Word)

1859+ then return $ Skip $ SplitOnSeqWordInit 0 state (0 :: Word)

1865 else do

1860 else do

1866 (rb, rhead) <- liftIO $ RB.new patLen

1861 (rb, rhead) <- liftIO $ RB.new patLen

1867− return $ Skip $ GO_KARP_RABIN_ACCUM 0 state rb rhead

1862+ skip $ SplitOnSeqKRInit 0 state rb rhead

1863+

1868− stepOuter _ (GO_YIELD x ns) = return $ Yield x ns

1864+ stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next

1869− stepOuter _ GO_DONE = return Stop

1870

1865

1871 ---------------------------

1866 ---------------------------

1872 -- Empty pattern

1867 -- Empty pattern

1873 ---------------------------

1868 ---------------------------

1874

1869

1875− stepOuter _ (GO_EMPTY_PAT_WITH s x) = do

1870+ stepOuter gst (SplitOnSeqEmpty st) = 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

1881 res <- step (adaptState gst) st

1871 res <- step (adaptState gst) st

1882− return

1883− $ case res of

1872+ case res of

1884− Yield x s -> Skip $ GO_EMPTY_PAT_WITH s x

1873+ Yield x s -> do

1885− Skip s -> Skip $ GO_EMPTY_PAT s

1874+ acc <- initial

1886− Stop -> Stop

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

1887

1886

1888 -----------------

1887 -----------------

1889 -- Single Pattern

1888 -- Single Pattern

1890 -----------------

1889 -----------------

1891

1890

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

1907 res <- step (adaptState gst) st

1892 res <- step (adaptState gst) st

1908 case res of

1893 case res of

1909− Yield x s -> return $ Skip $ GO_SINGLE_PAT_WITH fs s pat_ x

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)

1910− Skip s -> return $ Skip $ GO_SINGLE_PAT_NEXT fs s pat_

1903+ Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat

1911 Stop -> do

1904 Stop -> do

1912 r <- done fs

1905 r <- done fs

1913− return $ Skip $ GO_YIELD r GO_DONE

1906+ return $ Skip $ SplitOnSeqYield r SplitOnSeqDone

1914

1907

1915 ---------------------------

1908 ---------------------------

1916 -- Short Pattern - Shift Or

1909 -- Short Pattern - Shift Or

1917 ---------------------------

1910 ---------------------------

1918

1911

1919− -- XXX mask can be omitted in most of the places?

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 (GO_SHORT_PAT_ACCUM idx st wrd) = do

1920+ stepOuter gst (SplitOnSeqWordInit idx st wrd) = do

1921 res <- step (adaptState gst) st

1922 case res of

1923 Yield x s -> do

1924 let wrd1 = addToWord wrd x

1925− if idx == maxIndex

1925+ if idx /= maxIndex

1926− then if wrd1 .&. mask == pat

1926+ then skip $ SplitOnSeqWordInit (idx + 1) s wrd1

1927+ else do

1928+ fs <- initial

1929+ if wrd1 .&. wordMask == wordPat

1927− then do

1930+ then do

1928− r <- initial >>= done

1931+ r <- done fs

1929− return

1932+ let next = SplitOnSeqWordInit 0 s (0 :: Word)

1930− $ Skip

1933+ skip $ SplitOnSeqYield r next

1931− $ GO_YIELD r $ GO_SHORT_PAT_ACCUM 0 s (0 :: Word)

1934+ else skip $ SplitOnSeqWordLoop fs s wrd1

1932− else do

1935+ Skip s -> skip $ SplitOnSeqWordInit idx s wrd

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

1937 Stop -> do

1936 Stop -> do

1938− ini <- initial

1937+ fs <- initial

1939 if idx /= 0

1938 if idx /= 0

1940− then return $ Skip $ GO_SHORT_PAT_DRAIN idx ini wrd

1939+ then skip $ SplitOnSeqWordDone idx fs wrd

1941 else do

1940 else do

1942− r <- done ini

1941+ r <- done fs

1943− return $ Skip $ GO_YIELD r GO_DONE

1942+ skip $ SplitOnSeqYield r SplitOnSeqDone

1944− stepOuter _ (GO_SHORT_PAT_NEXT_WITH fs s x wrd) = do

1943+

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

1944+ stepOuter gst (SplitOnSeqWordLoop fs st wrd) = do

1960 res <- step (adaptState gst) st

1945 res <- step (adaptState gst) st

1961 case res of

1946 case res of

1962− Yield x s -> return $ Skip $ GO_SHORT_PAT_NEXT_WITH fs s x wrd

1947+ Yield x s -> do

1963− Skip s -> return $ Skip $ GO_SHORT_PAT_NEXT fs s wrd

1948+ let wrd1 = wordMask .&. addToWord wrd x

1964− Stop -> return $ Skip $ GO_SHORT_PAT_DRAIN patLen fs wrd

1949+ old = wrd `shiftR` (elemBits * (patLen - 1))

1965− -- XXX Check if this is correct? Consider n == 1, fromIntegral (mask

1950+ fs1 <- fstep fs (toEnum $ fromIntegral old)

1966− -- .&. word) is considering more than one element. We should change the mask

1951+ if wrd1 .&. wordMask == wordPat

1967− -- to ignore the previous elements.

1952+ then do

1968− stepOuter _ (GO_SHORT_PAT_DRAIN 0 fs _) = do

1969− r <- done fs

1953+ r <- done fs1

1970− return $ Skip $ GO_YIELD r GO_DONE

1954+ let next = SplitOnSeqWordInit 0 s (0 :: Word)

1971− stepOuter _ (GO_SHORT_PAT_DRAIN n fs wrd) = do

1955+ skip $ SplitOnSeqYield r next

1972− let old = wrd `shiftR` (elemBits * (n - 1))

1956+ else skip $ SplitOnSeqWordLoop fs1 s wrd1

1973− mask1 = (1 `shiftL` (elemBits * (n - 1))) - 1 :: Word

1957+ Skip s -> skip $ SplitOnSeqWordLoop fs s wrd

1974− wrd1 = mask1 .&. wrd

1958+ Stop -> skip $ SplitOnSeqWordDone patLen fs wrd

1975− fres <- fstep fs (toEnum $ fromIntegral old)

1976− return $ Skip $ GO_SHORT_PAT_DRAIN (n - 1) fres wrd1

1977

1959

1978 -------------------------------

1960 -------------------------------

1979 -- General Pattern - Karp Rabin

1961 -- General Pattern - Karp Rabin

1980 -------------------------------

1962 -------------------------------

1981

1963

1982− stepOuter gst (GO_KARP_RABIN_ACCUM idx st rb rh) = do

1964+ stepOuter gst (SplitOnSeqKRInit idx st rb rh) = do

1983 res <- step (adaptState gst) st

1965 res <- step (adaptState gst) st

1984 case res of

1966 case res of

1985 Yield x s -> do

1967 Yield x s -> do

@@ -1991,22 +1973,30 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =

1991 if ringHash == patHash

1973 if ringHash == patHash

1992 then do

1974 then do

1993 r <- initial >>= done

1975 r <- initial >>= done

1994− return

1976+ let next = SplitOnSeqKRInit 0 s rb rh1

1995− $ Skip $ GO_YIELD r $ GO_KARP_RABIN_ACCUM 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

1996 else do

1981 else do

1997− ini <- initial

1982+ fs <- initial

1998− return $ Skip $ GO_KARP_RABIN_NEXT ini s rb rh1 ringHash

1983+ skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash

1999− else return $ Skip $ GO_KARP_RABIN_ACCUM (idx + 1) s rb rh1

1984+ else skip $ SplitOnSeqKRInit (idx + 1) s rb rh1

2000− Skip s -> return $ Skip $ GO_KARP_RABIN_ACCUM idx s rb rh

1985+ Skip s -> skip $ SplitOnSeqKRInit idx s rb rh

2001 Stop -> do

1986 Stop -> do

2002− ini <- initial

1987+ fs <- initial

2003 let rh1 = RB.moveBy (0 - idx) rb (rh :: Ptr a)

1988 let rh1 = RB.moveBy (0 - idx) rb (rh :: Ptr a)

2004 if idx /= 0

1989 if idx /= 0

2005− then return $ Skip $ GO_KARP_RABIN_DRAIN idx ini rb rh1

1990+ then skip $ SplitOnSeqKRDone idx fs rb rh1

2006 else do

1991 else do

2007− r <- done ini

1992+ r <- done fs

2008− return $ Skip $ GO_YIELD r GO_DONE

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+ --

2009− stepOuter gst (GO_KARP_RABIN_NEXT fs0 st0 rb rh0 cksum0) =

1999+ stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =

2010 go SPEC fs0 st0 rh0 cksum0

2000 go SPEC fs0 st0 rh0 cksum0

2011

2001

2012 where

2002 where

@@ -2017,25 +2007,50 @@ splitOn patArr (Fold fstep initial done) (Stream step state) =

2017 Yield x s -> do

2007 Yield x s -> do

2018 old <- liftIO $ peek rh

2008 old <- liftIO $ peek rh

2019 let cksum1 = deltaCksum cksum old x

2009 let cksum1 = deltaCksum cksum old x

2020− fres <- fstep fs old

2010+ fs1 <- fstep fs old

2021 if (cksum1 == patHash)

2011 if (cksum1 == patHash)

2022 then do

2012 then do

2023− r <- done fres

2013+ r <- done fs1

2024− return $ Skip $ GO_YIELD r $ GO_KARP_RABIN_ACCUM 0 s rb rh

2014+ skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh

2025 else do

2015 else do

2026 rh1 <- liftIO (RB.unsafeInsert rb rh x)

2016 rh1 <- liftIO (RB.unsafeInsert rb rh x)

2027− go SPEC fres s rh1 cksum1

2017+ go SPEC fs1 s rh1 cksum1

2028 Skip s -> go SPEC fs s rh cksum

2018 Skip s -> go SPEC fs s rh cksum

2029− Stop -> return $ Skip $ GO_KARP_RABIN_DRAIN patLen fs rb rh

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+ -}

2030− stepOuter _ (GO_KARP_RABIN_DRAIN 0 fs _ _) = do

2046+ stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do

2031 r <- done fs

2047 r <- done fs

2032− return $ Skip $ GO_YIELD r GO_DONE

2048+ skip $ SplitOnSeqYield r SplitOnSeqDone

2033− stepOuter _ (GO_KARP_RABIN_DRAIN n fs rb rh) = do

2049+ stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do

2034 old <- liftIO $ peek rh

2050 old <- liftIO $ peek rh

2035 let rh1 = RB.advance rb rh

2051 let rh1 = RB.advance rb rh

2036− fres <- fstep fs old

2052+ fs1 <- fstep fs old

2037− return $ Skip $ GO_KARP_RABIN_DRAIN (n - 1) fres rb rh1

2053+ skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1

2038− stepOuter _ _ = undefined

2039

2054

2040data SplitOnState1 s a =

2055data SplitOnState1 s a =

2041 GO_START1

2056 GO_START1

src/Streamly/Memory/Ring.hs+2−2

@@ -54,8 +54,8 @@ import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

54-- structure. We should not leak out references to it for immutable use.

55--

56data Ring a = Ring

57− { ringStart :: !(ForeignPtr a) -- first address

57+ { ringStart :: {-# UNPACK #-} !(ForeignPtr a) -- first address

58− , ringBound :: !(Ptr a) -- first address beyond allocated memory

58+ , ringBound :: {-# UNPACK #-} !(Ptr a) -- first address beyond allocated memory

59 }

61-- | Create a new ringbuffer and return the ring buffer and the ringHead.