git.iko.soy/streamly
on master help sign in
Implement exception handling for unfolds
16de98cHarendra Kumar
src/Streamly/Internal/Data/Unfold.hs+143
compare with master
unifiedsplit
@@ -92,9 +92,18 @@ module Streamly.Internal.Data.Unfold
92 , concat
92 , concat
93 , concatMapM
93 , concatMapM
94 , outerProduct
94 , outerProduct
95+
96+ -- * Exceptions
97+ , before
98+ , after
99+ , onException
100+ , finally
101+ , bracket
102+ , handle
95 )
103 )
96where
104where
97
105
106+import Control.Exception (Exception)
98import Data.Void (Void)
107import Data.Void (Void)
99import GHC.Types (SPEC(..))
108import GHC.Types (SPEC(..))
100import Prelude hiding (concat, map, takeWhile, take, filter)
109import Prelude hiding (concat, map, takeWhile, take, filter)
@@ -106,6 +115,9 @@ import Streamly.Streams.StreamD.Type (pattern Stream)
106import Streamly.Internal.Data.Unfold.Types (Unfold(..))
115import Streamly.Internal.Data.Unfold.Types (Unfold(..))
107import Streamly.Internal.Data.Fold.Types (Fold(..))
116import Streamly.Internal.Data.Fold.Types (Fold(..))
108import Streamly.Internal.Data.SVar (defState)
117import Streamly.Internal.Data.SVar (defState)
118+import Control.Monad.Catch (MonadCatch)
119+
120+import qualified Control.Monad.Catch as MC
109
121
110-------------------------------------------------------------------------------
122-------------------------------------------------------------------------------
111-- Input operations
123-- Input operations
@@ -396,3 +408,134 @@ concatMapM f (Unfold step1 inject1) = Unfold step inject
396 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
408 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
397 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
409 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
398 Stop -> Skip (ConcatMapOuter ost)
410 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