1ddcfc4 · streamly · git.iko.soy

git.iko.soy/streamly

Fix zipWith(M) to work concurrently according to the stream type

https://github.com/composewell/streamly/issues/158

commit1ddcfc46345cd4e4c3115a955668a6d5158a64cd authorAdithya Kumar <adithya@composewell.com> dateFri 2 Apr 2021, 04:38 (6 years ago) parentd3cc553

4 files changed+53−26

compare with master

unifiedsplit

Changelog.md+3 benchmark/Streamly/Benchmark/Data/Stream/StreamK.hs+1−1 src/Streamly/Internal/Data/Stream/StreamK.hs+12−6 src/Streamly/Internal/Data/Stream/Zip.hs+37−19

Changelog.md+3

@@ -15,6 +15,9 @@

1515* `encodeLatin1` now silently truncates any character beyond 255 to incorrect

1616 characters in the input stream. Use `encodeLatin1'` to recover previous

1717 functionality.

18+* The zipping function in `Streamly.Prelude.zipWith` and

19+ `Streamly.Prelude.zipWithM` is now applied concurrently for concurrent

20+ streams.

1821

1922### Breaking changes

2023

benchmark/Streamly/Benchmark/Data/Stream/StreamK.hs+1−1

@@ -352,7 +352,7 @@ iterateDropWhileTrue streamLen iterStreamLen maxIters = iterateSource iterStream

352352-------------------------------------------------------------------------------

353353

354354{-# INLINE zipWith #-}

355−zipWith :: Monad m => Stream m Int -> m ()

355+zipWith :: S.MonadAsync m => Stream m Int -> m ()

356356zipWith src = drain $ S.zipWith (,) src src

357357

358358-------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamK.hs+12−6

@@ -995,34 +995,40 @@ mapMaybe f m = go m

995995-- Serial Zipping

996996------------------------------------------------------------------------------

997997

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

998+-- XXX We can probably implement zipWith in terms of zipWithM

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

1000+-- function is applied concurrently for concurrent streams.

9991001--

10001002-- @since 0.1.0

10011003{-# INLINABLE zipWith #-}

1002−zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c

1004+zipWith ::

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

10031006zipWith f = go

10041007 where

10051008 go mx my = mkStream $ \st yld sng stp -> do

10061009 let merge a ra =

10071010 let single2 b = sng (f a b)

1008− yield2 b rb = yld (f a b) (go ra rb)

1011+ runIt x = foldStream st yld sng stp x

1012+ yield2 b rb = runIt (return (f a b) `consM` go ra rb)

10091013 in foldStream (adaptState st) yield2 single2 stp my

10101014 let single1 a = merge a nil

10111015 yield1 = merge

10121016 foldStream (adaptState st) yield1 single1 stp mx

10131017

1014−-- | Zip two streams serially using a monadic zipping function.

1018+-- | Zip two streams serially using a monadic zipping function. The zipping

1019+-- function is applied concurrently for concurrent streams.

10151020--

10161021-- @since 0.1.0

10171022{-# INLINABLE zipWithM #-}

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

1023+zipWithM ::

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

10191025zipWithM f m1 m2 = go m1 m2

10201026 where

10211027 go mx my = mkStream $ \st yld sng stp -> do

10221028 let merge a ra =

10231029 let runIt x = foldStream st yld sng stp x

10241030 single2 b = f a b >>= sng

1025− yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)

1031+ yield2 b rb = runIt (f a b `consM` go ra rb)

10261032 in foldStream (adaptState st) yield2 single2 stp my

10271033 let single1 a = merge a nil

10281034 yield1 = merge

src/Streamly/Internal/Data/Stream/Zip.hs+37−19

@@ -35,6 +35,8 @@ module Streamly.Internal.Data.Stream.Zip

3535 )

3636where

3737

38+#include "inline.hs"

39+

3840import Control.Applicative (liftA2)

3941import Control.DeepSeq (NFData(..))

4042#if MIN_VERSION_deepseq(1,4,3)

@@ -56,18 +58,13 @@ import Streamly.Internal.Data.Stream.StreamK (IsStream(..), Stream)

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

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

5860

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

62+

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

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

61− (cmpBy, eqBy, foldl', foldr, fromList, toList, fromStreamS, toStreamS)

62−import qualified Streamly.Internal.Data.Stream.StreamK as K (repeat)

63−import qualified Streamly.Internal.Data.Stream.StreamK.Type as K

64−import qualified Streamly.Internal.Data.Stream.StreamD as D (zipWithM)

65−import qualified Streamly.Internal.Data.Stream.StreamD.Type as D

66−#ifdef USE_STREAMK_ONLY

67−import qualified Streamly.Internal.Data.Stream.StreamK as S (zipWith, zipWithM)

68−#else

69−import qualified Streamly.Internal.Data.Stream.StreamD as S (zipWith, zipWithM)

70−#endif

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

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

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

7168

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

7370

@@ -85,11 +82,21 @@ import Prelude hiding (map, repeat, zipWith, errorWithoutStackTrace)

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

8683--

8784-- @since 0.4.0

88−{-# INLINABLE zipWithM #-}

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

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

91−

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

85+{-# INLINE_EARLY zipWithM #-}

86+zipWithM ::

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

88+zipWithM f m1 m2 =

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

90+

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

92+{-# INLINE zipWithMSerial #-}

93+zipWithMSerial ::

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

95+zipWithMSerial f m1 m2 =

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

97+

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

99+-- function is applied concurrently for concurrent streams.

93100--

94101-- @

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

@@ -97,9 +104,17 @@ zipWithM f m1 m2 = P.fromStreamS $ S.zipWithM f (P.toStreamS m1) (P.toStreamS m2

97104-- @

98105--

99106-- @since 0.1.0

100−{-# INLINABLE zipWith #-}

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

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

107+{-# INLINE_EARLY zipWith #-}

108+zipWith ::

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

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

111+

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

113+{-# INLINE zipWithSerial #-}

114+zipWithSerial ::

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

116+zipWithSerial f m1 m2 =

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

103118

104119------------------------------------------------------------------------------

105120-- Parallel Zipping

@@ -204,7 +219,10 @@ instance Monad m => Functor (ZipSerialM m) where

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

205220 pure = ZipSerialM . K.repeat

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

207− (<*>) = zipWith id

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

223+ where

224+ m1_ = fromStream (toStream m1)

225+ m2_ = fromStream (toStream m2)

208226

209227FOLDABLE_INSTANCE(ZipSerialM)

210228TRAVERSABLE_INSTANCE(ZipSerialM)