-- |
-- Module      : Streamly.Internal.Data.StreamK.Transformer
-- Copyright   : (c) 2017 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.StreamK.Transformer
    (
      foldlT
    , foldrT

    , liftInner
    , evalStateT
    )
where

import Control.Monad.Trans.Class (MonadTrans(lift))
import Control.Monad.Trans.State.Strict (StateT)
import Streamly.Internal.Data.StreamK.Type
    (StreamK, nil, cons, uncons, concatEffect)

import qualified Control.Monad.Trans.State.Strict as State

-- | Lazy left fold to an arbitrary transformer monad.
{-# INLINE foldlT #-}
foldlT :: (Monad m, Monad (s m), MonadTrans s)
    => (s m b -> a -> s m b) -> s m b -> StreamK m a -> s m b
foldlT :: forall (m :: * -> *) (s :: (* -> *) -> * -> *) b a.
(Monad m, Monad (s m), MonadTrans s) =>
(s m b -> a -> s m b) -> s m b -> StreamK m a -> s m b
foldlT s m b -> a -> s m b
step = s m b -> StreamK m a -> s m b
go
  where
    go :: s m b -> StreamK m a -> s m b
go s m b
acc StreamK m a
m1 = do
        Maybe (a, StreamK m a)
res <- m (Maybe (a, StreamK m a)) -> s m (Maybe (a, StreamK m a))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (Maybe (a, StreamK m a)) -> s m (Maybe (a, StreamK m a)))
-> m (Maybe (a, StreamK m a)) -> s m (Maybe (a, StreamK m a))
forall a b. (a -> b) -> a -> b
$ StreamK m a -> m (Maybe (a, StreamK m a))
forall (m :: * -> *) a.
Applicative m =>
StreamK m a -> m (Maybe (a, StreamK m a))
uncons StreamK m a
m1
        case Maybe (a, StreamK m a)
res of
            Just (a
h, StreamK m a
t) -> s m b -> StreamK m a -> s m b
go (s m b -> a -> s m b
step s m b
acc a
h) StreamK m a
t
            Maybe (a, StreamK m a)
Nothing -> s m b
acc

-- | Right associative fold to an arbitrary transformer monad.
{-# INLINE foldrT #-}
foldrT :: (Monad m, Monad (s m), MonadTrans s)
    => (a -> s m b -> s m b) -> s m b -> StreamK m a -> s m b
foldrT :: forall (m :: * -> *) (s :: (* -> *) -> * -> *) a b.
(Monad m, Monad (s m), MonadTrans s) =>
(a -> s m b -> s m b) -> s m b -> StreamK m a -> s m b
foldrT a -> s m b -> s m b
step s m b
final = StreamK m a -> s m b
go
  where
    go :: StreamK m a -> s m b
go StreamK m a
m1 = do
        Maybe (a, StreamK m a)
res <- m (Maybe (a, StreamK m a)) -> s m (Maybe (a, StreamK m a))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (Maybe (a, StreamK m a)) -> s m (Maybe (a, StreamK m a)))
-> m (Maybe (a, StreamK m a)) -> s m (Maybe (a, StreamK m a))
forall a b. (a -> b) -> a -> b
$ StreamK m a -> m (Maybe (a, StreamK m a))
forall (m :: * -> *) a.
Applicative m =>
StreamK m a -> m (Maybe (a, StreamK m a))
uncons StreamK m a
m1
        case Maybe (a, StreamK m a)
res of
            Just (a
h, StreamK m a
t) -> a -> s m b -> s m b
step a
h (StreamK m a -> s m b
go StreamK m a
t)
            Maybe (a, StreamK m a)
Nothing -> s m b
final

------------------------------------------------------------------------------
-- Lifting inner monad
------------------------------------------------------------------------------

{-# INLINE evalStateT #-}
evalStateT :: Monad m => m s -> StreamK (StateT s m) a -> StreamK m a
evalStateT :: forall (m :: * -> *) s a.
Monad m =>
m s -> StreamK (StateT s m) a -> StreamK m a
evalStateT = m s -> StreamK (StateT s m) a -> StreamK m a
forall (m :: * -> *) s a.
Monad m =>
m s -> StreamK (StateT s m) a -> StreamK m a
go

    where

    go :: m a -> StreamK (StateT a m) a -> StreamK m a
go m a
st StreamK (StateT a m) a
m1 = m (StreamK m a) -> StreamK m a
forall (m :: * -> *) a. Monad m => m (StreamK m a) -> StreamK m a
concatEffect (m (StreamK m a) -> StreamK m a) -> m (StreamK m a) -> StreamK m a
forall a b. (a -> b) -> a -> b
$ ((Maybe (a, StreamK (StateT a m) a), a) -> StreamK m a)
-> m (Maybe (a, StreamK (StateT a m) a), a) -> m (StreamK m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe (a, StreamK (StateT a m) a), a) -> StreamK m a
f (m a
st m a
-> (a -> m (Maybe (a, StreamK (StateT a m) a), a))
-> m (Maybe (a, StreamK (StateT a m) a), a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= StateT a m (Maybe (a, StreamK (StateT a m) a))
-> a -> m (Maybe (a, StreamK (StateT a m) a), a)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT (StreamK (StateT a m) a
-> StateT a m (Maybe (a, StreamK (StateT a m) a))
forall (m :: * -> *) a.
Applicative m =>
StreamK m a -> m (Maybe (a, StreamK m a))
uncons StreamK (StateT a m) a
m1))

    f :: (Maybe (a, StreamK (StateT a m) a), a) -> StreamK m a
f (Maybe (a, StreamK (StateT a m) a)
res, a
s1) =
        case Maybe (a, StreamK (StateT a m) a)
res of
            Just (a
h, StreamK (StateT a m) a
t) -> a -> StreamK m a -> StreamK m a
forall a (m :: * -> *). a -> StreamK m a -> StreamK m a
cons a
h (m a -> StreamK (StateT a m) a -> StreamK m a
go (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
s1) StreamK (StateT a m) a
t)
            Maybe (a, StreamK (StateT a m) a)
Nothing -> StreamK m a
forall (m :: * -> *) a. StreamK m a
nil

{-# INLINE liftInner #-}
liftInner :: (Monad m, MonadTrans t, Monad (t m)) =>
    StreamK m a -> StreamK (t m) a
liftInner :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, MonadTrans t, Monad (t m)) =>
StreamK m a -> StreamK (t m) a
liftInner = StreamK m a -> StreamK (t m) a
forall {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadTrans t, Monad m, Monad (t m)) =>
StreamK m a -> StreamK (t m) a
go

    where

    go :: StreamK m a -> StreamK (t m) a
go StreamK m a
m1 = t m (StreamK (t m) a) -> StreamK (t m) a
forall (m :: * -> *) a. Monad m => m (StreamK m a) -> StreamK m a
concatEffect (t m (StreamK (t m) a) -> StreamK (t m) a)
-> t m (StreamK (t m) a) -> StreamK (t m) a
forall a b. (a -> b) -> a -> b
$ (Maybe (a, StreamK m a) -> StreamK (t m) a)
-> t m (Maybe (a, StreamK m a)) -> t m (StreamK (t m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe (a, StreamK m a) -> StreamK (t m) a
f (t m (Maybe (a, StreamK m a)) -> t m (StreamK (t m) a))
-> t m (Maybe (a, StreamK m a)) -> t m (StreamK (t m) a)
forall a b. (a -> b) -> a -> b
$ m (Maybe (a, StreamK m a)) -> t m (Maybe (a, StreamK m a))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (Maybe (a, StreamK m a)) -> t m (Maybe (a, StreamK m a)))
-> m (Maybe (a, StreamK m a)) -> t m (Maybe (a, StreamK m a))
forall a b. (a -> b) -> a -> b
$ StreamK m a -> m (Maybe (a, StreamK m a))
forall (m :: * -> *) a.
Applicative m =>
StreamK m a -> m (Maybe (a, StreamK m a))
uncons StreamK m a
m1

    f :: Maybe (a, StreamK m a) -> StreamK (t m) a
f Maybe (a, StreamK m a)
res =
        case Maybe (a, StreamK m a)
res of
            Just (a
h, StreamK m a
t) -> a -> StreamK (t m) a -> StreamK (t m) a
forall a (m :: * -> *). a -> StreamK m a -> StreamK m a
cons a
h (StreamK m a -> StreamK (t m) a
go StreamK m a
t)
            Maybe (a, StreamK m a)
Nothing -> StreamK (t m) a
forall (m :: * -> *) a. StreamK m a
nil