@@ -176,7 +176,7 @@ module Streamly.Internal.Data.Fold
176176 -- By elements
177177 , takeEndBy
178178 , takeEndBy_
179− -- , takeEndBySeq
179+ , takeEndBySeq
180180 {-
181181 , drop
182182 , dropWhile
@@ -285,16 +285,19 @@ module Streamly.Internal.Data.Fold
285285where
286286
287287#include "inline.hs"
288+#include "ArrayMacros.h"
288289
289290import Control.Monad.IO.Class (MonadIO(..))
290291import Data.Bifunctor (first)
292+import Data.Bits (shiftL, (.|.), (.&.))
291293import Data.Either (isLeft, isRight, fromLeft, fromRight)
292294import Data.Functor.Identity (Identity(..))
293295import Data.Int (Int64)
294296import Data.IORef (newIORef, readIORef, writeIORef)
295297import Data.Map.Strict (Map)
296298import Data.Maybe (isJust, fromJust)
297−import Foreign.Storable (Storable)
299+import Data.Word (Word32)
300+import Foreign.Storable (Storable(..))
298301import Streamly.Internal.Data.IsMap (IsMap(..))
299302import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
300303import Streamly.Internal.Data.Unfold.Type (Unfold(..))
@@ -303,10 +306,12 @@ import Streamly.Internal.Data.Stream.Serial (SerialT(..))
303306
304307import qualified Data.IntSet as IntSet
305308import qualified Data.Set as Set
306−import qualified Streamly.Internal.Data.IsMap as IsMap
307309import qualified Prelude
308310import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
311+import qualified Streamly.Internal.Data.Array.Foreign.Type as Array
312+import qualified Streamly.Internal.Data.IsMap as IsMap
309313import qualified Streamly.Internal.Data.Pipe.Type as Pipe
314+import qualified Streamly.Internal.Data.Ring.Foreign as Ring
310315import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD
311316
312317import Prelude hiding
@@ -327,6 +332,7 @@ import Streamly.Internal.Data.Fold.Type
327332-- >>> import qualified Streamly.Data.Fold as Fold
328333-- >>> import qualified Streamly.Internal.Data.Fold as Fold
329334-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold
335+-- >>> import qualified Streamly.Data.Array.Foreign as Array
330336-- >>> import qualified Streamly.Internal.Data.Parser as Parser
331337-- >>> import Streamly.Internal.Data.Stream.Serial (SerialT(..))
332338-- >>> import Data.IORef (newIORef, readIORef, writeIORef)
@@ -1339,21 +1345,160 @@ takeEndBy predicate (Fold fstep finitial fextract) =
13391345-- Binary splitting on a separator
13401346------------------------------------------------------------------------------
13411347
1342−{-
1343−-- | Find the first occurrence of the specified sequence in the input stream
1344−-- and break the input stream into two parts, the first part consisting of the
1345−-- stream before the sequence and the second part consisting of the sequence
1346−-- and the rest of the stream.
1348+data SplitOnSeqState acc a rb rh w ck =
1349+ SplitOnSeqEmpty !acc
1350+ | SplitOnSeqSingle !acc !a
1351+ | SplitOnSeqWord !acc !Int !w
1352+ | SplitOnSeqWordLoop !acc !w
1353+ | SplitOnSeqKR !acc !Int !rb !rh
1354+ | SplitOnSeqKRLoop !acc !ck !rb !rh
1355+
1356+-- XXX Need to add tests for takeEndBySeq, we have tests for takeEndBySeq_ .
1357+
1358+-- | Continue taking the input until the input sequence matches the supplied
1359+-- sequence, taking the supplied sequence as well. If the pattern is empty this
1360+-- acts as an identity fold.
13471361--
1348−-- > let breakOn_ pat xs = S.fold (S.breakOn pat FL.toList FL.toList) $ S.fromList xs
1362+-- >>> s = Stream.fromList "hello there. How are you?"
1363+-- >>> f = Fold.takeEndBySeq (Array.fromList "re") Fold.toList
1364+-- >>> Stream.fold f s
1365+-- "hello there"
13491366--
1350−-- >>> breakOn_ "dear" "Hello dear world!"
1351−-- > ("Hello ","dear world!")
1367+-- >>> Stream.toList $ Stream.foldMany f s
1368+-- ["hello there",". How are"," you?"]
13521369--
1353−{-# INLINE breakOn #-}
1354−breakOn :: Monad m => Array a -> Fold m a b -> Fold m a c -> Fold m a (b,c)
1355−breakOn pat f m = undefined
1356−-}
1370+-- /Pre-release/
1371+{-# INLINE takeEndBySeq #-}
1372+takeEndBySeq :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a) =>
1373+ Array.Array a
1374+ -> Fold m a b
1375+ -> Fold m a b
1376+takeEndBySeq patArr (Fold fstep finitial fextract) =
1377+ Fold step initial extract
1378+
1379+ where
1380+
1381+ patLen = Array.length patArr
1382+
1383+ initial = do
1384+ res <- finitial
1385+ case res of
1386+ Partial acc
1387+ | patLen == 0 ->
1388+ -- XXX Should we match nothing or everything on empty
1389+ -- pattern?
1390+ -- Done <$> fextract acc
1391+ return $ Partial $ SplitOnSeqEmpty acc
1392+ | patLen == 1 -> do
1393+ pat <- liftIO $ Array.unsafeIndexIO 0 patArr
1394+ return $ Partial $ SplitOnSeqSingle acc pat
1395+ | SIZE_OF(a) * patLen <= sizeOf (undefined :: Word) ->
1396+ return $ Partial $ SplitOnSeqWord acc 0 0
1397+ | otherwise -> do
1398+ (rb, rhead) <- liftIO $ Ring.new patLen
1399+ return $ Partial $ SplitOnSeqKR acc 0 rb rhead
1400+ Done b -> return $ Done b
1401+
1402+ -- Word pattern related
1403+ maxIndex = patLen - 1
1404+
1405+ elemBits = SIZE_OF(a) * 8
1406+
1407+ wordMask :: Word
1408+ wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1409+
1410+ wordPat :: Word
1411+ wordPat = wordMask .&. Array.foldl' addToWord 0 patArr
1412+
1413+ addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1414+
1415+ -- For Rabin-Karp search
1416+ k = 2891336453 :: Word32
1417+ coeff = k ^ patLen
1418+
1419+ addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1420+
1421+ deltaCksum cksum old new =
1422+ addCksum cksum new - coeff * fromIntegral (fromEnum old)
1423+
1424+ -- XXX shall we use a random starting hash or 1 instead of 0?
1425+ -- XXX Need to keep this cached across fold calls in foldmany
1426+ -- XXX We may need refold to inject the cached state instead of
1427+ -- initializing the state every time.
1428+ -- XXX Allocation of ring buffer should also be done once
1429+ patHash = Array.foldl' addCksum 0 patArr
1430+
1431+ step (SplitOnSeqEmpty s) x = do
1432+ res <- fstep s x
1433+ case res of
1434+ Partial s1 -> return $ Partial $ SplitOnSeqEmpty s1
1435+ Done b -> return $ Done b
1436+ step (SplitOnSeqSingle s pat) x = do
1437+ res <- fstep s x
1438+ case res of
1439+ Partial s1
1440+ | pat /= x -> return $ Partial $ SplitOnSeqSingle s1 pat
1441+ | otherwise -> Done <$> fextract s1
1442+ Done b -> return $ Done b
1443+ step (SplitOnSeqWord s idx wrd) x = do
1444+ res <- fstep s x
1445+ let wrd1 = addToWord wrd x
1446+ case res of
1447+ Partial s1
1448+ | idx == maxIndex -> do
1449+ if wrd1 .&. wordMask == wordPat
1450+ then Done <$> fextract s1
1451+ else return $ Partial $ SplitOnSeqWordLoop s1 wrd1
1452+ | otherwise ->
1453+ return $ Partial $ SplitOnSeqWord s1 (idx + 1) wrd1
1454+ Done b -> return $ Done b
1455+ step (SplitOnSeqWordLoop s wrd) x = do
1456+ res <- fstep s x
1457+ let wrd1 = addToWord wrd x
1458+ case res of
1459+ Partial s1
1460+ | wrd1 .&. wordMask == wordPat ->
1461+ Done <$> fextract s1
1462+ | otherwise ->
1463+ return $ Partial $ SplitOnSeqWordLoop s1 wrd1
1464+ Done b -> return $ Done b
1465+ step (SplitOnSeqKR s idx rb rh) x = do
1466+ res <- fstep s x
1467+ case res of
1468+ Partial s1 -> do
1469+ rh1 <- liftIO $ Ring.unsafeInsert rb rh x
1470+ if idx == maxIndex
1471+ then do
1472+ let fld = Ring.unsafeFoldRing (Ring.ringBound rb)
1473+ let !ringHash = fld addCksum 0 rb
1474+ if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
1475+ then Done <$> fextract s1
1476+ else return $ Partial $ SplitOnSeqKRLoop s1 ringHash rb rh1
1477+ else
1478+ return $ Partial $ SplitOnSeqKR s1 (idx + 1) rb rh1
1479+ Done b -> return $ Done b
1480+ step (SplitOnSeqKRLoop s cksum rb rh) x = do
1481+ res <- fstep s x
1482+ case res of
1483+ Partial s1 -> do
1484+ old <- liftIO $ peek rh
1485+ rh1 <- liftIO $ Ring.unsafeInsert rb rh x
1486+ let ringHash = deltaCksum cksum old x
1487+ if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
1488+ then Done <$> fextract s1
1489+ else return $ Partial $ SplitOnSeqKRLoop s1 ringHash rb rh1
1490+ Done b -> return $ Done b
1491+
1492+ extract state =
1493+ let st =
1494+ case state of
1495+ SplitOnSeqEmpty s -> s
1496+ SplitOnSeqSingle s _ -> s
1497+ SplitOnSeqWord s _ _ -> s
1498+ SplitOnSeqWordLoop s _ -> s
1499+ SplitOnSeqKR s _ _ _ -> s
1500+ SplitOnSeqKRLoop s _ _ _ -> s
1501+ in fextract st
13571502
13581503------------------------------------------------------------------------------
13591504-- Distributing