@@ -6,6 +6,7 @@
66{-# LANGUAGE RankNTypes #-}
77{-# LANGUAGE RecordWildCards #-}
88{-# LANGUAGE ScopedTypeVariables #-}
9+{-# LANGUAGE TupleSections #-}
910
1011#include "inline.hs"
1112
@@ -59,10 +60,13 @@ module Streamly.Internal.Data.Unfold
5960 Unfold
6061
6162 -- * Operations on Input
63+ , lmap
6264 , close
6365 , first
6466 , second
65− , lmap
67+ , discardFirst
68+ , discardSecond
69+ , swap
6670 -- coapply
6771 -- comonad
6872
@@ -74,9 +78,14 @@ module Streamly.Internal.Data.Unfold
7478 , fromStream
7579 , fromStream1
7680 , fromStream2
81+ , nilM
82+ , consM
83+ , effect
7784 , singleton
85+ , identity
7886 , replicateM
7987 , fromList
88+ , fromListM
8089 , enumerateFromStepIntegral
8190 , enumerateFromToIntegral
8291 , enumerateFromIntegral
@@ -121,6 +130,7 @@ import Streamly.Internal.Data.SVar (defState)
121130import Control.Monad.Catch (MonadCatch)
122131
123132import qualified Control.Monad.Catch as MC
133+import qualified Data.Tuple as Tuple
124134import qualified Streamly.Streams.StreamK as K
125135import qualified Streamly.Streams.StreamD as D
126136
@@ -128,13 +138,17 @@ import qualified Streamly.Streams.StreamD as D
128138-- Input operations
129139-------------------------------------------------------------------------------
130140
141+{-# INLINE_NORMAL lmap #-}
142+lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
143+lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
144+
131145-- | Supply the seed to an unfold closing the input end of the unfold.
132146--
133147-- /Internal/
134148--
135149{-# INLINE_NORMAL close #-}
136150close :: Unfold m a b -> a -> Unfold m Void b
137−close (Unfold ustep uinject) a = Unfold ustep (const $ uinject a)
151+close unf a = lmap (const a) unf
138152
139153-- XXX rename to closeFst/closeSnd
140154--
@@ -146,7 +160,7 @@ close (Unfold ustep uinject) a = Unfold ustep (const $ uinject a)
146160--
147161{-# INLINE_NORMAL first #-}
148162first :: Unfold m (a, b) c -> a -> Unfold m b c
149−first (Unfold ustep uinject) a = Unfold ustep (\b -> uinject (a, b))
163+first unf a = lmap (a, ) unf
150164
151165-- | Supply the second component of the tuple to an unfold that accepts a tuple
152166-- as a seed resulting in a fold that accepts the first component of the tuple
@@ -156,11 +170,19 @@ first (Unfold ustep uinject) a = Unfold ustep (\b -> uinject (a, b))
156170--
157171{-# INLINE_NORMAL second #-}
158172second :: Unfold m (a, b) c -> b -> Unfold m a c
159−second (Unfold ustep uinject) b = Unfold ustep (\a -> uinject (a, b))
173+second unf b = lmap (, b) unf
160174
161−{-# INLINE_NORMAL lmap #-}
162−lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
163−lmap f (Unfold ustep uinject) = Unfold ustep (\x -> uinject $ f x)
175+{-# INLINE_NORMAL discardFirst #-}
176+discardFirst :: Unfold m a b -> Unfold m (c, a) b
177+discardFirst = lmap snd
178+
179+{-# INLINE_NORMAL discardSecond #-}
180+discardSecond :: Unfold m a b -> Unfold m (a, c) b
181+discardSecond = lmap fst
182+
183+{-# INLINE_NORMAL swap #-}
184+swap :: Unfold m (a, c) b -> Unfold m (c, a) b
185+swap = lmap Tuple.swap
164186
165187-------------------------------------------------------------------------------
166188-- Output operations
@@ -241,20 +263,68 @@ fromStream2 f = Unfold streamStep (\(a, b) -> return $ D.toStreamD $ f a b)
241263-- Unfolds
242264-------------------------------------------------------------------------------
243265
244−-- | Identity unfold. Generates a singleton stream with the seed as the only
245−-- element in the stream.
266+-- | Lift a monadic function into an unfold generating a nil stream with a side
267+-- effect.
246268--
247−-- > singleton = replicateM 1
269+{-# INLINE nilM #-}
270+nilM :: Monad m => (a -> m c) -> Unfold m a b
271+nilM f = Unfold step return
272+ where
273+ {-# INLINE_LATE step #-}
274+ step x = f x >> return Stop
275+
276+-- | Prepend a monadic single element generator function to an 'Unfold'.
277+--
278+-- /Internal/
279+{-# INLINE_NORMAL consM #-}
280+consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
281+consM action unf = Unfold step inject
282+
283+ where
284+
285+ inject = return . Left
286+
287+ {-# INLINE_LATE step #-}
288+ step (Left a) = do
289+ action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
290+ step (Right (UnStream step1 st)) = do
291+ res <- step1 defState st
292+ case res of
293+ Yield x s -> return $ Yield x (Right (Stream step1 s))
294+ Skip s -> return $ Skip (Right (Stream step1 s))
295+ Stop -> return Stop
296+
297+-- | Lift a monadic effect into an unfold generating a singleton stream.
298+--
299+{-# INLINE effect #-}
300+effect :: Monad m => m b -> Unfold m Void b
301+effect eff = Unfold step inject
302+ where
303+ inject _ = return True
304+ {-# INLINE_LATE step #-}
305+ step True = eff >>= \r -> return $ Yield r False
306+ step False = return Stop
307+
308+-- | Lift a monadic function into an unfold generating a singleton stream.
248309--
249310{-# INLINE singleton #-}
250−singleton :: Monad m => Unfold m a a
251−singleton = Unfold step inject
311+singleton :: Monad m => (a -> m b) -> Unfold m a b
312+singleton f = Unfold step inject
252313 where
253314 inject x = return $ Just x
254315 {-# INLINE_LATE step #-}
255− step (Just x) = return $ Yield x Nothing
316+ step (Just x) = f x >>= \r -> return $ Yield r Nothing
256317 step Nothing = return Stop
257318
319+-- | Identity unfold. Generates a singleton stream with the seed as the only
320+-- element in the stream.
321+--
322+-- > identity = singleton return
323+--
324+{-# INLINE identity #-}
325+identity :: Monad m => Unfold m a a
326+identity = singleton return
327+
258328-- | Generates a stream replicating the seed @n@ times.
259329--
260330{-# INLINE replicateM #-}
@@ -278,6 +348,16 @@ fromList = Unfold step inject
278348 step (x:xs) = return $ Yield x xs
279349 step [] = return Stop
280350
351+-- | Convert a list of monadic values to a 'Stream'
352+{-# INLINE_LATE fromListM #-}
353+fromListM :: Monad m => Unfold m [m a] a
354+fromListM = Unfold step inject
355+ where
356+ inject x = return x
357+ {-# INLINE_LATE step #-}
358+ step (x:xs) = x >>= \r -> return $ Yield r xs
359+ step [] = return Stop
360+
281361-------------------------------------------------------------------------------
282362-- Filtering
283363-------------------------------------------------------------------------------
@@ -444,7 +524,7 @@ concatMapM f (Unfold step1 inject1) = Unfold step inject
444524 Skip s -> return $ Skip (ConcatMapOuter s)
445525 Stop -> return Stop
446526
447− step (ConcatMapInner ost (Stream istep ist)) = do
527+ step (ConcatMapInner ost (UnStream istep ist)) = do
448528 r <- istep defState ist
449529 return $ case r of
450530 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
@@ -455,6 +535,54 @@ concatMapM f (Unfold step1 inject1) = Unfold step inject
455535-- Exceptions
456536------------------------------------------------------------------------------
457537
538+-- MBD: If we add an (Unfold m e b) argument to it we can encompass "handle" as
539+-- well under this.
540+{-# INLINE_NORMAL gbracket #-}
541+gbracket
542+ :: Monad m
543+ => (a -> m c)
544+ -> (forall s. m s -> m (Either e s))
545+ -> (c -> m d)
546+ -> Unfold m (c, e) b
547+ -> Unfold m c b
548+ -> Unfold m a b
549+gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
550+ Unfold step inject
551+
552+ where
553+
554+ inject x = do
555+ r <- bef x
556+ s <- inject1 r
557+ return $ Right (s, r)
558+
559+ {-# INLINE_LATE step #-}
560+ step (Right (st, v)) = do
561+ res <- exc $ step1 st
562+ case res of
563+ Right r -> case r of
564+ Yield x s -> return $ Yield x (Right (s, v))
565+ Skip s -> return $ Skip (Right (s, v))
566+ Stop -> aft v >> return Stop
567+ Left e -> do
568+ r <- einject (v, e)
569+ return $ Skip (Left r)
570+ step (Left st) = do
571+ res <- estep st
572+ case res of
573+ Yield x s -> return $ Yield x (Left s)
574+ Skip s -> return $ Skip (Left s)
575+ Stop -> return Stop
576+
577+-- The custom implementation of "before" is slightly faster (5-7%) than
578+-- "_before". This is just to document and make sure that we can always use
579+-- gbracket to implement before. The same applies to other combinators as well.
580+--
581+{-# INLINE_NORMAL _before #-}
582+_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
583+_before action unf = gbracket (\x -> action x >> return x) (fmap Right)
584+ (\_ -> return ()) undefined unf
585+
458586-- | Run a side effect before the unfold yields its first element.
459587--
460588-- /Internal/
@@ -477,6 +605,10 @@ before action (Unfold step1 inject1) = Unfold step inject
477605 Skip s -> return $ Skip s
478606 Stop -> return Stop
479607
608+{-# INLINE_NORMAL _after #-}
609+_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
610+_after aft = gbracket return (fmap Right) aft undefined
611+
480612-- | Run a side effect whenever the unfold stops normally.
481613--
482614-- /Internal/
@@ -498,6 +630,13 @@ after action (Unfold step1 inject1) = Unfold step inject
498630 Skip s -> return $ Skip (s, v)
499631 Stop -> action v >> return Stop
500632
633+{-# INLINE_NORMAL _onException #-}
634+_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
635+_onException action unf =
636+ gbracket return MC.try
637+ (\_ -> return ())
638+ (nilM (\(a, (_ :: MC.SomeException)) -> action a)) unf
639+
501640-- | Run a side effect whenever the unfold aborts due to an exception.
502641--
503642-- /Internal/
@@ -519,6 +658,12 @@ onException action (Unfold step1 inject1) = Unfold step inject
519658 Skip s -> return $ Skip (s, v)
520659 Stop -> return Stop
521660
661+{-# INLINE_NORMAL _finally #-}
662+_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
663+_finally action unf =
664+ gbracket return MC.try action
665+ (nilM (\(a, (_ :: MC.SomeException)) -> action a)) unf
666+
522667-- | Run a side effect whenever the unfold stops normally or aborts due to an
523668-- exception.
524669--
@@ -541,6 +686,12 @@ finally action (Unfold step1 inject1) = Unfold step inject
541686 Skip s -> return $ Skip (s, v)
542687 Stop -> action v >> return Stop
543688
689+{-# INLINE_NORMAL _bracket #-}
690+_bracket :: MonadCatch m
691+ => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
692+_bracket bef aft unf =
693+ gbracket bef MC.try aft (nilM (\(a, (_ :: MC.SomeException)) -> aft a)) unf
694+
544695-- | @bracket before after between@ runs the @before@ action and then unfolds
545696-- its output using the @between@ unfold. When the @between@ unfold is done or
546697-- if an exception occurs then the @after@ action is run with the output of
@@ -567,26 +718,12 @@ bracket bef aft (Unfold step1 inject1) = Unfold step inject
567718 Skip s -> return $ Skip (s, v)
568719 Stop -> aft v >> return Stop
569720
570−-- | When unfolding an unfold if an exception occurs, unfold aborts
571−-- and the specified exception handler is run with the exception as argument.
721+-- | When unfolding if an exception occurs, unfold the exception using the
722+-- exception unfold.
572723--
573724-- /Internal/
574725{-# INLINE_NORMAL handle #-}
575726handle :: (MonadCatch m, Exception e)
576− => (e -> m b) -> Unfold m a b -> Unfold m a b
577−handle f (Unfold step1 inject1) = Unfold step inject
578−
579− where
580−
581− inject x = inject1 x >>= return . Just
582−
583− {-# INLINE_LATE step #-}
584− step (Just st) = do
585− res <- MC.try $ step1 st
586− case res of
587− Left e -> f e >>= \x -> return (Yield x Nothing)
588− Right r -> case r of
589− Yield x s -> return $ Yield x (Just s)
590− Skip s -> return $ Skip (Just s)
591− Stop -> return Stop
592− step Nothing = return Stop
727+ => Unfold m e b -> Unfold m a b -> Unfold m a b
728+handle exc unf =
729+ gbracket return MC.try (\_ -> return ()) (discardFirst exc) unf