on master help sign in
Add zipWith
commitd96aedee1fc431dfe1f7283e8a197e79a07aab64 authorHarendra Kumar <harendra@composewell.com> dateMon 24 Feb 2020, 22:36 (7 years ago) parent9fd164b
1 file changed+45−1
src/Streamly/Internal/Data/Unfold.hs+45−1
src/Streamly/Internal/Data/Unfold.hs+45−1
@@ -105,6 +105,10 @@ module Streamly.Internal.Data.Unfold
105105 , filter
106106 , filterM
107107
108+ -- * Zipping
109+ , zipWithM
110+ , zipWith
111+
108112 -- * Nesting
109113 , concat
110114 , concatMapM
@@ -130,7 +134,8 @@ import Control.Monad.IO.Class (MonadIO(..))
130134import Control.Monad.Trans.Control (MonadBaseControl)
131135import Data.Void (Void)
132136import GHC.Types (SPEC(..))
133−import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)
137+import Prelude
138+ hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith)
134139
135140import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
136141#if __GLASGOW_HASKELL__ < 800
@@ -520,6 +525,45 @@ enumerateFromToIntegral to =
520525enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
521526enumerateFromIntegral = enumerateFromToIntegral maxBound
522527
528+-------------------------------------------------------------------------------
529+-- Zipping
530+-------------------------------------------------------------------------------
531+
532+{-# INLINE_NORMAL zipWithM #-}
533+zipWithM :: Monad m
534+ => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
535+zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
536+
537+ where
538+
539+ inject (x, y) = do
540+ s1 <- inject1 x
541+ s2 <- inject2 y
542+ return (s1, s2, Nothing)
543+
544+ {-# INLINE_LATE step #-}
545+ step (s1, s2, Nothing) = do
546+ r <- step1 s1
547+ return $
548+ case r of
549+ Yield x s -> Skip (s, s2, Just x)
550+ Skip s -> Skip (s, s2, Nothing)
551+ Stop -> Stop
552+
553+ step (s1, s2, Just x) = do
554+ r <- step2 s2
555+ case r of
556+ Yield y s -> do
557+ z <- f x y
558+ return $ Yield z (s1, s, Nothing)
559+ Skip s -> return $ Skip (s1, s, Just x)
560+ Stop -> return Stop
561+
562+{-# INLINE zipWith #-}
563+zipWith :: Monad m
564+ => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
565+zipWith f = zipWithM (\a b -> return (f a b))
566+
523567-------------------------------------------------------------------------------
524568-- Nested
525569-------------------------------------------------------------------------------