#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Stream.StreamDK
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--

module Streamly.Internal.Data.Stream.StreamDK
    (
    -- * Stream Type

      Stream
    , Step (..)

    -- * Construction
    , nil
    , cons
    , consM
    , unfoldr
    , unfoldrM
    , replicateM

    -- * Folding
    , uncons
    , foldrS

    -- * Specific Folds
    , drain
    )
where

import Streamly.Internal.Data.Stream.StreamDK.Type (Stream(..), Step(..))

-------------------------------------------------------------------------------
-- Construction
-------------------------------------------------------------------------------

nil :: Monad m => Stream m a
nil :: Stream m a
nil = m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (m (Step m a) -> Stream m a) -> m (Step m a) -> Stream m a
forall a b. (a -> b) -> a -> b
$ Step m a -> m (Step m a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step m a
forall (m :: * -> *) a. Step m a
Stop

{-# INLINE_NORMAL cons #-}
cons :: Monad m => a -> Stream m a -> Stream m a
cons :: a -> Stream m a -> Stream m a
cons a
x Stream m a
xs = m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (m (Step m a) -> Stream m a) -> m (Step m a) -> Stream m a
forall a b. (a -> b) -> a -> b
$ Step m a -> m (Step m a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step m a -> m (Step m a)) -> Step m a -> m (Step m a)
forall a b. (a -> b) -> a -> b
$ a -> Stream m a -> Step m a
forall (m :: * -> *) a. a -> Stream m a -> Step m a
Yield a
x Stream m a
xs

consM :: Monad m => m a -> Stream m a -> Stream m a
consM :: m a -> Stream m a -> Stream m a
consM m a
eff Stream m a
xs = m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (m (Step m a) -> Stream m a) -> m (Step m a) -> Stream m a
forall a b. (a -> b) -> a -> b
$ m a
eff m a -> (a -> m (Step m a)) -> m (Step m a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> Step m a -> m (Step m a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step m a -> m (Step m a)) -> Step m a -> m (Step m a)
forall a b. (a -> b) -> a -> b
$ a -> Stream m a -> Step m a
forall (m :: * -> *) a. a -> Stream m a -> Step m a
Yield a
x Stream m a
xs

unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM :: (s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM s -> m (Maybe (a, s))
next s
state = m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (s -> m (Step m a)
step' s
state)
  where
    step' :: s -> m (Step m a)
step' s
st = do
        Maybe (a, s)
r <- s -> m (Maybe (a, s))
next s
st
        Step m a -> m (Step m a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step m a -> m (Step m a)) -> Step m a -> m (Step m a)
forall a b. (a -> b) -> a -> b
$ case Maybe (a, s)
r of
            Just (a
x, s
s) -> a -> Stream m a -> Step m a
forall (m :: * -> *) a. a -> Stream m a -> Step m a
Yield a
x (m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (s -> m (Step m a)
step' s
s))
            Maybe (a, s)
Nothing     -> Step m a
forall (m :: * -> *) a. Step m a
Stop
{-
unfoldrM next s0 = buildM $ \yld stp ->
    let go s = do
            r <- next s
            case r of
                Just (a, b) -> yld a (go b)
                Nothing -> stp
    in go s0
-}

{-# INLINE unfoldr #-}
unfoldr :: Monad m => (b -> Maybe (a, b)) -> b -> Stream m a
unfoldr :: (b -> Maybe (a, b)) -> b -> Stream m a
unfoldr b -> Maybe (a, b)
next b
s0 = (forall b. (a -> b -> b) -> b -> b) -> Stream m a
forall (m :: * -> *) a.
Monad m =>
(forall b. (a -> b -> b) -> b -> b) -> Stream m a
build ((forall b. (a -> b -> b) -> b -> b) -> Stream m a)
-> (forall b. (a -> b -> b) -> b -> b) -> Stream m a
forall a b. (a -> b) -> a -> b
$ \a -> b -> b
yld b
stp ->
    let go :: b -> b
go b
s =
            case b -> Maybe (a, b)
next b
s of
                Just (a
a, b
b) -> a -> b -> b
yld a
a (b -> b
go b
b)
                Maybe (a, b)
Nothing -> b
stp
    in b -> b
go b
s0

replicateM :: Monad m => Int -> a -> Stream m a
replicateM :: Int -> a -> Stream m a
replicateM Int
n a
x = m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (Int -> m (Step m a)
forall (m :: * -> *) t.
(Monad m, Ord t, Num t) =>
t -> m (Step m a)
step Int
n)
    where
    step :: t -> m (Step m a)
step t
i = Step m a -> m (Step m a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step m a -> m (Step m a)) -> Step m a -> m (Step m a)
forall a b. (a -> b) -> a -> b
$
        if t
i t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0
        then Step m a
forall (m :: * -> *) a. Step m a
Stop
        else a -> Stream m a -> Step m a
forall (m :: * -> *) a. a -> Stream m a -> Step m a
Yield a
x (m (Step m a) -> Stream m a
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (t -> m (Step m a)
step (t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
1)))

-------------------------------------------------------------------------------
-- Folding
-------------------------------------------------------------------------------

uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
uncons :: Stream m a -> m (Maybe (a, Stream m a))
uncons (Stream m (Step m a)
step) = do
    Step m a
r <- m (Step m a)
step
    Maybe (a, Stream m a) -> m (Maybe (a, Stream m a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (a, Stream m a) -> m (Maybe (a, Stream m a)))
-> Maybe (a, Stream m a) -> m (Maybe (a, Stream m a))
forall a b. (a -> b) -> a -> b
$ case Step m a
r of
        Yield a
x Stream m a
xs -> (a, Stream m a) -> Maybe (a, Stream m a)
forall a. a -> Maybe a
Just (a
x, Stream m a
xs)
        Step m a
Stop -> Maybe (a, Stream m a)
forall a. Maybe a
Nothing

-- | Lazy right associative fold to a stream.
{-# INLINE_NORMAL foldrS #-}
foldrS :: Monad m
       => (a -> Stream m b -> Stream m b)
       -> Stream m b
       -> Stream m a
       -> Stream m b
foldrS :: (a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrS a -> Stream m b -> Stream m b
f Stream m b
streamb = Stream m a -> Stream m b
go
    where
    go :: Stream m a -> Stream m b
go (Stream m (Step m a)
stepa) = m (Step m b) -> Stream m b
forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (m (Step m b) -> Stream m b) -> m (Step m b) -> Stream m b
forall a b. (a -> b) -> a -> b
$ do
        Step m a
r <- m (Step m a)
stepa
        case Step m a
r of
            Yield a
x Stream m a
xs -> let Stream m (Step m b)
step = a -> Stream m b -> Stream m b
f a
x (Stream m a -> Stream m b
go Stream m a
xs) in m (Step m b)
step
            Step m a
Stop -> let Stream m (Step m b)
step = Stream m b
streamb in m (Step m b)
step

{-# INLINE_LATE foldrM #-}
foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM :: (a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM a -> m b -> m b
fstep m b
acc = Stream m a -> m b
go
    where
    go :: Stream m a -> m b
go (Stream m (Step m a)
step) = do
        Step m a
r <- m (Step m a)
step
        case Step m a
r of
            Yield a
x Stream m a
xs -> a -> m b -> m b
fstep a
x (Stream m a -> m b
go Stream m a
xs)
            Step m a
Stop -> m b
acc

{-# INLINE_NORMAL build #-}
build :: Monad m
    => forall a. (forall b. (a -> b -> b) -> b -> b) -> Stream m a
build :: forall a. (forall b. (a -> b -> b) -> b -> b) -> Stream m a
build forall b. (a -> b -> b) -> b -> b
g = (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
forall b. (a -> b -> b) -> b -> b
g a -> Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
cons Stream m a
forall (m :: * -> *) a. Monad m => Stream m a
nil

{-# RULES
"foldrM/build"  forall k z (g :: forall b. (a -> b -> b) -> b -> b).
                foldrM k z (build g) = g k z #-}

{-
-- To fuse foldrM with unfoldrM we need the type m1 to be polymorphic such that
-- it is either Monad m or Stream m.  So that we can use cons/nil as well as
-- monadic construction function as its arguments.
--
{-# INLINE_NORMAL buildM #-}
buildM :: Monad m
    => forall a. (forall b. (a -> m1 b -> m1 b) -> m1 b -> m1 b) -> Stream m a
buildM g = g cons nil
-}

-------------------------------------------------------------------------------
-- Specific folds
-------------------------------------------------------------------------------

{-# INLINE drain #-}
drain :: Monad m => Stream m a -> m ()
drain :: Stream m a -> m ()
drain = (a -> m () -> m ()) -> m () -> Stream m a -> m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
_ m ()
xs -> m ()
xs) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
{-
drain (Stream step) = do
    r <- step
    case r of
        Yield _ next -> drain next
        Stop      -> return ()
        -}