@@ -417,13 +417,13 @@ concatMapM f (Unfold step1 inject1) = Unfold step inject
417417--
418418-- /Internal/
419419{-# INLINE_NORMAL before #-}
420−before :: Monad m => m c -> Unfold m a b -> Unfold m a b
420+before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
421421before action (Unfold step1 inject1) = Unfold step inject
422422
423423 where
424424
425425 inject x = do
426− _ <- action
426+ _ <- action x
427427 st <- inject1 x
428428 return st
429429
@@ -439,34 +439,42 @@ before action (Unfold step1 inject1) = Unfold step inject
439439--
440440-- /Internal/
441441{-# 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
442+after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
443+after action (Unfold step1 inject1) = Unfold step inject
444444
445445 where
446446
447+ inject x = do
448+ s <- inject1 x
449+ return (s, x)
450+
447451 {-# INLINE_LATE step #-}
448− step st = do
452+ step (st, v) = do
449453 res <- step1 st
450454 case res of
451− Yield x s -> return $ Yield x s
452− Skip s -> return $ Skip s
453− Stop -> action >> return Stop
455+ Yield x s -> return $ Yield x (s, v)
456+ Skip s -> return $ Skip (s, v)
457+ Stop -> action v >> return Stop
454458
455459-- | Run a side effect whenever the unfold aborts due to an exception.
456460--
457461-- /Internal/
458462{-# 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
463+onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
464+onException action (Unfold step1 inject1) = Unfold step inject
461465
462466 where
463467
468+ inject x = do
469+ s <- inject1 x
470+ return (s, x)
471+
464472 {-# INLINE_LATE step #-}
465− step st = do
466− res <- step1 st `MC.onException` action
473+ step (st, v) = do
474+ res <- step1 st `MC.onException` action v
467475 case res of
468− Yield x s -> return $ Yield x s
469− Skip s -> return $ Skip s
476+ Yield x s -> return $ Yield x (s, v)
477+ Skip s -> return $ Skip (s, v)
470478 Stop -> return Stop
471479
472480-- | Run a side effect whenever the unfold stops normally or aborts due to an
@@ -474,46 +482,47 @@ onException action (Unfold step1 inject1) = Unfold step inject1
474482--
475483-- /Internal/
476484{-# 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
485+finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
486+finally action (Unfold step1 inject1) = Unfold step inject
479487
480488 where
481489
490+ inject x = do
491+ s <- inject1 x
492+ return (s, x)
493+
482494 {-# INLINE_LATE step #-}
483− step st = do
484− res <- step1 st `MC.onException` action
495+ step (st, v) = do
496+ res <- step1 st `MC.onException` action v
485497 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.
498+ Yield x s -> return $ Yield x (s, v)
499+ Skip s -> return $ Skip (s, v)
500+ Stop -> action v >> return Stop
501+
502+-- | @bracket before after between@ runs the @before@ action and then unfolds
503+-- its output using the @between@ unfold. When the @between@ unfold is done or
504+-- if an exception occurs then the @after@ action is run with the output of
505+-- @before@ as argument.
494506--
495507-- /Internal/
496508{-# INLINE_NORMAL bracket #-}
497509bracket :: 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
510+ => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
511+bracket bef aft (Unfold step1 inject1) = Unfold step inject
500512
501513 where
502514
503515 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)
516+ r <- bef x
517+ s <- inject1 r
518+ return (s, r)
510519
511520 {-# INLINE_LATE step #-}
512− step (UnStream step2 state, v) = do
513− res <- step2 defState state `MC.onException` aft v
521+ step (st, v) = do
522+ res <- step1 st `MC.onException` aft v
514523 case res of
515− Yield x s -> return $ Yield x (Stream step2 s, v)
516− Skip s -> return $ Skip (Stream step2 s, v)
524+ Yield x s -> return $ Yield x (s, v)
525+ Skip s -> return $ Skip (s, v)
517526 Stop -> aft v >> return Stop
518527
519528-- | When unfolding an unfold if an exception occurs, unfold aborts