d99fb19 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Use a single yield point in scanlMAfter'

d99fb19Harendra Kumar5 years ago

core/src/Streamly/Internal/Data/Stream/StreamD/Transform.hs+21−10

compare with master

unifiedsplit

@@ -131,6 +131,7 @@ import Control.Monad.IO.Class (MonadIO(..))

131import Control.Monad.Trans.Class (MonadTrans(lift))

132import Data.IORef (newIORef, mkWeakIORef)

133import Data.Maybe (fromJust, isJust)

134+import Fusion.Plugin.Types (Fuse(..))

134import GHC.Types (SPEC(..))

135import GHC.Types (SPEC(..))

135import qualified Control.Monad.Catch as MC

136import qualified Control.Monad.Catch as MC

136

137

@@ -578,6 +579,12 @@ postscanlM' fstep begin (Stream step state) =

578postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

579postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

579postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)

580postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)

580

581

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+

581-- We can possibly have the "done" function as a Maybe to provide an option to

588-- We can possibly have the "done" function as a Maybe to provide an option to

582-- emit or not emit the accumulator when the stream stops.

589-- emit or not emit the accumulator when the stream stops.

583--

590--

@@ -587,21 +594,26 @@ postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)

587postscanlMAfter' :: Monad m

594postscanlMAfter' :: Monad m

588 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

595 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

589postscanlMAfter' fstep initial done (Stream step1 state1) = do

596postscanlMAfter' fstep initial done (Stream step1 state1) = do

590− Stream step (Just (state1, initial))

597+ Stream step (PScanAfterStep state1 initial)

591

598

592 where

599 where

593

600

594 {-# INLINE_LATE step #-}

601 {-# INLINE_LATE step #-}

595− step gst (Just (st, acc)) = do

602+ step gst (PScanAfterStep st acc) = do

596 r <- step1 (adaptState gst) st

603 r <- step1 (adaptState gst) st

597 case r of

604 case r of

598 Yield x s -> do

605 Yield x s -> do

599− old <- acc

606+ !old <- acc

600− y <- fstep old x

607+ !y <- fstep old x

601− y `seq` return (Yield y (Just (s, return y)))

608+ return (Skip $ PScanAfterYield y (PScanAfterStep s (return y)))

602− Skip s -> return $ Skip $ Just (s, acc)

609+ Skip s -> return $ Skip $ PScanAfterStep s acc

603− Stop -> acc >>= done >>= \res -> return (Yield res Nothing)

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

604− step _ Nothing = return Stop

616+ step _ PScanAfterStop = return Stop

605

617

606{-# INLINE_NORMAL postscanlM #-}

618{-# INLINE_NORMAL postscanlM #-}

607postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

619postscanlM :: 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

646scanlMAfter' :: Monad m

658scanlMAfter' :: Monad m

647 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

659 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

648scanlMAfter' fstep initial done s =

660scanlMAfter' 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

651

662

652{-# INLINE scanl' #-}

663{-# INLINE scanl' #-}

653scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

664scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b