1b00239 · streamly · git.iko.soy

git.iko.soy/streamly

Refactor, rename, add tree/graph traversal combinators

Simplify the implementation of iterateSmapMWith

commit1b0023993b1dd1b1f5d7565cd40bdcb3d5b87fcd authorHarendra Kumar <harendra@composewell.com> dateTue 27 Dec 2022, 15:40 (4 years ago) parentd6679cb

2 files changed+279−42

compare with master

unifiedsplit

core/src/Streamly/Internal/Data/Stream/Expand.hs+158−42 core/src/Streamly/Internal/Data/Stream/StreamD/Type.hs+121

core/src/Streamly/Internal/Data/Stream/Expand.hs+158−42

@@ -113,16 +113,24 @@ module Streamly.Internal.Data.Stream.Expand

113113 , bindWith

114114 , concatSmapMWith

115115

116− -- * ConcatPairsWith

116+ -- * MergeMapWith

117117 -- | See the notes about suitable merge functions in the 'concatMapWith'

118118 -- section.

119119 , mergeMapWith

120120

121− -- * IterateMap

121+ -- * Iterate

122122 -- | Map and flatten Trees of Streams

123− , iterateMapWith

124− , iterateSmapMWith

125− , iterateMapLeftsWith

123+ , iterateConcatMapWith

124+ , iterateMergeMapWith

125+

126+ , iterateConcatMapDFS

127+ , iterateConcatMapBFS

128+

129+ -- More experimental ops

130+ , iterateConcatMapBFSRev

131+ , iterateLeftsConcatMapWith

132+ , iterateConcatScanWith

133+ , iterateConcatScan

126134 , iterateUnfold

127135 )

128136where

@@ -130,7 +138,7 @@ where

130138#include "inline.hs"

131139

132140import Streamly.Internal.Data.Stream.Bottom

133− ( concatEffect, concatMapM, concatMap, smapM, fromPure, fromEffect

141+ ( concatEffect, concatMapM, concatMap, smapM, fromPure

134142 , zipWith, zipWithM)

135143import Streamly.Internal.Data.Stream.Type

136144 ( Stream, fromStreamD, fromStreamK, toStreamD, toStreamK

@@ -629,41 +637,128 @@ mergeMapWith par f m =

629637 (toStreamK m)

630638

631639------------------------------------------------------------------------------

632−-- IterateMap - Map and flatten Trees of Streams

640+-- iterateConcatMap - Map and flatten Trees of Streams

633641------------------------------------------------------------------------------

634642

635−-- | Like 'iterateM' but iterates after mapping a stream generator on the

636−-- output.

643+-- | Yield an input element in the output stream, map a stream generator on it

644+-- and repeat the process on the resulting stream. Resulting streams are

645+-- flattened using the 'concatMapWith' combinator. This can be used for a depth

646+-- first style (DFS) traversal of a tree like structure.

647+--

648+-- Example, list a directory tree using DFS:

649+--

650+-- >>> f = either Dir.readEitherPaths (const Stream.nil)

651+-- >>> input = Stream.fromPure (Left ".")

652+-- >>> ls = Stream.iterateConcatMapWith Stream.append f input

653+--

654+-- Note that 'iterateM' is a special case of 'iterateConcatMapWith':

655+--

656+-- >>> iterateM f = Stream.iterateConcatMapWith Stream.append (Stream.fromEffect . f) . Stream.fromEffect

657+--

658+-- /CPS/

659+--

660+-- /Pre-release/

661+--

662+{-# INLINE iterateConcatMapWith #-}

663+iterateConcatMapWith ::

664+ (Stream m a -> Stream m a -> Stream m a)

665+ -> (a -> Stream m a)

666+ -> Stream m a

667+ -> Stream m a

668+iterateConcatMapWith combine f = iterateStream

669+

670+ where

671+

672+ iterateStream = concatMapWith combine generate

673+

674+ generate x = fromPure x `combine` iterateStream (f x)

675+

676+-- | Traverse the stream in depth first style (DFS). Map each element in the

677+-- input stream to a stream and flatten, recursively map the resulting elements

678+-- as well to a stream and flatten until no more streams are generated.

679+--

680+-- Example, list a directory tree using DFS:

681+--

682+-- >>> f = either (Just . Dir.readEitherPaths) (const Nothing)

683+-- >>> input = Stream.fromPure (Left ".")

684+-- >>> ls = Stream.iterateConcatMapDFS f input

685+--

686+-- This is equivalent to using @iterateConcatMapWith Stream.append@.

687+--

688+-- /Pre-release/

689+{-# INLINE iterateConcatMapDFS #-}

690+iterateConcatMapDFS :: Monad m =>

691+ (a -> Maybe (Stream m a))

692+ -> Stream m a

693+ -> Stream m a

694+iterateConcatMapDFS f stream =

695+ fromStreamD

696+ $ D.iterateConcatMapDFS (fmap toStreamD . f ) (toStreamD stream)

697+

698+-- | Similar to 'iterateConcatMapDFS' except that it traverses the stream in

699+-- breadth first style (BFS). First, all the elements in the input stream are

700+-- emitted, and then their traversals are emitted.

637701--

638−-- Yield an input element in the output stream, map a stream generator on it

639−-- and then do the same on the resulting stream. This can be used for a depth

640−-- first traversal of a tree like structure.

702+-- Example, list a directory tree using BFS:

641703--

642−-- Note that 'iterateM' is a special case of 'iterateMapWith':

704+-- >>> f = either (Just . Dir.readEitherPaths) (const Nothing)

705+-- >>> input = Stream.fromPure (Left ".")

706+-- >>> ls = Stream.iterateConcatMapBFS f input

643707--

644−-- >>> iterateM f = Stream.iterateMapWith Stream.append (Stream.fromEffect . f) . Stream.fromEffect

708+-- /Pre-release/

709+{-# INLINE iterateConcatMapBFS #-}

710+iterateConcatMapBFS :: Monad m =>

711+ (a -> Maybe (Stream m a))

712+ -> Stream m a

713+ -> Stream m a

714+iterateConcatMapBFS f stream =

715+ fromStreamD

716+ $ D.iterateConcatMapBFS (fmap toStreamD . f ) (toStreamD stream)

717+

718+-- | Same as 'iterateConcatMapBFS' except that the traversal of the last

719+-- element on a level is emitted first and then going backwards up to the first

720+-- element (reversed ordering). This may be slightly faster than

721+-- 'iterateConcatMapBFS'.

722+--

723+{-# INLINE iterateConcatMapBFSRev #-}

724+iterateConcatMapBFSRev :: Monad m =>

725+ (a -> Maybe (Stream m a))

726+ -> Stream m a

727+ -> Stream m a

728+iterateConcatMapBFSRev f stream =

729+ fromStreamD

730+ $ D.iterateConcatMapBFSRev (fmap toStreamD . f ) (toStreamD stream)

731+

732+-- | Like 'iterateConcatMapWith' but uses the pairwise flattening combinator

733+-- 'mergeMapWith' for flattening the resulting streams. This can be used for a

734+-- balanced traversal of a tree like structure.

645735--

646−-- It can be used to traverse a tree structure. For example, to list a

647−-- directory tree:

736+-- Example, list a directory tree using balanced traversal:

648737--

738+-- >>> f = either Dir.readEitherPaths (const Stream.nil)

649739-- >>> input = Stream.fromPure (Left ".")

650−-- >>> f = either Dir.readEither (const Stream.nil)

651−-- >>> ls = Stream.iterateMapWith Stream.append f input

740+-- >>> ls = Stream.iterateMergeMapWith Stream.interleave f input

741+--

742+-- /CPS/

652743--

653744-- /Pre-release/

654745--

655−{-# INLINE iterateMapWith #-}

656−iterateMapWith ::

746+{-# INLINE iterateMergeMapWith #-}

747+iterateMergeMapWith ::

657748 (Stream m a -> Stream m a -> Stream m a)

658749 -> (a -> Stream m a)

659750 -> Stream m a

660751 -> Stream m a

661−iterateMapWith combine f = concatMapWith combine go

752+iterateMergeMapWith combine f = iterateStream

753+

662754 where

663− go x = fromPure x `combine` concatMapWith combine go (f x)

664755

665−-- | Same as @iterateMapWith Stream.serial@ but more efficient due to stream

666−-- fusion.

756+ iterateStream = mergeMapWith combine generate

757+

758+ generate x = fromPure x `combine` iterateStream (f x)

759+

760+-- | Same as @iterateConcatMapWith Stream.append@ but more efficient due to

761+-- stream fusion.

667762--

668763-- /Unimplemented/

669764{-# INLINE iterateUnfold #-}

@@ -689,27 +784,45 @@ iterateUnfold = undefined

689784--

690785-- /Pre-release/

691786--

692−{-# INLINE iterateSmapMWith #-}

693−iterateSmapMWith

787+{-# INLINE iterateConcatScanWith #-}

788+iterateConcatScanWith

694789 :: Monad m

695790 => (Stream m a -> Stream m a -> Stream m a)

696791 -> (b -> a -> m (b, Stream m a))

697792 -> m b

698793 -> Stream m a

699794 -> Stream m a

700−iterateSmapMWith combine f initial stream =

701− concatMap

702− (\b -> concatMapWith combine (go b) stream)

703− (fromEffect initial)

795+iterateConcatScanWith combine f initial stream =

796+ concatEffect $ do

797+ b <- initial

798+ iterateStream (b, stream)

704799

705800 where

706801

707− go b a = fromPure a `combine` feedback b a

802+ iterateStream (b, s) = pure $ concatMapWith combine (generate b) s

708803

709− feedback b a =

710− concatMap

711− (\(b1, s) -> concatMapWith combine (go b1) s)

712− (fromEffect $ f b a)

804+ generate b a = fromPure a `combine` feedback b a

805+

806+ feedback b a = concatEffect $ f b a >>= iterateStream

807+

808+-- Next stream is to be generated by the return value of the previous stream. A

809+-- general intuitive way of doing that could be to use an appending monad

810+-- instance for streams where the result of the previous stream is used to

811+-- generate the next one. In the first pass we can just emit the values in the

812+-- stream and keep building a buffered list/stream, once done we can then

813+-- process the buffered stream.

814+

815+{-# INLINE iterateConcatScan #-}

816+iterateConcatScan

817+ :: Monad m

818+ => (b -> a -> m b)

819+ -> (b -> m (Maybe (b, Stream m a)))

820+ -> b

821+ -> Stream m a

822+iterateConcatScan scanner generate initial =

823+ fromStreamD

824+ $ D.iterateConcatScan

825+ scanner (fmap (fmap (fmap toStreamD)) . generate) initial

713826

714827------------------------------------------------------------------------------

715828-- Either streams

@@ -731,24 +844,27 @@ concatMapEitherWith

731844concatMapEitherWith = undefined

732845-}

733846

847+-- XXX We should prefer using the Maybe stream returning signatures over this.

848+-- This API should perhaps be removed in favor of those.

849+

734850-- | In an 'Either' stream iterate on 'Left's. This is a special case of

735−-- 'iterateMapWith':

851+-- 'iterateConcatMapWith':

736852--

737−-- >>> iterateMapLeftsWith combine f = Stream.iterateMapWith combine (either f (const Stream.nil))

853+-- >>> iterateLeftsConcatMapWith combine f = Stream.iterateConcatMapWith combine (either f (const Stream.nil))

738854--

739855-- To traverse a directory tree:

740856--

741857-- >>> input = Stream.fromPure (Left ".")

742−-- >>> ls = Stream.iterateMapLeftsWith Stream.append Dir.readEither input

858+-- >>> ls = Stream.iterateLeftsConcatMapWith Stream.append Dir.readEither input

743859--

744860-- /Pre-release/

745861--

746−{-# INLINE iterateMapLeftsWith #-}

747−iterateMapLeftsWith

862+{-# INLINE iterateLeftsConcatMapWith #-}

863+iterateLeftsConcatMapWith

748864 :: (b ~ Either a c)

749865 => (Stream m b -> Stream m b -> Stream m b)

750866 -> (a -> Stream m b)

751867 -> Stream m b

752868 -> Stream m b

753−iterateMapLeftsWith combine f =

754− iterateMapWith combine (either f (const nil))

869+iterateLeftsConcatMapWith combine f =

870+ iterateConcatMapWith combine (either f (const nil))

core/src/Streamly/Internal/Data/Stream/StreamD/Type.hs+121

@@ -89,6 +89,11 @@ module Streamly.Internal.Data.Stream.StreamD.Type

8989 , concatMapM

9090 , concatEffect

9191

92+ , iterateConcatScan

93+ , iterateConcatMapDFS

94+ , iterateConcatMapBFS

95+ , iterateConcatMapBFSRev

96+

9297 , FoldMany (..) -- for inspection testing

9398 , FoldManyPost (..)

9499 , foldMany

@@ -885,6 +890,122 @@ instance Monad m => Monad (Stream m) where

885890 (>>) = (*>)

886891-}

887892

893+-- | Generate a stream from an initial state, scan and concat the stream,

894+-- generate a stream again from the final state of the previous scan and repeat

895+-- the process.

896+{-# INLINE_NORMAL iterateConcatScan #-}

897+iterateConcatScan :: Monad m =>

898+ (b -> a -> m b)

899+ -> (b -> m (Maybe (b, Stream m a)))

900+ -> b

901+ -> Stream m a

902+iterateConcatScan scanner generate initial = Stream step (Left initial)

903+

904+ where

905+

906+ {-# INLINE_LATE step #-}

907+ step _ (Left acc) = do

908+ r <- generate acc

909+ case r of

910+ Nothing -> return Stop

911+ Just v -> return $ Skip (Right v)

912+

913+ step gst (Right (st, UnStream inner_step inner_st)) = do

914+ r <- inner_step (adaptState gst) inner_st

915+ case r of

916+ Yield b inner_s -> do

917+ acc <- scanner st b

918+ return $ Yield b (Right (acc, Stream inner_step inner_s))

919+ Skip inner_s ->

920+ return $ Skip (Right (st, Stream inner_step inner_s))

921+ Stop -> return $ Skip (Left st)

922+

923+-- Note: The iterate function returns a Maybe Stream instead of returning a nil

924+-- stream for indicating a leaf node. This is to optimize so that we do not

925+-- have to store any state. This makes the stored state proportional to the

926+-- number of non-leaf nodes rather than total number of nodes.

927+

928+-- | This function may be slightly faster than iterateConcatMapBFS because it

929+-- traverses the elements on a level in reverse order, therefore, does not have

930+-- to reverse the list stroing those.

931+{-# INLINE_NORMAL iterateConcatMapBFSRev #-}

932+iterateConcatMapBFSRev :: Monad m =>

933+ (a -> Maybe (Stream m a))

934+ -> Stream m a

935+ -> Stream m a

936+iterateConcatMapBFSRev f stream = Stream step (stream, [])

937+

938+ where

939+

940+ {-# INLINE_LATE step #-}

941+ step gst (UnStream step1 st, xs) = do

942+ r <- step1 (adaptState gst) st

943+ case r of

944+ Yield a s -> do

945+ let xs1 =

946+ case f a of

947+ Nothing -> xs

948+ Just x -> x:xs

949+ return $ Yield a (Stream step1 s, xs1)

950+ Skip s -> return $ Skip (Stream step1 s, xs)

951+ Stop ->

952+ case xs of

953+ (y:ys) -> return $ Skip (y, ys)

954+ [] -> return Stop

955+

956+{-# INLINE_NORMAL iterateConcatMapBFS #-}

957+iterateConcatMapBFS :: Monad m =>

958+ (a -> Maybe (Stream m a))

959+ -> Stream m a

960+ -> Stream m a

961+iterateConcatMapBFS f stream = Stream step (stream, [], [])

962+

963+ where

964+

965+ {-# INLINE_LATE step #-}

966+ step gst (UnStream step1 st, xs, ys) = do

967+ r <- step1 (adaptState gst) st

968+ case r of

969+ Yield a s -> do

970+ let ys1 =

971+ case f a of

972+ Nothing -> ys

973+ Just y -> y:ys

974+ return $ Yield a (Stream step1 s, xs, ys1)

975+ Skip s -> return $ Skip (Stream step1 s, xs, ys)

976+ Stop ->

977+ case xs of

978+ (x:xs1) -> return $ Skip (x, xs1, ys)

979+ [] ->

980+ case reverse ys of

981+ (x:xs1) -> return $ Skip (x, xs1, [])

982+ [] -> return Stop

983+

984+{-# INLINE_NORMAL iterateConcatMapDFS #-}

985+iterateConcatMapDFS :: Monad m =>

986+ (a -> Maybe (Stream m a))

987+ -> Stream m a

988+ -> Stream m a

989+iterateConcatMapDFS f stream = Stream step (stream, [])

990+

991+ where

992+

993+ {-# INLINE_LATE step #-}

994+ step gst (UnStream step1 st, xs) = do

995+ r <- step1 (adaptState gst) st

996+ case r of

997+ Yield a s -> do

998+ let st1 =

999+ case f a of

1000+ Nothing -> (Stream step1 s, xs)

1001+ Just x -> (x, Stream step1 s:xs)

1002+ return $ Yield a st1

1003+ Skip s -> return $ Skip (Stream step1 s, xs)

1004+ Stop ->

1005+ case xs of

1006+ (y:ys) -> return $ Skip (y, ys)

1007+ [] -> return Stop

1008+

8881009------------------------------------------------------------------------------

8891010-- Grouping/Splitting

8901011------------------------------------------------------------------------------