1b00239 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor, rename, add tree/graph traversal combinators

1b00239Harendra Kumar4 years ago

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

compare with master

unifiedsplit

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

113 , bindWith

114 , concatSmapMWith

115

116− -- * ConcatPairsWith

116+ -- * MergeMapWith

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

118 -- section.

119 , mergeMapWith

120

121− -- * IterateMap

121+ -- * Iterate

122 -- | Map and flatten Trees of Streams

123− , iterateMapWith

123+ , iterateConcatMapWith

124− , iterateSmapMWith

124+ , iterateMergeMapWith

125− , iterateMapLeftsWith

125+

126+ , iterateConcatMapDFS

127+ , iterateConcatMapBFS

128+

129+ -- More experimental ops

130+ , iterateConcatMapBFSRev

131+ , iterateLeftsConcatMapWith

132+ , iterateConcatScanWith

133+ , iterateConcatScan

126 , iterateUnfold

134 , iterateUnfold

127 )

135 )

128where

136where

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

130#include "inline.hs"

138#include "inline.hs"

131

139

132import Streamly.Internal.Data.Stream.Bottom

140import Streamly.Internal.Data.Stream.Bottom

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

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

134 , zipWith, zipWithM)

142 , zipWith, zipWithM)

135import Streamly.Internal.Data.Stream.Type

143import Streamly.Internal.Data.Stream.Type

136 ( Stream, fromStreamD, fromStreamK, toStreamD, toStreamK

144 ( Stream, fromStreamD, fromStreamK, toStreamD, toStreamK

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

629 (toStreamK m)

637 (toStreamK m)

630

638

631------------------------------------------------------------------------------

639------------------------------------------------------------------------------

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

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

633------------------------------------------------------------------------------

641------------------------------------------------------------------------------

634

642

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

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

636−-- output.

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.

637--

701--

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

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

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.

641--

703--

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

643--

707--

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.

645--

735--

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:

648--

737--

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

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

739-- >>> 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/

652--

743--

653-- /Pre-release/

744-- /Pre-release/

654--

745--

655−{-# INLINE iterateMapWith #-}

746+{-# INLINE iterateMergeMapWith #-}

656−iterateMapWith ::

747+iterateMergeMapWith ::

657 (Stream m a -> Stream m a -> Stream m a)

748 (Stream m a -> Stream m a -> Stream m a)

658 -> (a -> Stream m a)

749 -> (a -> Stream m a)

659 -> Stream m a

750 -> Stream m a

660 -> Stream m a

751 -> Stream m a

661−iterateMapWith combine f = concatMapWith combine go

752+iterateMergeMapWith combine f = iterateStream

753+

662 where

754 where

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

664

755

756+ iterateStream = mergeMapWith combine generate

757+

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

759+

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

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

666−-- fusion.

761+-- stream fusion.

667--

762--

668-- /Unimplemented/

763-- /Unimplemented/

669{-# INLINE iterateUnfold #-}

764{-# INLINE iterateUnfold #-}

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

689--

784--

690-- /Pre-release/

785-- /Pre-release/

691--

786--

692−{-# INLINE iterateSmapMWith #-}

787+{-# INLINE iterateConcatScanWith #-}

693−iterateSmapMWith

788+iterateConcatScanWith

694 :: Monad m

789 :: Monad m

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

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

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

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

697 -> m b

792 -> m b

698 -> Stream m a

793 -> Stream m a

699 -> Stream m a

794 -> Stream m a

700−iterateSmapMWith combine f initial stream =

795+iterateConcatScanWith combine f initial stream =

701− concatMap

796+ concatEffect $ do

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

797+ b <- initial

703− (fromEffect initial)

798+ iterateStream (b, stream)

704

799

705 where

800 where

706

801

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

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

708

803

709− feedback b a =

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

710− concatMap

805+

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

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

712− (fromEffect $ f b a)

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

713

826

714------------------------------------------------------------------------------

827------------------------------------------------------------------------------

715-- Either streams

828-- Either streams

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

731concatMapEitherWith = undefined

844concatMapEitherWith = undefined

732-}

845-}

733

846

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+

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

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

735−-- 'iterateMapWith':

851+-- 'iterateConcatMapWith':

736--

852--

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

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

738--

854--

739-- To traverse a directory tree:

855-- To traverse a directory tree:

740--

856--

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

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

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

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

743--

859--

744-- /Pre-release/

860-- /Pre-release/

745--

861--

746−{-# INLINE iterateMapLeftsWith #-}

862+{-# INLINE iterateLeftsConcatMapWith #-}

747−iterateMapLeftsWith

863+iterateLeftsConcatMapWith

748 :: (b ~ Either a c)

864 :: (b ~ Either a c)

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

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

750 -> (a -> Stream m b)

866 -> (a -> Stream m b)

751 -> Stream m b

867 -> Stream m b

752 -> Stream m b

868 -> Stream m b

753−iterateMapLeftsWith combine f =

869+iterateLeftsConcatMapWith combine f =

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

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

89 , concatMapM

90 , concatEffect

92+ , iterateConcatScan

93+ , iterateConcatMapDFS

94+ , iterateConcatMapBFS

95+ , iterateConcatMapBFSRev

96+

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

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

93 , FoldManyPost (..)

98 , FoldManyPost (..)

94 , foldMany

99 , foldMany

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

885 (>>) = (*>)

890 (>>) = (*>)

886-}

891-}

887

892

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+

888------------------------------------------------------------------------------

1009------------------------------------------------------------------------------

889-- Grouping/Splitting

1010-- Grouping/Splitting

890------------------------------------------------------------------------------

1011------------------------------------------------------------------------------