@@ -1,17 +1,29 @@
1#include "inline.hs"
1#include "inline.hs"
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
11−module Streamly.Internal.Data.Stream.StreamD.SplitGroup
13+module Streamly.Internal.Data.Stream.StreamD.Nesting
13 -- ** Grouping
25 -- ** Grouping
14− groupsOf2
26+ , groupsOf2
15 , groupsBy
27 , groupsBy
16 , groupsRollingBy
28 , groupsRollingBy
@@ -23,31 +35,164 @@ module Streamly.Internal.Data.Stream.StreamD.SplitGroup
24 , splitInnerBy
36 , splitInnerBy
25 , splitInnerBySuffix
37 , splitInnerBySuffix
39+ -- ** Flattening nested streams
42+ , ConcatMapUState (..)
44+ , ConcatUnfoldInterleaveState (..)
45+ , concatUnfoldInterleave
46+ , concatUnfoldRoundrobin
49+ , InterleaveState(..)
54+ , roundRobin -- interleaveFair?/ParallelFair
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 )
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(..))
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(..))
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
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
83+import Prelude hiding (concatMap, mapM, zipWith)
85+------------------------------------------------------------------------------
87+------------------------------------------------------------------------------
89+{-# INLINE_NORMAL indexed #-}
90+indexed :: Monad m => Stream m a -> Stream m (Int, a)
91+indexed (Stream step state) = Stream step' (state, 0)
93+ {-# INLINE_LATE step' #-}
94+ step' gst (st, i) = i `seq` do
95+ r <- step (adaptState gst) st
97+ Yield x s -> return $ Yield (i, x) (s, i+1)
98+ Skip s -> return $ Skip (s, i)
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)
105+ {-# INLINE_LATE step' #-}
106+ step' gst (st, i) = i `seq` do
107+ r <- step (adaptState gst) st
109+ Yield x s -> let i' = i - 1
110+ in return $ Yield (i, x) (s, i')
111+ Skip s -> return $ Skip (s, i)
114+{-# INLINE_NORMAL zipWithM #-}
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)
119+ {-# INLINE_LATE step #-}
120+ step gst (sa, sb, Nothing) = do
121+ r <- stepa (adaptState gst) sa
124+ Yield x sa' -> Skip (sa', sb, Just x)
125+ Skip sa' -> Skip (sa', sb, Nothing)
128+ step gst (sa, sb, Just x) = do
129+ r <- stepb (adaptState gst) sb
133+ return $ Yield z (sa, sb', Nothing)
134+ Skip sb' -> return $ Skip (sa, sb', Just x)
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 #-}
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))
146+------------------------------------------------------------------------------
148+------------------------------------------------------------------------------
150+{-# INLINE_NORMAL mergeByM #-}
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)
157+ {-# INLINE_LATE step #-}
159+ -- one of the values is missing, and the corresponding stream is running
160+ step gst (Just sa, sb, Nothing, b) = do
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)
167+ step gst (sa, Just sb, a, Nothing) = do
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)
174+ -- both the values are available
175+ step _ (sa, sb, Just a, Just b) = do
177+ return $ case res of
178+ GT -> Yield b (sa, sb, Just a, Nothing)
179+ _ -> Yield a (sa, sb, Nothing, Just b)
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)
185+ step _ (sa, Nothing, Just a, Nothing) =
186+ return $ Yield a (sa, Nothing, Nothing, Nothing)
188+ step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
190+{-# INLINE mergeBy #-}
193+ => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
194+mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
51------------------------------------------------------------------------------
196------------------------------------------------------------------------------
52-- Grouping/Splitting
197-- Grouping/Splitting
53------------------------------------------------------------------------------
198------------------------------------------------------------------------------
@@ -1003,3 +1148,718 @@ splitInnerBySuffix splitter joiner (Stream step1 state1) =
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
1152+-------------------------------------------------------------------------------
1153+-- Stream transformations using Unfolds
1154+-------------------------------------------------------------------------------
1156+-- Define a unique structure to use in inspection testing
1157+data ConcatMapUState o i =
1159+ | ConcatMapUInner o i
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
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
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)
1175+ {-# INLINE_LATE step #-}
1176+ step gst (ConcatMapUOuter o) = do
1177+ r <- ostep (adaptState gst) o
1181+ i `seq` return (Skip (ConcatMapUInner o' i))
1182+ Skip o' -> return $ Skip (ConcatMapUOuter o')
1183+ Stop -> return $ Stop
1185+ step _ (ConcatMapUInner o i) = do
1188+ Yield x i' -> Yield x (ConcatMapUInner o i')
1189+ Skip i' -> Skip (ConcatMapUInner o i')
1190+ Stop -> Skip (ConcatMapUOuter o)
1192+data ConcatUnfoldInterleaveState o i =
1193+ ConcatUnfoldInterleaveOuter o [i]
1194+ | ConcatUnfoldInterleaveInner o [i]
1195+ | ConcatUnfoldInterleaveInnerL [i] [i]
1196+ | ConcatUnfoldInterleaveInnerR [i] [i]
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.
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.
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.
1211+-- Ideally, we need some scheduling bias to inner streams vs outer stream.
1212+-- Maybe we can configure the behavior.
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 [])
1219+ {-# INLINE_LATE step #-}
1220+ step gst (ConcatUnfoldInterleaveOuter o ls) = do
1221+ r <- ostep (adaptState gst) o
1225+ i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
1226+ Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
1227+ Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1229+ step _ (ConcatUnfoldInterleaveInner _ []) = undefined
1230+ step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
1233+ Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
1234+ Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
1235+ Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
1237+ step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
1238+ step _ (ConcatUnfoldInterleaveInnerL [] rs) =
1239+ return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
1241+ step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
1244+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
1245+ Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
1246+ Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
1248+ step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
1249+ step _ (ConcatUnfoldInterleaveInnerR ls []) =
1250+ return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1252+ step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
1255+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
1256+ Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
1257+ Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
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.
1263+-- This could be inefficient if the tasks are too small.
1265+-- Compared to concatUnfoldInterleave this one switches streams on Skips.
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 [])
1272+ {-# INLINE_LATE step #-}
1273+ step gst (ConcatUnfoldInterleaveOuter o ls) = do
1274+ r <- ostep (adaptState gst) o
1278+ i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
1279+ Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
1280+ Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1282+ step _ (ConcatUnfoldInterleaveInner o []) =
1283+ return $ Skip (ConcatUnfoldInterleaveOuter o [])
1285+ step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
1288+ Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
1289+ Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
1290+ Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
1292+ step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
1293+ step _ (ConcatUnfoldInterleaveInnerL [] rs) =
1294+ return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
1296+ step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
1299+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
1300+ Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
1301+ Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
1303+ step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
1304+ step _ (ConcatUnfoldInterleaveInnerR ls []) =
1305+ return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
1307+ step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
1310+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
1311+ Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
1312+ Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
1314+data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
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.
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)
1327+ {-# INLINE_LATE step #-}
1328+ step gst (AppendFirst st) = do
1331+ Yield a s -> Yield a (AppendFirst s)
1332+ Skip s -> Skip (AppendFirst s)
1333+ Stop -> Skip (AppendSecond state2)
1335+ step gst (AppendSecond st) = do
1338+ Yield a s -> Yield a (AppendSecond s)
1339+ Skip s -> Skip (AppendSecond s)
1342+data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
1343+ | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
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)
1352+ {-# INLINE_LATE step #-}
1353+ step gst (InterleaveFirst st1 st2) = do
1356+ Yield a s -> Yield a (InterleaveSecond s st2)
1357+ Skip s -> Skip (InterleaveFirst s st2)
1358+ Stop -> Skip (InterleaveSecondOnly st2)
1360+ step gst (InterleaveSecond st1 st2) = do
1363+ Yield a s -> Yield a (InterleaveFirst st1 s)
1364+ Skip s -> Skip (InterleaveSecond st1 s)
1365+ Stop -> Skip (InterleaveFirstOnly st1)
1367+ step gst (InterleaveFirstOnly st1) = do
1370+ Yield a s -> Yield a (InterleaveFirstOnly s)
1371+ Skip s -> Skip (InterleaveFirstOnly s)
1374+ step gst (InterleaveSecondOnly st2) = do
1377+ Yield a s -> Yield a (InterleaveSecondOnly s)
1378+ Skip s -> Skip (InterleaveSecondOnly s)
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)
1388+ {-# INLINE_LATE step #-}
1389+ step gst (InterleaveFirst st1 st2) = do
1392+ Yield a s -> Yield a (InterleaveSecond s st2)
1393+ Skip s -> Skip (InterleaveFirst s st2)
1396+ step gst (InterleaveSecond st1 st2) = do
1399+ Yield a s -> Yield a (InterleaveFirst st1 s)
1400+ Skip s -> Skip (InterleaveSecond st1 s)
1403+ step _ (InterleaveFirstOnly _) = undefined
1404+ step _ (InterleaveSecondOnly _) = undefined
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)
1413+ {-# INLINE_LATE step #-}
1414+ step gst (InterleaveFirst st1 st2) = do
1417+ Yield a s -> Yield a (InterleaveSecond s st2)
1418+ Skip s -> Skip (InterleaveFirst s st2)
1421+ step gst (InterleaveSecond st1 st2) = do
1424+ Yield a s -> Yield a (InterleaveFirst st1 s)
1425+ Skip s -> Skip (InterleaveSecond st1 s)
1426+ Stop -> Skip (InterleaveFirstOnly st1)
1428+ step gst (InterleaveFirstOnly st1) = do
1431+ Yield a s -> Yield a (InterleaveFirstOnly s)
1432+ Skip s -> Skip (InterleaveFirstOnly s)
1435+ step _ (InterleaveSecondOnly _) = undefined
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
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)
1451+ {-# INLINE_LATE step #-}
1452+ step gst (InterleaveInfixFirst st1 st2) = do
1455+ Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
1456+ Skip s -> Skip (InterleaveInfixFirst s st2)
1459+ step gst (InterleaveInfixSecondBuf st1 st2) = do
1462+ Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
1463+ Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
1464+ Stop -> Skip (InterleaveInfixFirstOnly st1)
1466+ step gst (InterleaveInfixSecondYield st1 st2 x) = do
1469+ Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
1470+ Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
1473+ step _ (InterleaveInfixFirstYield st1 st2 x) = do
1474+ return $ Yield x (InterleaveInfixSecondBuf st1 st2)
1476+ step gst (InterleaveInfixFirstOnly st1) = do
1479+ Yield a s -> Yield a (InterleaveInfixFirstOnly s)
1480+ Skip s -> Skip (InterleaveInfixFirstOnly s)
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)
1490+ {-# INLINE_LATE step #-}
1491+ step gst (InterleaveFirst st1 st2) = do
1494+ Yield a s -> Yield a (InterleaveSecond s st2)
1495+ Skip s -> Skip (InterleaveSecond s st2)
1496+ Stop -> Skip (InterleaveSecondOnly st2)
1498+ step gst (InterleaveSecond st1 st2) = do
1501+ Yield a s -> Yield a (InterleaveFirst st1 s)
1502+ Skip s -> Skip (InterleaveFirst st1 s)
1503+ Stop -> Skip (InterleaveFirstOnly st1)
1505+ step gst (InterleaveSecondOnly st2) = do
1508+ Yield a s -> Yield a (InterleaveSecondOnly s)
1509+ Skip s -> Skip (InterleaveSecondOnly s)
1512+ step gst (InterleaveFirstOnly st1) = do
1515+ Yield a s -> Yield a (InterleaveFirstOnly s)
1516+ Skip s -> Skip (InterleaveFirstOnly s)
1519+data ICUState s1 s2 i1 i2 =
1524+ | ICUFirstInner s1 s2 i1
1525+ | ICUSecondInner s1 s2 i2
1526+ | ICUFirstOnlyInner s1 i1
1527+ | ICUSecondOnlyInner s2 i2
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.
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, ...]
1539+{-# INLINE_NORMAL gintercalateSuffix #-}
1542+ => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
1544+ (Unfold istep1 inject1) (Stream step1 state1)
1545+ (Unfold istep2 inject2) (Stream step2 state2) =
1546+ Stream step (ICUFirst state1 state2)
1550+ {-# INLINE_LATE step #-}
1551+ step gst (ICUFirst s1 s2) = do
1552+ r <- step1 (adaptState gst) s1
1556+ i `seq` return (Skip (ICUFirstInner s s2 i))
1557+ Skip s -> return $ Skip (ICUFirst s s2)
1558+ Stop -> return Stop
1560+ step gst (ICUFirstOnly s1) = do
1561+ r <- step1 (adaptState gst) s1
1565+ i `seq` return (Skip (ICUFirstOnlyInner s i))
1566+ Skip s -> return $ Skip (ICUFirstOnly s)
1567+ Stop -> return Stop
1569+ step _ (ICUFirstInner s1 s2 i1) = do
1572+ Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
1573+ Skip i' -> Skip (ICUFirstInner s1 s2 i')
1574+ Stop -> Skip (ICUSecond s1 s2)
1576+ step _ (ICUFirstOnlyInner s1 i1) = do
1579+ Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
1580+ Skip i' -> Skip (ICUFirstOnlyInner s1 i')
1581+ Stop -> Skip (ICUFirstOnly s1)
1583+ step gst (ICUSecond s1 s2) = do
1584+ r <- step2 (adaptState gst) s2
1588+ i `seq` return (Skip (ICUSecondInner s1 s i))
1589+ Skip s -> return $ Skip (ICUSecond s1 s)
1590+ Stop -> return $ Skip (ICUFirstOnly s1)
1592+ step _ (ICUSecondInner s1 s2 i2) = do
1595+ Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
1596+ Skip i' -> Skip (ICUSecondInner s1 s2 i')
1597+ Stop -> Skip (ICUFirst s1 s2)
1599+ step _ (ICUSecondOnly _s2) = undefined
1600+ step _ (ICUSecondOnlyInner _s2 _i2) = undefined
1602+{-# ANN type InterposeSuffixState Fuse #-}
1603+data InterposeSuffixState s1 i1 =
1604+ InterposeSuffixFirst s1
1605+ -- | InterposeSuffixFirstYield s1 i1
1606+ | InterposeSuffixFirstInner s1 i1
1607+ | InterposeSuffixSecond s1
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 #-}
1617+ => m c -> Unfold m b c -> Stream m b -> Stream m c
1620+ (Unfold istep1 inject1) (Stream step1 state1) =
1621+ Stream step (InterposeSuffixFirst state1)
1625+ {-# INLINE_LATE step #-}
1626+ step gst (InterposeSuffixFirst s1) = do
1627+ r <- step1 (adaptState gst) s1
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
1637+ step _ (InterposeSuffixFirstYield s1 i1) = do
1640+ Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
1641+ Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
1642+ Stop -> Skip (InterposeSuffixFirst s1)
1645+ step _ (InterposeSuffixFirstInner s1 i1) = do
1648+ Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
1649+ Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
1650+ Stop -> Skip (InterposeSuffixSecond s1)
1652+ step _ (InterposeSuffixSecond s1) = do
1654+ return $ Yield r (InterposeSuffixFirst s1)
1656+data ICALState s1 s2 i1 i2 a =
1658+ -- | ICALFirstYield s1 s2 i1
1659+ | ICALFirstInner s1 s2 i1
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
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.
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, ...]
1679+{-# INLINE_NORMAL gintercalate #-}
1682+ => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
1684+ (Unfold istep1 inject1) (Stream step1 state1)
1685+ (Unfold istep2 inject2) (Stream step2 state2) =
1686+ Stream step (ICALFirst state1 state2)
1690+ {-# INLINE_LATE step #-}
1691+ step gst (ICALFirst s1 s2) = do
1692+ r <- step1 (adaptState gst) s1
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
1702+ step _ (ICALFirstYield s1 s2 i1) = do
1705+ Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
1706+ Skip i' -> Skip (ICALFirstYield s1 s2 i')
1707+ Stop -> Skip (ICALFirst s1 s2)
1710+ step _ (ICALFirstInner s1 s2 i1) = do
1713+ Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
1714+ Skip i' -> Skip (ICALFirstInner s1 s2 i')
1715+ Stop -> Skip (ICALSecondInject s1 s2)
1717+ step gst (ICALFirstOnly s1) = do
1718+ r <- step1 (adaptState gst) s1
1722+ i `seq` return (Skip (ICALFirstOnlyInner s i))
1723+ Skip s -> return $ Skip (ICALFirstOnly s)
1724+ Stop -> return Stop
1726+ step _ (ICALFirstOnlyInner s1 i1) = do
1729+ Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
1730+ Skip i' -> Skip (ICALFirstOnlyInner s1 i')
1731+ Stop -> Skip (ICALFirstOnly s1)
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
1742+ i `seq` return (Skip (ICALFirstInject s1 s i))
1743+ Skip s -> return $ Skip (ICALSecondInject s1 s)
1744+ Stop -> return $ Skip (ICALFirstOnly s1)
1746+ step gst (ICALFirstInject s1 s2 i2) = do
1747+ r <- step1 (adaptState gst) s1
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
1757+ step _ (ICALFirstBuf s1 s2 i1 i2) = do
1760+ Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
1761+ Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
1764+ step _ (ICALSecondInner s1 s2 i1 i2 v) = do
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)
1772+ step _ (ICALSecondInner s1 s2 i1 i2) = do
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)
1781+ step _ (ICALFirstResume s1 s2 i1 i2 x) = do
1782+ return $ Yield x (ICALFirstInner s1 s2 i1 i2)
1785+{-# ANN type InterposeState Fuse #-}
1786+data InterposeState s1 i1 a =
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
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
1802+ (Unfold istep1 inject1) (Stream step1 state1) =
1803+ Stream step (InterposeFirst state1)
1807+ {-# INLINE_LATE step #-}
1808+ step gst (InterposeFirst s1) = do
1809+ r <- step1 (adaptState gst) s1
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
1819+ step _ (InterposeFirstYield s1 i1) = do
1822+ Yield x i' -> Yield x (InterposeFirstInner s1 i')
1823+ Skip i' -> Skip (InterposeFirstYield s1 i')
1824+ Stop -> Skip (InterposeFirst s1)
1827+ step _ (InterposeFirstInner s1 i1) = do
1830+ Yield x i' -> Yield x (InterposeFirstInner s1 i')
1831+ Skip i' -> Skip (InterposeFirstInner s1 i')
1832+ Stop -> Skip (InterposeFirstInject s1)
1834+ step gst (InterposeFirstInject s1) = do
1835+ r <- step1 (adaptState gst) s1
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
1845+ step _ (InterposeFirstBuf s1 i1) = do
1848+ Yield x i' -> Skip (InterposeSecondYield s1 i' x)
1849+ Skip i' -> Skip (InterposeFirstBuf s1 i')
1854+ step _ (InterposeSecondYield s1 i1 v) = do
1856+ return $ Yield r (InterposeFirstResume s1 i1 v)
1858+ step _ (InterposeSecondYield s1 i1) = do
1860+ return $ Yield r (InterposeFirstInner s1 i1)
1863+ step _ (InterposeFirstResume s1 i1 v) = do
1864+ return $ Yield v (InterposeFirstInner s1 i1)