-- |
-- Module      : Streamly.Internal.Data.Fold.Step
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.Fold.Step
    (
    -- * Types
      Step (..)

    , mapMStep
    , chainStepM
    )
where

import Data.Bifunctor (Bifunctor(..))
import Fusion.Plugin.Types (Fuse(..))

------------------------------------------------------------------------------
-- Step of a fold
------------------------------------------------------------------------------

-- The Step functor around b allows expressing early termination like a right
-- fold. Traditional list right folds use function composition and laziness to
-- terminate early whereas we use data constructors. It allows stream fusion in
-- contrast to the foldr/build fusion when composing with functions.

-- | Represents the result of the @step@ of a 'Fold'.  'Partial' returns an
-- intermediate state of the fold, the fold step can be called again with the
-- state or the driver can use @extract@ on the state to get the result out.
-- 'Done' returns the final result and the fold cannot be driven further.
--
-- /Pre-release/
--
{-# ANN type Step Fuse #-}
data Step s b
    = Partial !s
    | Done !b

-- | 'first' maps over 'Partial' and 'second' maps over 'Done'.
--
instance Bifunctor Step where
    {-# INLINE bimap #-}
    bimap :: forall a b c d. (a -> b) -> (c -> d) -> Step a c -> Step b d
bimap a -> b
f c -> d
_ (Partial a
a) = forall s b. s -> Step s b
Partial (a -> b
f a
a)
    bimap a -> b
_ c -> d
g (Done c
b) = forall s b. b -> Step s b
Done (c -> d
g c
b)

    {-# INLINE first #-}
    first :: forall a b c. (a -> b) -> Step a c -> Step b c
first a -> b
f (Partial a
a) = forall s b. s -> Step s b
Partial (a -> b
f a
a)
    first a -> b
_ (Done c
x) = forall s b. b -> Step s b
Done c
x

    {-# INLINE second #-}
    second :: forall b c a. (b -> c) -> Step a b -> Step a c
second b -> c
_ (Partial a
x) = forall s b. s -> Step s b
Partial a
x
    second b -> c
f (Done b
a) = forall s b. b -> Step s b
Done (b -> c
f b
a)

-- | 'fmap' maps over 'Done'.
--
-- @
-- fmap = 'second'
-- @
--
instance Functor (Step s) where
    {-# INLINE fmap #-}
    fmap :: forall a b. (a -> b) -> Step s a -> Step s b
fmap = forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second

-- | Map a monadic function over the result @b@ in @Step s b@.
--
-- /Internal/
{-# INLINE mapMStep #-}
mapMStep :: Applicative m => (a -> m b) -> Step s a -> m (Step s b)
mapMStep :: forall (m :: * -> *) a b s.
Applicative m =>
(a -> m b) -> Step s a -> m (Step s b)
mapMStep a -> m b
f Step s a
res =
    case Step s a
res of
        Partial s
s -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
Partial s
s
        Done a
b -> forall s b. b -> Step s b
Done forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m b
f a
b

-- | If 'Partial' then map the state, if 'Done' then call the next step.
{-# INLINE chainStepM #-}
chainStepM :: Applicative m =>
    (s1 -> m s2) -> (a -> m (Step s2 b)) -> Step s1 a -> m (Step s2 b)
chainStepM :: forall (m :: * -> *) s1 s2 a b.
Applicative m =>
(s1 -> m s2) -> (a -> m (Step s2 b)) -> Step s1 a -> m (Step s2 b)
chainStepM s1 -> m s2
f a -> m (Step s2 b)
_ (Partial s1
s) = forall s b. s -> Step s b
Partial forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s1 -> m s2
f s1
s
chainStepM s1 -> m s2
_ a -> m (Step s2 b)
g (Done a
b) = a -> m (Step s2 b)
g a
b