@@ -131,6 +131,7 @@ import Control.Monad.IO.Class (MonadIO(..))
131131import Control.Monad.Trans.Class (MonadTrans(lift))
132132import Data.IORef (newIORef, mkWeakIORef)
133133import Data.Maybe (fromJust, isJust)
134+import Fusion.Plugin.Types (Fuse(..))
134135import GHC.Types (SPEC(..))
135136import qualified Control.Monad.Catch as MC
136137
@@ -578,6 +579,12 @@ postscanlM' fstep begin (Stream step state) =
578579postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
579580postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
580581
582+{-# ANN type PScanAfterState Fuse #-}
583+data PScanAfterState m st acc =
584+ PScanAfterStep st (m acc)
585+ | PScanAfterYield acc (PScanAfterState m st acc)
586+ | PScanAfterStop
587+
581588-- We can possibly have the "done" function as a Maybe to provide an option to
582589-- emit or not emit the accumulator when the stream stops.
583590--
@@ -587,21 +594,26 @@ postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
587594postscanlMAfter' :: Monad m
588595 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
589596postscanlMAfter' fstep initial done (Stream step1 state1) = do
590− Stream step (Just (state1, initial))
597+ Stream step (PScanAfterStep state1 initial)
591598
592599 where
593600
594601 {-# INLINE_LATE step #-}
595− step gst (Just (st, acc)) = do
602+ step gst (PScanAfterStep st acc) = do
596603 r <- step1 (adaptState gst) st
597604 case r of
598605 Yield x s -> do
599− old <- acc
600− y <- fstep old x
601− y `seq` return (Yield y (Just (s, return y)))
602− Skip s -> return $ Skip $ Just (s, acc)
603− Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
604− step _ Nothing = return Stop
606+ !old <- acc
607+ !y <- fstep old x
608+ return (Skip $ PScanAfterYield y (PScanAfterStep s (return y)))
609+ Skip s -> return $ Skip $ PScanAfterStep s acc
610+ -- Strictness is important for fusion
611+ Stop -> do
612+ !v <- acc
613+ !res <- done v
614+ return (Skip $ PScanAfterYield res PScanAfterStop)
615+ step _ (PScanAfterYield acc next) = return $ Yield acc next
616+ step _ PScanAfterStop = return Stop
605617
606618{-# INLINE_NORMAL postscanlM #-}
607619postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
@@ -646,8 +658,7 @@ scanlM' fstep begin (Stream step state) = Stream step' Nothing
646658scanlMAfter' :: Monad m
647659 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
648660scanlMAfter' fstep initial done s =
649− (initial >>= \x -> x `seq` return x) `consM`
650− postscanlMAfter' fstep initial done s
661+ initial `consM` postscanlMAfter' fstep initial done s
651662
652663{-# INLINE scanl' #-}
653664scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b