git.iko.soy/streamly
on master help sign in
Add zipWith
d96aedeHarendra Kumar
src/Streamly/Internal/Data/Unfold.hs+45−1
compare with master
unifiedsplit
@@ -105,6 +105,10 @@ module Streamly.Internal.Data.Unfold
105 , filter
105 , filter
106 , filterM
106 , filterM
107
107
108+ -- * Zipping
109+ , zipWithM
110+ , zipWith
111+
108 -- * Nesting
112 -- * Nesting
109 , concat
113 , concat
110 , concatMapM
114 , concatMapM
@@ -130,7 +134,8 @@ import Control.Monad.IO.Class (MonadIO(..))
130import Control.Monad.Trans.Control (MonadBaseControl)
134import Control.Monad.Trans.Control (MonadBaseControl)
131import Data.Void (Void)
135import Data.Void (Void)
132import GHC.Types (SPEC(..))
136import GHC.Types (SPEC(..))
137+import Prelude
133−import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)
138+ hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith)
134
139
135import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
140import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
136#if __GLASGOW_HASKELL__ < 800
141#if __GLASGOW_HASKELL__ < 800
@@ -520,6 +525,45 @@ enumerateFromToIntegral to =
520enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
525enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
521enumerateFromIntegral = enumerateFromToIntegral maxBound
526enumerateFromIntegral = enumerateFromToIntegral maxBound
522
527
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+
523-------------------------------------------------------------------------------
567-------------------------------------------------------------------------------
524-- Nested
568-- Nested
525-------------------------------------------------------------------------------
569-------------------------------------------------------------------------------