git.iko.soy/streamly
on 09c955e help sign in
streamly/src/Streamly/Streams/Zip.hs
248 lines 8.0 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE ConstraintKinds #-}
3{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE FlexibleInstances #-}
5{-# LANGUAGE GeneralizedNewtypeDeriving#-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE StandaloneDeriving #-}
9{-# LANGUAGE UndecidableInstances #-} -- XXX
11-- |
12-- Module : Streamly.Streams.Zip
13-- Copyright : (c) 2017 Harendra Kumar
14--
15-- License : BSD3
16-- Maintainer : harendra.kumar@gmail.com
17-- Stability : experimental
18-- Portability : GHC
19--
20--
21module Streamly.Streams.Zip
22 (
23 zipWith
24 , zipWithM
25 , zipAsyncWith
26 , zipAsyncWithM
28 , ZipSerialM
29 , ZipSerial
30 , ZipStream -- deprecated
31 , zipSerially
32 , zipping -- deprecated
34 , ZipAsyncM
35 , ZipAsync
36 , zipAsyncly
37 , zippingAsync -- deprecated
38 )
39where
41import Data.Semigroup (Semigroup(..))
42import Prelude hiding (map, repeat, zipWith)
44import Streamly.Streams.StreamK (IsStream(..), Stream(..))
45import Streamly.Streams.Async (mkAsync')
46import Streamly.Streams.Serial (map)
47import Streamly.SVar (MonadAsync, rstState)
49import qualified Streamly.Streams.StreamK as K
51#include "Instances.hs"
53------------------------------------------------------------------------------
54-- Serial Zipping
55------------------------------------------------------------------------------
57{-# INLINE zipWithS #-}
58zipWithS :: (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
59zipWithS f m1 m2 = go m1 m2
60 where
61 go mx my = Stream $ \st stp sng yld -> do
62 let merge a ra =
63 let single2 b = sng (f a b)
64 yield2 b rb = yld (f a b) (go ra rb)
65 in unStream my (rstState st) stp single2 yield2
66 let single1 a = merge a K.nil
67 yield1 a ra = merge a ra
68 unStream mx (rstState st) stp single1 yield1
70-- | Zip two streams serially using a pure zipping function.
71--
72-- @since 0.1.0
73{-# INLINABLE zipWith #-}
74zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
75zipWith f m1 m2 = fromStream $ zipWithS f (toStream m1) (toStream m2)
77-- | Zip two streams serially using a monadic zipping function.
78--
79-- @since 0.1.0
80zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
81zipWithM f m1 m2 = fromStream $ go (toStream m1) (toStream m2)
82 where
83 go mx my = Stream $ \st stp sng yld -> do
84 let merge a ra =
85 let runIt x = unStream x (rstState st) stp sng yld
86 single2 b = f a b >>= sng
87 yield2 b rb = f a b >>= \x -> runIt (x `K.cons` go ra rb)
88 in unStream my (rstState st) stp single2 yield2
89 let single1 a = merge a K.nil
90 yield1 a ra = merge a ra
91 unStream mx (rstState st) stp single1 yield1
93------------------------------------------------------------------------------
94-- Serially Zipping Streams
95------------------------------------------------------------------------------
97-- | The applicative instance of 'ZipSerialM' zips a number of streams serially
98-- i.e. it produces one element from each stream serially and then zips all
99-- those elements.
100--
101-- @
102-- main = (toList . 'zipSerially' $ (,,) \<$\> s1 \<*\> s2 \<*\> s3) >>= print
103-- where s1 = fromFoldable [1, 2]
104-- s2 = fromFoldable [3, 4]
105-- s3 = fromFoldable [5, 6]
106-- @
107-- @
108-- [(1,3,5),(2,4,6)]
109-- @
110--
111-- The 'Semigroup' instance of this type works the same way as that of
112-- 'SerialT'.
113--
114-- @since 0.2.0
115newtype ZipSerialM m a = ZipSerialM {getZipSerialM :: Stream m a}
116 deriving (Semigroup, Monoid)
118-- |
119-- @since 0.1.0
120{-# DEPRECATED ZipStream "Please use 'ZipSerialM' instead." #-}
121type ZipStream = ZipSerialM
123-- | An IO stream whose applicative instance zips streams serially.
124--
125-- @since 0.2.0
126type ZipSerial a = ZipSerialM IO a
128-- | Fix the type of a polymorphic stream as 'ZipSerialM'.
129--
130-- @since 0.2.0
131zipSerially :: IsStream t => ZipSerialM m a -> t m a
132zipSerially = K.adapt
134-- | Same as 'zipSerially'.
135--
136-- @since 0.1.0
137{-# DEPRECATED zipping "Please use zipSerially instead." #-}
138zipping :: IsStream t => ZipSerialM m a -> t m a
139zipping = zipSerially
141instance IsStream ZipSerialM where
142 toStream = getZipSerialM
143 fromStream = ZipSerialM
145 {-# INLINE consM #-}
146 {-# SPECIALIZE consM :: IO a -> ZipSerialM IO a -> ZipSerialM IO a #-}
147 consM :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a
148 consM m r = fromStream $ K.consMSerial m (toStream r)
150 {-# INLINE (|:) #-}
151 {-# SPECIALIZE (|:) :: IO a -> ZipSerialM IO a -> ZipSerialM IO a #-}
152 (|:) :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a
153 m |: r = fromStream $ K.consMSerial m (toStream r)
155instance Monad m => Functor (ZipSerialM m) where
156 fmap = map
158instance Monad m => Applicative (ZipSerialM m) where
159 pure = ZipSerialM . K.repeat
160 m1 <*> m2 = fromStream $ zipWith id (toStream m1) (toStream m2)
162------------------------------------------------------------------------------
163-- Parallel Zipping
164------------------------------------------------------------------------------
166-- | Zip two streams concurrently (i.e. both the elements being zipped are
167-- generated concurrently) using a pure zipping function.
168--
169-- @since 0.1.0
170zipAsyncWith :: (IsStream t, MonadAsync m)
171 => (a -> b -> c) -> t m a -> t m b -> t m c
172zipAsyncWith f m1 m2 = fromStream $ Stream $ \st stp sng yld -> do
173 ma <- mkAsync' (rstState st) m1
174 mb <- mkAsync' (rstState st) m2
175 unStream (toStream (zipWith f ma mb)) (rstState st) stp sng yld
177-- | Zip two streams asyncly (i.e. both the elements being zipped are generated
178-- concurrently) using a monadic zipping function.
179--
180-- @since 0.4.0
181zipAsyncWithM :: (IsStream t, MonadAsync m)
182 => (a -> b -> m c) -> t m a -> t m b -> t m c
183zipAsyncWithM f m1 m2 = fromStream $ Stream $ \st stp sng yld -> do
184 ma <- mkAsync' (rstState st) m1
185 mb <- mkAsync' (rstState st) m2
186 unStream (toStream (zipWithM f ma mb)) (rstState st) stp sng yld
188------------------------------------------------------------------------------
189-- Parallely Zipping Streams
190------------------------------------------------------------------------------
191--
192-- | Like 'ZipSerialM' but zips in parallel, it generates all the elements to
193-- be zipped concurrently.
194--
195-- @
196-- main = (toList . 'zipAsyncly' $ (,,) \<$\> s1 \<*\> s2 \<*\> s3) >>= print
197-- where s1 = fromFoldable [1, 2]
198-- s2 = fromFoldable [3, 4]
199-- s3 = fromFoldable [5, 6]
200-- @
201-- @
202-- [(1,3,5),(2,4,6)]
203-- @
204--
205-- The 'Semigroup' instance of this type works the same way as that of
206-- 'SerialT'.
207--
208-- @since 0.2.0
209newtype ZipAsyncM m a = ZipAsyncM {getZipAsyncM :: Stream m a}
210 deriving (Semigroup, Monoid)
212-- | An IO stream whose applicative instance zips streams wAsyncly.
213--
214-- @since 0.2.0
215type ZipAsync a = ZipAsyncM IO a
217-- | Fix the type of a polymorphic stream as 'ZipAsyncM'.
218--
219-- @since 0.2.0
220zipAsyncly :: IsStream t => ZipAsyncM m a -> t m a
221zipAsyncly = K.adapt
223-- | Same as 'zipAsyncly'.
224--
225-- @since 0.1.0
226{-# DEPRECATED zippingAsync "Please use zipAsyncly instead." #-}
227zippingAsync :: IsStream t => ZipAsyncM m a -> t m a
228zippingAsync = zipAsyncly
229instance IsStream ZipAsyncM where
230 toStream = getZipAsyncM
231 fromStream = ZipAsyncM
233 {-# INLINE consM #-}
234 {-# SPECIALIZE consM :: IO a -> ZipAsyncM IO a -> ZipAsyncM IO a #-}
235 consM :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a
236 consM m r = fromStream $ K.consMSerial m (toStream r)
238 {-# INLINE (|:) #-}
239 {-# SPECIALIZE (|:) :: IO a -> ZipAsyncM IO a -> ZipAsyncM IO a #-}
240 (|:) :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a
241 m |: r = fromStream $ K.consMSerial m (toStream r)
243instance Monad m => Functor (ZipAsyncM m) where
244 fmap = map
246instance MonadAsync m => Applicative (ZipAsyncM m) where
247 pure = ZipAsyncM . K.repeat
248 m1 <*> m2 = zipAsyncWith id m1 m2