git.iko.soy/streamly
on master help sign in
Add takeEndBySeq_ fold
d5b3e37Harendra Kumar
core/src/Streamly/Internal/Data/Fold.hs+159−1
compare with master
unifiedsplit
@@ -177,6 +177,7 @@ module Streamly.Internal.Data.Fold
177 , takeEndBy
177 , takeEndBy
178 , takeEndBy_
178 , takeEndBy_
179 , takeEndBySeq
179 , takeEndBySeq
180+ , takeEndBySeq_
180 {-
181 {-
181 , drop
182 , drop
182 , dropWhile
183 , dropWhile
@@ -289,7 +290,7 @@ where
289
290
290import Control.Monad.IO.Class (MonadIO(..))
291import Control.Monad.IO.Class (MonadIO(..))
291import Data.Bifunctor (first)
292import Data.Bifunctor (first)
292−import Data.Bits (shiftL, (.|.), (.&.))
293+import Data.Bits (shiftL, shiftR, (.|.), (.&.))
293import Data.Either (isLeft, isRight, fromLeft, fromRight)
294import Data.Either (isLeft, isRight, fromLeft, fromRight)
294import Data.Functor.Identity (Identity(..))
295import Data.Functor.Identity (Identity(..))
295import Data.Int (Int64)
296import Data.Int (Int64)
@@ -1500,6 +1501,163 @@ takeEndBySeq patArr (Fold fstep finitial fextract) =
1500 SplitOnSeqKRLoop s _ _ _ -> s
1501 SplitOnSeqKRLoop s _ _ _ -> s
1501 in fextract st
1502 in fextract st
1502
1503
1504+-- | Like 'takeEndBySeq' but discards the matched sequence.
1505+--
1506+-- /Pre-release/
1507+--
1508+{-# INLINE takeEndBySeq_ #-}
1509+takeEndBySeq_ :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a) =>
1510+ Array.Array a
1511+ -> Fold m a b
1512+ -> Fold m a b
1513+takeEndBySeq_ patArr (Fold fstep finitial fextract) =
1514+ Fold step initial extract
1515+
1516+ where
1517+
1518+ patLen = Array.length patArr
1519+
1520+ initial = do
1521+ res <- finitial
1522+ case res of
1523+ Partial acc
1524+ | patLen == 0 ->
1525+ -- XXX Should we match nothing or everything on empty
1526+ -- pattern?
1527+ -- Done <$> fextract acc
1528+ return $ Partial $ SplitOnSeqEmpty acc
1529+ | patLen == 1 -> do
1530+ pat <- liftIO $ Array.unsafeIndexIO 0 patArr
1531+ return $ Partial $ SplitOnSeqSingle acc pat
1532+ -- XXX Need to add tests for this case
1533+ | SIZE_OF(a) * patLen <= sizeOf (undefined :: Word) ->
1534+ return $ Partial $ SplitOnSeqWord acc 0 0
1535+ | otherwise -> do
1536+ (rb, rhead) <- liftIO $ Ring.new patLen
1537+ return $ Partial $ SplitOnSeqKR acc 0 rb rhead
1538+ Done b -> return $ Done b
1539+
1540+ -- Word pattern related
1541+ maxIndex = patLen - 1
1542+
1543+ elemBits = SIZE_OF(a) * 8
1544+
1545+ wordMask :: Word
1546+ wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1547+
1548+ elemMask :: Word
1549+ elemMask = (1 `shiftL` elemBits) - 1
1550+
1551+ wordPat :: Word
1552+ wordPat = wordMask .&. Array.foldl' addToWord 0 patArr
1553+
1554+ addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1555+
1556+ -- For Rabin-Karp search
1557+ k = 2891336453 :: Word32
1558+ coeff = k ^ patLen
1559+
1560+ addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1561+
1562+ deltaCksum cksum old new =
1563+ addCksum cksum new - coeff * fromIntegral (fromEnum old)
1564+
1565+ -- XXX shall we use a random starting hash or 1 instead of 0?
1566+ -- XXX Need to keep this cached across fold calls in foldMany
1567+ -- XXX We may need refold to inject the cached state instead of
1568+ -- initializing the state every time.
1569+ -- XXX Allocation of ring buffer should also be done once
1570+ patHash = Array.foldl' addCksum 0 patArr
1571+
1572+ step (SplitOnSeqEmpty s) x = do
1573+ res <- fstep s x
1574+ case res of
1575+ Partial s1 -> return $ Partial $ SplitOnSeqEmpty s1
1576+ Done b -> return $ Done b
1577+ step (SplitOnSeqSingle s pat) x = do
1578+ if pat /= x
1579+ then do
1580+ res <- fstep s x
1581+ case res of
1582+ Partial s1 -> return $ Partial $ SplitOnSeqSingle s1 pat
1583+ Done b -> return $ Done b
1584+ else Done <$> fextract s
1585+ step (SplitOnSeqWord s idx wrd) x = do
1586+ let wrd1 = addToWord wrd x
1587+ if idx == maxIndex
1588+ then do
1589+ if wrd1 .&. wordMask == wordPat
1590+ then Done <$> fextract s
1591+ else return $ Partial $ SplitOnSeqWordLoop s wrd1
1592+ else return $ Partial $ SplitOnSeqWord s (idx + 1) wrd1
1593+ step (SplitOnSeqWordLoop s wrd) x = do
1594+ let wrd1 = addToWord wrd x
1595+ old = (wordMask .&. wrd)
1596+ `shiftR` (elemBits * (patLen - 1))
1597+ res <- fstep s (toEnum $ fromIntegral old)
1598+ case res of
1599+ Partial s1
1600+ | wrd1 .&. wordMask == wordPat ->
1601+ Done <$> fextract s1
1602+ | otherwise ->
1603+ return $ Partial $ SplitOnSeqWordLoop s1 wrd1
1604+ Done b -> return $ Done b
1605+ step (SplitOnSeqKR s idx rb rh) x = do
1606+ rh1 <- liftIO $ Ring.unsafeInsert rb rh x
1607+ if idx == maxIndex
1608+ then do
1609+ let fld = Ring.unsafeFoldRing (Ring.ringBound rb)
1610+ let !ringHash = fld addCksum 0 rb
1611+ if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
1612+ then Done <$> fextract s
1613+ else return $ Partial $ SplitOnSeqKRLoop s ringHash rb rh1
1614+ else return $ Partial $ SplitOnSeqKR s (idx + 1) rb rh1
1615+ step (SplitOnSeqKRLoop s cksum rb rh) x = do
1616+ old <- liftIO $ peek rh
1617+ res <- fstep s old
1618+ case res of
1619+ Partial s1 -> do
1620+ rh1 <- liftIO $ Ring.unsafeInsert rb rh x
1621+ let ringHash = deltaCksum cksum old x
1622+ if ringHash == patHash && Ring.unsafeEqArray rb rh1 patArr
1623+ then Done <$> fextract s1
1624+ else return $ Partial $ SplitOnSeqKRLoop s1 ringHash rb rh1
1625+ Done b -> return $ Done b
1626+
1627+ -- XXX extract should return backtrack count as well. If the fold
1628+ -- terminates early inside extract, we may still have buffered data
1629+ -- remaining which will be lost if we do not communicate that to the
1630+ -- driver.
1631+ extract state = do
1632+ let consumeWord s n wrd = do
1633+ if n == 0
1634+ then fextract s
1635+ else do
1636+ let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1637+ r <- fstep s (toEnum $ fromIntegral old)
1638+ case r of
1639+ Partial s1 -> consumeWord s1 (n - 1) wrd
1640+ Done b -> return b
1641+
1642+ let consumeRing s n rb rh =
1643+ if n == 0
1644+ then fextract s
1645+ else do
1646+ old <- liftIO $ peek rh
1647+ let rh1 = Ring.advance rb rh
1648+ r <- fstep s old
1649+ case r of
1650+ Partial s1 -> consumeRing s1 (n - 1) rb rh1
1651+ Done b -> return b
1652+
1653+ case state of
1654+ SplitOnSeqEmpty s -> fextract s
1655+ SplitOnSeqSingle s _ -> fextract s
1656+ SplitOnSeqWord s idx wrd -> consumeWord s idx wrd
1657+ SplitOnSeqWordLoop s wrd -> consumeWord s patLen wrd
1658+ SplitOnSeqKR s idx rb _ -> consumeRing s idx rb (Ring.startOf rb)
1659+ SplitOnSeqKRLoop s _ rb rh -> consumeRing s patLen rb rh
1660+
1503------------------------------------------------------------------------------
1661------------------------------------------------------------------------------
1504-- Distributing
1662-- Distributing
1505------------------------------------------------------------------------------
1663------------------------------------------------------------------------------