-- |
-- 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 :: (a -> m c) -> Producer m a b
nilM a -> m c
f = (a -> m (Step a b)) -> (a -> m a) -> (a -> m a) -> Producer m a b
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)
forall s a. a -> m (Step s a)
step a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a -> m a
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 m c -> m (Step s a) -> m (Step s a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step s a -> m (Step s a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step s a
forall s a. Step s a
Stop

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

{-# INLINE unfoldrM #-}
unfoldrM :: Monad m => (a -> m (Maybe (b, a))) -> Producer m a b
unfoldrM :: (a -> m (Maybe (b, a))) -> Producer m a b
unfoldrM a -> m (Maybe (b, a))
next = (a -> m (Step a b)) -> (a -> m a) -> (a -> m a) -> Producer m a b
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 a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a -> m a
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
        Step a b -> m (Step a b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step a b -> m (Step a b)) -> Step a b -> m (Step a b)
forall a b. (a -> b) -> a -> b
$ case Maybe (b, a)
r of
            Just (b
x, a
s) -> b -> a -> Step a b
forall s a. a -> s -> Step s a
Yield b
x a
s
            Maybe (b, a)
Nothing -> Step a b
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 :: Producer m [a] a
fromList = ([a] -> m (Step [a] a))
-> ([a] -> m [a]) -> ([a] -> m [a]) -> Producer m [a] a
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] a)
forall (m :: * -> *) a. Monad m => [a] -> m (Step [a] a)
step [a] -> m [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a] -> m [a]
forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    {-# INLINE_LATE step #-}
    step :: [a] -> m (Step [a] a)
step (a
x:[a]
xs) = Step [a] a -> m (Step [a] a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step [a] a -> m (Step [a] a)) -> Step [a] a -> m (Step [a] a)
forall a b. (a -> b) -> a -> b
$ a -> [a] -> Step [a] a
forall s a. a -> s -> Step s a
Yield a
x [a]
xs
    step [] = Step [a] a -> m (Step [a] a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step [a] a
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 :: (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) =
    (s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
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 (c -> m s) -> (a -> c) -> a -> m s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> c
f) ((c -> a) -> m c -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap c -> a
g (m c -> m a) -> (s -> m c) -> s -> m a
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 :: (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) = (s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
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 (a -> m s) -> (a -> a) -> a -> m s
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 :: (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) = (s -> m (Step s c)) -> (a -> m s) -> (s -> m a) -> Producer m a c
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 = (Step s b -> Step s c) -> m (Step s b) -> m (Step s c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((b -> c) -> Step s b -> Step s c
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 :: (a -> b) -> Producer m a a -> Producer m a b
fmap = (a -> b) -> Producer m a a -> Producer m a b
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 :: 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) =
    (NestedLoop s s -> m (Step (NestedLoop s s) c))
-> (NestedLoop a b -> m (NestedLoop s s))
-> (NestedLoop s s -> m (NestedLoop a b))
-> Producer m (NestedLoop a b) c
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
        NestedLoop s s -> m (NestedLoop s s)
forall (m :: * -> *) a. Monad m => a -> m a
return (NestedLoop s s -> m (NestedLoop s s))
-> NestedLoop s s -> m (NestedLoop s s)
forall a b. (a -> b) -> a -> b
$ s -> NestedLoop s s
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
        NestedLoop s s -> m (NestedLoop s s)
forall (m :: * -> *) a. Monad m => a -> m a
return (NestedLoop s s -> m (NestedLoop s s))
-> NestedLoop s s -> m (NestedLoop s s)
forall a b. (a -> b) -> a -> b
$ s -> s -> NestedLoop s s
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
                Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c))
-> Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall a b. (a -> b) -> a -> b
$ NestedLoop s s -> Step (NestedLoop s s) c
forall s a. s -> Step s a
Skip (s -> s -> NestedLoop s s
forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
s s
innerSt)
            Skip s
s    -> Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c))
-> Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall a b. (a -> b) -> a -> b
$ NestedLoop s s -> Step (NestedLoop s s) c
forall s a. s -> Step s a
Skip (s -> NestedLoop s s
forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop s
s)
            Step s b
Stop      -> Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (NestedLoop s s) c
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
        Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c))
-> Step (NestedLoop s s) c -> m (Step (NestedLoop s s) c)
forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
s -> c -> NestedLoop s s -> Step (NestedLoop s s) c
forall s a. a -> s -> Step s a
Yield c
x (s -> s -> NestedLoop s s
forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
ost s
s)
            Skip s
s    -> NestedLoop s s -> Step (NestedLoop s s) c
forall s a. s -> Step s a
Skip (s -> s -> NestedLoop s s
forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop s
ost s
s)
            Step s c
Stop      -> NestedLoop s s -> Step (NestedLoop s s) c
forall s a. s -> Step s a
Skip (s -> NestedLoop s s
forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop s
ost)

    extract :: NestedLoop s s -> m (NestedLoop a b)
extract (OuterLoop s
s1) = a -> NestedLoop a b
forall s1 s2. s1 -> NestedLoop s1 s2
OuterLoop (a -> NestedLoop a b) -> m a -> m (NestedLoop a b)
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
        NestedLoop a b -> m (NestedLoop a b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b -> NestedLoop a b
forall s1 s2. s1 -> s2 -> NestedLoop s1 s2
InnerLoop a
r1 b
r2)