git.iko.soy/streamly
on 5949f4c help sign in
streamly/src/Streamly/Internal/Data/Stream/Zip.hs
273 lines 8.4 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 TypeFamilies #-}
9{-# LANGUAGE UndecidableInstances #-} -- XXX
11-- |
12-- Module : Streamly.Internal.Data.Stream.Zip
13-- Copyright : (c) 2017 Harendra Kumar
14--
15-- License : BSD3
16-- Maintainer : streamly@composewell.com
17-- Stability : experimental
18-- Portability : GHC
19--
20--
21module Streamly.Internal.Data.Stream.Zip
22 (
23 ZipSerialM
24 , ZipSerial
25 , zipSerially
27 , ZipAsyncM
28 , ZipAsync
29 , zipAsyncly
31 , zipWith
32 , zipWithM
33 , zipAsyncWith
34 , zipAsyncWithM
36 -- * Deprecated
37 , ZipStream
38 , zipping
39 , zippingAsync
40 )
41where
43import Control.Applicative (liftA2)
44import Control.DeepSeq (NFData(..))
45#if MIN_VERSION_deepseq(1,4,3)
46import Control.DeepSeq (NFData1(..))
47#endif
48import Data.Functor.Identity (Identity, runIdentity)
49import Data.Foldable (fold)
50#if __GLASGOW_HASKELL__ < 808
51import Data.Semigroup (Semigroup(..))
52#endif
53import GHC.Exts (IsList(..), IsString(..))
54import Text.Read (Lexeme(Ident), lexP, parens, prec, readPrec, readListPrec,
55 readListPrecDefault)
56import Prelude hiding (map, repeat, zipWith)
58import Streamly.Internal.Data.Stream.StreamK (IsStream(..), Stream)
59import Streamly.Internal.Data.SVar (MonadAsync, adaptState)
61import qualified Streamly.Internal.Data.Stream.Prelude as P
62import qualified Streamly.Internal.Data.Stream.StreamK as K
63import qualified Streamly.Internal.Data.Stream.StreamD as D
65#ifdef USE_STREAMK_ONLY
66import qualified Streamly.Internal.Data.Stream.StreamK as S
67#else
68import qualified Streamly.Internal.Data.Stream.StreamD as S
69#endif
71#include "Instances.hs"
73-- | Like 'zipWith' but using a monadic zipping function.
74--
75-- @since 0.4.0
76{-# INLINABLE zipWithM #-}
77zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
78zipWithM f m1 m2 = P.fromStreamS $ S.zipWithM f (P.toStreamS m1) (P.toStreamS m2)
80-- | Zip two streams serially using a pure zipping function.
81--
82-- @
83-- > S.toList $ S.zipWith (+) (S.fromList [1,2,3]) (S.fromList [4,5,6])
84-- [5,7,9]
85-- @
86--
87-- @since 0.1.0
88{-# INLINABLE zipWith #-}
89zipWith :: (IsStream t, Monad m) => (a -> b -> c) -> t m a -> t m b -> t m c
90zipWith f m1 m2 = P.fromStreamS $ S.zipWith f (P.toStreamS m1) (P.toStreamS m2)
92------------------------------------------------------------------------------
93-- Parallel Zipping
94------------------------------------------------------------------------------
96-- The CPS version and the direct version of zipAsyncWithM below seem to have
97-- identical performance. However, we need to use the StreamD version of
98-- mkParallel which uses a direct implementation of fromSVar. In comparison to
99-- CPS version of fromSVar the direct version gives a 2x improvement.
101-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped
102-- are generated concurrently.
103--
104-- @since 0.4.0
105{-# INLINE zipAsyncWithM #-}
106zipAsyncWithM :: (IsStream t, MonadAsync m)
107 => (a -> b -> m c) -> t m a -> t m b -> t m c
108zipAsyncWithM f m1 m2 = K.mkStream $ \st stp sng yld -> do
109 ma <- D.mkParallel (adaptState st) (D.toStreamD m1)
110 mb <- D.mkParallel (adaptState st) (D.toStreamD m2)
111 K.foldStream st stp sng yld $ zipWithM f (D.fromStreamD ma)
112 (D.fromStreamD mb)
113{-
114zipAsyncWithM f m1 m2 =
115 fromStreamD $ D.zipAsyncWithM f (toStreamD m1) (toStreamD m2)
116-}
118-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped
119-- are generated concurrently.
120--
121-- @since 0.1.0
122{-# INLINE zipAsyncWith #-}
123zipAsyncWith :: (IsStream t, MonadAsync m)
124 => (a -> b -> c) -> t m a -> t m b -> t m c
125zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))
127------------------------------------------------------------------------------
128-- Serially Zipping Streams
129------------------------------------------------------------------------------
131-- | The applicative instance of 'ZipSerialM' zips a number of streams serially
132-- i.e. it produces one element from each stream serially and then zips all
133-- those elements.
134--
135-- @
136-- main = (toList . 'zipSerially' $ (,,) \<$\> s1 \<*\> s2 \<*\> s3) >>= print
137-- where s1 = fromFoldable [1, 2]
138-- s2 = fromFoldable [3, 4]
139-- s3 = fromFoldable [5, 6]
140-- @
141-- @
142-- [(1,3,5),(2,4,6)]
143-- @
144--
145-- The 'Semigroup' instance of this type works the same way as that of
146-- 'SerialT'.
147--
148-- @since 0.2.0
149newtype ZipSerialM m a = ZipSerialM {getZipSerialM :: Stream m a}
150 deriving (Semigroup, Monoid)
152-- |
153-- @since 0.1.0
154{-# DEPRECATED ZipStream "Please use 'ZipSerialM' instead." #-}
155type ZipStream = ZipSerialM
157-- | An IO stream whose applicative instance zips streams serially.
158--
159-- @since 0.2.0
160type ZipSerial = ZipSerialM IO
162-- | Fix the type of a polymorphic stream as 'ZipSerialM'.
163--
164-- @since 0.2.0
165zipSerially :: IsStream t => ZipSerialM m a -> t m a
166zipSerially = K.adapt
168-- | Same as 'zipSerially'.
169--
170-- @since 0.1.0
171{-# DEPRECATED zipping "Please use zipSerially instead." #-}
172zipping :: IsStream t => ZipSerialM m a -> t m a
173zipping = zipSerially
175consMZip :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a
176consMZip m ms = fromStream $ K.consMStream m (toStream ms)
178instance IsStream ZipSerialM where
179 toStream = getZipSerialM
180 fromStream = ZipSerialM
182 {-# INLINE consM #-}
183 {-# SPECIALIZE consM :: IO a -> ZipSerialM IO a -> ZipSerialM IO a #-}
184 consM :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a
185 consM = consMZip
187 {-# INLINE (|:) #-}
188 {-# SPECIALIZE (|:) :: IO a -> ZipSerialM IO a -> ZipSerialM IO a #-}
189 (|:) :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a
190 (|:) = consMZip
192LIST_INSTANCES(ZipSerialM)
193NFDATA1_INSTANCE(ZipSerialM)
195instance Monad m => Functor (ZipSerialM m) where
196 {-# INLINE fmap #-}
197 fmap f (ZipSerialM m) = D.fromStreamD $ D.mapM (return . f) $ D.toStreamD m
199instance Monad m => Applicative (ZipSerialM m) where
200 pure = ZipSerialM . K.repeat
201 {-# INLINE (<*>) #-}
202 (<*>) = zipWith id
204FOLDABLE_INSTANCE(ZipSerialM)
205TRAVERSABLE_INSTANCE(ZipSerialM)
207------------------------------------------------------------------------------
208-- Parallely Zipping Streams
209------------------------------------------------------------------------------
210--
211-- | Like 'ZipSerialM' but zips in parallel, it generates all the elements to
212-- be zipped concurrently.
213--
214-- @
215-- main = (toList . 'zipAsyncly' $ (,,) \<$\> s1 \<*\> s2 \<*\> s3) >>= print
216-- where s1 = fromFoldable [1, 2]
217-- s2 = fromFoldable [3, 4]
218-- s3 = fromFoldable [5, 6]
219-- @
220-- @
221-- [(1,3,5),(2,4,6)]
222-- @
223--
224-- The 'Semigroup' instance of this type works the same way as that of
225-- 'SerialT'.
226--
227-- @since 0.2.0
228newtype ZipAsyncM m a = ZipAsyncM {getZipAsyncM :: Stream m a}
229 deriving (Semigroup, Monoid)
231-- | An IO stream whose applicative instance zips streams wAsyncly.
232--
233-- @since 0.2.0
234type ZipAsync = ZipAsyncM IO
236-- | Fix the type of a polymorphic stream as 'ZipAsyncM'.
237--
238-- @since 0.2.0
239zipAsyncly :: IsStream t => ZipAsyncM m a -> t m a
240zipAsyncly = K.adapt
242-- | Same as 'zipAsyncly'.
243--
244-- @since 0.1.0
245{-# DEPRECATED zippingAsync "Please use zipAsyncly instead." #-}
246zippingAsync :: IsStream t => ZipAsyncM m a -> t m a
247zippingAsync = zipAsyncly
249consMZipAsync :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a
250consMZipAsync m ms = fromStream $ K.consMStream m (toStream ms)
252instance IsStream ZipAsyncM where
253 toStream = getZipAsyncM
254 fromStream = ZipAsyncM
256 {-# INLINE consM #-}
257 {-# SPECIALIZE consM :: IO a -> ZipAsyncM IO a -> ZipAsyncM IO a #-}
258 consM :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a
259 consM = consMZipAsync
261 {-# INLINE (|:) #-}
262 {-# SPECIALIZE (|:) :: IO a -> ZipAsyncM IO a -> ZipAsyncM IO a #-}
263 (|:) :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a
264 (|:) = consMZipAsync
266instance Monad m => Functor (ZipAsyncM m) where
267 {-# INLINE fmap #-}
268 fmap f (ZipAsyncM m) = D.fromStreamD $ D.mapM (return . f) $ D.toStreamD m
270instance MonadAsync m => Applicative (ZipAsyncM m) where
271 pure = ZipAsyncM . K.repeat
272 {-# INLINE (<*>) #-}
273 m1 <*> m2 = zipAsyncWith id m1 m2