git.iko.soy/streamly
on master help sign in
Add iterateUnfoldMany and friends
83bca76Harendra Kumar
core/src/Streamly/Internal/Data/Stream/Expand.hs+31−7
compare with master
unifiedsplit
@@ -120,6 +120,10 @@ module Streamly.Internal.Data.Stream.Expand
120
120
121 -- * Iterate
121 -- * Iterate
122 -- | Map and flatten Trees of Streams
122 -- | Map and flatten Trees of Streams
123+ , iterateUnfoldManyDFS
124+ , iterateUnfoldManyBFS
125+ , iterateUnfoldManyBFSRev
126+
123 , iterateConcatMapWith
127 , iterateConcatMapWith
124 , iterateMergeMapWith
128 , iterateMergeMapWith
125
129
@@ -131,7 +135,6 @@ module Streamly.Internal.Data.Stream.Expand
131 , iterateLeftsConcatMapWith
135 , iterateLeftsConcatMapWith
132 , iterateConcatScanWith
136 , iterateConcatScanWith
133 , iterateConcatScan
137 , iterateConcatScan
134− , iterateUnfold
135 )
138 )
136where
139where
137
140
@@ -757,13 +760,34 @@ iterateMergeMapWith combine f = iterateStream
757
760
758 generate x = fromPure x `combine` iterateStream (f x)
761 generate x = fromPure x `combine` iterateStream (f x)
759
762
760−-- | Same as @iterateConcatMapWith Stream.append@ but more efficient due to
763+-- | Same as @iterateConcatMapDFS@ but more efficient due to stream fusion.
761−-- stream fusion.
762--
764--
763−-- /Unimplemented/
765+-- Example, list a directory tree using DFS:
764−{-# INLINE iterateUnfold #-}
766+--
767+-- >>> f = Unfold.either Dir.eitherReaderPaths Unfold.nil
768+-- >>> input = Stream.fromPure (Left ".")
769+-- >>> ls = Stream.iterateUnfoldManyDFS f input
770+--
771+-- /Pre-release/
772+{-# INLINE iterateUnfoldManyDFS #-}
765−iterateUnfold :: Unfold m a a -> Stream m a -> Stream m a
773+iterateUnfoldManyDFS :: Monad m => Unfold m a a -> Stream m a -> Stream m a
766−iterateUnfold = undefined
774+iterateUnfoldManyDFS u = fromStreamD . D.iterateUnfoldManyDFS u . toStreamD
775+
776+-- | Like 'iterateUnfoldManyDFS' but uses breadth first style traversal.
777+--
778+-- /Pre-release/
779+{-# INLINE iterateUnfoldManyBFS #-}
780+iterateUnfoldManyBFS :: Monad m => Unfold m a a -> Stream m a -> Stream m a
781+iterateUnfoldManyBFS u = fromStreamD . D.iterateUnfoldManyBFS u . toStreamD
782+
783+-- | Like 'iterateUnfoldManyBFS' but processes the children in reverse order,
784+-- therefore, may be slightly faster.
785+--
786+-- /Pre-release/
787+{-# INLINE iterateUnfoldManyBFSRev #-}
788+iterateUnfoldManyBFSRev :: Monad m => Unfold m a a -> Stream m a -> Stream m a
789+iterateUnfoldManyBFSRev u =
790+ fromStreamD . D.iterateUnfoldManyBFSRev u . toStreamD
767
791
768------------------------------------------------------------------------------
792------------------------------------------------------------------------------
769-- Flattening Graphs
793-- Flattening Graphs
core/src/Streamly/Internal/Data/Stream/StreamD/Type.hs+124
@@ -89,6 +89,10 @@ module Streamly.Internal.Data.Stream.StreamD.Type
89 , concatMapM
89 , concatMapM
90 , concatEffect
90 , concatEffect
91
91
92+ , iterateUnfoldManyDFS
93+ , iterateUnfoldManyBFS
94+ , iterateUnfoldManyBFSRev
95+
92 , iterateConcatScan
96 , iterateConcatScan
93 , iterateConcatMapDFS
97 , iterateConcatMapDFS
94 , iterateConcatMapBFS
98 , iterateConcatMapBFS
@@ -1006,6 +1010,126 @@ iterateConcatMapDFS f stream = Stream step (stream, [])
1006 (y:ys) -> return $ Skip (y, ys)
1010 (y:ys) -> return $ Skip (y, ys)
1007 [] -> return Stop
1011 [] -> return Stop
1008
1012
1013+{-# ANN type IterateUnfoldState Fuse #-}
1014+data IterateUnfoldState o i =
1015+ IterateUnfoldOuter o
1016+ | IterateUnfoldInner o i [i]
1017+
1018+{-# INLINE_NORMAL iterateUnfoldManyDFS #-}
1019+iterateUnfoldManyDFS :: Monad m =>
1020+ Unfold m a a
1021+ -> Stream m a
1022+ -> Stream m a
1023+iterateUnfoldManyDFS (Unfold istep inject) (Stream ostep ost) =
1024+ Stream step (IterateUnfoldOuter ost)
1025+
1026+ where
1027+
1028+ {-# INLINE_LATE step #-}
1029+ step gst (IterateUnfoldOuter o) = do
1030+ r <- ostep (adaptState gst) o
1031+ case r of
1032+ Yield a s -> do
1033+ i <- inject a
1034+ i `seq` return (Yield a (IterateUnfoldInner s i []))
1035+ Skip s -> return $ Skip (IterateUnfoldOuter s)
1036+ Stop -> return Stop
1037+
1038+ step _ (IterateUnfoldInner o i ii) = do
1039+ r <- istep i
1040+ case r of
1041+ Yield x s -> do
1042+ i1 <- inject x
1043+ i1 `seq` return $ Yield x (IterateUnfoldInner o i1 (s:ii))
1044+ Skip s -> return $ Skip (IterateUnfoldInner o s ii)
1045+ Stop ->
1046+ case ii of
1047+ (y:ys) -> return $ Skip (IterateUnfoldInner o y ys)
1048+ [] -> return $ Skip (IterateUnfoldOuter o)
1049+
1050+{-# ANN type IterateUnfoldBFSRevState Fuse #-}
1051+data IterateUnfoldBFSRevState o i =
1052+ IterateUnfoldBFSRevOuter o [i]
1053+ | IterateUnfoldBFSRevInner i [i]
1054+
1055+{-# INLINE_NORMAL iterateUnfoldManyBFSRev #-}
1056+iterateUnfoldManyBFSRev :: Monad m =>
1057+ Unfold m a a
1058+ -> Stream m a
1059+ -> Stream m a
1060+iterateUnfoldManyBFSRev (Unfold istep inject) (Stream ostep ost) =
1061+ Stream step (IterateUnfoldBFSRevOuter ost [])
1062+
1063+ where
1064+
1065+ {-# INLINE_LATE step #-}
1066+ step gst (IterateUnfoldBFSRevOuter o ii) = do
1067+ r <- ostep (adaptState gst) o
1068+ case r of
1069+ Yield a s -> do
1070+ i <- inject a
1071+ i `seq` return (Yield a (IterateUnfoldBFSRevOuter s (i:ii)))
1072+ Skip s -> return $ Skip (IterateUnfoldBFSRevOuter s ii)
1073+ Stop ->
1074+ case ii of
1075+ (y:ys) -> return $ Skip (IterateUnfoldBFSRevInner y ys)
1076+ [] -> return Stop
1077+
1078+ step _ (IterateUnfoldBFSRevInner i ii) = do
1079+ r <- istep i
1080+ case r of
1081+ Yield x s -> do
1082+ i1 <- inject x
1083+ i1 `seq` return $ Yield x (IterateUnfoldBFSRevInner s (i1:ii))
1084+ Skip s -> return $ Skip (IterateUnfoldBFSRevInner s ii)
1085+ Stop ->
1086+ case ii of
1087+ (y:ys) -> return $ Skip (IterateUnfoldBFSRevInner y ys)
1088+ [] -> return Stop
1089+
1090+{-# ANN type IterateUnfoldBFSState Fuse #-}
1091+data IterateUnfoldBFSState o i =
1092+ IterateUnfoldBFSOuter o [i]
1093+ | IterateUnfoldBFSInner i [i] [i]
1094+
1095+{-# INLINE_NORMAL iterateUnfoldManyBFS #-}
1096+iterateUnfoldManyBFS :: Monad m =>
1097+ Unfold m a a
1098+ -> Stream m a
1099+ -> Stream m a
1100+iterateUnfoldManyBFS (Unfold istep inject) (Stream ostep ost) =
1101+ Stream step (IterateUnfoldBFSOuter ost [])
1102+
1103+ where
1104+
1105+ {-# INLINE_LATE step #-}
1106+ step gst (IterateUnfoldBFSOuter o rii) = do
1107+ r <- ostep (adaptState gst) o
1108+ case r of
1109+ Yield a s -> do
1110+ i <- inject a
1111+ i `seq` return (Yield a (IterateUnfoldBFSOuter s (i:rii)))
1112+ Skip s -> return $ Skip (IterateUnfoldBFSOuter s rii)
1113+ Stop ->
1114+ case reverse rii of
1115+ (y:ys) -> return $ Skip (IterateUnfoldBFSInner y ys [])
1116+ [] -> return Stop
1117+
1118+ step _ (IterateUnfoldBFSInner i ii rii) = do
1119+ r <- istep i
1120+ case r of
1121+ Yield x s -> do
1122+ i1 <- inject x
1123+ i1 `seq` return $ Yield x (IterateUnfoldBFSInner s ii (i1:rii))
1124+ Skip s -> return $ Skip (IterateUnfoldBFSInner s ii rii)
1125+ Stop ->
1126+ case ii of
1127+ (y:ys) -> return $ Skip (IterateUnfoldBFSInner y ys rii)
1128+ [] ->
1129+ case reverse rii of
1130+ (y:ys) -> return $ Skip (IterateUnfoldBFSInner y ys [])
1131+ [] -> return Stop
1132+
1009------------------------------------------------------------------------------
1133------------------------------------------------------------------------------
1010-- Grouping/Splitting
1134-- Grouping/Splitting
1011------------------------------------------------------------------------------
1135------------------------------------------------------------------------------