e433d17 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add takeEndBySeq fold

e433d17Harendra Kumar4 years ago

core/src/Streamly/Internal/Data/Fold.hs+160−15

compare with master

unifiedsplit

@@ -176,7 +176,7 @@ module Streamly.Internal.Data.Fold

176 -- By elements

177 , takeEndBy

178 , takeEndBy_

179− -- , takeEndBySeq

179+ , takeEndBySeq

180 {-

181 , drop

182 , dropWhile

@@ -285,16 +285,19 @@ module Streamly.Internal.Data.Fold

285where

286

287#include "inline.hs"

288+#include "ArrayMacros.h"

288

289

289import Control.Monad.IO.Class (MonadIO(..))

290import Control.Monad.IO.Class (MonadIO(..))

290import Data.Bifunctor (first)

291import Data.Bifunctor (first)

292+import Data.Bits (shiftL, (.|.), (.&.))

291import Data.Either (isLeft, isRight, fromLeft, fromRight)

293import Data.Either (isLeft, isRight, fromLeft, fromRight)

292import Data.Functor.Identity (Identity(..))

294import Data.Functor.Identity (Identity(..))

293import Data.Int (Int64)

295import Data.Int (Int64)

294import Data.IORef (newIORef, readIORef, writeIORef)

296import Data.IORef (newIORef, readIORef, writeIORef)

295import Data.Map.Strict (Map)

297import Data.Map.Strict (Map)

296import Data.Maybe (isJust, fromJust)

298import Data.Maybe (isJust, fromJust)

299+import Data.Word (Word32)

297−import Foreign.Storable (Storable)

300+import Foreign.Storable (Storable(..))

298import Streamly.Internal.Data.IsMap (IsMap(..))

301import Streamly.Internal.Data.IsMap (IsMap(..))

299import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

302import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

300import Streamly.Internal.Data.Unfold.Type (Unfold(..))

303import Streamly.Internal.Data.Unfold.Type (Unfold(..))

@@ -303,10 +306,12 @@ import Streamly.Internal.Data.Stream.Serial (SerialT(..))

303

306

304import qualified Data.IntSet as IntSet

307import qualified Data.IntSet as IntSet

305import qualified Data.Set as Set

308import qualified Data.Set as Set

306−import qualified Streamly.Internal.Data.IsMap as IsMap

307import qualified Prelude

309import qualified Prelude

308import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA

310import 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

309import qualified Streamly.Internal.Data.Pipe.Type as Pipe

313import qualified Streamly.Internal.Data.Pipe.Type as Pipe

314+import qualified Streamly.Internal.Data.Ring.Foreign as Ring

310import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

315import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

311

316

312import Prelude hiding

317import Prelude hiding

@@ -327,6 +332,7 @@ import Streamly.Internal.Data.Fold.Type

327-- >>> import qualified Streamly.Data.Fold as Fold

332-- >>> import qualified Streamly.Data.Fold as Fold

328-- >>> import qualified Streamly.Internal.Data.Fold as Fold

333-- >>> import qualified Streamly.Internal.Data.Fold as Fold

329-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold

334-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold

335+-- >>> import qualified Streamly.Data.Array.Foreign as Array

330-- >>> import qualified Streamly.Internal.Data.Parser as Parser

336-- >>> import qualified Streamly.Internal.Data.Parser as Parser

331-- >>> import Streamly.Internal.Data.Stream.Serial (SerialT(..))

337-- >>> import Streamly.Internal.Data.Stream.Serial (SerialT(..))

332-- >>> import Data.IORef (newIORef, readIORef, writeIORef)

338-- >>> import Data.IORef (newIORef, readIORef, writeIORef)

@@ -1339,21 +1345,160 @@ takeEndBy predicate (Fold fstep finitial fextract) =

1339-- Binary splitting on a separator

1345-- Binary splitting on a separator

1340------------------------------------------------------------------------------

1346------------------------------------------------------------------------------

1341

1347

1342−{-

1348+data SplitOnSeqState acc a rb rh w ck =

1343−-- | Find the first occurrence of the specified sequence in the input stream

1349+ SplitOnSeqEmpty !acc

1344−-- and break the input stream into two parts, the first part consisting of the

1350+ | SplitOnSeqSingle !acc !a

1345−-- stream before the sequence and the second part consisting of the sequence

1351+ | SplitOnSeqWord !acc !Int !w

1346−-- and the rest of the stream.

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.

1347--

1361--

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"

1349--

1366--

1350−-- >>> breakOn_ "dear" "Hello dear world!"

1367+-- >>> Stream.toList $ Stream.foldMany f s

1351−-- > ("Hello ","dear world!")

1368+-- ["hello there",". How are"," you?"]

1352--

1369--

1353−{-# INLINE breakOn #-}

1370+-- /Pre-release/

1354−breakOn :: Monad m => Array a -> Fold m a b -> Fold m a c -> Fold m a (b,c)

1371+{-# INLINE takeEndBySeq #-}

1355−breakOn pat f m = undefined

1372+takeEndBySeq :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a) =>

1356−-}

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

1357

1502

1358------------------------------------------------------------------------------

1503------------------------------------------------------------------------------

1359-- Distributing

1504-- Distributing