mirror of
https://github.com/composewell/streamly.git
synced 2024-09-11 19:47:10 +03:00
Add a scan for folds using "Scan" type
This commit is contained in:
parent
fe2dba682c
commit
db1cee542a
@ -39,6 +39,7 @@ import qualified Streamly.Internal.Data.Fold as FL
|
||||
import qualified Streamly.Data.Fold.Prelude as Fold
|
||||
import qualified Streamly.Internal.Data.Unfold as Unfold
|
||||
import qualified Streamly.Internal.Data.Pipe as Pipe
|
||||
import qualified Streamly.Internal.Data.Scan as Scan
|
||||
import qualified Streamly.Internal.Data.Stream as Stream
|
||||
|
||||
import Test.Tasty.Bench
|
||||
@ -388,9 +389,16 @@ o_1_space_serial_transformation value =
|
||||
value
|
||||
"pipe-mapM"
|
||||
(Stream.fold
|
||||
(FL.transform
|
||||
(FL.pipe
|
||||
(Pipe.mapM (\x -> return $ x + 1))
|
||||
FL.drain))
|
||||
, benchIOSink
|
||||
value
|
||||
"fold-runScan"
|
||||
(Stream.fold
|
||||
(FL.runScan
|
||||
(Scan.mapM (\x -> return $ x + 1))
|
||||
FL.drain))
|
||||
, benchIOSink
|
||||
value
|
||||
"fold-scan"
|
||||
|
@ -136,6 +136,7 @@ module Streamly.Internal.Data.Fold.Combinators
|
||||
-- ** Scanning Input
|
||||
, scan
|
||||
, scanMany
|
||||
, runScan
|
||||
, pipe
|
||||
, indexed
|
||||
|
||||
@ -231,6 +232,7 @@ import Foreign.Storable (Storable, peek)
|
||||
import Streamly.Internal.Data.MutArray.Type (MutArray(..))
|
||||
import Streamly.Internal.Data.Maybe.Strict (Maybe'(..), toMaybe)
|
||||
import Streamly.Internal.Data.Pipe.Type (Pipe (..))
|
||||
import Streamly.Internal.Data.Scan (Scan (..))
|
||||
import Streamly.Internal.Data.Unbox (Unbox, sizeOf)
|
||||
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
|
||||
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
|
||||
@ -241,6 +243,7 @@ import qualified Streamly.Internal.Data.MutArray.Type as MA
|
||||
import qualified Streamly.Internal.Data.Array.Type as Array
|
||||
import qualified Streamly.Internal.Data.Fold.Window as Fold
|
||||
import qualified Streamly.Internal.Data.Pipe.Type as Pipe
|
||||
import qualified Streamly.Internal.Data.Scan as Scan
|
||||
import qualified Streamly.Internal.Data.Ring as Ring
|
||||
import qualified Streamly.Internal.Data.Stream.Type as StreamD
|
||||
|
||||
@ -477,6 +480,50 @@ pipe (Pipe consume produce pinitial) (Fold fstep finitial fextract ffinal) =
|
||||
|
||||
final (Tuple' _ fs) = ffinal fs
|
||||
|
||||
{-# INLINE runScanWith #-}
|
||||
runScanWith :: Monad m => Bool -> Scan m a b -> Fold m b c -> Fold m a c
|
||||
runScanWith isMany
|
||||
(Scan stepL initialL)
|
||||
(Fold stepR initialR extractR finalR) =
|
||||
Fold step initial extract final
|
||||
|
||||
where
|
||||
|
||||
step (sL, sR) x = do
|
||||
rL <- stepL sL x
|
||||
case rL of
|
||||
Scan.Yield sL1 b -> do
|
||||
rR <- stepR sR b
|
||||
case rR of
|
||||
Partial sR1 -> return $ Partial (sL1, sR1)
|
||||
Done bR -> return (Done bR)
|
||||
Scan.Skip sL1 -> return $ Partial (sL1, sR)
|
||||
-- XXX We have dropped the input.
|
||||
-- XXX Need same behavior for Stop in Fold so that the driver can
|
||||
-- consistently assume it is dropped.
|
||||
Scan.Stop ->
|
||||
if isMany
|
||||
then return $ Partial (initialL, sR)
|
||||
else Done <$> finalR sR
|
||||
|
||||
initial = do
|
||||
r <- initialR
|
||||
case r of
|
||||
Partial sR -> return $ Partial (initialL, sR)
|
||||
Done b -> return $ Done b
|
||||
|
||||
extract = extractR . snd
|
||||
|
||||
final = finalR . snd
|
||||
|
||||
-- | Scan the input of a 'Fold' to change it in a stateful manner using a
|
||||
-- 'Scan'. The scan stops as soon as the fold terminates.
|
||||
--
|
||||
-- /Pre-release/
|
||||
{-# INLINE runScan #-}
|
||||
runScan :: Monad m => Scan m a b -> Fold m b c -> Fold m a c
|
||||
runScan = runScanWith False
|
||||
|
||||
{-# INLINE scanWith #-}
|
||||
scanWith :: Monad m => Bool -> Fold m a b -> Fold m b c -> Fold m a c
|
||||
scanWith isMany
|
||||
@ -495,6 +542,10 @@ scanWith isMany
|
||||
case rR of
|
||||
Partial sR1 ->
|
||||
if isMany
|
||||
-- XXX recursive call. If initialL returns Done then it
|
||||
-- will not terminate. In that case we should return
|
||||
-- error in the beginning itself. And we should remove
|
||||
-- this recursion, assuming it won't return Done.
|
||||
then runStep initialL sR1
|
||||
else Done <$> finalR sR1
|
||||
Done bR -> return $ Done bR
|
||||
|
Loading…
Reference in New Issue
Block a user