@@ -1,7 +1,7 @@
11#include "inline.hs"
22
33-- |
4−-- Module : Streamly.Internal.Data.Stream.StreamD.Flatten
4+-- Module : Streamly.Internal.Data.Stream.StreamD.Combine
55-- Copyright : (c) 2018 Composewell Technologies
66-- (c) Roman Leshchinskiy 2008-2010
77-- (c) The University of Glasgow, 2009
@@ -10,10 +10,20 @@
1010-- Stability : experimental
1111-- Portability : GHC
1212
13−module Streamly.Internal.Data.Stream.StreamD.Flatten
13+module Streamly.Internal.Data.Stream.StreamD.Combine
1414 (
15+ -- * Zipping
16+ indexed
17+ , indexedR
18+ , zipWith
19+ , zipWithM
20+
21+ -- * Merging
22+ , mergeBy
23+ , mergeByM
24+
1525 -- ** Flattening nested streams
16− concatMapM
26+ , concatMapM
1727 , concatMap
1828 , ConcatMapUState (..)
1929 , concatMapU
@@ -35,6 +45,9 @@ module Streamly.Internal.Data.Stream.StreamD.Flatten
3545 )
3646where
3747
48+#if __GLASGOW_HASKELL__ >= 801
49+import Data.Functor.Identity ( Identity )
50+#endif
3851import Fusion.Plugin.Types (Fuse(..))
3952import Streamly.Internal.Data.Unfold.Types (Unfold(..))
4053
@@ -48,6 +61,117 @@ import Prelude hiding
4861import Streamly.Internal.Data.Stream.StreamD.Type
4962import Streamly.Internal.Data.SVar
5063
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+
51175-------------------------------------------------------------------------------
52176-- Stream transformations using Unfolds
53177-------------------------------------------------------------------------------