@@ -1,17 +1,29 @@
11#include "inline.hs"
22
33-- |
4−-- Module : Streamly.Internal.Data.Stream.StreamD.SplitGroup
5−-- Copyright : (c) 2020 Composewell Technologies
6−-- License : BSD-3-Clause
4+-- Module : Streamly.Internal.Data.Stream.StreamD.Nesting
5+-- Copyright : (c) 2018 Composewell Technologies
6+-- (c) Roman Leshchinskiy 2008-2010
7+-- (c) The University of Glasgow, 2009
8+-- License : BSD3
79-- Maintainer : streamly@composewell.com
810-- Stability : experimental
911-- Portability : GHC
1012
11−module Streamly.Internal.Data.Stream.StreamD.SplitGroup
13+module Streamly.Internal.Data.Stream.StreamD.Nesting
1214 (
15+ -- * Zipping
16+ indexed
17+ , indexedR
18+ , zipWith
19+ , zipWithM
20+
21+ -- * Merging
22+ , mergeBy
23+ , mergeByM
24+
1325 -- ** Grouping
14− groupsOf2
26+ , groupsOf2
1527 , groupsBy
1628 , groupsRollingBy
1729
@@ -23,31 +35,164 @@ module Streamly.Internal.Data.Stream.StreamD.SplitGroup
2335
2436 , splitInnerBy
2537 , 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
2659 )
2760where
2861
2962import Control.Monad.IO.Class (MonadIO(..))
3063import Data.Bits (shiftR, shiftL, (.|.), (.&.))
64+#if __GLASGOW_HASKELL__ >= 801
65+import Data.Functor.Identity ( Identity )
66+#endif
3167import Data.Word (Word32)
3268import Foreign.Storable (Storable(..))
33−import GHC.Types (SPEC(..))
3469import Fusion.Plugin.Types (Fuse(..))
70+import GHC.Types (SPEC(..))
71+
3572import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
3673import Streamly.Internal.Data.Fold.Types (Fold(..))
74+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
3775
3876import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
3977import qualified Streamly.Internal.Data.Fold as FL
4078import qualified Streamly.Memory.Ring as RB
4179
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)
4880import Streamly.Internal.Data.Stream.StreamD.Type
4981import Streamly.Internal.Data.SVar
5082
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+
51196------------------------------------------------------------------------------
52197-- Grouping/Splitting
53198------------------------------------------------------------------------------
@@ -1003,3 +1148,718 @@ splitInnerBySuffix splitter joiner (Stream step1 state1) =
10031148
10041149 step _ (SplitYielding x next) = return $ Yield x next
10051150 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+ -}