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

git.iko.soy/streamly

on edbb7b8 help sign in

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

301 lines 9.1 kB history raw

1{-# LANGUAGE UndecidableInstances #-}

3-- |

4-- Module : Streamly.Internal.Data.Stream.Zip

6--

7-- License : BSD3

8-- Maintainer : streamly@composewell.com

9-- Stability : experimental

10-- Portability : GHC

11--

12-- To run examples in this module:

13--

14-- >>> import qualified Streamly.Prelude as Stream

15--

16module Streamly.Internal.Data.Stream.Zip

17 (

18 ZipSerialM

19 , ZipSerial

20 , fromZipSerial

22 , ZipAsyncM

23 , ZipAsync

24 , fromZipAsync

26 , zipWith

27 , zipWithM

28 , zipAsyncWith

29 , zipAsyncWithM

31 -- * Deprecated

32 , ZipStream

33 , zipping

34 , zippingAsync

35 )

36where

38#include "inline.hs"

40import Control.Applicative (liftA2)

41import Control.DeepSeq (NFData(..))

42#if MIN_VERSION_deepseq(1,4,3)

43import Control.DeepSeq (NFData1(..))

44#endif

45import Data.Foldable (Foldable(foldl'), fold)

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

47import Data.Maybe (fromMaybe)

48import Data.Semigroup (Endo(..))

49#if __GLASGOW_HASKELL__ < 808

50import Data.Semigroup (Semigroup(..))

51#endif

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

53import Text.Read

54 ( Lexeme(Ident), lexP, parens, prec, readPrec, readListPrec

55 , readListPrecDefault)

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

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

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

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

61import Streamly.Internal.Data.Stream.Serial (SerialT)

63import qualified Streamly.Internal.Data.Stream.Parallel as Par

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

65 (cmpBy, eqBy, foldl', foldr, fromList, toList)

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

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

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

71#include "Instances.hs"

73-- $setup

74-- >>> import Control.Concurrent (threadDelay)

75-- >>> :{

76-- delay n = do

77-- threadDelay (n * 1000000) -- sleep for n seconds

78-- putStrLn (show n ++ " sec") -- print "n sec"

79-- return n -- IO Int

80-- :}

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

83--

84-- @since 0.4.0

85{-# INLINE_EARLY zipWithM #-}

86zipWithM ::

87 (IsStream t, MonadAsync m) => (a -> b -> m c) -> t m a -> t m b -> t m c

88zipWithM f m1 m2 =

89 K.fromStream $ K.zipWithM f (K.toStream m1) (K.toStream m2)

91{-# RULES "zipWithM serial" zipWithM = zipWithMSerial #-}

92{-# INLINE zipWithMSerial #-}

93zipWithMSerial ::

94 Monad m => (a -> b -> m c) -> SerialT m a -> SerialT m b -> SerialT m c

95zipWithMSerial f m1 m2 =

96 D.fromStreamD $ D.zipWithM f (D.toStreamD m1) (D.toStreamD m2)

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

99-- function is applied concurrently for concurrent streams.

100--

101-- @

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

103-- [5,7,9]

104-- @

105--

106-- @since 0.1.0

107{-# INLINE_EARLY zipWith #-}

108zipWith ::

109 (IsStream t, MonadAsync m) => (a -> b -> c) -> t m a -> t m b -> t m c

110zipWith f m1 m2 = K.fromStream $ K.zipWith f (K.toStream m1) (K.toStream m2)

111

112{-# RULES "zipWith serial" zipWith = zipWithSerial #-}

113{-# INLINE zipWithSerial #-}

114zipWithSerial ::

115 Monad m => (a -> b -> c) -> SerialT m a -> SerialT m b -> SerialT m c

116zipWithSerial f m1 m2 =

117 D.fromStreamD $ D.zipWith f (D.toStreamD m1) (D.toStreamD m2)

118

119------------------------------------------------------------------------------

120-- Parallel Zipping

121------------------------------------------------------------------------------

122

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

124-- are generated concurrently.

125--

126-- @since 0.4.0

127{-# INLINE zipAsyncWithM #-}

128zipAsyncWithM :: (IsStream t, MonadAsync m)

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

130zipAsyncWithM f m1 m2 = D.fromStreamD $

131 D.zipWithM f (Par.mkParallelD $ D.toStreamD m1)

132 (Par.mkParallelD $ D.toStreamD m2)

133

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

135-- are generated concurrently.

136--

137-- @since 0.1.0

138{-# INLINE zipAsyncWith #-}

139zipAsyncWith :: (IsStream t, MonadAsync m)

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

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

142

143------------------------------------------------------------------------------

144-- Serially Zipping Streams

145------------------------------------------------------------------------------

146

147-- | For 'ZipSerialM' streams:

148--

149-- @

150-- (<>) = 'Streamly.Prelude.serial'

151-- (<*>) = 'Streamly.Prelude.serial.zipWith' id

152-- @

153--

154-- Applicative evaluates the streams being zipped serially:

155--

156-- >>> s1 = Stream.fromFoldable [1, 2]

157-- >>> s2 = Stream.fromFoldable [3, 4]

158-- >>> s3 = Stream.fromFoldable [5, 6]

159-- >>> Stream.toList $ Stream.fromZipSerial $ (,,) <$> s1 <*> s2 <*> s3

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

161--

162-- /Since: 0.2.0 ("Streamly")/

163--

164-- @since 0.8.0

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

166 deriving (Semigroup, Monoid)

167

168-- |

169-- @since 0.1.0

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

171type ZipStream = ZipSerialM

172

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

174--

175-- /Since: 0.2.0 ("Streamly")/

176--

177-- @since 0.8.0

178type ZipSerial = ZipSerialM IO

179

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

181--

182-- /Since: 0.2.0 ("Streamly")/

183--

184-- @since 0.8.0

185fromZipSerial :: IsStream t => ZipSerialM m a -> t m a

186fromZipSerial = K.adapt

187

188-- | Same as 'fromZipSerial'.

189--

190-- @since 0.1.0

191{-# DEPRECATED zipping "Please use fromZipSerial instead." #-}

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

193zipping = fromZipSerial

194

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

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

197

198instance IsStream ZipSerialM where

199 toStream = getZipSerialM

200 fromStream = ZipSerialM

201

202 {-# INLINE consM #-}

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

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

205 consM = consMZip

206

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

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

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

210 (|:) = consMZip

211

212LIST_INSTANCES(ZipSerialM)

213NFDATA1_INSTANCE(ZipSerialM)

214

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

216 {-# INLINE fmap #-}

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

218

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

220 pure = ZipSerialM . K.repeat

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

222 m1 <*> m2 = fromStream $ toStream $ zipWithSerial id m1_ m2_

223 where

224 m1_ = fromStream (toStream m1)

225 m2_ = fromStream (toStream m2)

226

227FOLDABLE_INSTANCE(ZipSerialM)

228TRAVERSABLE_INSTANCE(ZipSerialM)

229

230------------------------------------------------------------------------------

231-- Parallely Zipping Streams

232------------------------------------------------------------------------------

233--

234-- | For 'ZipAsyncM' streams:

235--

236-- @

237-- (<>) = 'Streamly.Prelude.serial'

238-- (<*>) = 'Streamly.Prelude.serial.zipAsyncWith' id

239-- @

240--

241-- Applicative evaluates the streams being zipped concurrently, the following

242-- would take half the time that it would take in serial zipping:

243--

244-- >>> s = Stream.fromFoldableM $ Prelude.map delay [1, 1, 1]

245-- >>> Stream.toList $ Stream.fromZipAsync $ (,) <$> s <*> s

246-- ...

247-- [(1,1),(1,1),(1,1)]

248--

249-- /Since: 0.2.0 ("Streamly")/

250--

251-- @since 0.8.0

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

253 deriving (Semigroup, Monoid)

254

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

256--

257-- /Since: 0.2.0 ("Streamly")/

258--

259-- @since 0.8.0

260type ZipAsync = ZipAsyncM IO

261

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

263--

264-- /Since: 0.2.0 ("Streamly")/

265--

266-- @since 0.8.0

267fromZipAsync :: IsStream t => ZipAsyncM m a -> t m a

268fromZipAsync = K.adapt

269

270-- | Same as 'fromZipAsync'.

271--

272-- @since 0.1.0

273{-# DEPRECATED zippingAsync "Please use fromZipAsync instead." #-}

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

275zippingAsync = fromZipAsync

276

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

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

279

280instance IsStream ZipAsyncM where

281 toStream = getZipAsyncM

282 fromStream = ZipAsyncM

283

284 {-# INLINE consM #-}

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

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

287 consM = consMZipAsync

288

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

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

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

292 (|:) = consMZipAsync

293

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

295 {-# INLINE fmap #-}

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

297

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

299 pure = ZipAsyncM . K.repeat

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

301 m1 <*> m2 = zipAsyncWith id m1 m2