a46c7bf · streamly · git.iko.soy

git.iko.soy/streamly

Make the type of Consumer similar to Fold

commita46c7bf6a9661f2c9f57da4e65562089812ef897 authorHarendra Kumar <harendra@composewell.com> dateFri 8 Oct 2021, 23:27 (5 years ago) parent9738b4b

6 files changed+262−113

compare with master

unifiedsplit

src/Streamly/Internal/Data/Consumer/Type.hs+91−11 src/Streamly/Internal/Data/Fold/Type.hs+100−31 src/Streamly/Internal/Data/Stream/IsStream/Reduce.hs+9−12 src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs+1−1 src/Streamly/Internal/Data/Stream/StreamD/Type.hs+40−45 src/Streamly/Internal/FileSystem/Handle.hs+21−13

src/Streamly/Internal/Data/Consumer/Type.hs+91−11

@@ -6,35 +6,115 @@

66-- Stability : experimental

77-- Portability : GHC

88--

9−-- Consumer is to Fold as Unfold is to Stream.

9+-- The 'Fold' type embeds a default initial value, therefore, it is like a

10+-- 'Monoid' whereas the 'Consumer' type has to be supplied with an initial

11+-- value, therefore, it is more like a 'Semigroup' operation.

12+--

13+-- See the file splitting example in the @streamly-examples@ repository for an

14+-- application of the 'Consumer' type. The 'Fold' type does not perform as well

15+-- in this situation.

16+--

17+-- 'Consumer' type is to 'Fold' as 'Unfold' type is to 'Stream'. Like 'Unfold'

18+-- provides better optimizaiton than stream in nested operations similarly

19+-- 'Consumer' provides better optimization than 'Fold'.

1020--

1121module Streamly.Internal.Data.Consumer.Type

1222 (

1323 -- * Types

1424 Consumer (..)

1525

26+ -- Combinators

27+ , lmapM

28+

1629 -- Consumers

1730 , drainBy

31+ , take

1832 )

1933where

2034

21−import Control.Monad (void)

35+import Fusion.Plugin.Types (Fuse(..))

36+import Streamly.Internal.Data.Fold.Step (Step(..))

2237

23−-- All folds in the Fold module should be implemented as Refolds.

38+import Prelude hiding (take)

39+

40+-- All folds in the Fold module should be implemented using Consumers.

2441--

25−-- | Experimental type to provide a side input to the fold for generating the

26−-- initial state. For example, if we have to fold chunks of a stream and write

27−-- each chunk to a different file, then we can generate the file name using a

28−-- monadic action. This is a generalized version of 'Fold'.

42+-- | Like 'Fold' except that the initial state of the accmulator can be

43+-- generated using a dynamically supplied input. This affords better stream

44+-- fusion optimization in nested fold operations where the initial fold state

45+-- is determined based on a dynamic value.

2946--

3047-- /Internal/

3148data Consumer m c a b =

3249 -- | @Fold @ @ step @ @ inject @ @ extract@

33− forall s. Consumer (s -> a -> m s) (c -> m s) (s -> m b)

50+ forall s. Consumer (s -> a -> m (Step s b)) (c -> m (Step s b)) (s -> m b)

51+

52+------------------------------------------------------------------------------

53+-- Mapping on input

54+------------------------------------------------------------------------------

55+

56+-- | @lmapM f fold@ maps the monadic function @f@ on the input of the fold.

57+--

58+-- /Internal/

59+{-# INLINE lmapM #-}

60+lmapM :: Monad m => (a -> m b) -> Consumer m c b r -> Consumer m c a r

61+lmapM f (Consumer step inject extract) = Consumer step1 inject extract

62+

63+ where

64+

65+ step1 x a = f a >>= step x

66+

67+------------------------------------------------------------------------------

68+-- Consumers

69+------------------------------------------------------------------------------

3470

3571-- |

3672--

3773-- /Internal/

38−{-# INLINABLE drainBy #-}

39−drainBy :: Monad m => (a -> m b) -> Consumer m c a ()

40−drainBy f = Consumer (const (void . f)) (\_ -> return ()) return

74+{-# INLINE drainBy #-}

75+drainBy :: Monad m => (c -> a -> m b) -> Consumer m c a ()

76+drainBy f = Consumer step inject extract

77+

78+ where

79+

80+ inject = return . Partial

81+

82+ step c a = f c a >> return (Partial c)

83+

84+ extract _ = return ()

85+

86+-- Required to fuse "take" with "many" in "chunksOf", for ghc-9.x

87+{-# ANN type Tuple'Fused Fuse #-}

88+data Tuple'Fused a b = Tuple'Fused !a !b deriving Show

89+

90+-- | Take at most @n@ input elements and fold them using the supplied fold. A

91+-- negative count is treated as 0.

92+--

93+-- /Internal/

94+{-# INLINE take #-}

95+take :: Monad m => Int -> Consumer m x a b -> Consumer m x a b

96+take n (Consumer fstep finject fextract) = Consumer step inject extract

97+

98+ where

99+

100+ inject x = do

101+ res <- finject x

102+ case res of

103+ Partial s ->

104+ if n > 0

105+ then return $ Partial $ Tuple'Fused 0 s

106+ else Done <$> fextract s

107+ Done b -> return $ Done b

108+

109+ step (Tuple'Fused i r) a = do

110+ res <- fstep r a

111+ case res of

112+ Partial sres -> do

113+ let i1 = i + 1

114+ s1 = Tuple'Fused i1 sres

115+ if i1 < n

116+ then return $ Partial s1

117+ else Done <$> fextract sres

118+ Done bres -> return $ Done bres

119+

120+ extract (Tuple'Fused _ r) = fextract r

src/Streamly/Internal/Data/Fold/Type.hs+100−31

@@ -209,6 +209,7 @@ module Streamly.Internal.Data.Fold.Type

209209 -- * Folds

210210 , fromPure

211211 , fromEffect

212+ , fromConsumer

212213 , drain

213214 , toList

214215 , toStreamK

@@ -248,6 +249,8 @@ module Streamly.Internal.Data.Fold.Type

248249 , ManyState

249250 , many

250251 , manyPost

252+ , consumeMany

253+ , consumeMany1

251254 , chunksOf

252255

253256 -- ** Nesting

@@ -258,10 +261,6 @@ module Streamly.Internal.Data.Fold.Type

258261 , snoc

259262 , duplicate

260263 , finish

261−

262− -- * Consumer

263− , simplify

264− , chunksOf2

265264 )

266265where

267266

@@ -271,7 +270,7 @@ import Data.Maybe (isJust, fromJust)

271270import Fusion.Plugin.Types (Fuse(..))

272271import Streamly.Internal.Data.Fold.Step (Step(..), mapMStep, chainStepM)

273272import Streamly.Internal.Data.Maybe.Strict (Maybe'(..), toMaybe)

274−import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))

273+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

275274import Streamly.Internal.Data.Consumer.Type (Consumer(..))

276275

277276import qualified Streamly.Internal.Data.Stream.StreamK.Type as K

@@ -505,18 +504,18 @@ mkFoldM_ :: Monad m => (b -> a -> m (Step b b)) -> m (Step b b) -> Fold m a b

505504mkFoldM_ step initial = mkFoldM step initial return

506505

507506------------------------------------------------------------------------------

508−-- Fold2

507+-- Consumer

509508------------------------------------------------------------------------------

510509

511510-- This is similar to how we run an Unfold to generate a Stream. A Fold is like

512511-- a Stream and a Fold2 is like an Unfold.

513512--

514−-- | Convert more general type 'Fold2' into a simpler type 'Fold'

513+-- | Make a fold from a consumer.

515514--

516515-- /Internal/

517−simplify :: Functor m => Consumer m c a b -> c -> Fold m a b

518−simplify (Consumer step inject extract) c =

519− Fold (\x a -> Partial <$> step x a) (Partial <$> inject c) extract

516+fromConsumer :: Consumer m c a b -> c -> Fold m a b

517+fromConsumer (Consumer step inject extract) c =

518+ Fold step (inject c) extract

520519

521520------------------------------------------------------------------------------

522521-- Basic Folds

@@ -1312,32 +1311,102 @@ manyPost (Fold sstep sinitial sextract) (Fold cstep cinitial cextract) =

13121311chunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c

13131312chunksOf n split = many (take n split)

13141313

1315−-- |

1314+-- | Like 'many' but uses a 'Consumer' for collecting.

1315+--

1316+{-# INLINE consumeMany #-}

1317+consumeMany :: Monad m => Fold m a b -> Consumer m x b c -> Consumer m x a c

1318+consumeMany (Fold sstep sinitial sextract) (Consumer cstep cinject cextract) =

1319+ Consumer step inject extract

1320+

1321+ where

1322+

1323+ -- cs = collect state

1324+ -- ss = split state

1325+ -- cres = collect state result

1326+ -- sres = split state result

1327+ -- cb = collect done

1328+ -- sb = split done

1329+

1330+ -- Caution! There is mutual recursion here, inlining the right functions is

1331+ -- important.

1332+

1333+ {-# INLINE split #-}

1334+ split cs f sres =

1335+ case sres of

1336+ Partial ss -> return $ Partial $ Tuple' cs (f ss)

1337+ Done sb -> cstep cs sb >>= collect

1338+

1339+ {-# INLINE collect #-}

1340+ collect cres =

1341+ case cres of

1342+ Partial cs -> sinitial >>= split cs Left

1343+ Done cb -> return $ Done cb

1344+

1345+ inject x = cinject x >>= collect

1346+

1347+ {-# INLINE step_ #-}

1348+ step_ ss cs a = sstep ss a >>= split cs Right

1349+

1350+ {-# INLINE step #-}

1351+ step (Tuple' cs (Left ss)) a = step_ ss cs a

1352+ step (Tuple' cs (Right ss)) a = step_ ss cs a

1353+

1354+ -- Do not extract the split fold if no item was consumed.

1355+ extract (Tuple' cs (Left _)) = cextract cs

1356+ extract (Tuple' cs (Right ss )) = do

1357+ cres <- sextract ss >>= cstep cs

1358+ case cres of

1359+ Partial s -> cextract s

1360+ Done b -> return b

1361+

1362+{-# ANN type ConsumeManyState Fuse #-}

1363+data ConsumeManyState x cs ss = ConsumeMany x cs (Either ss ss)

1364+

1365+-- | Like 'many' but uses a 'Consumer' for splitting.

13161366--

13171367-- /Internal/

1318−{-# INLINE chunksOf2 #-}

1319−chunksOf2 :: Monad m =>

1320− Int -> Fold m a b -> Consumer m x b c -> Consumer m x a c

1321−chunksOf2 n (Fold step1 initial1 extract1) (Consumer step2 inject2 extract2) =

1322− Consumer step' inject' extract'

1368+{-# INLINE consumeMany1 #-}

1369+consumeMany1 :: Monad m => Consumer m x a b -> Fold m b c -> Consumer m x a c

1370+consumeMany1 (Consumer sstep sinject sextract) (Fold cstep cinitial cextract) =

1371+ Consumer step inject extract

13231372

13241373 where

13251374

1326− loopUntilPartial s = do

1327− res <- initial1

1328− case res of

1329− Partial fs -> return $ Tuple3' 0 fs s

1330− Done _ -> loopUntilPartial s

1375+ -- cs = collect state

1376+ -- ss = split state

1377+ -- cres = collect state result

1378+ -- sres = split state result

1379+ -- cb = collect done

1380+ -- sb = split done

1381+

1382+ -- Caution! There is mutual recursion here, inlining the right functions is

1383+ -- important.

13311384

1332− inject' x = inject2 x >>= loopUntilPartial

1385+ {-# INLINE split #-}

1386+ split x cs f sres =

1387+ case sres of

1388+ Partial ss -> return $ Partial $ ConsumeMany x cs (f ss)

1389+ Done sb -> cstep cs sb >>= collect x

13331390

1334− step' (Tuple3' i r1 r2) a =

1335− if i < n

1336− then do

1337− res <- step1 r1 a

1338− case res of

1339− Partial s -> return $ Tuple3' (i + 1) s r2

1340− Done b -> step2 r2 b >>= loopUntilPartial

1341− else extract1 r1 >>= step2 r2 >>= loopUntilPartial

1391+ {-# INLINE collect #-}

1392+ collect x cres =

1393+ case cres of

1394+ Partial cs -> sinject x >>= split x cs Left

1395+ Done cb -> return $ Done cb

1396+

1397+ inject x = cinitial >>= collect x

13421398

1343− extract' (Tuple3' _ r1 r2) = extract1 r1 >>= step2 r2 >>= extract2

1399+ {-# INLINE step_ #-}

1400+ step_ x ss cs a = sstep ss a >>= split x cs Right

1401+

1402+ {-# INLINE step #-}

1403+ step (ConsumeMany x cs (Left ss)) a = step_ x ss cs a

1404+ step (ConsumeMany x cs (Right ss)) a = step_ x ss cs a

1405+

1406+ -- Do not extract the split fold if no item was consumed.

1407+ extract (ConsumeMany _ cs (Left _)) = cextract cs

1408+ extract (ConsumeMany _ cs (Right ss )) = do

1409+ cres <- sextract ss >>= cstep cs

1410+ case cres of

1411+ Partial s -> cextract s

1412+ Done b -> return b

src/Streamly/Internal/Data/Stream/IsStream/Reduce.hs+9−12

@@ -29,6 +29,7 @@ module Streamly.Internal.Data.Stream.IsStream.Reduce

2929 -- | Apply folds on a stream.

3030 , foldMany

3131 , foldManyPost

32+ , consumeMany

3233 , foldSequence

3334 , foldIterateM

3435

@@ -138,9 +139,6 @@ module Streamly.Internal.Data.Stream.IsStream.Reduce

138139 -- * Nested splitting

139140 , splitInnerBy

140141 , splitInnerBySuffix

141−

142− -- * Fold2

143− , chunksOf2

144142 )

145143where

146144

@@ -296,6 +294,14 @@ foldMany

296294 -> t m b

297295foldMany f m = fromStreamD $ D.foldMany f (toStreamD m)

298296

297+-- | Like 'foldMany' but using the 'Consumer' type instead of 'Fold'.

298+--

299+-- /Pre-release/

300+{-# INLINE consumeMany #-}

301+consumeMany :: (IsStream t, Monad m) =>

302+ Consumer m c a b -> m c -> t m a -> t m b

303+consumeMany f action = fromStreamD . D.consumeMany f action . toStreamD

304+

299305-- | Apply a stream of folds to an input stream and emit the results in the

300306-- output stream.

301307--

@@ -992,15 +998,6 @@ chunksOf

992998 => Int -> Fold m a b -> t m a -> t m b

993999chunksOf n f = fromStreamD . D.chunksOf n f . toStreamD

9941000

995−-- |

996−--

997−-- /Pre-release/

998−{-# INLINE chunksOf2 #-}

999−chunksOf2

1000− :: (IsStream t, Monad m)

1001− => Int -> m c -> Consumer m c a b -> t m a -> t m b

1002−chunksOf2 n action f m = fromStreamD $ D.groupsOf2 n action f (toStreamD m)

1003−

10041001-- | @arraysOf n stream@ groups the elements in the input stream into arrays of

10051002-- @n@ elements each.

10061003--

src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs+1−1

@@ -113,6 +113,7 @@ module Streamly.Internal.Data.Stream.StreamD.Nesting

113113 -- | Apply folds on a stream.

114114 , foldMany

115115 , foldIterateM

116+ , consumeMany

116117

117118 -- ** Parsing

118119 -- | Parsing is opposite to flattening. 'parseMany' is dual to concatMap or

@@ -126,7 +127,6 @@ module Streamly.Internal.Data.Stream.StreamD.Nesting

126127 -- ** Grouping

127128 -- | Group segments of a stream and fold. Special case of parsing.

128129 , chunksOf

129− , groupsOf2

130130 , groupsBy

131131 , groupsRollingBy

132132

src/Streamly/Internal/Data/Stream/StreamD/Type.hs+40−45

@@ -83,13 +83,12 @@ module Streamly.Internal.Data.Stream.StreamD.Type

8383 , FoldManyPost (..)

8484 , foldMany

8585 , foldManyPost

86− , groupsOf2

86+ , consumeMany

8787 , chunksOf

8888 )

8989where

9090

9191import Control.Applicative (liftA2)

92−import Control.Monad (when)

9392import Control.Monad.Catch (MonadThrow, throwM)

9493import Control.Monad.Trans.Class (lift, MonadTrans)

9594import Data.Functor (($>))

@@ -922,55 +921,51 @@ foldMany (Fold fstep initial extract) (Stream step state) =

922921chunksOf :: Monad m => Int -> Fold m a b -> Stream m a -> Stream m b

923922chunksOf n f = foldMany (FL.take n f)

924923

925−data GroupState2 s fs

926− = GroupStart2 s

927− | GroupBuffer2 s fs Int

928− | GroupYield2 fs (GroupState2 s fs)

929− | GroupFinish2

930−

931−{-# INLINE_NORMAL groupsOf2 #-}

932−groupsOf2

933− :: Monad m

934− => Int

935− -> m c

936− -> Consumer m c a b

937− -> Stream m a

938− -> Stream m b

939−groupsOf2 n input (Consumer fstep inject extract) (Stream step state) =

940− n `seq` Stream step' (GroupStart2 state)

924+-- Keep the argument order consistent with consumeIterateM

925+--

926+-- | Like 'foldMany' but for the 'Consumer' type.

927+--

928+-- /Internal/

929+{-# INLINE_NORMAL consumeMany #-}

930+consumeMany :: Monad m => Consumer m x a b -> m x -> Stream m a -> Stream m b

931+consumeMany (Consumer fstep inject extract) action (Stream step state) =

932+ Stream step' (FoldManyStart state)

941933

942934 where

943935

936+ {-# INLINE consume #-}

937+ consume x s fs = do

938+ res <- fstep fs x

939+ return

940+ $ Skip

941+ $ case res of

942+ FL.Done b -> FoldManyYield b (FoldManyStart s)

943+ FL.Partial ps -> FoldManyLoop s ps

944+

944945 {-# INLINE_LATE step' #-}

945− step' _ (GroupStart2 st) = do

946− -- XXX shall we use the Natural type instead? Need to check performance

947− -- implications.

948− when (n <= 0) $

949− -- XXX we can pass the module string from the higher level API

950− error $ "Streamly.Internal.Data.Stream.StreamD.Type.groupsOf: the size of "

951− ++ "groups [" ++ show n ++ "] must be a natural number"

952− -- fs = fold state

953− fs <- input >>= inject

954− return $ Skip (GroupBuffer2 st fs 0)

955−

956− step' gst (GroupBuffer2 st fs i) = do

946+ step' _ (FoldManyStart st) = do

947+ r <- action >>= inject

948+ return

949+ $ Skip

950+ $ case r of

951+ FL.Done b -> FoldManyYield b (FoldManyStart st)

952+ FL.Partial fs -> FoldManyFirst fs st

953+ step' gst (FoldManyFirst fs st) = do

957954 r <- step (adaptState gst) st

958955 case r of

959− Yield x s -> do

960− !fs' <- fstep fs x

961− let i' = i + 1

962− return $

963− if i' >= n

964− then Skip (GroupYield2 fs' (GroupStart2 s))

965− else Skip (GroupBuffer2 s fs' i')

966− Skip s -> return $ Skip (GroupBuffer2 s fs i)

967− Stop -> return $ Skip (GroupYield2 fs GroupFinish2)

968−

969− step' _ (GroupYield2 fs next) = do

970− r <- extract fs

971− return $ Yield r next

972−

973− step' _ GroupFinish2 = return Stop

956+ Yield x s -> consume x s fs

957+ Skip s -> return $ Skip (FoldManyFirst fs s)

958+ Stop -> return Stop

959+ step' gst (FoldManyLoop st fs) = do

960+ r <- step (adaptState gst) st

961+ case r of

962+ Yield x s -> consume x s fs

963+ Skip s -> return $ Skip (FoldManyLoop s fs)

964+ Stop -> do

965+ b <- extract fs

966+ return $ Skip (FoldManyYield b FoldManyDone)

967+ step' _ (FoldManyYield b next) = return $ Yield b next

968+ step' _ FoldManyDone = return Stop

974969

975970------------------------------------------------------------------------------

976971-- Other instances

src/Streamly/Internal/FileSystem/Handle.hs+21−13

@@ -45,7 +45,7 @@ module Streamly.Internal.FileSystem.Handle

4545 -- * Byte Stream Write

4646 -- Byte stream write (Folds)

4747 , write

48− , write2

48+ , consumer

4949 -- , writeUtf8

5050 -- , writeUtf8ByLines

5151 -- , writeByFrames

@@ -125,6 +125,7 @@ import Streamly.Internal.System.IO (defaultChunkSize)

125125

126126import qualified Streamly.Internal.Data.Array.Foreign as A

127127import qualified Streamly.Internal.Data.Array.Stream.Foreign as AS

128+import qualified Streamly.Internal.Data.Consumer.Type as Consumer

128129import qualified Streamly.Internal.Data.Fold as FL

129130import qualified Streamly.Internal.Data.Fold.Type as FL

130131import qualified Streamly.Internal.Data.Stream.IsStream as S

@@ -456,12 +457,12 @@ putBytes = putBytesWithBufferOf defaultChunkSize

456457writeChunks :: (MonadIO m, Storable a) => Handle -> Fold m (Array a) ()

457458writeChunks h = FL.drainBy (putChunk h)

458459

459−{-# INLINE writeChunks2 #-}

460−writeChunks2 :: (MonadIO m, Storable a) => Consumer m Handle (Array a) ()

461−writeChunks2 =

462− Consumer (\h arr -> liftIO $ putChunk h arr >> return h)

463− return

464− (\_ -> return ())

460+-- | Like writeChunks but uses the experimental 'Consumer' API.

461+--

462+-- /Internal/

463+{-# INLINE consumeChunks #-}

464+consumeChunks :: (MonadIO m, Storable a) => Consumer m Handle (Array a) ()

465+consumeChunks = Consumer.drainBy putChunk

465466

466467-- XXX lpackArraysChunksOf should be written idiomatically

467468--

@@ -498,9 +499,13 @@ writeChunksWithBufferOf n h = lpackArraysChunksOf n (writeChunks h)

498499writeWithBufferOf :: MonadIO m => Int -> Handle -> Fold m Word8 ()

499500writeWithBufferOf n h = FL.chunksOf n (writeNUnsafe n) (writeChunks h)

500501

501−{-# INLINE writeWithBufferOf2 #-}

502−writeWithBufferOf2 :: MonadIO m => Int -> Consumer m Handle Word8 ()

503−writeWithBufferOf2 n = FL.chunksOf2 n (writeNUnsafe n) writeChunks2

502+-- | Like 'writeWithBufferOf' but uses the experimental 'Consumer' API.

503+--

504+-- /Internal/

505+{-# INLINE consumerWithBufferOf #-}

506+consumerWithBufferOf :: MonadIO m => Int -> Consumer m Handle Word8 ()

507+consumerWithBufferOf n =

508+ FL.consumeMany (FL.take n $ writeNUnsafe n) consumeChunks

504509

505510-- | Write a byte stream to a file handle. Accumulates the input in chunks of

506511-- up to 'Streamly.Internal.Data.Array.Foreign.Type.defaultChunkSize' before writing

@@ -513,9 +518,12 @@ writeWithBufferOf2 n = FL.chunksOf2 n (writeNUnsafe n) writeChunks2

513518write :: MonadIO m => Handle -> Fold m Word8 ()

514519write = writeWithBufferOf defaultChunkSize

515520

516−{-# INLINE write2 #-}

517−write2 :: MonadIO m => Consumer m Handle Word8 ()

518−write2 = writeWithBufferOf2 defaultChunkSize

521+-- | Like 'write' but uses the experimental 'Consumer' API.

522+--

523+-- /Internal/

524+{-# INLINE consumer #-}

525+consumer :: MonadIO m => Consumer m Handle Word8 ()

526+consumer = consumerWithBufferOf defaultChunkSize

519527

520528{-

521529{-# INLINE write #-}