@@ -257,6 +257,7 @@ module Streamly.Streams.StreamD
257257 , the
258258
259259 -- * Exceptions
260+ , gbracket
260261 , before
261262 , after
262263 , bracket
@@ -273,7 +274,7 @@ module Streamly.Streams.StreamD
273274 )
274275where
275276
276−import Control.Exception (Exception)
277+import Control.Exception (Exception, SomeException)
277278import Control.Monad (void)
278279import Control.Monad.Catch (MonadCatch)
279280import Control.Monad.IO.Class (MonadIO(..))
@@ -2256,6 +2257,52 @@ interleaveFstThenConcat
22562257-- Exceptions
22572258------------------------------------------------------------------------------
22582259
2260+data GbracketState s1 s2 v
2261+ = GBracketInit
2262+ | GBracketNormal s1 v
2263+ | GBracketException s2
2264+
2265+-- | The most general bracketing and exception combinator. All other
2266+-- combinators can be expressed in terms of this combinator. This can also be
2267+-- used for cases which are not covered by the standard combinators.
2268+--
2269+-- /Internal/
2270+--
2271+{-# INLINE_NORMAL gbracket #-}
2272+gbracket
2273+ :: Monad m
2274+ => m c -- ^ before
2275+ -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2276+ -> (c -> m d) -- ^ after, on normal stop
2277+ -> (c -> e -> Stream m b) -- ^ on exception
2278+ -> (c -> Stream m b) -- ^ stream generator
2279+ -> Stream m b
2280+gbracket bef exc aft fexc fnormal =
2281+ Stream step GBracketInit
2282+
2283+ where
2284+
2285+ {-# INLINE_LATE step #-}
2286+ step _ GBracketInit = do
2287+ r <- bef
2288+ return $ Skip $ GBracketNormal (fnormal r) r
2289+
2290+ step gst (GBracketNormal (UnStream step1 st) v) = do
2291+ res <- exc $ step1 gst st
2292+ case res of
2293+ Right r -> case r of
2294+ Yield x s ->
2295+ return $ Yield x (GBracketNormal (Stream step1 s) v)
2296+ Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2297+ Stop -> aft v >> return Stop
2298+ Left e -> return $ Skip (GBracketException (fexc v e))
2299+ step gst (GBracketException (UnStream step1 st)) = do
2300+ res <- step1 gst st
2301+ case res of
2302+ Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2303+ Skip s -> return $ Skip (GBracketException (Stream step1 s))
2304+ Stop -> return Stop
2305+
22592306-- | Run a side effect before the stream yields its first element.
22602307{-# INLINE_NORMAL before #-}
22612308before :: Monad m => m b -> Stream m a -> Stream m a
@@ -2301,7 +2348,14 @@ after action (Stream step state) = Stream step' state
23012348-- exception is not caught, simply rethrown.
23022349{-# INLINE_NORMAL onException #-}
23032350onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2304−onException action (Stream step state) = Stream step' state
2351+onException action str =
2352+ gbracket (return ()) MC.try return
2353+ (\_ (_ :: MC.SomeException) -> nilM $ action)
2354+ (\_ -> str)
2355+
2356+{-# INLINE_NORMAL _onException #-}
2357+_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2358+_onException action (Stream step state) = Stream step' state
23052359
23062360 where
23072361
@@ -2313,13 +2367,6 @@ onException action (Stream step state) = Stream step' state
23132367 Skip s -> return $ Skip s
23142368 Stop -> return Stop
23152369
2316−-- | Run a side effect whenever the stream stops normally or aborts due to an
2317−-- exception.
2318−{-# INLINE finally #-}
2319−finally :: MonadCatch m => m b -> Stream m a -> Stream m a
2320−-- finally action xs = after action $ onException action xs
2321−finally action xs = bracket (return ()) (\_ -> action) (const xs)
2322−
23232370-- XXX bracket is like concatMap, it generates a stream and then flattens it.
23242371-- Like concatMap it has 10x worse performance compared to linear fused
23252372-- compositions.
@@ -2330,41 +2377,70 @@ finally action xs = bracket (return ()) (\_ -> action) (const xs)
23302377-- an exception.
23312378{-# INLINE_NORMAL bracket #-}
23322379bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2333−bracket bef aft bet = Stream step' Nothing
2380+bracket bef aft bet =
2381+ gbracket bef MC.try aft (\a (_ :: SomeException) -> nilM $ aft a) bet
2382+
2383+data BracketState s v = BracketInit | BracketRun s v
2384+
2385+{-# INLINE_NORMAL _bracket #-}
2386+_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2387+_bracket bef aft bet = Stream step' BracketInit
23342388
23352389 where
23362390
23372391 {-# INLINE_LATE step' #-}
2338− step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
2392+ step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
23392393
23402394 -- NOTE: It is important to use UnStream instead of the Stream pattern
23412395 -- here, otherwise we get huge perf degradation, see note in concatMap.
2342− step' gst (Just (UnStream step state, v)) = do
2343− res <- step gst state `MC.onException` aft v
2396+ step' gst (BracketRun (UnStream step state) v) = do
2397+ -- res <- step gst state `MC.onException` aft v
2398+ res <- MC.try $ step gst state
23442399 case res of
2345− Yield x s -> return $ Yield x (Just (Stream step s, v))
2346− Skip s -> return $ Skip (Just (Stream step s, v))
2347− Stop -> aft v >> return Stop
2400+ Left (_ :: SomeException) -> aft v >> return Stop
2401+ Right r -> case r of
2402+ Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
2403+ Skip s -> return $ Skip (BracketRun (Stream step s) v)
2404+ Stop -> aft v >> return Stop
2405+
2406+-- | Run a side effect whenever the stream stops normally or aborts due to an
2407+-- exception.
2408+{-# INLINE finally #-}
2409+finally :: MonadCatch m => m b -> Stream m a -> Stream m a
2410+-- finally action xs = after action $ onException action xs
2411+finally action xs = bracket (return ()) (\_ -> action) (const xs)
23482412
23492413-- | When evaluating a stream if an exception occurs, stream evaluation aborts
23502414-- and the specified exception handler is run with the exception as argument.
23512415{-# INLINE_NORMAL handle #-}
2352−handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
2353−handle f (Stream step state) = Stream step' (Just state)
2416+handle :: (MonadCatch m, Exception e)
2417+ => (e -> Stream m a) -> Stream m a -> Stream m a
2418+handle f str =
2419+ gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
2420+
2421+{-# INLINE_NORMAL _handle #-}
2422+_handle :: (MonadCatch m, Exception e)
2423+ => (e -> Stream m a) -> Stream m a -> Stream m a
2424+_handle f (Stream step state) = Stream step' (Left state)
23542425
23552426 where
23562427
23572428 {-# INLINE_LATE step' #-}
2358− step' gst (Just st) = do
2429+ step' gst (Left st) = do
23592430 res <- MC.try $ step gst st
23602431 case res of
2361− Left e -> f e >>= \x -> return (Yield x Nothing)
2432+ Left e -> return $ Skip $ Right (f e)
23622433 Right r -> case r of
2363− Yield x s -> return $ Yield x (Just s)
2364− Skip s -> return $ Skip (Just s)
2434+ Yield x s -> return $ Yield x (Left s)
2435+ Skip s -> return $ Skip (Left s)
23652436 Stop -> return Stop
23662437
2367− step' _ Nothing = return Stop
2438+ step' gst (Right (UnStream step1 st)) = do
2439+ res <- step1 gst st
2440+ case res of
2441+ Yield x s -> return $ Yield x (Right (Stream step1 s))
2442+ Skip s -> return $ Skip (Right (Stream step1 s))
2443+ Stop -> return Stop
23682444
23692445-------------------------------------------------------------------------------
23702446-- General transformation