@@ -112,16 +112,23 @@ module Streamly.Internal.Data.Unfold
112112
113113 -- * Exceptions
114114 , gbracket
115+ , gbracketIO
115116 , before
116117 , after
118+ , afterIO
117119 , onException
118120 , finally
121+ , finallyIO
119122 , bracket
123+ , bracketIO
120124 , handle
121125 )
122126where
123127
124128import Control.Exception (Exception)
129+import Control.Monad.IO.Class (MonadIO(..))
130+import Control.Monad.Trans.Control (MonadBaseControl)
131+import Data.IORef (readIORef, writeIORef)
125132import Data.Void (Void)
126133import GHC.Types (SPEC(..))
127134import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)
@@ -132,7 +139,7 @@ import Streamly.Internal.Data.Stream.StreamD.Type (pattern Stream)
132139#endif
133140import Streamly.Internal.Data.Unfold.Types (Unfold(..))
134141import Streamly.Internal.Data.Fold.Types (Fold(..))
135−import Streamly.Internal.Data.SVar (defState)
142+import Streamly.Internal.Data.SVar (defState, MonadAsync)
136143import Control.Monad.Catch (MonadCatch)
137144
138145import qualified Prelude
@@ -653,6 +660,53 @@ gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
653660 Skip s -> return $ Skip (Left s)
654661 Stop -> return Stop
655662
663+-- | The most general bracketing and exception combinator. All other
664+-- combinators can be expressed in terms of this combinator. This can also be
665+-- used for cases which are not covered by the standard combinators.
666+--
667+-- /Internal/
668+--
669+{-# INLINE_NORMAL gbracketIO #-}
670+gbracketIO
671+ :: (MonadIO m, MonadBaseControl IO m)
672+ => (a -> m c) -- ^ before
673+ -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
674+ -> (c -> m d) -- ^ after, on normal stop, or GC
675+ -> Unfold m (c, e) b -- ^ on exception
676+ -> Unfold m c b -- ^ unfold to run
677+ -> Unfold m a b
678+gbracketIO bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
679+ Unfold step inject
680+
681+ where
682+
683+ inject x = do
684+ r <- bef x
685+ ref <- D.newFinalizedIORef (aft r)
686+ s <- inject1 r
687+ return $ Right (s, r, ref)
688+
689+ {-# INLINE_LATE step #-}
690+ step (Right (st, v, ref)) = do
691+ res <- exc $ step1 st
692+ case res of
693+ Right r -> case r of
694+ Yield x s -> return $ Yield x (Right (s, v, ref))
695+ Skip s -> return $ Skip (Right (s, v, ref))
696+ Stop -> do
697+ liftIO $ writeIORef ref Nothing
698+ aft v >> return Stop
699+ Left e -> do
700+ liftIO (writeIORef ref Nothing)
701+ r <- einject (v, e)
702+ return $ Skip (Left r)
703+ step (Left st) = do
704+ res <- estep st
705+ case res of
706+ Yield x s -> return $ Yield x (Left s)
707+ Skip s -> return $ Skip (Left s)
708+ Stop -> return Stop
709+
656710-- The custom implementation of "before" is slightly faster (5-7%) than
657711-- "_before". This is just to document and make sure that we can always use
658712-- gbracket to implement before. The same applies to other combinators as well.
@@ -690,6 +744,10 @@ _after aft = gbracket return (fmap Right) aft undefined
690744
691745-- | Run a side effect whenever the unfold stops normally.
692746--
747+-- Prefer afterIO over this as the @after@ action in this combinator is not
748+-- executed if the unfold is partially evaluated lazily and then garbage
749+-- collected.
750+--
693751-- /Internal/
694752{-# INLINE_NORMAL after #-}
695753after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
@@ -709,6 +767,33 @@ after action (Unfold step1 inject1) = Unfold step inject
709767 Skip s -> return $ Skip (s, v)
710768 Stop -> action v >> return Stop
711769
770+-- | Run a side effect whenever the unfold stops normally
771+-- or is garbage collected after a partial lazy evaluation.
772+--
773+-- /Internal/
774+{-# INLINE_NORMAL afterIO #-}
775+afterIO :: (MonadIO m, MonadBaseControl IO m)
776+ => (a -> m c) -> Unfold m a b -> Unfold m a b
777+afterIO action (Unfold step1 inject1) = Unfold step inject
778+
779+ where
780+
781+ inject x = do
782+ s <- inject1 x
783+ ref <- D.newFinalizedIORef (action x)
784+ return (s, ref)
785+
786+ {-# INLINE_LATE step #-}
787+ step (st, ref) = do
788+ res <- step1 st
789+ case res of
790+ Yield x s -> return $ Yield x (s, ref)
791+ Skip s -> return $ Skip (s, ref)
792+ Stop -> liftIO $ do
793+ Just f <- readIORef ref
794+ writeIORef ref Nothing
795+ f >> return Stop
796+
712797{-# INLINE_NORMAL _onException #-}
713798_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
714799_onException action unf =
@@ -746,6 +831,10 @@ _finally action unf =
746831-- | Run a side effect whenever the unfold stops normally or aborts due to an
747832-- exception.
748833--
834+-- Prefer finallyIO over this as the @after@ action in this combinator is not
835+-- executed if the unfold is partially evaluated lazily and then garbage
836+-- collected.
837+--
749838-- /Internal/
750839{-# INLINE_NORMAL finally #-}
751840finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
@@ -765,6 +854,32 @@ finally action (Unfold step1 inject1) = Unfold step inject
765854 Skip s -> return $ Skip (s, v)
766855 Stop -> action v >> return Stop
767856
857+-- | Run a side effect whenever the unfold stops normally, aborts due to an
858+-- exception or if it is garbage collected after a partial lazy evaluation.
859+--
860+-- /Internal/
861+{-# INLINE_NORMAL finallyIO #-}
862+finallyIO :: (MonadAsync m, MonadCatch m)
863+ => (a -> m c) -> Unfold m a b -> Unfold m a b
864+finallyIO action (Unfold step1 inject1) = Unfold step inject
865+
866+ where
867+
868+ inject x = do
869+ s <- inject1 x
870+ ref <- D.newFinalizedIORef (action x)
871+ return (s, x, ref)
872+
873+ {-# INLINE_LATE step #-}
874+ step (st, v, ref) = do
875+ res <- step1 st `MC.onException` action v
876+ case res of
877+ Yield x s -> return $ Yield x (s, v, ref)
878+ Skip s -> return $ Skip (s, v, ref)
879+ Stop -> do
880+ liftIO (writeIORef ref Nothing)
881+ action v >> return Stop
882+
768883{-# INLINE_NORMAL _bracket #-}
769884_bracket :: MonadCatch m
770885 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
@@ -777,6 +892,10 @@ _bracket bef aft unf =
777892-- if an exception occurs then the @after@ action is run with the output of
778893-- @before@ as argument.
779894--
895+-- Prefer bracketIO over this as the @after@ action in this combinator is not
896+-- executed if the unfold is partially evaluated lazily and then garbage
897+-- collected.
898+--
780899-- /Internal/
781900{-# INLINE_NORMAL bracket #-}
782901bracket :: MonadCatch m
@@ -798,6 +917,36 @@ bracket bef aft (Unfold step1 inject1) = Unfold step inject
798917 Skip s -> return $ Skip (s, v)
799918 Stop -> aft v >> return Stop
800919
920+-- | @bracket before after between@ runs the @before@ action and then unfolds
921+-- its output using the @between@ unfold. When the @between@ unfold is done or
922+-- if an exception occurs then the @after@ action is run with the output of
923+-- @before@ as argument. The after action is also executed if the unfold is
924+-- paritally evaluated and then garbage collected.
925+--
926+-- /Internal/
927+{-# INLINE_NORMAL bracketIO #-}
928+bracketIO :: (MonadAsync m, MonadCatch m)
929+ => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
930+bracketIO bef aft (Unfold step1 inject1) = Unfold step inject
931+
932+ where
933+
934+ inject x = do
935+ r <- bef x
936+ s <- inject1 r
937+ ref <- D.newFinalizedIORef (aft r)
938+ return (s, r, ref)
939+
940+ {-# INLINE_LATE step #-}
941+ step (st, v, ref) = do
942+ res <- step1 st `MC.onException` aft v
943+ case res of
944+ Yield x s -> return $ Yield x (s, v, ref)
945+ Skip s -> return $ Skip (s, v, ref)
946+ Stop -> do
947+ liftIO (writeIORef ref Nothing)
948+ aft v >> return Stop
949+
801950-- | When unfolding if an exception occurs, unfold the exception using the
802951-- exception unfold supplied as the first argument to 'handle'.
803952--