src/Streamly/Internal/Data/Stream/Zip.hs · streamly

git.iko.soy/streamly

on 70fb992 help sign in

streamly/src/Streamly/Internal/Data/Stream/Zip.hs

266 lines 8.1 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

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.Foldable (Foldable(foldl'), fold)

49import Data.Functor.Identity (Identity(..), runIdentity)

50import Data.Maybe (fromMaybe)

51import Data.Semigroup (Endo(..))

52#if __GLASGOW_HASKELL__ < 808

53import Data.Semigroup (Semigroup(..))

54#endif

55import GHC.Exts (IsList(..), IsString(..))

56import Text.Read (Lexeme(Ident), lexP, parens, prec, readPrec, readListPrec,

57 readListPrecDefault)

58import Prelude hiding (map, repeat, zipWith, errorWithoutStackTrace)

60import Streamly.Internal.BaseCompat ((#.), errorWithoutStackTrace)

61import Streamly.Internal.Data.Stream.StreamK (IsStream(..), Stream)

62import Streamly.Internal.Data.Strict (Maybe'(..), toMaybe)

63import Streamly.Internal.Data.SVar (MonadAsync)

65import qualified Streamly.Internal.Data.Stream.Prelude as P

66import qualified Streamly.Internal.Data.Stream.StreamK as K

67import qualified Streamly.Internal.Data.Stream.StreamD as D

69#ifdef USE_STREAMK_ONLY

70import qualified Streamly.Internal.Data.Stream.StreamK as S

71#else

72import qualified Streamly.Internal.Data.Stream.StreamD as S

73#endif

75#include "Instances.hs"

77-- | Like 'zipWith' but using a monadic zipping function.

78--

79-- @since 0.4.0

80{-# INLINABLE zipWithM #-}

81zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c

82zipWithM f m1 m2 = P.fromStreamS $ S.zipWithM f (P.toStreamS m1) (P.toStreamS m2)

84-- | Zip two streams serially using a pure zipping function.

85--

86-- @

87-- > S.toList $ S.zipWith (+) (S.fromList [1,2,3]) (S.fromList [4,5,6])

88-- [5,7,9]

89-- @

90--

91-- @since 0.1.0

92{-# INLINABLE zipWith #-}

93zipWith :: (IsStream t, Monad m) => (a -> b -> c) -> t m a -> t m b -> t m c

94zipWith f m1 m2 = P.fromStreamS $ S.zipWith f (P.toStreamS m1) (P.toStreamS m2)

96------------------------------------------------------------------------------

97-- Parallel Zipping

98------------------------------------------------------------------------------

100-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped

101-- are generated concurrently.

102--

103-- @since 0.4.0

104{-# INLINE zipAsyncWithM #-}

105zipAsyncWithM :: (IsStream t, MonadAsync m)

106 => (a -> b -> m c) -> t m a -> t m b -> t m c

107zipAsyncWithM f m1 m2 = D.fromStreamD $

108 D.zipWithM f (D.mkParallelD $ D.toStreamD m1)

109 (D.mkParallelD $ D.toStreamD m2)

110

111-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped

112-- are generated concurrently.

113--

114-- @since 0.1.0

115{-# INLINE zipAsyncWith #-}

116zipAsyncWith :: (IsStream t, MonadAsync m)

117 => (a -> b -> c) -> t m a -> t m b -> t m c

118zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))

119

120------------------------------------------------------------------------------

121-- Serially Zipping Streams

122------------------------------------------------------------------------------

123

124-- | The applicative instance of 'ZipSerialM' zips a number of streams serially

125-- i.e. it produces one element from each stream serially and then zips all

126-- those elements.

127--

128-- @

129-- main = (toList . 'zipSerially' $ (,,) \<$\> s1 \<*\> s2 \<*\> s3) >>= print

130-- where s1 = fromFoldable [1, 2]

131-- s2 = fromFoldable [3, 4]

132-- s3 = fromFoldable [5, 6]

133-- @

134-- @

135-- [(1,3,5),(2,4,6)]

136-- @

137--

138-- The 'Semigroup' instance of this type works the same way as that of

139-- 'SerialT'.

140--

141-- @since 0.2.0

142newtype ZipSerialM m a = ZipSerialM {getZipSerialM :: Stream m a}

143 deriving (Semigroup, Monoid)

144

145-- |

146-- @since 0.1.0

147{-# DEPRECATED ZipStream "Please use 'ZipSerialM' instead." #-}

148type ZipStream = ZipSerialM

149

150-- | An IO stream whose applicative instance zips streams serially.

151--

152-- @since 0.2.0

153type ZipSerial = ZipSerialM IO

154

155-- | Fix the type of a polymorphic stream as 'ZipSerialM'.

156--

157-- @since 0.2.0

158zipSerially :: IsStream t => ZipSerialM m a -> t m a

159zipSerially = K.adapt

160

161-- | Same as 'zipSerially'.

162--

163-- @since 0.1.0

164{-# DEPRECATED zipping "Please use zipSerially instead." #-}

165zipping :: IsStream t => ZipSerialM m a -> t m a

166zipping = zipSerially

167

168consMZip :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a

169consMZip m ms = fromStream $ K.consMStream m (toStream ms)

170

171instance IsStream ZipSerialM where

172 toStream = getZipSerialM

173 fromStream = ZipSerialM

174

175 {-# INLINE consM #-}

176 {-# SPECIALIZE consM :: IO a -> ZipSerialM IO a -> ZipSerialM IO a #-}

177 consM :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a

178 consM = consMZip

179

180 {-# INLINE (|:) #-}

181 {-# SPECIALIZE (|:) :: IO a -> ZipSerialM IO a -> ZipSerialM IO a #-}

182 (|:) :: Monad m => m a -> ZipSerialM m a -> ZipSerialM m a

183 (|:) = consMZip

184

185LIST_INSTANCES(ZipSerialM)

186NFDATA1_INSTANCE(ZipSerialM)

187

188instance Monad m => Functor (ZipSerialM m) where

189 {-# INLINE fmap #-}

190 fmap f (ZipSerialM m) = D.fromStreamD $ D.mapM (return . f) $ D.toStreamD m

191

192instance Monad m => Applicative (ZipSerialM m) where

193 pure = ZipSerialM . K.repeat

194 {-# INLINE (<*>) #-}

195 (<*>) = zipWith id

196

197FOLDABLE_INSTANCE(ZipSerialM)

198TRAVERSABLE_INSTANCE(ZipSerialM)

199

200------------------------------------------------------------------------------

201-- Parallely Zipping Streams

202------------------------------------------------------------------------------

203--

204-- | Like 'ZipSerialM' but zips in parallel, it generates all the elements to

205-- be zipped concurrently.

206--

207-- @

208-- main = (toList . 'zipAsyncly' $ (,,) \<$\> s1 \<*\> s2 \<*\> s3) >>= print

209-- where s1 = fromFoldable [1, 2]

210-- s2 = fromFoldable [3, 4]

211-- s3 = fromFoldable [5, 6]

212-- @

213-- @

214-- [(1,3,5),(2,4,6)]

215-- @

216--

217-- The 'Semigroup' instance of this type works the same way as that of

218-- 'SerialT'.

219--

220-- @since 0.2.0

221newtype ZipAsyncM m a = ZipAsyncM {getZipAsyncM :: Stream m a}

222 deriving (Semigroup, Monoid)

223

224-- | An IO stream whose applicative instance zips streams wAsyncly.

225--

226-- @since 0.2.0

227type ZipAsync = ZipAsyncM IO

228

229-- | Fix the type of a polymorphic stream as 'ZipAsyncM'.

230--

231-- @since 0.2.0

232zipAsyncly :: IsStream t => ZipAsyncM m a -> t m a

233zipAsyncly = K.adapt

234

235-- | Same as 'zipAsyncly'.

236--

237-- @since 0.1.0

238{-# DEPRECATED zippingAsync "Please use zipAsyncly instead." #-}

239zippingAsync :: IsStream t => ZipAsyncM m a -> t m a

240zippingAsync = zipAsyncly

241

242consMZipAsync :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a

243consMZipAsync m ms = fromStream $ K.consMStream m (toStream ms)

244

245instance IsStream ZipAsyncM where

246 toStream = getZipAsyncM

247 fromStream = ZipAsyncM

248

249 {-# INLINE consM #-}

250 {-# SPECIALIZE consM :: IO a -> ZipAsyncM IO a -> ZipAsyncM IO a #-}

251 consM :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a

252 consM = consMZipAsync

253

254 {-# INLINE (|:) #-}

255 {-# SPECIALIZE (|:) :: IO a -> ZipAsyncM IO a -> ZipAsyncM IO a #-}

256 (|:) :: Monad m => m a -> ZipAsyncM m a -> ZipAsyncM m a

257 (|:) = consMZipAsync

258

259instance Monad m => Functor (ZipAsyncM m) where

260 {-# INLINE fmap #-}

261 fmap f (ZipAsyncM m) = D.fromStreamD $ D.mapM (return . f) $ D.toStreamD m

262

263instance MonadAsync m => Applicative (ZipAsyncM m) where

264 pure = ZipAsyncM . K.repeat

265 {-# INLINE (<*>) #-}

266 m1 <*> m2 = zipAsyncWith id m1 m2