-- |
-- 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 :: forall a b. (a -> b) -> Builder s m a -> Builder s m b
fmap a -> b
f (Builder s -> m (s, a)
step1) = forall s (m :: * -> *) a. (s -> m (s, a)) -> Builder s m a
Builder (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) 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 :: forall (m :: * -> *) b s. Applicative m => b -> Builder s m b
fromPure b
b = forall s (m :: * -> *) a. (s -> m (s, a)) -> Builder s m a
Builder (\s
s -> 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 :: 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 (x, a)
stepL) (Builder x -> m (x, b)
stepR) = 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
        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 :: forall a. a -> Builder a m a
pure = 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
(<*>) = forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Builder x m a -> Builder x m b -> Builder x m c
sequenceWith forall a. a -> a
id

    {-# INLINE (*>) #-}
    *> :: forall a 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 (forall a b. a -> b -> a
const 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 = forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) (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 = 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 = 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
            forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
s2, b
y)