#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 :: forall (m :: * -> *) a. Monad m => Stream m a
nil = forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a. Step m a
Stop

{-# INLINE_NORMAL cons #-}
cons :: Monad m => a -> Stream m a -> Stream m a
cons :: forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
cons a
x Stream m a
xs = forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ 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 :: forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
consM m a
eff Stream m a
xs = forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream forall a b. (a -> b) -> a -> b
$ m a
eff forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ 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 :: forall (m :: * -> *) s a.
Monad m =>
(s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM s -> m (Maybe (a, s))
next s
state = 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
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (a, s)
r of
            Just (a
x, s
s) -> forall (m :: * -> *) a. a -> Stream m a -> Step m a
Yield a
x (forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (s -> m (Step m a)
step' s
s))
            Maybe (a, s)
Nothing     -> 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 :: forall (m :: * -> *) b a.
Monad m =>
(b -> Maybe (a, b)) -> b -> Stream m a
unfoldr b -> Maybe (a, b)
next b
s0 = forall (m :: * -> *) a.
Monad m =>
(forall b. (a -> b -> b) -> b -> b) -> Stream m a
build 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 :: forall (m :: * -> *) a. Monad m => Int -> a -> Stream m a
replicateM Int
n a
x = forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (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 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        if t
i forall a. Ord a => a -> a -> Bool
<= t
0
        then forall (m :: * -> *) a. Step m a
Stop
        else forall (m :: * -> *) a. a -> Stream m a -> Step m a
Yield a
x (forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream (t -> m (Step m a)
step (t
i forall a. Num a => a -> a -> a
- t
1)))

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

uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
uncons :: forall (m :: * -> *) a.
Monad m =>
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
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step m a
r of
        Yield a
x Stream m a
xs -> forall a. a -> Maybe a
Just (a
x, Stream m a
xs)
        Step m a
Stop -> 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 :: forall (m :: * -> *) a b.
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
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) = forall (m :: * -> *) a. m (Step m a) -> Stream m a
Stream 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 :: forall (m :: * -> *) a b.
Monad m =>
(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 (m :: * -> *) a.
Monad m =>
(forall b. (a -> b -> b) -> b -> b) -> Stream m a
build forall b. (a -> b -> b) -> b -> b
g = forall b. (a -> b -> b) -> b -> b
g forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
cons 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 :: forall (m :: * -> *) a. Monad m => Stream m a -> m ()
drain = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
_ m ()
xs -> m ()
xs) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
{-
drain (Stream step) = do
    r <- step
    case r of
        Yield _ next -> drain next
        Stop      -> return ()
        -}