@@ -10,6 +10,10 @@ module Streamly.Internal.Data.Unfold.Types
1010 ( Unfold (..)
1111 , lmap
1212 , map
13+ , const
14+ , apSequence
15+ , apDiscardSnd
16+ , cross
1317 )
1418where
1519
@@ -17,7 +21,7 @@ where
1721
1822import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
1923
20−import Prelude hiding (map)
24+import Prelude hiding (const, map)
2125
2226------------------------------------------------------------------------------
2327-- Monadic Unfolds
@@ -63,3 +67,86 @@ map f (Unfold ustep uinject) = Unfold step uinject
6367instance Functor m => Functor (Unfold m a) where
6468 {-# INLINE fmap #-}
6569 fmap = map
70+
71+------------------------------------------------------------------------------
72+-- Applicative
73+------------------------------------------------------------------------------
74+
75+{-# INLINE const #-}
76+const :: Applicative m => m b -> Unfold m a b
77+const m = Unfold step inject
78+
79+ where
80+
81+ inject _ = pure False
82+
83+ step False = (`Yield` True) <$> m
84+ step True = pure Stop
85+
86+-- | Outer product discarding the first element.
87+--
88+-- /Unimplemented/
89+--
90+{-# INLINE_NORMAL apSequence #-}
91+apSequence :: -- Monad m =>
92+ Unfold m a b -> Unfold m a c -> Unfold m a c
93+apSequence (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
94+
95+-- | Outer product discarding the second element.
96+--
97+-- /Unimplemented/
98+--
99+{-# INLINE_NORMAL apDiscardSnd #-}
100+apDiscardSnd :: -- Monad m =>
101+ Unfold m a b -> Unfold m a c -> Unfold m a b
102+apDiscardSnd (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
103+
104+data Cross a s1 b s2 = CrossOuter a s1 | CrossInner a s1 b s2
105+
106+-- | Create a cross product (vector product or cartesian product) of the
107+-- output streams of two unfolds.
108+--
109+{-# INLINE_NORMAL cross #-}
110+cross :: Monad m => Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
111+cross (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
112+
113+ where
114+
115+ inject a = do
116+ s1 <- inject1 a
117+ return $ CrossOuter a s1
118+
119+ {-# INLINE_LATE step #-}
120+ step (CrossOuter a s1) = do
121+ r <- step1 s1
122+ case r of
123+ Yield b s -> do
124+ s2 <- inject2 a
125+ return $ Skip (CrossInner a s b s2)
126+ Skip s -> return $ Skip (CrossOuter a s)
127+ Stop -> return Stop
128+
129+ step (CrossInner a s1 b s2) = do
130+ r <- step2 s2
131+ return $ case r of
132+ Yield c s -> Yield (b, c) (CrossInner a s1 b s)
133+ Skip s -> Skip (CrossInner a s1 b s)
134+ Stop -> Skip (CrossOuter a s1)
135+
136+-- | Example:
137+--
138+-- >>> Stream.toList $ Stream.unfold ((,) <$> Unfold.lmap fst Unfold.fromList <*> Unfold.lmap snd Unfold.fromList) ([1,2],[3,4])
139+-- [(1,3),(1,4),(2,3),(2,4)]
140+--
141+instance Monad m => Applicative (Unfold m a) where
142+ {-# INLINE pure #-}
143+ pure = const . return
144+
145+ {-# INLINE (<*>) #-}
146+ u1 <*> u2 = fmap (\(a, b) -> a b) (cross u1 u2)
147+
148+ -- {-# INLINE (*>) #-}
149+ -- (*>) = apSequence
150+
151+ -- {-# INLINE (<*) #-}
152+ -- (<*) = apDiscardSnd