@@ -92,9 +92,18 @@ module Streamly.Internal.Data.Unfold
9292 , concat
9393 , concatMapM
9494 , outerProduct
95+
96+ -- * Exceptions
97+ , before
98+ , after
99+ , onException
100+ , finally
101+ , bracket
102+ , handle
95103 )
96104where
97105
106+import Control.Exception (Exception)
98107import Data.Void (Void)
99108import GHC.Types (SPEC(..))
100109import Prelude hiding (concat, map, takeWhile, take, filter)
@@ -106,6 +115,9 @@ import Streamly.Streams.StreamD.Type (pattern Stream)
106115import Streamly.Internal.Data.Unfold.Types (Unfold(..))
107116import Streamly.Internal.Data.Fold.Types (Fold(..))
108117import Streamly.Internal.Data.SVar (defState)
118+import Control.Monad.Catch (MonadCatch)
119+
120+import qualified Control.Monad.Catch as MC
109121
110122-------------------------------------------------------------------------------
111123-- Input operations
@@ -396,3 +408,134 @@ concatMapM f (Unfold step1 inject1) = Unfold step inject
396408 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
397409 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
398410 Stop -> Skip (ConcatMapOuter ost)
411+
412+------------------------------------------------------------------------------
413+-- Exceptions
414+------------------------------------------------------------------------------
415+
416+-- | Run a side effect before the unfold yields its first element.
417+--
418+-- /Internal/
419+{-# INLINE_NORMAL before #-}
420+before :: Monad m => m c -> Unfold m a b -> Unfold m a b
421+before action (Unfold step1 inject1) = Unfold step inject
422+
423+ where
424+
425+ inject x = do
426+ _ <- action
427+ st <- inject1 x
428+ return st
429+
430+ {-# INLINE_LATE step #-}
431+ step st = do
432+ res <- step1 st
433+ case res of
434+ Yield x s -> return $ Yield x s
435+ Skip s -> return $ Skip s
436+ Stop -> return Stop
437+
438+-- | Run a side effect whenever the unfold stops normally.
439+--
440+-- /Internal/
441+{-# INLINE_NORMAL after #-}
442+after :: Monad m => m c -> Unfold m a b -> Unfold m a b
443+after action (Unfold step1 inject1) = Unfold step inject1
444+
445+ where
446+
447+ {-# INLINE_LATE step #-}
448+ step st = do
449+ res <- step1 st
450+ case res of
451+ Yield x s -> return $ Yield x s
452+ Skip s -> return $ Skip s
453+ Stop -> action >> return Stop
454+
455+-- | Run a side effect whenever the unfold aborts due to an exception.
456+--
457+-- /Internal/
458+{-# INLINE_NORMAL onException #-}
459+onException :: MonadCatch m => m c -> Unfold m a b -> Unfold m a b
460+onException action (Unfold step1 inject1) = Unfold step inject1
461+
462+ where
463+
464+ {-# INLINE_LATE step #-}
465+ step st = do
466+ res <- step1 st `MC.onException` action
467+ case res of
468+ Yield x s -> return $ Yield x s
469+ Skip s -> return $ Skip s
470+ Stop -> return Stop
471+
472+-- | Run a side effect whenever the unfold stops normally or aborts due to an
473+-- exception.
474+--
475+-- /Internal/
476+{-# INLINE_NORMAL finally #-}
477+finally :: MonadCatch m => m c -> Unfold m a b -> Unfold m a b
478+finally action (Unfold step1 inject1) = Unfold step inject1
479+
480+ where
481+
482+ {-# INLINE_LATE step #-}
483+ step st = do
484+ res <- step1 st `MC.onException` action
485+ case res of
486+ Yield x s -> return $ Yield x s
487+ Skip s -> return $ Skip s
488+ Stop -> action >> return Stop
489+
490+-- | @bracket before after between@ runs the @before@ action and then runs the
491+-- @between@ function using the output of @before@ as its argument. When the
492+-- @between@ unfold is done or if an exception occurs while running @between@
493+-- the @after@ action is run with the output of @before@ as argument.
494+--
495+-- /Internal/
496+{-# INLINE_NORMAL bracket #-}
497+bracket :: MonadCatch m
498+ => m c -> (c -> m d) -> (c -> Unfold m a b) -> Unfold m a b
499+bracket bef aft bet = Unfold step inject
500+
501+ where
502+
503+ inject x = do
504+ r <- bef
505+ go (bet r) r
506+ where
507+ go (Unfold step2 inject2) r = do
508+ st <- inject2 x
509+ return (Stream (\_ -> step2) st, r)
510+
511+ {-# INLINE_LATE step #-}
512+ step (Stream step2 state, v) = do
513+ res <- step2 defState state `MC.onException` aft v
514+ case res of
515+ Yield x s -> return $ Yield x (Stream step2 s, v)
516+ Skip s -> return $ Skip (Stream step2 s, v)
517+ Stop -> aft v >> return Stop
518+
519+-- | When unfolding an unfold if an exception occurs, unfold aborts
520+-- and the specified exception handler is run with the exception as argument.
521+--
522+-- /Internal/
523+{-# INLINE_NORMAL handle #-}
524+handle :: (MonadCatch m, Exception e)
525+ => (e -> m b) -> Unfold m a b -> Unfold m a b
526+handle f (Unfold step1 inject1) = Unfold step inject
527+
528+ where
529+
530+ inject x = inject1 x >>= return . Just
531+
532+ {-# INLINE_LATE step #-}
533+ step (Just st) = do
534+ res <- MC.try $ step1 st
535+ case res of
536+ Left e -> f e >>= \x -> return (Yield x Nothing)
537+ Right r -> case r of
538+ Yield x s -> return $ Yield x (Just s)
539+ Skip s -> return $ Skip (Just s)
540+ Stop -> return Stop
541+ step Nothing = return Stop