-- |
-- Module      : Streamly.Internal.Data.Builder
-- Copyright   : (c) 2022 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.Builder
    (
    -- * Imports
    -- $setup

    -- * Types
      Builder (..)
    )
where

#if !MIN_VERSION_base(4,18,0)
import Control.Applicative (liftA2)
#endif

------------------------------------------------------------------------------
-- The Builder type
------------------------------------------------------------------------------

-- | A simple stateful function composing monad that chains state passing
-- functions. This can be considered as a simplified version of the State monad
-- or even a Fold. Unlike fold the step function is one-shot and not called in
-- a loop.
newtype Builder s m a =
  Builder (s -> m (a, s))

-- | Maps a function on the output of the fold (the type @b@).
instance Functor m => Functor (Builder s m) where
    {-# INLINE fmap #-}
    fmap :: forall a b. (a -> b) -> Builder s m a -> Builder s m b
fmap a -> b
f (Builder s -> m (a, s)
step1) = (s -> m (b, s)) -> Builder s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> Builder s m a
Builder (((a, s) -> (b, s)) -> m (a, s) -> m (b, s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (a
a, s
s) -> (a -> b
f a
a, s
s)) (m (a, s) -> m (b, s)) -> (s -> m (a, s)) -> s -> m (b, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> m (a, s)
step1)

{-# INLINE fromPure #-}
fromPure :: Applicative m => b -> Builder s m b
fromPure :: forall (m :: * -> *) b s. Applicative m => b -> Builder s m b
fromPure b
b = (s -> m (b, s)) -> Builder s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> Builder s m a
Builder (\s
s -> (b, s) -> m (b, s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
b, s
s))

-- | Chain the actions and zip the outputs.
{-# INLINE sequenceWith #-}
sequenceWith :: Monad m =>
    (a -> b -> c) -> Builder x m a -> Builder x m b -> Builder x m c
sequenceWith :: forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Builder x m a -> Builder x m b -> Builder x m c
sequenceWith a -> b -> c
func (Builder x -> m (a, x)
stepL) (Builder x -> m (b, x)
stepR) = (x -> m (c, x)) -> Builder x m c
forall s (m :: * -> *) a. (s -> m (a, s)) -> Builder s m a
Builder x -> m (c, x)
step

    where

    step :: x -> m (c, x)
step x
s = do
        (a
x, x
s1) <- x -> m (a, x)
stepL x
s
        (b
y, x
s2) <- x -> m (b, x)
stepR x
s1
        (c, x) -> m (c, x)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> b -> c
func a
x b
y, x
s2)

instance Monad m => Applicative (Builder a m) where
    {-# INLINE pure #-}
    pure :: forall a. a -> Builder a m a
pure = a -> Builder a m a
forall (m :: * -> *) b s. Applicative m => b -> Builder s m b
fromPure

    {-# INLINE (<*>) #-}
    <*> :: forall a b. Builder a m (a -> b) -> Builder a m a -> Builder a m b
(<*>) = ((a -> b) -> a -> b)
-> Builder a m (a -> b) -> Builder a m a -> Builder a m b
forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Builder x m a -> Builder x m b -> Builder x m c
sequenceWith (a -> b) -> a -> b
forall a. a -> a
id

    {-# INLINE (*>) #-}
    *> :: forall a b. Builder a m a -> Builder a m b -> Builder a m b
(*>) = (a -> b -> b) -> Builder a m a -> Builder a m b -> Builder a m b
forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Builder x m a -> Builder x m b -> Builder x m c
sequenceWith ((b -> b) -> a -> b -> b
forall a b. a -> b -> a
const b -> b
forall a. a -> a
id)

    {-# INLINE liftA2 #-}
    liftA2 :: forall a b c.
(a -> b -> c) -> Builder a m a -> Builder a m b -> Builder a m c
liftA2 a -> b -> c
f Builder a m a
x = Builder a m (b -> c) -> Builder a m b -> Builder a m c
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) ((a -> b -> c) -> Builder a m a -> Builder a m (b -> c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b -> c
f Builder a m a
x)

instance Monad m => Monad (Builder a m) where
    {-# INLINE return #-}
    return :: forall a. a -> Builder a m a
return = a -> Builder a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    {-# INLINE (>>=) #-}
    (Builder a -> m (a, a)
stepL) >>= :: forall a b. Builder a m a -> (a -> Builder a m b) -> Builder a m b
>>= a -> Builder a m b
f = (a -> m (b, a)) -> Builder a m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> Builder s m a
Builder a -> m (b, a)
step

        where

        step :: a -> m (b, a)
step a
s = do
            (a
x, a
s1) <- a -> m (a, a)
stepL a
s
            let Builder a -> m (b, a)
stepR = a -> Builder a m b
f a
x
            (b
y, a
s2) <- a -> m (b, a)
stepR a
s1
            (b, a) -> m (b, a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
y, a
s2)