Change signatures of zipWith/zipWithM

benchmark/Streamly/Benchmark/Data/Unfold.hs

@@ -384,7 +384,7 @@ dropWhileMFalse size start =
 {-# INLINE zipWith #-}
 zipWith :: Monad m => Int -> Int -> m ()
 zipWith size start =
-    drainProductDefault (size + start) (UF.zipWith (+)) (start, start + 1)
+    drainProductDefault (size + start) (UF.zipWith (+)) start
 
 {-# INLINE zipWithM #-}
 zipWithM :: Monad m => Int -> Int -> m ()
@@ -392,12 +392,12 @@ zipWithM size start =
     drainProductDefault
         (size + start)
         (UF.zipWithM (\a b -> return $ a + b))
-        (start, start + 1)
+        start
 
 {-# INLINE teeZipWith #-}
 teeZipWith :: Monad m => Int -> Int -> m ()
 teeZipWith size start =
-    drainProductDefault (size + start) (UF.teeZipWith (+)) start
+    drainProductDefault (size + start) (UF.zipWith (+)) start
 
 -------------------------------------------------------------------------------
 -- Applicative

src/Streamly/Data/Unfold.hs

@@ -94,7 +94,6 @@ module Streamly.Data.Unfold
     , dropWhileM
 
     -- ** Zipping
-    , zipWithM
     , zipWith
 
     -- ** Cross Product

src/Streamly/Internal/Data/Unfold.hs

@@ -159,7 +159,6 @@ module Streamly.Internal.Data.Unfold
     -- ** Zipping
     , zipWithM
     , zipWith
-    , teeZipWith
 
     -- ** Cross product
     , crossWithM
@@ -881,24 +880,6 @@ fromProducer = Unfold step (return . FromSVarRead)
 
     step (FromSVarInit _) = undefined
 
--------------------------------------------------------------------------------
--- Zipping
--------------------------------------------------------------------------------
-
--- | Distribute the input to two unfolds and then zip the outputs to a single
--- stream.
---
--- @
--- S.mapM_ print $ S.unfoldMany (UF.teeZipWith (,) UF.identity (UF.singleton sqrt)) $ S.fromList [1..10]
--- @
---
--- /Pre-release/
---
-{-# INLINE_NORMAL teeZipWith #-}
-teeZipWith :: Monad m
-    => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c
-teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2
-
 ------------------------------------------------------------------------------
 -- Exceptions
 ------------------------------------------------------------------------------

src/Streamly/Internal/Data/Unfold/Type.hs

@@ -54,8 +54,8 @@ import Prelude hiding (const, map, concatMap, zipWith)
 -- $setup
 -- >>> :m
 -- >>> import qualified Streamly.Prelude as Stream
--- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream
--- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
+-- >>> import qualified Streamly.Data.Fold as Fold
+-- >>> import qualified Streamly.Data.Unfold as Unfold
 
 ------------------------------------------------------------------------------
 -- Monadic Unfolds
@@ -125,7 +125,7 @@ yieldM m = Unfold step inject
 
 -- | Discards the unfold input and always returns the argument of 'yield'.
 yield :: Applicative m => b -> Unfold m a b
-yield = yieldM . pure
+yield = yieldM Prelude.. pure
 
 -- | Outer product discarding the first element.
 --
@@ -218,7 +218,7 @@ apply u1 u2 = fmap (\(a, b) -> a b) (cross u1 u2)
 --
 instance Monad m => Applicative (Unfold m a) where
     {-# INLINE pure #-}
-    pure = const Prelude.. return
+    pure = yield
 
     {-# INLINE (<*>) #-}
     (<*>) = apply
@@ -314,8 +314,6 @@ instance Monad m => Monad (Unfold m a) where
 -- Category
 -------------------------------------------------------------------------------
 
--- XXX change it to yieldM or change yieldM in Prelude to singletonM
---
 -- | Lift a monadic function into an unfold generating a singleton stream.
 --
 {-# INLINE singletonM #-}
@@ -397,20 +395,25 @@ instance Monad m => Category (Unfold m) where
 -}
 
 -------------------------------------------------------------------------------
--- Arrow
+-- Zipping
 -------------------------------------------------------------------------------
 
--- | Stops as soon as any of the unfolds stops.
+-- | Distribute the input to two unfolds and then zip the outputs to a single
+-- stream using a monadic zip function.
+--
+-- Stops as soon as any of the unfolds stops.
+--
+-- /Pre-release/
 {-# INLINE_NORMAL zipWithM #-}
 zipWithM :: Monad m
-    => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
+    => (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
 zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
 
     where
 
-    inject (x, y) = do
+    inject x = do
         s1 <- inject1 x
-        s2 <- inject2 y
+        s2 <- inject2 x
         return (s1, s2, Nothing)
 
     {-# INLINE_LATE step #-}
@@ -431,23 +434,27 @@ zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
             Skip s -> return $ Skip (s1, s, Just x)
             Stop   -> return Stop
 
--- | Divide the input into two unfolds and then zip the outputs to a single
--- stream.
+-- | Like 'zipWithM' but with a pure zip function.
 --
--- @
---   S.mapM_ print
--- $ S.unfoldMany (UF.zipWith (,) UF.identity (UF.singleton sqrt))
--- $ S.map (\x -> (x,x))
--- $ S.fromList [1..10]
--- @
+-- >>> square = fmap (\x -> x * x) Unfold.fromList
+-- >>> cube = fmap (\x -> x * x * x) Unfold.fromList
+-- >>> u = Unfold.zipWith (,) square cube
+-- >>> Unfold.fold u Fold.toList [1..5]
+-- [(1,1),(4,8),(9,27),(16,64),(25,125)]
+--
+-- > zipWith f = zipWithM (\a b -> return $ f a b)
 --
 -- /Pre-release/
 --
 {-# INLINE zipWith #-}
 zipWith :: Monad m
-    => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
+    => (b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
 zipWith f = zipWithM (\a b -> return (f a b))
 
+-------------------------------------------------------------------------------
+-- Arrow
+-------------------------------------------------------------------------------
+
 {-
 -- XXX There are multiple ways of combining the outputs of two unfolds, we
 -- could zip, merge, append and more. What is the preferred way for Arrow
@@ -464,5 +471,5 @@ instance Monad m => Arrow (Unfold m) where
     arr = singleton
 
     {-# INLINE (***) #-}
-    (***) = zipWith (,)
+    u1 *** u2 = zipWith (,) (lmap fst u1) (lmap snd u2)
 -}

test/Streamly/Test/Data/Unfold.hs

@@ -319,7 +319,7 @@ zipWithM =
                   unf2 = UF.enumerateFromToIntegral 20
                   fA = applyFun2 f :: Int -> Int -> Int
                   fM a b = modify (+ 1) >> return (fA a b)
-                  unf = UF.zipWithM fM unf1 unf2
+                  unf = UF.zipWithM fM (UF.lmap fst unf1) (UF.lmap snd unf2)
                   lst = Prelude.zipWith fA [1 .. 10] [1 .. 20]
                in testUnfoldMD unf (1, 1) 0 10 lst