-- |
-- 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

import Control.Applicative (liftA2)

------------------------------------------------------------------------------
-- 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 (s, a))

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

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

-- | 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 :: (a -> b -> c) -> Builder x m a -> Builder x m b -> Builder x m c
sequenceWith a -> b -> c
func (Builder x -> m (x, a)
stepL) (Builder x -> m (x, b)
stepR) = (x -> m (x, c)) -> Builder x m c
forall s (m :: * -> *) a. (s -> m (s, a)) -> Builder s m a
Builder x -> m (x, c)
step

    where

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

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

    {-# INLINE (<*>) #-}
    <*> :: 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 (*>) #-}
    *> :: 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 :: (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 :: 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) >>= :: Builder a m a -> (a -> Builder a m b) -> Builder a m b
>>= a -> Builder a m b
f = (a -> m (a, b)) -> Builder a m b
forall s (m :: * -> *) a. (s -> m (s, a)) -> Builder s m a
Builder a -> m (a, b)
step

        where

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