Add Applicative instance for Unfold

benchmark/Streamly/Benchmark/Data/Unfold.hs

@@ -416,9 +416,12 @@ concatMapM size start =
 
     unfoldOut = UF.enumerateFromToIntegral (start + sizeOuter)
 
-{-# INLINE _ap #-}
-_ap :: Int -> Int -> m ()
-_ap = undefined
+{-# INLINE toNullAp #-}
+toNullAp :: Monad m => Int -> Int -> m ()
+toNullAp linearCount start =
+    let end = start + Nested.nestedCount2 linearCount
+        s = Nested.source end
+    in UF.fold ((+) <$> s <*> s) FL.drain start
 
 {-# INLINE _apDiscardFst #-}
 _apDiscardFst :: Int -> Int -> m ()
@@ -527,7 +530,8 @@ o_1_space_nested :: Int -> [Benchmark]
 o_1_space_nested size =
     [ bgroup
           "outer-product"
-          [ benchIO "toNull" $ Nested.toNull size
+          [ benchIO "toNullAp" $ toNullAp size
+          , benchIO "toNull" $ Nested.toNull size
           , benchIO "toNull3" $ Nested.toNull3 size
           , benchIO "concat" $ Nested.concat size
           , benchIO "breakAfterSome" $ Nested.breakAfterSome size

src/Streamly/Internal/Data/Unfold.hs

@@ -125,9 +125,6 @@ module Streamly.Internal.Data.Unfold
     , concat
     , concatMapM
     , outerProduct
-    , ap
-    , apDiscardFst
-    , apDiscardSnd
 
     -- * Exceptions
     , gbracket_
@@ -159,7 +156,7 @@ import Streamly.Internal.Data.IOFinalizer
     (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
 import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
 import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
-import Streamly.Internal.Data.Unfold.Types (Unfold(..), lmap, map)
+import Streamly.Internal.Data.Unfold.Types (Unfold(..), lmap, map, const)
 import System.Mem (performMajorGC)
 
 import qualified Prelude
@@ -429,13 +426,6 @@ singleton f = singletonM $ return . f
 identity :: Monad m => Unfold m a a
 identity = singletonM return
 
-const :: Monad m => m b -> Unfold m a b
-const m = Unfold step inject
-    where
-    inject _ = return ()
-    step () = m >>= \r -> return $ Yield r ()
-
-
 -- | Convert a list of pure values to a 'Stream'
 {-# INLINE_LATE fromList #-}
 fromList :: Monad m => Unfold m [a] a
@@ -952,6 +942,8 @@ concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
 data OuterProductState s1 s2 sy x y =
     OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
 
+-- XXX this can be written in terms of "cross".
+--
 -- | Create an outer product (vector product or cartesian product) of the
 -- output streams of two unfolds.
 --
@@ -981,34 +973,6 @@ outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
             Skip s    -> Skip (OuterProductInner ost sy s x)
             Stop      -> Skip (OuterProductOuter ost sy)
 
--- Special cases of outer product
--- | Outer product with a function application.
---
--- /Unimplemented/
---
-{-# INLINE_NORMAL ap #-}
-ap :: -- Monad m =>
-    Unfold m a (b -> c) -> Unfold m d b -> Unfold m (a, d) c
-ap (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
-
--- | Outer product discarding the first element.
---
--- /Unimplemented/
---
-{-# INLINE_NORMAL apDiscardFst #-}
-apDiscardFst :: -- Monad m =>
-    Unfold m a b -> Unfold m c d -> Unfold m (a, c) d
-apDiscardFst (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
-
--- | Outer product discarding the second element.
---
--- /Unimplemented/
---
-{-# INLINE_NORMAL apDiscardSnd #-}
-apDiscardSnd :: -- Monad m =>
-    Unfold m a b -> Unfold m c d -> Unfold m (a, c) b
-apDiscardSnd (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
-
 -- XXX This can be used to implement a Monad instance for "Unfold m ()".
 
 data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2

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

@@ -10,6 +10,10 @@ module Streamly.Internal.Data.Unfold.Types
     ( Unfold (..)
     , lmap
     , map
+    , const
+    , apSequence
+    , apDiscardSnd
+    , cross
     )
 where
 
@@ -17,7 +21,7 @@ where
 
 import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
 
-import Prelude hiding (map)
+import Prelude hiding (const, map)
 
 ------------------------------------------------------------------------------
 -- Monadic Unfolds
@@ -63,3 +67,86 @@ map f (Unfold ustep uinject) = Unfold step uinject
 instance Functor m => Functor (Unfold m a) where
     {-# INLINE fmap #-}
     fmap = map
+
+------------------------------------------------------------------------------
+-- Applicative
+------------------------------------------------------------------------------
+
+{-# INLINE const #-}
+const :: Applicative m => m b -> Unfold m a b
+const m = Unfold step inject
+
+    where
+
+    inject _ = pure False
+
+    step False = (`Yield` True) <$> m
+    step True = pure Stop
+
+-- | Outer product discarding the first element.
+--
+-- /Unimplemented/
+--
+{-# INLINE_NORMAL apSequence #-}
+apSequence :: -- Monad m =>
+    Unfold m a b -> Unfold m a c -> Unfold m a c
+apSequence (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
+
+-- | Outer product discarding the second element.
+--
+-- /Unimplemented/
+--
+{-# INLINE_NORMAL apDiscardSnd #-}
+apDiscardSnd :: -- Monad m =>
+    Unfold m a b -> Unfold m a c -> Unfold m a b
+apDiscardSnd (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
+
+data Cross a s1 b s2 = CrossOuter a s1 | CrossInner a s1 b s2
+
+-- | Create a cross product (vector product or cartesian product) of the
+-- output streams of two unfolds.
+--
+{-# INLINE_NORMAL cross #-}
+cross :: Monad m => Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
+cross (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
+
+    where
+
+    inject a = do
+        s1 <- inject1 a
+        return $ CrossOuter a s1
+
+    {-# INLINE_LATE step #-}
+    step (CrossOuter a s1) = do
+        r <- step1 s1
+        case r of
+            Yield b s -> do
+                s2 <- inject2 a
+                return $ Skip (CrossInner a s b s2)
+            Skip s    -> return $ Skip (CrossOuter a s)
+            Stop      -> return Stop
+
+    step (CrossInner a s1 b s2) = do
+        r <- step2 s2
+        return $ case r of
+            Yield c s -> Yield (b, c) (CrossInner a s1 b s)
+            Skip s    -> Skip (CrossInner a s1 b s)
+            Stop      -> Skip (CrossOuter a s1)
+
+-- | Example:
+--
+-- >>> Stream.toList $ Stream.unfold ((,) <$> Unfold.lmap fst Unfold.fromList <*> Unfold.lmap snd Unfold.fromList) ([1,2],[3,4])
+-- [(1,3),(1,4),(2,3),(2,4)]
+--
+instance Monad m => Applicative (Unfold m a) where
+    {-# INLINE pure #-}
+    pure = const . return
+
+    {-# INLINE (<*>) #-}
+    u1 <*> u2 = fmap (\(a, b) -> a b) (cross u1 u2)
+
+    -- {-# INLINE (*>) #-}
+    -- (*>) = apSequence
+
+    -- {-# INLINE (<*) #-}
+    -- (<*) = apDiscardSnd

test/Streamly/Test/Data/Unfold.hs

@@ -146,8 +146,8 @@ singletonM =
 
 const :: Bool
 const =
-    let unf = UF.take 10 $ UF.const (modify (+ 1) >> get)
-     in testUnfoldMD unf (0 :: Int) 0 10 [1 .. 10]
+    let unf = UF.const (modify (+ 1) >> get)
+     in testUnfoldMD unf (0 :: Int) 0 1 [1]
 
 unfoldrM :: Property
 unfoldrM =