-- |
-- Module      : Streamly.Internal.Data.Producer.Type
-- Copyright   : (c) 2021 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- See "Streamly.Internal.Data.Producer" for introduction.
--

module Streamly.Internal.Data.Producer.Type
    (
    -- * Type
    Producer (..)

    -- * Producers
    , nil
    , nilM
    , unfoldrM
    , fromList

    -- * Mapping
    , translate
    , lmap

    -- * Nesting
    , NestedLoop (..)
    , concat
    )
where

#include "inline.hs"

import Fusion.Plugin.Types (Fuse(..))
import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
import Prelude hiding (concat, map)

------------------------------------------------------------------------------
-- Type
------------------------------------------------------------------------------

-- | A @Producer m a b@ is a generator of a stream of values of type @b@ from a
-- seed of type 'a' in 'Monad' @m@.
--
-- /Pre-release/

data Producer m a b =
    -- | @Producer step inject extract@
    forall s. Producer (s -> m (Step s b)) (a -> m s) (s -> m a)

------------------------------------------------------------------------------
-- Producers
------------------------------------------------------------------------------

{-# INLINE nilM #-}
nilM :: Monad m => (a -> m c) -> Producer m a b
nilM :: forall (m :: * -> *) a c b. Monad m => (a -> m c) -> Producer m a b
nilM a -> m c
f = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer forall {s} {a}. a -> m (Step s a)
step forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    {-# INLINE_LATE step #-}
    step :: a -> m (Step s a)
step a
x = a -> m c
f a
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE nil #-}
nil :: Monad m => Producer m a b
nil :: forall (m :: * -> *) a b. Monad m => Producer m a b
nil = forall (m :: * -> *) a c b. Monad m => (a -> m c) -> Producer m a b
nilM (\a
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ())

{-# INLINE unfoldrM #-}
unfoldrM :: Monad m => (a -> m (Maybe (b, a))) -> Producer m a b
unfoldrM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m (Maybe (b, a))) -> Producer m a b
unfoldrM a -> m (Maybe (b, a))
next = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer a -> m (Step a b)
step forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    {-# INLINE_LATE step #-}
    step :: a -> m (Step a b)
step a
st = do
        Maybe (b, a)
r <- a -> m (Maybe (b, a))
next a
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (b, a)
r of
            Just (b
x, a
s) -> forall s a. a -> s -> Step s a
Yield b
x a
s
            Maybe (b, a)
Nothing -> forall s a. Step s a
Stop

-- | Convert a list of pure values to a 'Stream'
--
-- /Pre-release/
{-# INLINE_LATE fromList #-}
fromList :: Monad m => Producer m [a] a
fromList :: forall (m :: * -> *) a. Monad m => Producer m [a] a
fromList = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer forall {m :: * -> *} {a}. Monad m => [a] -> m (Step [a] a)
step forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    {-# INLINE_LATE step #-}
    step :: [a] -> m (Step [a] a)
step (a
x:[a]
xs) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x [a]
xs
    step [] = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Mapping
------------------------------------------------------------------------------

-- | Interconvert the producer between two interconvertible input types.
--
-- /Pre-release/
{-# INLINE_NORMAL translate #-}
translate :: Functor m =>
    (a -> c) -> (c -> a) -> Producer m c b -> Producer m a b
translate :: forall (m :: * -> *) a c b.
Functor m =>
(a -> c) -> (c -> a) -> Producer m c b -> Producer m a b
translate a -> c
f c -> a
g (Producer s -> m (Step s b)
step c -> m s
inject s -> m c
extract) =
    forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer s -> m (Step s b)
step (c -> m s
inject forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> c
f) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap c -> a
g forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> m c
extract)

-- | Map the producer input to another value of the same type.
--
-- /Pre-release/
{-# INLINE_NORMAL lmap #-}
lmap :: (a -> a) -> Producer m a b -> Producer m a b
lmap :: forall a (m :: * -> *) b.
(a -> a) -> Producer m a b -> Producer m a b
lmap a -> a
f (Producer s -> m (Step s b)
step a -> m s
inject s -> m a
extract) = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer s -> m (Step s b)
step (a -> m s
inject forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f) s -> m a
extract

------------------------------------------------------------------------------
-- Functor
------------------------------------------------------------------------------

-- | Map a function on the output of the producer (the type @b@).
--
-- /Pre-release/
{-# INLINE_NORMAL map #-}
map :: Functor m => (b -> c) -> Producer m a b -> Producer m a c
map :: forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Producer m a b -> Producer m a c
map b -> c
f (Producer s -> m (Step s b)
ustep a -> m s
uinject s -> m a
uextract) = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer s -> m (Step s c)
step a -> m s
uinject s -> m a
uextract

    where

    {-# INLINE_LATE step #-}
    step :: s -> m (Step s c)
step s
st = 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 b -> c
f) (s -> m (Step s b)
ustep s
st)

-- | Maps a function on the output of the producer (the type @b@).
instance Functor m => Functor (Producer m a) where
    {-# INLINE fmap #-}
    fmap :: forall a b. (a -> b) -> Producer m a a -> Producer m a b
fmap = forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Producer m a b -> Producer m a c
map

------------------------------------------------------------------------------
-- Nesting
------------------------------------------------------------------------------

-- | State representing a nested loop.
{-# ANN type NestedLoop Fuse #-}
data NestedLoop s1 s2 = OuterLoop s1 | InnerLoop s1 s2

-- | Apply the second unfold to each output element of the first unfold and
-- flatten the output in a single stream.
--
-- /Pre-release/
--
{-# INLINE_NORMAL concat #-}
concat :: Monad m =>
    Producer m a b -> Producer m b c -> Producer m (NestedLoop a b) c
concat :: forall (m :: * -> *) a b c.
Monad m =>
Producer m a b -> Producer m b c -> Producer m (NestedLoop a b) c
concat (Producer s -> m (Step s b)
step1 a -> m s
inject1 s -> m a
extract1) (Producer s -> m (Step s c)
step2 b -> m s
inject2 s -> m b
extract2) =
    forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer NestedLoop s s -> m (Step (NestedLoop s s) c)
step NestedLoop a b -> m (NestedLoop s s)
inject NestedLoop s s -> m (NestedLoop a b)
extract

    where

    inject :: NestedLoop a b -> m (NestedLoop s s)
inject (OuterLoop a
x) = do
        s
s <- a -> m s
inject1 a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop s
s
    inject (InnerLoop a
x b
y) = do
        s
s1 <- a -> m s
inject1 a
x
        s
s2 <- b -> m s
inject2 b
y
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
s1 s
s2

    {-# INLINE_LATE step #-}
    step :: NestedLoop s s -> m (Step (NestedLoop s s) c)
step (OuterLoop s
st) = do
        Step s b
r <- s -> m (Step s b)
step1 s
st
        case Step s b
r of
            Yield b
x s
s -> do
                s
innerSt <- b -> m s
inject2 b
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
s s
innerSt)
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop s
s)
            Step s b
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step (InnerLoop s
ost s
ist) = do
        Step s c
r <- s -> m (Step s c)
step2 s
ist
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
s -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
ost s
s)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
ost s
s)
            Step s c
Stop      -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop s
ost)

    extract :: NestedLoop s s -> m (NestedLoop a b)
extract (OuterLoop s
s1) = forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m a
extract1 s
s1
    extract (InnerLoop s
s1 s
s2) = do
        a
r1 <- s -> m a
extract1 s
s1
        b
r2 <- s -> m b
extract2 s
s2
        forall (m :: * -> *) a. Monad m => a -> m a
return (forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop a
r1 b
r2)