on master help sign in
Merge SplitGroup and Combine modules to create StreamD.Nesting.
30a474bPranay Sashank
src/Streamly/Internal/Data/Stream/StreamD.hs+2−2
@@ -311,6 +311,6 @@ import Streamly.Internal.Data.Stream.StreamD.Type
311import Streamly.Internal.Data.Stream.StreamD.Generate
311import Streamly.Internal.Data.Stream.StreamD.Generate
312import Streamly.Internal.Data.Stream.StreamD.Eliminate
312import Streamly.Internal.Data.Stream.StreamD.Eliminate
313import Streamly.Internal.Data.Stream.StreamD.Exceptions
313import Streamly.Internal.Data.Stream.StreamD.Exceptions
314−import Streamly.Internal.Data.Stream.StreamD.Combine
314+import Streamly.Internal.Data.Stream.StreamD.Nesting
315−import Streamly.Internal.Data.Stream.StreamD.SplitGroup
315+
316import Streamly.Internal.Data.Stream.StreamD.Transform
316import Streamly.Internal.Data.Stream.StreamD.Transform
src/Streamly/Internal/Data/Stream/StreamD/Combine.hsdeleted−888
@@ -1,888 +0,0 @@
1−#include "inline.hs"
2−
3−-- |
4−-- Module : Streamly.Internal.Data.Stream.StreamD.Combine
5−-- Copyright : (c) 2018 Composewell Technologies
6−-- (c) Roman Leshchinskiy 2008-2010
7−-- (c) The University of Glasgow, 2009
8−-- License : BSD3
9−-- Maintainer : streamly@composewell.com
10−-- Stability : experimental
11−-- Portability : GHC
12−
13−module Streamly.Internal.Data.Stream.StreamD.Combine
14− (
15− -- * Zipping
16− indexed
17− , indexedR
18− , zipWith
19− , zipWithM
20−
21− -- * Merging
22− , mergeBy
23− , mergeByM
24−
25− -- ** Flattening nested streams
26− , concatMapM
27− , concatMap
28− , ConcatMapUState (..)
29− , concatMapU
30− , ConcatUnfoldInterleaveState (..)
31− , concatUnfoldInterleave
32− , concatUnfoldRoundrobin
33− , AppendState(..)
34− , append
35− , InterleaveState(..)
36− , interleave
37− , interleaveMin
38− , interleaveSuffix
39− , interleaveInfix
40− , roundRobin -- interleaveFair?/ParallelFair
41− , gintercalateSuffix
42− , interposeSuffix
43− , gintercalate
44− , interpose
45− )
46−where
47−
48−#if __GLASGOW_HASKELL__ >= 801
49−import Data.Functor.Identity ( Identity )
50−#endif
51−import Fusion.Plugin.Types (Fuse(..))
52−import Streamly.Internal.Data.Unfold.Types (Unfold(..))
53−
54−
55−import Prelude hiding
56− ( map, mapM, mapM_, repeat, foldr, last, take, filter
57− , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
58− , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
59− , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
60− , reverse, iterate, splitAt)
61−import Streamly.Internal.Data.Stream.StreamD.Type
62−import Streamly.Internal.Data.SVar
63−
64−------------------------------------------------------------------------------
65−-- Zipping
66−------------------------------------------------------------------------------
67−
68−{-# INLINE_NORMAL indexed #-}
69−indexed :: Monad m => Stream m a -> Stream m (Int, a)
70−indexed (Stream step state) = Stream step' (state, 0)
71− where
72− {-# INLINE_LATE step' #-}
73− step' gst (st, i) = i `seq` do
74− r <- step (adaptState gst) st
75− case r of
76− Yield x s -> return $ Yield (i, x) (s, i+1)
77− Skip s -> return $ Skip (s, i)
78− Stop -> return Stop
79−
80−{-# INLINE_NORMAL indexedR #-}
81−indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
82−indexedR m (Stream step state) = Stream step' (state, m)
83− where
84− {-# INLINE_LATE step' #-}
85− step' gst (st, i) = i `seq` do
86− r <- step (adaptState gst) st
87− case r of
88− Yield x s -> let i' = i - 1
89− in return $ Yield (i, x) (s, i')
90− Skip s -> return $ Skip (s, i)
91− Stop -> return Stop
92−
93−{-# INLINE_NORMAL zipWithM #-}
94−zipWithM :: Monad m
95− => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
96−zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
97− where
98− {-# INLINE_LATE step #-}
99− step gst (sa, sb, Nothing) = do
100− r <- stepa (adaptState gst) sa
101− return $
102− case r of
103− Yield x sa' -> Skip (sa', sb, Just x)
104− Skip sa' -> Skip (sa', sb, Nothing)
105− Stop -> Stop
106−
107− step gst (sa, sb, Just x) = do
108− r <- stepb (adaptState gst) sb
109− case r of
110− Yield y sb' -> do
111− z <- f x y
112− return $ Yield z (sa, sb', Nothing)
113− Skip sb' -> return $ Skip (sa, sb', Just x)
114− Stop -> return Stop
115−
116−#if __GLASGOW_HASKELL__ >= 801
117−{-# RULES "zipWithM xs xs"
118− forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
119−#endif
120−
121−{-# INLINE zipWith #-}
122−zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
123−zipWith f = zipWithM (\a b -> return (f a b))
124−
125−------------------------------------------------------------------------------
126−-- Merging
127−------------------------------------------------------------------------------
128−
129−{-# INLINE_NORMAL mergeByM #-}
130−mergeByM
131− :: (Monad m)
132− => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
133−mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
134− Stream step (Just ta, Just tb, Nothing, Nothing)
135− where
136− {-# INLINE_LATE step #-}
137−
138− -- one of the values is missing, and the corresponding stream is running
139− step gst (Just sa, sb, Nothing, b) = do
140− r <- stepa gst sa
141− return $ case r of
142− Yield a sa' -> Skip (Just sa', sb, Just a, b)
143− Skip sa' -> Skip (Just sa', sb, Nothing, b)
144− Stop -> Skip (Nothing, sb, Nothing, b)
145−
146− step gst (sa, Just sb, a, Nothing) = do
147− r <- stepb gst sb
148− return $ case r of
149− Yield b sb' -> Skip (sa, Just sb', a, Just b)
150− Skip sb' -> Skip (sa, Just sb', a, Nothing)
151− Stop -> Skip (sa, Nothing, a, Nothing)
152−
153− -- both the values are available
154− step _ (sa, sb, Just a, Just b) = do
155− res <- cmp a b
156− return $ case res of
157− GT -> Yield b (sa, sb, Just a, Nothing)
158− _ -> Yield a (sa, sb, Nothing, Just b)
159−
160− -- one of the values is missing, corresponding stream is done
161− step _ (Nothing, sb, Nothing, Just b) =
162− return $ Yield b (Nothing, sb, Nothing, Nothing)
163−
164− step _ (sa, Nothing, Just a, Nothing) =
165− return $ Yield a (sa, Nothing, Nothing, Nothing)
166−
167− step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
168−
169−{-# INLINE mergeBy #-}
170−mergeBy
171− :: (Monad m)
172− => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
173−mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
174−
175−-------------------------------------------------------------------------------
176−-- Stream transformations using Unfolds
177−-------------------------------------------------------------------------------
178−
179−-- Define a unique structure to use in inspection testing
180−data ConcatMapUState o i =
181− ConcatMapUOuter o
182− | ConcatMapUInner o i
183−
184−-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
185−-- to streams and then flattens the generated streams into a single output
186−-- stream.
187−
188−-- This is like 'concatMap' but uses an unfold with an explicit state to
189−-- generate the stream instead of a 'Stream' type generator. This allows better
190−-- optimization via fusion. This can be many times more efficient than
191−-- 'concatMap'.
192−
193−{-# INLINE_NORMAL concatMapU #-}
194−concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
195−concatMapU (Unfold istep inject) (Stream ostep ost) =
196− Stream step (ConcatMapUOuter ost)
197− where
198− {-# INLINE_LATE step #-}
199− step gst (ConcatMapUOuter o) = do
200− r <- ostep (adaptState gst) o
201− case r of
202− Yield a o' -> do
203− i <- inject a
204− i `seq` return (Skip (ConcatMapUInner o' i))
205− Skip o' -> return $ Skip (ConcatMapUOuter o')
206− Stop -> return $ Stop
207−
208− step _ (ConcatMapUInner o i) = do
209− r <- istep i
210− return $ case r of
211− Yield x i' -> Yield x (ConcatMapUInner o i')
212− Skip i' -> Skip (ConcatMapUInner o i')
213− Stop -> Skip (ConcatMapUOuter o)
214−
215−data ConcatUnfoldInterleaveState o i =
216− ConcatUnfoldInterleaveOuter o [i]
217− | ConcatUnfoldInterleaveInner o [i]
218− | ConcatUnfoldInterleaveInnerL [i] [i]
219− | ConcatUnfoldInterleaveInnerR [i] [i]
220−
221−-- XXX use arrays to store state instead of lists.
222−-- XXX In general we can use different scheduling strategies e.g. how to
223−-- schedule the outer vs inner loop or assigning weights to different streams
224−-- or outer and inner loops.
225−
226−-- After a yield, switch to the next stream. Do not switch streams on Skip.
227−-- Yield from outer stream switches to the inner stream.
228−--
229−-- There are two choices here, (1) exhaust the outer stream first and then
230−-- start yielding from the inner streams, this is much simpler to implement,
231−-- (2) yield at least one element from an inner stream before going back to
232−-- outer stream and opening the next stream from it.
233−--
234−-- Ideally, we need some scheduling bias to inner streams vs outer stream.
235−-- Maybe we can configure the behavior.
236−--
237−{-# INLINE_NORMAL concatUnfoldInterleave #-}
238−concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
239−concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
240− Stream step (ConcatUnfoldInterleaveOuter ost [])
241− where
242− {-# INLINE_LATE step #-}
243− step gst (ConcatUnfoldInterleaveOuter o ls) = do
244− r <- ostep (adaptState gst) o
245− case r of
246− Yield a o' -> do
247− i <- inject a
248− i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
249− Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
250− Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
251−
252− step _ (ConcatUnfoldInterleaveInner _ []) = undefined
253− step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
254− r <- istep st
255− return $ case r of
256− Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
257− Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
258− Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
259−
260− step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
261− step _ (ConcatUnfoldInterleaveInnerL [] rs) =
262− return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
263−
264− step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
265− r <- istep st
266− return $ case r of
267− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
268− Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
269− Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
270−
271− step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
272− step _ (ConcatUnfoldInterleaveInnerR ls []) =
273− return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
274−
275− step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
276− r <- istep st
277− return $ case r of
278− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
279− Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
280− Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
281−
282−-- XXX In general we can use different scheduling strategies e.g. how to
283−-- schedule the outer vs inner loop or assigning weights to different streams
284−-- or outer and inner loops.
285−--
286−-- This could be inefficient if the tasks are too small.
287−--
288−-- Compared to concatUnfoldInterleave this one switches streams on Skips.
289−--
290−{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
291−concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
292−concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
293− Stream step (ConcatUnfoldInterleaveOuter ost [])
294− where
295− {-# INLINE_LATE step #-}
296− step gst (ConcatUnfoldInterleaveOuter o ls) = do
297− r <- ostep (adaptState gst) o
298− case r of
299− Yield a o' -> do
300− i <- inject a
301− i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
302− Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
303− Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
304−
305− step _ (ConcatUnfoldInterleaveInner o []) =
306− return $ Skip (ConcatUnfoldInterleaveOuter o [])
307−
308− step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
309− r <- istep st
310− return $ case r of
311− Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
312− Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
313− Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
314−
315− step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
316− step _ (ConcatUnfoldInterleaveInnerL [] rs) =
317− return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
318−
319− step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
320− r <- istep st
321− return $ case r of
322− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
323− Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
324− Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
325−
326− step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
327− step _ (ConcatUnfoldInterleaveInnerR ls []) =
328− return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
329−
330− step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
331− r <- istep st
332− return $ case r of
333− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
334− Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
335− Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
336−
337−data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
338−
339−-- Note that this could be much faster compared to the CPS stream. However, as
340−-- the number of streams being composed increases this may become expensive.
341−-- Need to see where the breaking point is between the two.
342−--
343−{-# INLINE_NORMAL append #-}
344−append :: Monad m => Stream m a -> Stream m a -> Stream m a
345−append (Stream step1 state1) (Stream step2 state2) =
346− Stream step (AppendFirst state1)
347−
348− where
349−
350− {-# INLINE_LATE step #-}
351− step gst (AppendFirst st) = do
352− r <- step1 gst st
353− return $ case r of
354− Yield a s -> Yield a (AppendFirst s)
355− Skip s -> Skip (AppendFirst s)
356− Stop -> Skip (AppendSecond state2)
357−
358− step gst (AppendSecond st) = do
359− r <- step2 gst st
360− return $ case r of
361− Yield a s -> Yield a (AppendSecond s)
362− Skip s -> Skip (AppendSecond s)
363− Stop -> Stop
364−
365−data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
366− | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
367−
368−{-# INLINE_NORMAL interleave #-}
369−interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
370−interleave (Stream step1 state1) (Stream step2 state2) =
371− Stream step (InterleaveFirst state1 state2)
372−
373− where
374−
375− {-# INLINE_LATE step #-}
376− step gst (InterleaveFirst st1 st2) = do
377− r <- step1 gst st1
378− return $ case r of
379− Yield a s -> Yield a (InterleaveSecond s st2)
380− Skip s -> Skip (InterleaveFirst s st2)
381− Stop -> Skip (InterleaveSecondOnly st2)
382−
383− step gst (InterleaveSecond st1 st2) = do
384− r <- step2 gst st2
385− return $ case r of
386− Yield a s -> Yield a (InterleaveFirst st1 s)
387− Skip s -> Skip (InterleaveSecond st1 s)
388− Stop -> Skip (InterleaveFirstOnly st1)
389−
390− step gst (InterleaveFirstOnly st1) = do
391− r <- step1 gst st1
392− return $ case r of
393− Yield a s -> Yield a (InterleaveFirstOnly s)
394− Skip s -> Skip (InterleaveFirstOnly s)
395− Stop -> Stop
396−
397− step gst (InterleaveSecondOnly st2) = do
398− r <- step2 gst st2
399− return $ case r of
400− Yield a s -> Yield a (InterleaveSecondOnly s)
401− Skip s -> Skip (InterleaveSecondOnly s)
402− Stop -> Stop
403−
404−{-# INLINE_NORMAL interleaveMin #-}
405−interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
406−interleaveMin (Stream step1 state1) (Stream step2 state2) =
407− Stream step (InterleaveFirst state1 state2)
408−
409− where
410−
411− {-# INLINE_LATE step #-}
412− step gst (InterleaveFirst st1 st2) = do
413− r <- step1 gst st1
414− return $ case r of
415− Yield a s -> Yield a (InterleaveSecond s st2)
416− Skip s -> Skip (InterleaveFirst s st2)
417− Stop -> Stop
418−
419− step gst (InterleaveSecond st1 st2) = do
420− r <- step2 gst st2
421− return $ case r of
422− Yield a s -> Yield a (InterleaveFirst st1 s)
423− Skip s -> Skip (InterleaveSecond st1 s)
424− Stop -> Stop
425−
426− step _ (InterleaveFirstOnly _) = undefined
427− step _ (InterleaveSecondOnly _) = undefined
428−
429−{-# INLINE_NORMAL interleaveSuffix #-}
430−interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
431−interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
432− Stream step (InterleaveFirst state1 state2)
433−
434− where
435−
436− {-# INLINE_LATE step #-}
437− step gst (InterleaveFirst st1 st2) = do
438− r <- step1 gst st1
439− return $ case r of
440− Yield a s -> Yield a (InterleaveSecond s st2)
441− Skip s -> Skip (InterleaveFirst s st2)
442− Stop -> Stop
443−
444− step gst (InterleaveSecond st1 st2) = do
445− r <- step2 gst st2
446− return $ case r of
447− Yield a s -> Yield a (InterleaveFirst st1 s)
448− Skip s -> Skip (InterleaveSecond st1 s)
449− Stop -> Skip (InterleaveFirstOnly st1)
450−
451− step gst (InterleaveFirstOnly st1) = do
452− r <- step1 gst st1
453− return $ case r of
454− Yield a s -> Yield a (InterleaveFirstOnly s)
455− Skip s -> Skip (InterleaveFirstOnly s)
456− Stop -> Stop
457−
458− step _ (InterleaveSecondOnly _) = undefined
459−
460−data InterleaveInfixState s1 s2 a
461− = InterleaveInfixFirst s1 s2
462− | InterleaveInfixSecondBuf s1 s2
463− | InterleaveInfixSecondYield s1 s2 a
464− | InterleaveInfixFirstYield s1 s2 a
465− | InterleaveInfixFirstOnly s1
466−
467−{-# INLINE_NORMAL interleaveInfix #-}
468−interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
469−interleaveInfix (Stream step1 state1) (Stream step2 state2) =
470− Stream step (InterleaveInfixFirst state1 state2)
471−
472− where
473−
474− {-# INLINE_LATE step #-}
475− step gst (InterleaveInfixFirst st1 st2) = do
476− r <- step1 gst st1
477− return $ case r of
478− Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
479− Skip s -> Skip (InterleaveInfixFirst s st2)
480− Stop -> Stop
481−
482− step gst (InterleaveInfixSecondBuf st1 st2) = do
483− r <- step2 gst st2
484− return $ case r of
485− Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
486− Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
487− Stop -> Skip (InterleaveInfixFirstOnly st1)
488−
489− step gst (InterleaveInfixSecondYield st1 st2 x) = do
490− r <- step1 gst st1
491− return $ case r of
492− Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
493− Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
494− Stop -> Stop
495−
496− step _ (InterleaveInfixFirstYield st1 st2 x) = do
497− return $ Yield x (InterleaveInfixSecondBuf st1 st2)
498−
499− step gst (InterleaveInfixFirstOnly st1) = do
500− r <- step1 gst st1
501− return $ case r of
502− Yield a s -> Yield a (InterleaveInfixFirstOnly s)
503− Skip s -> Skip (InterleaveInfixFirstOnly s)
504− Stop -> Stop
505−
506−{-# INLINE_NORMAL roundRobin #-}
507−roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
508−roundRobin (Stream step1 state1) (Stream step2 state2) =
509− Stream step (InterleaveFirst state1 state2)
510−
511− where
512−
513− {-# INLINE_LATE step #-}
514− step gst (InterleaveFirst st1 st2) = do
515− r <- step1 gst st1
516− return $ case r of
517− Yield a s -> Yield a (InterleaveSecond s st2)
518− Skip s -> Skip (InterleaveSecond s st2)
519− Stop -> Skip (InterleaveSecondOnly st2)
520−
521− step gst (InterleaveSecond st1 st2) = do
522− r <- step2 gst st2
523− return $ case r of
524− Yield a s -> Yield a (InterleaveFirst st1 s)
525− Skip s -> Skip (InterleaveFirst st1 s)
526− Stop -> Skip (InterleaveFirstOnly st1)
527−
528− step gst (InterleaveSecondOnly st2) = do
529− r <- step2 gst st2
530− return $ case r of
531− Yield a s -> Yield a (InterleaveSecondOnly s)
532− Skip s -> Skip (InterleaveSecondOnly s)
533− Stop -> Stop
534−
535− step gst (InterleaveFirstOnly st1) = do
536− r <- step1 gst st1
537− return $ case r of
538− Yield a s -> Yield a (InterleaveFirstOnly s)
539− Skip s -> Skip (InterleaveFirstOnly s)
540− Stop -> Stop
541−
542−data ICUState s1 s2 i1 i2 =
543− ICUFirst s1 s2
544− | ICUSecond s1 s2
545− | ICUSecondOnly s2
546− | ICUFirstOnly s1
547− | ICUFirstInner s1 s2 i1
548− | ICUSecondInner s1 s2 i2
549− | ICUFirstOnlyInner s1 i1
550− | ICUSecondOnlyInner s2 i2
551−
552−-- | Interleave streams (full streams, not the elements) unfolded from two
553−-- input streams and concat. Stop when the first stream stops. If the second
554−-- stream ends before the first one then first stream still keeps running alone
555−-- without any interleaving with the second stream.
556−--
557−-- [a1, a2, ... an] [b1, b2 ...]
558−-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
559−-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
560−-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
561−--
562−{-# INLINE_NORMAL gintercalateSuffix #-}
563−gintercalateSuffix
564− :: Monad m
565− => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
566−gintercalateSuffix
567− (Unfold istep1 inject1) (Stream step1 state1)
568− (Unfold istep2 inject2) (Stream step2 state2) =
569− Stream step (ICUFirst state1 state2)
570−
571− where
572−
573− {-# INLINE_LATE step #-}
574− step gst (ICUFirst s1 s2) = do
575− r <- step1 (adaptState gst) s1
576− case r of
577− Yield a s -> do
578− i <- inject1 a
579− i `seq` return (Skip (ICUFirstInner s s2 i))
580− Skip s -> return $ Skip (ICUFirst s s2)
581− Stop -> return Stop
582−
583− step gst (ICUFirstOnly s1) = do
584− r <- step1 (adaptState gst) s1
585− case r of
586− Yield a s -> do
587− i <- inject1 a
588− i `seq` return (Skip (ICUFirstOnlyInner s i))
589− Skip s -> return $ Skip (ICUFirstOnly s)
590− Stop -> return Stop
591−
592− step _ (ICUFirstInner s1 s2 i1) = do
593− r <- istep1 i1
594− return $ case r of
595− Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
596− Skip i' -> Skip (ICUFirstInner s1 s2 i')
597− Stop -> Skip (ICUSecond s1 s2)
598−
599− step _ (ICUFirstOnlyInner s1 i1) = do
600− r <- istep1 i1
601− return $ case r of
602− Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
603− Skip i' -> Skip (ICUFirstOnlyInner s1 i')
604− Stop -> Skip (ICUFirstOnly s1)
605−
606− step gst (ICUSecond s1 s2) = do
607− r <- step2 (adaptState gst) s2
608− case r of
609− Yield a s -> do
610− i <- inject2 a
611− i `seq` return (Skip (ICUSecondInner s1 s i))
612− Skip s -> return $ Skip (ICUSecond s1 s)
613− Stop -> return $ Skip (ICUFirstOnly s1)
614−
615− step _ (ICUSecondInner s1 s2 i2) = do
616− r <- istep2 i2
617− return $ case r of
618− Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
619− Skip i' -> Skip (ICUSecondInner s1 s2 i')
620− Stop -> Skip (ICUFirst s1 s2)
621−
622− step _ (ICUSecondOnly _s2) = undefined
623− step _ (ICUSecondOnlyInner _s2 _i2) = undefined
624−
625−{-# ANN type InterposeSuffixState Fuse #-}
626−data InterposeSuffixState s1 i1 =
627− InterposeSuffixFirst s1
628− -- | InterposeSuffixFirstYield s1 i1
629− | InterposeSuffixFirstInner s1 i1
630− | InterposeSuffixSecond s1
631−
632−-- Note that if an unfolded layer turns out to be nil we still emit the
633−-- separator effect. An alternate behavior could be to emit the separator
634−-- effect only if at least one element has been yielded by the unfolding.
635−-- However, that becomes a bit complicated, so we have chosen the former
636−-- behvaior for now.
637−{-# INLINE_NORMAL interposeSuffix #-}
638−interposeSuffix
639− :: Monad m
640− => m c -> Unfold m b c -> Stream m b -> Stream m c
641−interposeSuffix
642− action
643− (Unfold istep1 inject1) (Stream step1 state1) =
644− Stream step (InterposeSuffixFirst state1)
645−
646− where
647−
648− {-# INLINE_LATE step #-}
649− step gst (InterposeSuffixFirst s1) = do
650− r <- step1 (adaptState gst) s1
651− case r of
652− Yield a s -> do
653− i <- inject1 a
654− i `seq` return (Skip (InterposeSuffixFirstInner s i))
655− -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
656− Skip s -> return $ Skip (InterposeSuffixFirst s)
657− Stop -> return Stop
658−
659− {-
660− step _ (InterposeSuffixFirstYield s1 i1) = do
661− r <- istep1 i1
662− return $ case r of
663− Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
664− Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
665− Stop -> Skip (InterposeSuffixFirst s1)
666− -}
667−
668− step _ (InterposeSuffixFirstInner s1 i1) = do
669− r <- istep1 i1
670− return $ case r of
671− Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
672− Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
673− Stop -> Skip (InterposeSuffixSecond s1)
674−
675− step _ (InterposeSuffixSecond s1) = do
676− r <- action
677− return $ Yield r (InterposeSuffixFirst s1)
678−
679−data ICALState s1 s2 i1 i2 a =
680− ICALFirst s1 s2
681− -- | ICALFirstYield s1 s2 i1
682− | ICALFirstInner s1 s2 i1
683− | ICALFirstOnly s1
684− | ICALFirstOnlyInner s1 i1
685− | ICALSecondInject s1 s2
686− | ICALFirstInject s1 s2 i2
687− -- | ICALFirstBuf s1 s2 i1 i2
688− | ICALSecondInner s1 s2 i1 i2
689− -- -- | ICALSecondInner s1 s2 i1 i2 a
690− -- -- | ICALFirstResume s1 s2 i1 i2 a
691−
692−-- | Interleave streams (full streams, not the elements) unfolded from two
693−-- input streams and concat. Stop when the first stream stops. If the second
694−-- stream ends before the first one then first stream still keeps running alone
695−-- without any interleaving with the second stream.
696−--
697−-- [a1, a2, ... an] [b1, b2 ...]
698−-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
699−-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
700−-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
701−--
702−{-# INLINE_NORMAL gintercalate #-}
703−gintercalate
704− :: Monad m
705− => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
706−gintercalate
707− (Unfold istep1 inject1) (Stream step1 state1)
708− (Unfold istep2 inject2) (Stream step2 state2) =
709− Stream step (ICALFirst state1 state2)
710−
711− where
712−
713− {-# INLINE_LATE step #-}
714− step gst (ICALFirst s1 s2) = do
715− r <- step1 (adaptState gst) s1
716− case r of
717− Yield a s -> do
718− i <- inject1 a
719− i `seq` return (Skip (ICALFirstInner s s2 i))
720− -- i `seq` return (Skip (ICALFirstYield s s2 i))
721− Skip s -> return $ Skip (ICALFirst s s2)
722− Stop -> return Stop
723−
724− {-
725− step _ (ICALFirstYield s1 s2 i1) = do
726− r <- istep1 i1
727− return $ case r of
728− Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
729− Skip i' -> Skip (ICALFirstYield s1 s2 i')
730− Stop -> Skip (ICALFirst s1 s2)
731− -}
732−
733− step _ (ICALFirstInner s1 s2 i1) = do
734− r <- istep1 i1
735− return $ case r of
736− Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
737− Skip i' -> Skip (ICALFirstInner s1 s2 i')
738− Stop -> Skip (ICALSecondInject s1 s2)
739−
740− step gst (ICALFirstOnly s1) = do
741− r <- step1 (adaptState gst) s1
742− case r of
743− Yield a s -> do
744− i <- inject1 a
745− i `seq` return (Skip (ICALFirstOnlyInner s i))
746− Skip s -> return $ Skip (ICALFirstOnly s)
747− Stop -> return Stop
748−
749− step _ (ICALFirstOnlyInner s1 i1) = do
750− r <- istep1 i1
751− return $ case r of
752− Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
753− Skip i' -> Skip (ICALFirstOnlyInner s1 i')
754− Stop -> Skip (ICALFirstOnly s1)
755−
756− -- We inject the second stream even before checking if the first stream
757− -- would yield any more elements. There is no clear choice whether we
758− -- should do this before or after that. Doing it after may make the state
759− -- machine a bit simpler though.
760− step gst (ICALSecondInject s1 s2) = do
761− r <- step2 (adaptState gst) s2
762− case r of
763− Yield a s -> do
764− i <- inject2 a
765− i `seq` return (Skip (ICALFirstInject s1 s i))
766− Skip s -> return $ Skip (ICALSecondInject s1 s)
767− Stop -> return $ Skip (ICALFirstOnly s1)
768−
769− step gst (ICALFirstInject s1 s2 i2) = do
770− r <- step1 (adaptState gst) s1
771− case r of
772− Yield a s -> do
773− i <- inject1 a
774− i `seq` return (Skip (ICALSecondInner s s2 i i2))
775− -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
776− Skip s -> return $ Skip (ICALFirstInject s s2 i2)
777− Stop -> return Stop
778−
779− {-
780− step _ (ICALFirstBuf s1 s2 i1 i2) = do
781− r <- istep1 i1
782− return $ case r of
783− Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
784− Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
785− Stop -> Stop
786−
787− step _ (ICALSecondInner s1 s2 i1 i2 v) = do
788− r <- istep2 i2
789− return $ case r of
790− Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
791− Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
792− Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
793− -}
794−
795− step _ (ICALSecondInner s1 s2 i1 i2) = do
796− r <- istep2 i2
797− return $ case r of
798− Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
799− Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
800− Stop -> Skip (ICALFirstInner s1 s2 i1)
801− -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
802−
803− {-
804− step _ (ICALFirstResume s1 s2 i1 i2 x) = do
805− return $ Yield x (ICALFirstInner s1 s2 i1 i2)
806− -}
807−
808−{-# ANN type InterposeState Fuse #-}
809−data InterposeState s1 i1 a =
810− InterposeFirst s1
811− -- | InterposeFirstYield s1 i1
812− | InterposeFirstInner s1 i1
813− | InterposeFirstInject s1
814− -- | InterposeFirstBuf s1 i1
815− | InterposeSecondYield s1 i1
816− -- -- | InterposeSecondYield s1 i1 a
817− -- -- | InterposeFirstResume s1 i1 a
818−
819−-- Note that this only interposes the pure values, we may run many effects to
820−-- generate those values as some effects may not generate anything (Skip).
821−{-# INLINE_NORMAL interpose #-}
822−interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
823−interpose
824− action
825− (Unfold istep1 inject1) (Stream step1 state1) =
826− Stream step (InterposeFirst state1)
827−
828− where
829−
830− {-# INLINE_LATE step #-}
831− step gst (InterposeFirst s1) = do
832− r <- step1 (adaptState gst) s1
833− case r of
834− Yield a s -> do
835− i <- inject1 a
836− i `seq` return (Skip (InterposeFirstInner s i))
837− -- i `seq` return (Skip (InterposeFirstYield s i))
838− Skip s -> return $ Skip (InterposeFirst s)
839− Stop -> return Stop
840−
841− {-
842− step _ (InterposeFirstYield s1 i1) = do
843− r <- istep1 i1
844− return $ case r of
845− Yield x i' -> Yield x (InterposeFirstInner s1 i')
846− Skip i' -> Skip (InterposeFirstYield s1 i')
847− Stop -> Skip (InterposeFirst s1)
848− -}
849−
850− step _ (InterposeFirstInner s1 i1) = do
851− r <- istep1 i1
852− return $ case r of
853− Yield x i' -> Yield x (InterposeFirstInner s1 i')
854− Skip i' -> Skip (InterposeFirstInner s1 i')
855− Stop -> Skip (InterposeFirstInject s1)
856−
857− step gst (InterposeFirstInject s1) = do
858− r <- step1 (adaptState gst) s1
859− case r of
860− Yield a s -> do
861− i <- inject1 a
862− -- i `seq` return (Skip (InterposeFirstBuf s i))
863− i `seq` return (Skip (InterposeSecondYield s i))
864− Skip s -> return $ Skip (InterposeFirstInject s)
865− Stop -> return Stop
866−
867− {-
868− step _ (InterposeFirstBuf s1 i1) = do
869− r <- istep1 i1
870− return $ case r of
871− Yield x i' -> Skip (InterposeSecondYield s1 i' x)
872− Skip i' -> Skip (InterposeFirstBuf s1 i')
873− Stop -> Stop
874− -}
875−
876− {-
877− step _ (InterposeSecondYield s1 i1 v) = do
878− r <- action
879− return $ Yield r (InterposeFirstResume s1 i1 v)
880− -}
881− step _ (InterposeSecondYield s1 i1) = do
882− r <- action
883− return $ Yield r (InterposeFirstInner s1 i1)
884−
885− {-
886− step _ (InterposeFirstResume s1 i1 v) = do
887− return $ Yield v (InterposeFirstInner s1 i1)
888− -}
src/Streamly/Internal/Data/Stream/StreamD/SplitGroup.hs → src/Streamly/Internal/Data/Stream/StreamD/Nesting.hsrenamed+872−12
@@ -1,17 +1,29 @@
1#include "inline.hs"
1#include "inline.hs"
2
2
3-- |
3-- |
4−-- Module : Streamly.Internal.Data.Stream.StreamD.SplitGroup
4+-- Module : Streamly.Internal.Data.Stream.StreamD.Nesting
5−-- Copyright : (c) 2020 Composewell Technologies
5+-- Copyright : (c) 2018 Composewell Technologies
6+-- (c) Roman Leshchinskiy 2008-2010
7+-- (c) The University of Glasgow, 2009
6−-- License : BSD-3-Clause
8+-- License : BSD3
7-- Maintainer : streamly@composewell.com
9-- Maintainer : streamly@composewell.com
8-- Stability : experimental
10-- Stability : experimental
9-- Portability : GHC
11-- Portability : GHC
10
12
11−module Streamly.Internal.Data.Stream.StreamD.SplitGroup
13+module Streamly.Internal.Data.Stream.StreamD.Nesting
12 (
14 (
15+ -- * Zipping
16+ indexed
17+ , indexedR
18+ , zipWith
19+ , zipWithM
20+
21+ -- * Merging
22+ , mergeBy
23+ , mergeByM
24+
13 -- ** Grouping
25 -- ** Grouping
14− groupsOf2
26+ , groupsOf2
15 , groupsBy
27 , groupsBy
16 , groupsRollingBy
28 , groupsRollingBy
17
29
@@ -23,31 +35,164 @@ module Streamly.Internal.Data.Stream.StreamD.SplitGroup
23
35
24 , splitInnerBy
36 , splitInnerBy
25 , splitInnerBySuffix
37 , splitInnerBySuffix
38+
39+ -- ** Flattening nested streams
40+ , concatMapM
41+ , concatMap
42+ , ConcatMapUState (..)
43+ , concatMapU
44+ , ConcatUnfoldInterleaveState (..)
45+ , concatUnfoldInterleave
46+ , concatUnfoldRoundrobin
47+ , AppendState(..)
48+ , append
49+ , InterleaveState(..)
50+ , interleave
51+ , interleaveMin
52+ , interleaveSuffix
53+ , interleaveInfix
54+ , roundRobin -- interleaveFair?/ParallelFair
55+ , gintercalateSuffix
56+ , interposeSuffix
57+ , gintercalate
58+ , interpose
26 )
59 )
27where
60where
28
61
29import Control.Monad.IO.Class (MonadIO(..))
62import Control.Monad.IO.Class (MonadIO(..))
30import Data.Bits (shiftR, shiftL, (.|.), (.&.))
63import Data.Bits (shiftR, shiftL, (.|.), (.&.))
64+#if __GLASGOW_HASKELL__ >= 801
65+import Data.Functor.Identity ( Identity )
66+#endif
31import Data.Word (Word32)
67import Data.Word (Word32)
32import Foreign.Storable (Storable(..))
68import Foreign.Storable (Storable(..))
33−import GHC.Types (SPEC(..))
34import Fusion.Plugin.Types (Fuse(..))
69import Fusion.Plugin.Types (Fuse(..))
70+import GHC.Types (SPEC(..))
71+
35import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
72import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
36import Streamly.Internal.Data.Fold.Types (Fold(..))
73import Streamly.Internal.Data.Fold.Types (Fold(..))
74+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
37
75
38import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
76import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
39import qualified Streamly.Internal.Data.Fold as FL
77import qualified Streamly.Internal.Data.Fold as FL
40import qualified Streamly.Memory.Ring as RB
78import qualified Streamly.Memory.Ring as RB
41
79
42−import Prelude hiding
43− ( map, mapM, mapM_, repeat, foldr, last, take, filter
44− , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
45− , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
46− , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
47− , reverse, iterate, splitAt)
48import Streamly.Internal.Data.Stream.StreamD.Type
80import Streamly.Internal.Data.Stream.StreamD.Type
49import Streamly.Internal.Data.SVar
81import Streamly.Internal.Data.SVar
50
82
83+import Prelude hiding (concatMap, mapM, zipWith)
84+
85+------------------------------------------------------------------------------
86+-- Zipping
87+------------------------------------------------------------------------------
88+
89+{-# INLINE_NORMAL indexed #-}
90+indexed :: Monad m => Stream m a -> Stream m (Int, a)
91+indexed (Stream step state) = Stream step' (state, 0)
92+ where
93+ {-# INLINE_LATE step' #-}
94+ step' gst (st, i) = i `seq` do
95+ r <- step (adaptState gst) st
96+ case r of
97+ Yield x s -> return $ Yield (i, x) (s, i+1)
98+ Skip s -> return $ Skip (s, i)
99+ Stop -> return Stop
100+
101+{-# INLINE_NORMAL indexedR #-}
102+indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
103+indexedR m (Stream step state) = Stream step' (state, m)
104+ where
105+ {-# INLINE_LATE step' #-}
106+ step' gst (st, i) = i `seq` do
107+ r <- step (adaptState gst) st
108+ case r of
109+ Yield x s -> let i' = i - 1
110+ in return $ Yield (i, x) (s, i')
111+ Skip s -> return $ Skip (s, i)
112+ Stop -> return Stop
113+
114+{-# INLINE_NORMAL zipWithM #-}
115+zipWithM :: Monad m
116+ => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
117+zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
118+ where
119+ {-# INLINE_LATE step #-}
120+ step gst (sa, sb, Nothing) = do
121+ r <- stepa (adaptState gst) sa
122+ return $
123+ case r of
124+ Yield x sa' -> Skip (sa', sb, Just x)
125+ Skip sa' -> Skip (sa', sb, Nothing)
126+ Stop -> Stop
127+
128+ step gst (sa, sb, Just x) = do
129+ r <- stepb (adaptState gst) sb
130+ case r of
131+ Yield y sb' -> do
132+ z <- f x y
133+ return $ Yield z (sa, sb', Nothing)
134+ Skip sb' -> return $ Skip (sa, sb', Just x)
135+ Stop -> return Stop
136+
137+#if __GLASGOW_HASKELL__ >= 801
138+{-# RULES "zipWithM xs xs"
139+ forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
140+#endif
141+
142+{-# INLINE zipWith #-}
143+zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
144+zipWith f = zipWithM (\a b -> return (f a b))
145+
146+------------------------------------------------------------------------------
147+-- Merging
148+------------------------------------------------------------------------------
149+
150+{-# INLINE_NORMAL mergeByM #-}
151+mergeByM
152+ :: (Monad m)
153+ => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
154+mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
155+ Stream step (Just ta, Just tb, Nothing, Nothing)
156+ where
157+ {-# INLINE_LATE step #-}
158+
159+ -- one of the values is missing, and the corresponding stream is running
160+ step gst (Just sa, sb, Nothing, b) = do
161+ r <- stepa gst sa
162+ return $ case r of
163+ Yield a sa' -> Skip (Just sa', sb, Just a, b)
164+ Skip sa' -> Skip (Just sa', sb, Nothing, b)
165+ Stop -> Skip (Nothing, sb, Nothing, b)
166+
167+ step gst (sa, Just sb, a, Nothing) = do
168+ r <- stepb gst sb
169+ return $ case r of
170+ Yield b sb' -> Skip (sa, Just sb', a, Just b)
171+ Skip sb' -> Skip (sa, Just sb', a, Nothing)
172+ Stop -> Skip (sa, Nothing, a, Nothing)
173+
174+ -- both the values are available
175+ step _ (sa, sb, Just a, Just b) = do
176+ res <- cmp a b
177+ return $ case res of
178+ GT -> Yield b (sa, sb, Just a, Nothing)
179+ _ -> Yield a (sa, sb, Nothing, Just b)
180+
181+ -- one of the values is missing, corresponding stream is done
182+ step _ (Nothing, sb, Nothing, Just b) =
183+ return $ Yield b (Nothing, sb, Nothing, Nothing)
184+
185+ step _ (sa, Nothing, Just a, Nothing) =
186+ return $ Yield a (sa, Nothing, Nothing, Nothing)
187+
188+ step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
189+
190+{-# INLINE mergeBy #-}
191+mergeBy
192+ :: (Monad m)
193+ => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
194+mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
195+
51------------------------------------------------------------------------------
196------------------------------------------------------------------------------
52-- Grouping/Splitting
197-- Grouping/Splitting
53------------------------------------------------------------------------------
198------------------------------------------------------------------------------
@@ -1003,3 +1148,718 @@ splitInnerBySuffix splitter joiner (Stream step1 state1) =
1003
1148
1004 step _ (SplitYielding x next) = return $ Yield x next
1149 step _ (SplitYielding x next) = return $ Yield x next
1005 step _ SplitFinishing = return Stop
1150 step _ SplitFinishing = return Stop
1151+
1152+-------------------------------------------------------------------------------
1153+-- Stream transformations using Unfolds
1154+-------------------------------------------------------------------------------
1155+
1156+-- Define a unique structure to use in inspection testing
1157+data ConcatMapUState o i =
1158+ ConcatMapUOuter o
1159+ | ConcatMapUInner o i
1160+
1161+-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
1162+-- to streams and then flattens the generated streams into a single output
1163+-- stream.
1164+
1165+-- This is like 'concatMap' but uses an unfold with an explicit state to
1166+-- generate the stream instead of a 'Stream' type generator. This allows better
1167+-- optimization via fusion. This can be many times more efficient than
1168+-- 'concatMap'.
1169+
1170+{-# INLINE_NORMAL concatMapU #-}
1171+concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
1172+concatMapU (Unfold istep inject) (Stream ostep ost) =
1173+ Stream step (ConcatMapUOuter ost)
1174+ where
1175+ {-# INLINE_LATE step #-}
1176+ step gst (ConcatMapUOuter o) = do
1177+ r <- ostep (adaptState gst) o
1178+ case r of
1179+ Yield a o' -> do
1180+ i <- inject a
1181+ i `seq` return (Skip (ConcatMapUInner o' i))
1182+ Skip o' -> return $ Skip (ConcatMapUOuter o')
1183+ Stop -> return $ Stop
1184+
1185+ step _ (ConcatMapUInner o i) = do
1186+ r <- istep i
1187+ return $ case r of
1188+ Yield x i' -> Yield x (ConcatMapUInner o i')
1189+ Skip i' -> Skip (ConcatMapUInner o i')
1190+ Stop -> Skip (ConcatMapUOuter o)
1191+
1192+data ConcatUnfoldInterleaveState o i =
1193+ ConcatUnfoldInterleaveOuter o [i]
1194+ | ConcatUnfoldInterleaveInner o [i]
1195+ | ConcatUnfoldInterleaveInnerL [i] [i]
1196+ | ConcatUnfoldInterleaveInnerR [i] [i]
1197+
1198+-- XXX use arrays to store state instead of lists.
1199+-- XXX In general we can use different scheduling strategies e.g. how to
1200+-- schedule the outer vs inner loop or assigning weights to different streams
1201+-- or outer and inner loops.
1202+
1203+-- After a yield, switch to the next stream. Do not switch streams on Skip.
1204+-- Yield from outer stream switches to the inner stream.
1205+--
1206+-- There are two choices here, (1) exhaust the outer stream first and then
1207+-- start yielding from the inner streams, this is much simpler to implement,
1208+-- (2) yield at least one element from an inner stream before going back to
1209+-- outer stream and opening the next stream from it.
1210+--
1211+-- Ideally, we need some scheduling bias to inner streams vs outer stream.
1212+-- Maybe we can configure the behavior.
1213+--
1214+{-# INLINE_NORMAL concatUnfoldInterleave #-}
1215+concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
1216+concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
1217+ Stream step (ConcatUnfoldInterleaveOuter ost [])
1218+ where
1219+ {-# INLINE_LATE step #-}
1220+ step gst (ConcatUnfoldInterleaveOuter o ls) = do
1221+ r <- ostep (adaptState gst) o
1222+ case r of
1223+ Yield a o' -> do
1224+ i <- inject a
1225+ i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
1226+ Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
1227+ Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1228+
1229+ step _ (ConcatUnfoldInterleaveInner _ []) = undefined
1230+ step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
1231+ r <- istep st
1232+ return $ case r of
1233+ Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
1234+ Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
1235+ Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
1236+
1237+ step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
1238+ step _ (ConcatUnfoldInterleaveInnerL [] rs) =
1239+ return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
1240+
1241+ step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
1242+ r <- istep st
1243+ return $ case r of
1244+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
1245+ Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
1246+ Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
1247+
1248+ step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
1249+ step _ (ConcatUnfoldInterleaveInnerR ls []) =
1250+ return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1251+
1252+ step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
1253+ r <- istep st
1254+ return $ case r of
1255+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
1256+ Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
1257+ Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
1258+
1259+-- XXX In general we can use different scheduling strategies e.g. how to
1260+-- schedule the outer vs inner loop or assigning weights to different streams
1261+-- or outer and inner loops.
1262+--
1263+-- This could be inefficient if the tasks are too small.
1264+--
1265+-- Compared to concatUnfoldInterleave this one switches streams on Skips.
1266+--
1267+{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
1268+concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
1269+concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
1270+ Stream step (ConcatUnfoldInterleaveOuter ost [])
1271+ where
1272+ {-# INLINE_LATE step #-}
1273+ step gst (ConcatUnfoldInterleaveOuter o ls) = do
1274+ r <- ostep (adaptState gst) o
1275+ case r of
1276+ Yield a o' -> do
1277+ i <- inject a
1278+ i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
1279+ Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
1280+ Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1281+
1282+ step _ (ConcatUnfoldInterleaveInner o []) =
1283+ return $ Skip (ConcatUnfoldInterleaveOuter o [])
1284+
1285+ step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
1286+ r <- istep st
1287+ return $ case r of
1288+ Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
1289+ Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
1290+ Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
1291+
1292+ step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
1293+ step _ (ConcatUnfoldInterleaveInnerL [] rs) =
1294+ return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
1295+
1296+ step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
1297+ r <- istep st
1298+ return $ case r of
1299+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
1300+ Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
1301+ Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
1302+
1303+ step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
1304+ step _ (ConcatUnfoldInterleaveInnerR ls []) =
1305+ return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1306+
1307+ step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
1308+ r <- istep st
1309+ return $ case r of
1310+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
1311+ Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
1312+ Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
1313+
1314+data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
1315+
1316+-- Note that this could be much faster compared to the CPS stream. However, as
1317+-- the number of streams being composed increases this may become expensive.
1318+-- Need to see where the breaking point is between the two.
1319+--
1320+{-# INLINE_NORMAL append #-}
1321+append :: Monad m => Stream m a -> Stream m a -> Stream m a
1322+append (Stream step1 state1) (Stream step2 state2) =
1323+ Stream step (AppendFirst state1)
1324+
1325+ where
1326+
1327+ {-# INLINE_LATE step #-}
1328+ step gst (AppendFirst st) = do
1329+ r <- step1 gst st
1330+ return $ case r of
1331+ Yield a s -> Yield a (AppendFirst s)
1332+ Skip s -> Skip (AppendFirst s)
1333+ Stop -> Skip (AppendSecond state2)
1334+
1335+ step gst (AppendSecond st) = do
1336+ r <- step2 gst st
1337+ return $ case r of
1338+ Yield a s -> Yield a (AppendSecond s)
1339+ Skip s -> Skip (AppendSecond s)
1340+ Stop -> Stop
1341+
1342+data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
1343+ | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
1344+
1345+{-# INLINE_NORMAL interleave #-}
1346+interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
1347+interleave (Stream step1 state1) (Stream step2 state2) =
1348+ Stream step (InterleaveFirst state1 state2)
1349+
1350+ where
1351+
1352+ {-# INLINE_LATE step #-}
1353+ step gst (InterleaveFirst st1 st2) = do
1354+ r <- step1 gst st1
1355+ return $ case r of
1356+ Yield a s -> Yield a (InterleaveSecond s st2)
1357+ Skip s -> Skip (InterleaveFirst s st2)
1358+ Stop -> Skip (InterleaveSecondOnly st2)
1359+
1360+ step gst (InterleaveSecond st1 st2) = do
1361+ r <- step2 gst st2
1362+ return $ case r of
1363+ Yield a s -> Yield a (InterleaveFirst st1 s)
1364+ Skip s -> Skip (InterleaveSecond st1 s)
1365+ Stop -> Skip (InterleaveFirstOnly st1)
1366+
1367+ step gst (InterleaveFirstOnly st1) = do
1368+ r <- step1 gst st1
1369+ return $ case r of
1370+ Yield a s -> Yield a (InterleaveFirstOnly s)
1371+ Skip s -> Skip (InterleaveFirstOnly s)
1372+ Stop -> Stop
1373+
1374+ step gst (InterleaveSecondOnly st2) = do
1375+ r <- step2 gst st2
1376+ return $ case r of
1377+ Yield a s -> Yield a (InterleaveSecondOnly s)
1378+ Skip s -> Skip (InterleaveSecondOnly s)
1379+ Stop -> Stop
1380+
1381+{-# INLINE_NORMAL interleaveMin #-}
1382+interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
1383+interleaveMin (Stream step1 state1) (Stream step2 state2) =
1384+ Stream step (InterleaveFirst state1 state2)
1385+
1386+ where
1387+
1388+ {-# INLINE_LATE step #-}
1389+ step gst (InterleaveFirst st1 st2) = do
1390+ r <- step1 gst st1
1391+ return $ case r of
1392+ Yield a s -> Yield a (InterleaveSecond s st2)
1393+ Skip s -> Skip (InterleaveFirst s st2)
1394+ Stop -> Stop
1395+
1396+ step gst (InterleaveSecond st1 st2) = do
1397+ r <- step2 gst st2
1398+ return $ case r of
1399+ Yield a s -> Yield a (InterleaveFirst st1 s)
1400+ Skip s -> Skip (InterleaveSecond st1 s)
1401+ Stop -> Stop
1402+
1403+ step _ (InterleaveFirstOnly _) = undefined
1404+ step _ (InterleaveSecondOnly _) = undefined
1405+
1406+{-# INLINE_NORMAL interleaveSuffix #-}
1407+interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
1408+interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
1409+ Stream step (InterleaveFirst state1 state2)
1410+
1411+ where
1412+
1413+ {-# INLINE_LATE step #-}
1414+ step gst (InterleaveFirst st1 st2) = do
1415+ r <- step1 gst st1
1416+ return $ case r of
1417+ Yield a s -> Yield a (InterleaveSecond s st2)
1418+ Skip s -> Skip (InterleaveFirst s st2)
1419+ Stop -> Stop
1420+
1421+ step gst (InterleaveSecond st1 st2) = do
1422+ r <- step2 gst st2
1423+ return $ case r of
1424+ Yield a s -> Yield a (InterleaveFirst st1 s)
1425+ Skip s -> Skip (InterleaveSecond st1 s)
1426+ Stop -> Skip (InterleaveFirstOnly st1)
1427+
1428+ step gst (InterleaveFirstOnly st1) = do
1429+ r <- step1 gst st1
1430+ return $ case r of
1431+ Yield a s -> Yield a (InterleaveFirstOnly s)
1432+ Skip s -> Skip (InterleaveFirstOnly s)
1433+ Stop -> Stop
1434+
1435+ step _ (InterleaveSecondOnly _) = undefined
1436+
1437+data InterleaveInfixState s1 s2 a
1438+ = InterleaveInfixFirst s1 s2
1439+ | InterleaveInfixSecondBuf s1 s2
1440+ | InterleaveInfixSecondYield s1 s2 a
1441+ | InterleaveInfixFirstYield s1 s2 a
1442+ | InterleaveInfixFirstOnly s1
1443+
1444+{-# INLINE_NORMAL interleaveInfix #-}
1445+interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
1446+interleaveInfix (Stream step1 state1) (Stream step2 state2) =
1447+ Stream step (InterleaveInfixFirst state1 state2)
1448+
1449+ where
1450+
1451+ {-# INLINE_LATE step #-}
1452+ step gst (InterleaveInfixFirst st1 st2) = do
1453+ r <- step1 gst st1
1454+ return $ case r of
1455+ Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
1456+ Skip s -> Skip (InterleaveInfixFirst s st2)
1457+ Stop -> Stop
1458+
1459+ step gst (InterleaveInfixSecondBuf st1 st2) = do
1460+ r <- step2 gst st2
1461+ return $ case r of
1462+ Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
1463+ Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
1464+ Stop -> Skip (InterleaveInfixFirstOnly st1)
1465+
1466+ step gst (InterleaveInfixSecondYield st1 st2 x) = do
1467+ r <- step1 gst st1
1468+ return $ case r of
1469+ Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
1470+ Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
1471+ Stop -> Stop
1472+
1473+ step _ (InterleaveInfixFirstYield st1 st2 x) = do
1474+ return $ Yield x (InterleaveInfixSecondBuf st1 st2)
1475+
1476+ step gst (InterleaveInfixFirstOnly st1) = do
1477+ r <- step1 gst st1
1478+ return $ case r of
1479+ Yield a s -> Yield a (InterleaveInfixFirstOnly s)
1480+ Skip s -> Skip (InterleaveInfixFirstOnly s)
1481+ Stop -> Stop
1482+
1483+{-# INLINE_NORMAL roundRobin #-}
1484+roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
1485+roundRobin (Stream step1 state1) (Stream step2 state2) =
1486+ Stream step (InterleaveFirst state1 state2)
1487+
1488+ where
1489+
1490+ {-# INLINE_LATE step #-}
1491+ step gst (InterleaveFirst st1 st2) = do
1492+ r <- step1 gst st1
1493+ return $ case r of
1494+ Yield a s -> Yield a (InterleaveSecond s st2)
1495+ Skip s -> Skip (InterleaveSecond s st2)
1496+ Stop -> Skip (InterleaveSecondOnly st2)
1497+
1498+ step gst (InterleaveSecond st1 st2) = do
1499+ r <- step2 gst st2
1500+ return $ case r of
1501+ Yield a s -> Yield a (InterleaveFirst st1 s)
1502+ Skip s -> Skip (InterleaveFirst st1 s)
1503+ Stop -> Skip (InterleaveFirstOnly st1)
1504+
1505+ step gst (InterleaveSecondOnly st2) = do
1506+ r <- step2 gst st2
1507+ return $ case r of
1508+ Yield a s -> Yield a (InterleaveSecondOnly s)
1509+ Skip s -> Skip (InterleaveSecondOnly s)
1510+ Stop -> Stop
1511+
1512+ step gst (InterleaveFirstOnly st1) = do
1513+ r <- step1 gst st1
1514+ return $ case r of
1515+ Yield a s -> Yield a (InterleaveFirstOnly s)
1516+ Skip s -> Skip (InterleaveFirstOnly s)
1517+ Stop -> Stop
1518+
1519+data ICUState s1 s2 i1 i2 =
1520+ ICUFirst s1 s2
1521+ | ICUSecond s1 s2
1522+ | ICUSecondOnly s2
1523+ | ICUFirstOnly s1
1524+ | ICUFirstInner s1 s2 i1
1525+ | ICUSecondInner s1 s2 i2
1526+ | ICUFirstOnlyInner s1 i1
1527+ | ICUSecondOnlyInner s2 i2
1528+
1529+-- | Interleave streams (full streams, not the elements) unfolded from two
1530+-- input streams and concat. Stop when the first stream stops. If the second
1531+-- stream ends before the first one then first stream still keeps running alone
1532+-- without any interleaving with the second stream.
1533+--
1534+-- [a1, a2, ... an] [b1, b2 ...]
1535+-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
1536+-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
1537+-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
1538+--
1539+{-# INLINE_NORMAL gintercalateSuffix #-}
1540+gintercalateSuffix
1541+ :: Monad m
1542+ => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
1543+gintercalateSuffix
1544+ (Unfold istep1 inject1) (Stream step1 state1)
1545+ (Unfold istep2 inject2) (Stream step2 state2) =
1546+ Stream step (ICUFirst state1 state2)
1547+
1548+ where
1549+
1550+ {-# INLINE_LATE step #-}
1551+ step gst (ICUFirst s1 s2) = do
1552+ r <- step1 (adaptState gst) s1
1553+ case r of
1554+ Yield a s -> do
1555+ i <- inject1 a
1556+ i `seq` return (Skip (ICUFirstInner s s2 i))
1557+ Skip s -> return $ Skip (ICUFirst s s2)
1558+ Stop -> return Stop
1559+
1560+ step gst (ICUFirstOnly s1) = do
1561+ r <- step1 (adaptState gst) s1
1562+ case r of
1563+ Yield a s -> do
1564+ i <- inject1 a
1565+ i `seq` return (Skip (ICUFirstOnlyInner s i))
1566+ Skip s -> return $ Skip (ICUFirstOnly s)
1567+ Stop -> return Stop
1568+
1569+ step _ (ICUFirstInner s1 s2 i1) = do
1570+ r <- istep1 i1
1571+ return $ case r of
1572+ Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
1573+ Skip i' -> Skip (ICUFirstInner s1 s2 i')
1574+ Stop -> Skip (ICUSecond s1 s2)
1575+
1576+ step _ (ICUFirstOnlyInner s1 i1) = do
1577+ r <- istep1 i1
1578+ return $ case r of
1579+ Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
1580+ Skip i' -> Skip (ICUFirstOnlyInner s1 i')
1581+ Stop -> Skip (ICUFirstOnly s1)
1582+
1583+ step gst (ICUSecond s1 s2) = do
1584+ r <- step2 (adaptState gst) s2
1585+ case r of
1586+ Yield a s -> do
1587+ i <- inject2 a
1588+ i `seq` return (Skip (ICUSecondInner s1 s i))
1589+ Skip s -> return $ Skip (ICUSecond s1 s)
1590+ Stop -> return $ Skip (ICUFirstOnly s1)
1591+
1592+ step _ (ICUSecondInner s1 s2 i2) = do
1593+ r <- istep2 i2
1594+ return $ case r of
1595+ Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
1596+ Skip i' -> Skip (ICUSecondInner s1 s2 i')
1597+ Stop -> Skip (ICUFirst s1 s2)
1598+
1599+ step _ (ICUSecondOnly _s2) = undefined
1600+ step _ (ICUSecondOnlyInner _s2 _i2) = undefined
1601+
1602+{-# ANN type InterposeSuffixState Fuse #-}
1603+data InterposeSuffixState s1 i1 =
1604+ InterposeSuffixFirst s1
1605+ -- | InterposeSuffixFirstYield s1 i1
1606+ | InterposeSuffixFirstInner s1 i1
1607+ | InterposeSuffixSecond s1
1608+
1609+-- Note that if an unfolded layer turns out to be nil we still emit the
1610+-- separator effect. An alternate behavior could be to emit the separator
1611+-- effect only if at least one element has been yielded by the unfolding.
1612+-- However, that becomes a bit complicated, so we have chosen the former
1613+-- behvaior for now.
1614+{-# INLINE_NORMAL interposeSuffix #-}
1615+interposeSuffix
1616+ :: Monad m
1617+ => m c -> Unfold m b c -> Stream m b -> Stream m c
1618+interposeSuffix
1619+ action
1620+ (Unfold istep1 inject1) (Stream step1 state1) =
1621+ Stream step (InterposeSuffixFirst state1)
1622+
1623+ where
1624+
1625+ {-# INLINE_LATE step #-}
1626+ step gst (InterposeSuffixFirst s1) = do
1627+ r <- step1 (adaptState gst) s1
1628+ case r of
1629+ Yield a s -> do
1630+ i <- inject1 a
1631+ i `seq` return (Skip (InterposeSuffixFirstInner s i))
1632+ -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
1633+ Skip s -> return $ Skip (InterposeSuffixFirst s)
1634+ Stop -> return Stop
1635+
1636+ {-
1637+ step _ (InterposeSuffixFirstYield s1 i1) = do
1638+ r <- istep1 i1
1639+ return $ case r of
1640+ Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
1641+ Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
1642+ Stop -> Skip (InterposeSuffixFirst s1)
1643+ -}
1644+
1645+ step _ (InterposeSuffixFirstInner s1 i1) = do
1646+ r <- istep1 i1
1647+ return $ case r of
1648+ Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
1649+ Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
1650+ Stop -> Skip (InterposeSuffixSecond s1)
1651+
1652+ step _ (InterposeSuffixSecond s1) = do
1653+ r <- action
1654+ return $ Yield r (InterposeSuffixFirst s1)
1655+
1656+data ICALState s1 s2 i1 i2 a =
1657+ ICALFirst s1 s2
1658+ -- | ICALFirstYield s1 s2 i1
1659+ | ICALFirstInner s1 s2 i1
1660+ | ICALFirstOnly s1
1661+ | ICALFirstOnlyInner s1 i1
1662+ | ICALSecondInject s1 s2
1663+ | ICALFirstInject s1 s2 i2
1664+ -- | ICALFirstBuf s1 s2 i1 i2
1665+ | ICALSecondInner s1 s2 i1 i2
1666+ -- -- | ICALSecondInner s1 s2 i1 i2 a
1667+ -- -- | ICALFirstResume s1 s2 i1 i2 a
1668+
1669+-- | Interleave streams (full streams, not the elements) unfolded from two
1670+-- input streams and concat. Stop when the first stream stops. If the second
1671+-- stream ends before the first one then first stream still keeps running alone
1672+-- without any interleaving with the second stream.
1673+--
1674+-- [a1, a2, ... an] [b1, b2 ...]
1675+-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
1676+-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
1677+-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
1678+--
1679+{-# INLINE_NORMAL gintercalate #-}
1680+gintercalate
1681+ :: Monad m
1682+ => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
1683+gintercalate
1684+ (Unfold istep1 inject1) (Stream step1 state1)
1685+ (Unfold istep2 inject2) (Stream step2 state2) =
1686+ Stream step (ICALFirst state1 state2)
1687+
1688+ where
1689+
1690+ {-# INLINE_LATE step #-}
1691+ step gst (ICALFirst s1 s2) = do
1692+ r <- step1 (adaptState gst) s1
1693+ case r of
1694+ Yield a s -> do
1695+ i <- inject1 a
1696+ i `seq` return (Skip (ICALFirstInner s s2 i))
1697+ -- i `seq` return (Skip (ICALFirstYield s s2 i))
1698+ Skip s -> return $ Skip (ICALFirst s s2)
1699+ Stop -> return Stop
1700+
1701+ {-
1702+ step _ (ICALFirstYield s1 s2 i1) = do
1703+ r <- istep1 i1
1704+ return $ case r of
1705+ Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
1706+ Skip i' -> Skip (ICALFirstYield s1 s2 i')
1707+ Stop -> Skip (ICALFirst s1 s2)
1708+ -}
1709+
1710+ step _ (ICALFirstInner s1 s2 i1) = do
1711+ r <- istep1 i1
1712+ return $ case r of
1713+ Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
1714+ Skip i' -> Skip (ICALFirstInner s1 s2 i')
1715+ Stop -> Skip (ICALSecondInject s1 s2)
1716+
1717+ step gst (ICALFirstOnly s1) = do
1718+ r <- step1 (adaptState gst) s1
1719+ case r of
1720+ Yield a s -> do
1721+ i <- inject1 a
1722+ i `seq` return (Skip (ICALFirstOnlyInner s i))
1723+ Skip s -> return $ Skip (ICALFirstOnly s)
1724+ Stop -> return Stop
1725+
1726+ step _ (ICALFirstOnlyInner s1 i1) = do
1727+ r <- istep1 i1
1728+ return $ case r of
1729+ Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
1730+ Skip i' -> Skip (ICALFirstOnlyInner s1 i')
1731+ Stop -> Skip (ICALFirstOnly s1)
1732+
1733+ -- We inject the second stream even before checking if the first stream
1734+ -- would yield any more elements. There is no clear choice whether we
1735+ -- should do this before or after that. Doing it after may make the state
1736+ -- machine a bit simpler though.
1737+ step gst (ICALSecondInject s1 s2) = do
1738+ r <- step2 (adaptState gst) s2
1739+ case r of
1740+ Yield a s -> do
1741+ i <- inject2 a
1742+ i `seq` return (Skip (ICALFirstInject s1 s i))
1743+ Skip s -> return $ Skip (ICALSecondInject s1 s)
1744+ Stop -> return $ Skip (ICALFirstOnly s1)
1745+
1746+ step gst (ICALFirstInject s1 s2 i2) = do
1747+ r <- step1 (adaptState gst) s1
1748+ case r of
1749+ Yield a s -> do
1750+ i <- inject1 a
1751+ i `seq` return (Skip (ICALSecondInner s s2 i i2))
1752+ -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
1753+ Skip s -> return $ Skip (ICALFirstInject s s2 i2)
1754+ Stop -> return Stop
1755+
1756+ {-
1757+ step _ (ICALFirstBuf s1 s2 i1 i2) = do
1758+ r <- istep1 i1
1759+ return $ case r of
1760+ Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
1761+ Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
1762+ Stop -> Stop
1763+
1764+ step _ (ICALSecondInner s1 s2 i1 i2 v) = do
1765+ r <- istep2 i2
1766+ return $ case r of
1767+ Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
1768+ Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
1769+ Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
1770+ -}
1771+
1772+ step _ (ICALSecondInner s1 s2 i1 i2) = do
1773+ r <- istep2 i2
1774+ return $ case r of
1775+ Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
1776+ Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
1777+ Stop -> Skip (ICALFirstInner s1 s2 i1)
1778+ -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
1779+
1780+ {-
1781+ step _ (ICALFirstResume s1 s2 i1 i2 x) = do
1782+ return $ Yield x (ICALFirstInner s1 s2 i1 i2)
1783+ -}
1784+
1785+{-# ANN type InterposeState Fuse #-}
1786+data InterposeState s1 i1 a =
1787+ InterposeFirst s1
1788+ -- | InterposeFirstYield s1 i1
1789+ | InterposeFirstInner s1 i1
1790+ | InterposeFirstInject s1
1791+ -- | InterposeFirstBuf s1 i1
1792+ | InterposeSecondYield s1 i1
1793+ -- -- | InterposeSecondYield s1 i1 a
1794+ -- -- | InterposeFirstResume s1 i1 a
1795+
1796+-- Note that this only interposes the pure values, we may run many effects to
1797+-- generate those values as some effects may not generate anything (Skip).
1798+{-# INLINE_NORMAL interpose #-}
1799+interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
1800+interpose
1801+ action
1802+ (Unfold istep1 inject1) (Stream step1 state1) =
1803+ Stream step (InterposeFirst state1)
1804+
1805+ where
1806+
1807+ {-# INLINE_LATE step #-}
1808+ step gst (InterposeFirst s1) = do
1809+ r <- step1 (adaptState gst) s1
1810+ case r of
1811+ Yield a s -> do
1812+ i <- inject1 a
1813+ i `seq` return (Skip (InterposeFirstInner s i))
1814+ -- i `seq` return (Skip (InterposeFirstYield s i))
1815+ Skip s -> return $ Skip (InterposeFirst s)
1816+ Stop -> return Stop
1817+
1818+ {-
1819+ step _ (InterposeFirstYield s1 i1) = do
1820+ r <- istep1 i1
1821+ return $ case r of
1822+ Yield x i' -> Yield x (InterposeFirstInner s1 i')
1823+ Skip i' -> Skip (InterposeFirstYield s1 i')
1824+ Stop -> Skip (InterposeFirst s1)
1825+ -}
1826+
1827+ step _ (InterposeFirstInner s1 i1) = do
1828+ r <- istep1 i1
1829+ return $ case r of
1830+ Yield x i' -> Yield x (InterposeFirstInner s1 i')
1831+ Skip i' -> Skip (InterposeFirstInner s1 i')
1832+ Stop -> Skip (InterposeFirstInject s1)
1833+
1834+ step gst (InterposeFirstInject s1) = do
1835+ r <- step1 (adaptState gst) s1
1836+ case r of
1837+ Yield a s -> do
1838+ i <- inject1 a
1839+ -- i `seq` return (Skip (InterposeFirstBuf s i))
1840+ i `seq` return (Skip (InterposeSecondYield s i))
1841+ Skip s -> return $ Skip (InterposeFirstInject s)
1842+ Stop -> return Stop
1843+
1844+ {-
1845+ step _ (InterposeFirstBuf s1 i1) = do
1846+ r <- istep1 i1
1847+ return $ case r of
1848+ Yield x i' -> Skip (InterposeSecondYield s1 i' x)
1849+ Skip i' -> Skip (InterposeFirstBuf s1 i')
1850+ Stop -> Stop
1851+ -}
1852+
1853+ {-
1854+ step _ (InterposeSecondYield s1 i1 v) = do
1855+ r <- action
1856+ return $ Yield r (InterposeFirstResume s1 i1 v)
1857+ -}
1858+ step _ (InterposeSecondYield s1 i1) = do
1859+ r <- action
1860+ return $ Yield r (InterposeFirstInner s1 i1)
1861+
1862+ {-
1863+ step _ (InterposeFirstResume s1 i1 v) = do
1864+ return $ Yield v (InterposeFirstInner s1 i1)
1865+ -}
streamly.cabal+1−2
@@ -453,9 +453,8 @@ library
453 , Streamly.Internal.Data.Stream.StreamD.Type
453 , Streamly.Internal.Data.Stream.StreamD.Type
454 , Streamly.Internal.Data.Stream.StreamD
454 , Streamly.Internal.Data.Stream.StreamD
455 , Streamly.Internal.Data.Stream.StreamD.Generate
455 , Streamly.Internal.Data.Stream.StreamD.Generate
456− , Streamly.Internal.Data.Stream.StreamD.SplitGroup
457 , Streamly.Internal.Data.Stream.StreamD.Eliminate
456 , Streamly.Internal.Data.Stream.StreamD.Eliminate
458− , Streamly.Internal.Data.Stream.StreamD.Combine
457+ , Streamly.Internal.Data.Stream.StreamD.Nesting
459 , Streamly.Internal.Data.Stream.StreamD.Transform
458 , Streamly.Internal.Data.Stream.StreamD.Transform
460 , Streamly.Internal.Data.Stream.StreamD.Exceptions
459 , Streamly.Internal.Data.Stream.StreamD.Exceptions
461 , Streamly.Internal.Data.Stream.StreamDK.Type
460 , Streamly.Internal.Data.Stream.StreamDK.Type