-- |
-- Module      : Streamly.Internal.Data.Stream.StreamD.Generate
-- Copyright   : (c) 2020 Composewell Technologies and Contributors
--               (c) Roman Leshchinskiy 2008-2010
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Prefer unfolds ("Streamly.Internal.Data.Unfold") over the combinators in
-- this module. They are more powerful and efficient as they can be transformed
-- and composed on the input side efficiently and they can fuse in nested
-- operations (e.g.  unfoldMany). All the combinators in this module can be
-- expressed using unfolds with the same efficiency.
--
-- Operations in this module that are not in "Streamly.Internal.Data.Unfold":
-- generate, times, fromPrimIORef.
--
-- We should plan to replace this module with "Streamly.Internal.Data.Unfold"
-- in future.

-- A few combinators in this module have been adapted from the vector package
-- (c) Roman Leshchinskiy. See the notes in specific combinators.
--
module Streamly.Internal.Data.Stream.StreamD.Generate
  (
    -- * Primitives
      nil
    , nilM
    , cons
    , consM

    -- * From 'Unfold'
    , unfold

    -- * Unfolding
    , unfoldr
    , unfoldrM

    -- * From Values
    , fromPure
    , fromEffect
    , repeat
    , repeatM
    , replicate
    , replicateM

    -- * Enumeration
    , enumerateFromStepIntegral
    , enumerateFromIntegral
    , enumerateFromThenIntegral
    , enumerateFromToIntegral
    , enumerateFromThenToIntegral

    , enumerateFromStepNum
    , enumerateFromNum
    , enumerateFromThenNum
    , enumerateFromToFractional
    , enumerateFromThenToFractional

    -- * Time Enumeration
    , times

    -- * From Generators
    -- | Generate a monadic stream from a seed.
    , fromIndices
    , fromIndicesM
    , generate
    , generateM

    -- * Iteration
    , iterate
    , iterateM

    -- * From Containers
    -- | Transform an input structure into a stream.

    -- Note: Direct style stream does not support @fromFoldable@.
    , fromList
    , fromListM

    -- * Conversions
    , fromStreamK
    , toStreamK
    )
where

#include "inline.hs"

import Control.Monad.IO.Class (MonadIO(..))
import Streamly.Internal.Control.Concurrent (MonadAsync)
import Streamly.Internal.Data.Time.Clock
    (Clock(Monotonic), asyncClock, readClock)
import Streamly.Internal.Data.Time.Units
    (toAbsTime, AbsTime, toRelTime64, RelTime64)

import Prelude hiding (iterate, repeat, replicate, takeWhile)
import Streamly.Internal.Data.Stream.StreamD.Type

------------------------------------------------------------------------------
-- Primitives
------------------------------------------------------------------------------

-- | An empty 'Stream'.
{-# INLINE_NORMAL nil #-}
nil :: Monad m => Stream m a
nil :: Stream m a
nil = (State Stream m a -> () -> m (Step () a)) -> () -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> Step () a -> m (Step () a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step () a
forall s a. Step s a
Stop) ()

-- XXX implement in terms of consM?
-- cons x = consM (return x)
--
-- | Can fuse but has O(n^2) complexity.
{-# 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 State Stream m a -> s -> m (Step s a)
step s
state) = (State Stream m a -> Maybe s -> m (Step (Maybe s) a))
-> Maybe s -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 Maybe s
forall a. Maybe a
Nothing
    where
    {-# INLINE_LATE step1 #-}
    step1 :: State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 State Stream m a
_ Maybe s
Nothing   = Step (Maybe s) a -> m (Step (Maybe s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe s) a -> m (Step (Maybe s) a))
-> Step (Maybe s) a -> m (Step (Maybe s) a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe s -> Step (Maybe s) a
forall s a. a -> s -> Step s a
Yield a
x (s -> Maybe s
forall a. a -> Maybe a
Just s
state)
    step1 State Stream m a
gst (Just s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        Step (Maybe s) a -> m (Step (Maybe s) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe s) a -> m (Step (Maybe s) a))
-> Step (Maybe s) a -> m (Step (Maybe s) a)
forall a b. (a -> b) -> a -> b
$
          case Step s a
r of
            Yield a
a s
s -> a -> Maybe s -> Step (Maybe s) a
forall s a. a -> s -> Step s a
Yield a
a (s -> Maybe s
forall a. a -> Maybe a
Just s
s)
            Skip  s
s   -> Maybe s -> Step (Maybe s) a
forall s a. s -> Step s a
Skip (s -> Maybe s
forall a. a -> Maybe a
Just s
s)
            Step s a
Stop      -> Step (Maybe s) a
forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Unfolding
------------------------------------------------------------------------------

-- Adapted from vector package
{-# INLINE_NORMAL unfoldrM #-}
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 = (State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
forall p. p -> s -> m (Step s a)
step
  where
    {-# INLINE_LATE step #-}
    step :: p -> s -> m (Step s a)
step p
_ s
st = do
        Maybe (a, s)
r <- s -> m (Maybe (a, s))
next s
st
        Step s a -> m (Step s a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step s a -> m (Step s a)) -> Step s a -> m (Step s a)
forall a b. (a -> b) -> a -> b
$ case Maybe (a, s)
r of
            Just (a
x, s
s) -> a -> s -> Step s a
forall s a. a -> s -> Step s a
Yield a
x s
s
            Maybe (a, s)
Nothing     -> Step s a
forall s a. Step s a
Stop

{-# INLINE_LATE unfoldr #-}
unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
unfoldr :: (s -> Maybe (a, s)) -> s -> Stream m a
unfoldr s -> Maybe (a, s)
f = (s -> m (Maybe (a, s))) -> s -> Stream m a
forall (m :: * -> *) s a.
Monad m =>
(s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM (Maybe (a, s) -> m (Maybe (a, s))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (a, s) -> m (Maybe (a, s)))
-> (s -> Maybe (a, s)) -> s -> m (Maybe (a, s))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Maybe (a, s)
f)

------------------------------------------------------------------------------
-- From values
------------------------------------------------------------------------------

{-# INLINE_NORMAL repeatM #-}
repeatM :: Monad m => m a -> Stream m a
repeatM :: m a -> Stream m a
repeatM m a
x = (State Stream m a -> () -> m (Step () a)) -> () -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> m a
x m a -> (a -> m (Step () a)) -> m (Step () a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> Step () a -> m (Step () a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step () a -> m (Step () a)) -> Step () a -> m (Step () a)
forall a b. (a -> b) -> a -> b
$ a -> () -> Step () a
forall s a. a -> s -> Step s a
Yield a
r ()) ()

{-# INLINE_NORMAL repeat #-}
repeat :: Monad m => a -> Stream m a
repeat :: a -> Stream m a
repeat a
x = (State Stream m a -> () -> m (Step () a)) -> () -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> Step () a -> m (Step () a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step () a -> m (Step () a)) -> Step () a -> m (Step () a)
forall a b. (a -> b) -> a -> b
$ a -> () -> Step () a
forall s a. a -> s -> Step s a
Yield a
x ()) ()

-- Adapted from the vector package
{-# INLINE_NORMAL replicateM #-}
replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
replicateM :: Int -> m a -> Stream m a
replicateM Int
n m a
p = (State Stream m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Int -> m (Step Int a)
forall p. p -> Int -> m (Step Int a)
step Int
n
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ (Int
i :: Int)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step Int a
forall s a. Step s a
Stop
      | Bool
otherwise = do
          a
x <- m a
p
          Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> m (Step Int a)) -> Step Int a -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$ a -> Int -> Step Int a
forall s a. a -> s -> Step s a
Yield a
x (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL replicate #-}
replicate :: Monad m => Int -> a -> Stream m a
replicate :: Int -> a -> Stream m a
replicate Int
n a
x = Int -> m a -> Stream m a
forall (m :: * -> *) a. Monad m => Int -> m a -> Stream m a
replicateM Int
n (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x)

------------------------------------------------------------------------------
-- Enumeration of Num
------------------------------------------------------------------------------

-- | For floating point numbers if the increment is less than the precision then
-- it just gets lost. Therefore we cannot always increment it correctly by just
-- repeated addition.
-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
-- 9007199254740992 + 2     :: Double => 9.007199254740994e15
--
-- Instead we accumulate the increment counter and compute the increment
-- every time before adding it to the starting number.
--
-- This works for Integrals as well as floating point numbers, but
-- enumerateFromStepIntegral is faster for integrals.
{-# INLINE_NORMAL enumerateFromStepNum #-}
enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum :: a -> a -> Stream m a
enumerateFromStepNum a
from a
stride = (State Stream m a -> a -> m (Step a a)) -> a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> a -> m (Step a a)
forall (m :: * -> *) p. Monad m => p -> a -> m (Step a a)
step a
0
    where
    {-# INLINE_LATE step #-}
    step :: p -> a -> m (Step a a)
step p
_ !a
i = 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 -> a -> Step a a) -> a -> a -> Step a a
forall a b. (a -> b) -> a -> b
$! (a
from a -> a -> a
forall a. Num a => a -> a -> a
+ a
i a -> a -> a
forall a. Num a => a -> a -> a
* a
stride)) (a -> Step a a) -> a -> Step a a
forall a b. (a -> b) -> a -> b
$! (a
i a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)

{-# INLINE_NORMAL enumerateFromNum #-}
enumerateFromNum :: (Monad m, Num a) => a -> Stream m a
enumerateFromNum :: a -> Stream m a
enumerateFromNum a
from = a -> a -> Stream m a
forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from a
1

{-# INLINE_NORMAL enumerateFromThenNum #-}
enumerateFromThenNum :: (Monad m, Num a) => a -> a -> Stream m a
enumerateFromThenNum :: a -> a -> Stream m a
enumerateFromThenNum a
from a
next = a -> a -> Stream m a
forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from (a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from)

------------------------------------------------------------------------------
-- Enumeration of Integrals
------------------------------------------------------------------------------

data EnumState a = EnumInit | EnumYield a a a | EnumStop

{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
enumerateFromThenToIntegralUp
    :: (Monad m, Integral a)
    => a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp :: a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp a
from a
next a
to = (State Stream m a -> EnumState a -> m (Step (EnumState a) a))
-> EnumState a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> EnumState a -> m (Step (EnumState a) a)
forall (m :: * -> *) p.
Monad m =>
p -> EnumState a -> m (Step (EnumState a) a)
step EnumState a
forall a. EnumState a
EnumInit
    where
    {-# INLINE_LATE step #-}
    step :: p -> EnumState a -> m (Step (EnumState a) a)
step p
_ EnumState a
EnumInit =
        Step (EnumState a) a -> m (Step (EnumState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (EnumState a) a -> m (Step (EnumState a) a))
-> Step (EnumState a) a -> m (Step (EnumState a) a)
forall a b. (a -> b) -> a -> b
$
            if a
to a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
next
            then if a
to a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
from
                 then Step (EnumState a) a
forall s a. Step s a
Stop
                 else a -> EnumState a -> Step (EnumState a) a
forall s a. a -> s -> Step s a
Yield a
from EnumState a
forall a. EnumState a
EnumStop
            else -- from <= next <= to
                let stride :: a
stride = a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from
                in EnumState a -> Step (EnumState a) a
forall s a. s -> Step s a
Skip (EnumState a -> Step (EnumState a) a)
-> EnumState a -> Step (EnumState a) a
forall a b. (a -> b) -> a -> b
$ a -> a -> a -> EnumState a
forall a. a -> a -> a -> EnumState a
EnumYield a
from a
stride (a
to a -> a -> a
forall a. Num a => a -> a -> a
- a
stride)

    step p
_ (EnumYield a
x a
stride a
toMinus) =
        Step (EnumState a) a -> m (Step (EnumState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (EnumState a) a -> m (Step (EnumState a) a))
-> Step (EnumState a) a -> m (Step (EnumState a) a)
forall a b. (a -> b) -> a -> b
$
            if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
toMinus
            then a -> EnumState a -> Step (EnumState a) a
forall s a. a -> s -> Step s a
Yield a
x EnumState a
forall a. EnumState a
EnumStop
            else a -> EnumState a -> Step (EnumState a) a
forall s a. a -> s -> Step s a
Yield a
x (EnumState a -> Step (EnumState a) a)
-> EnumState a -> Step (EnumState a) a
forall a b. (a -> b) -> a -> b
$ a -> a -> a -> EnumState a
forall a. a -> a -> a -> EnumState a
EnumYield (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
stride) a
stride a
toMinus

    step p
_ EnumState a
EnumStop = Step (EnumState a) a -> m (Step (EnumState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (EnumState a) a
forall s a. Step s a
Stop

{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
enumerateFromThenToIntegralDn
    :: (Monad m, Integral a)
    => a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn :: a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn a
from a
next a
to = (State Stream m a -> EnumState a -> m (Step (EnumState a) a))
-> EnumState a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> EnumState a -> m (Step (EnumState a) a)
forall (m :: * -> *) p.
Monad m =>
p -> EnumState a -> m (Step (EnumState a) a)
step EnumState a
forall a. EnumState a
EnumInit
    where
    {-# INLINE_LATE step #-}
    step :: p -> EnumState a -> m (Step (EnumState a) a)
step p
_ EnumState a
EnumInit =
        Step (EnumState a) a -> m (Step (EnumState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (EnumState a) a -> m (Step (EnumState a) a))
-> Step (EnumState a) a -> m (Step (EnumState a) a)
forall a b. (a -> b) -> a -> b
$ if a
to a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
next
            then if a
to a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
from
                 then Step (EnumState a) a
forall s a. Step s a
Stop
                 else a -> EnumState a -> Step (EnumState a) a
forall s a. a -> s -> Step s a
Yield a
from EnumState a
forall a. EnumState a
EnumStop
            else -- from >= next >= to
                let stride :: a
stride = a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from
                in EnumState a -> Step (EnumState a) a
forall s a. s -> Step s a
Skip (EnumState a -> Step (EnumState a) a)
-> EnumState a -> Step (EnumState a) a
forall a b. (a -> b) -> a -> b
$ a -> a -> a -> EnumState a
forall a. a -> a -> a -> EnumState a
EnumYield a
from a
stride (a
to a -> a -> a
forall a. Num a => a -> a -> a
- a
stride)

    step p
_ (EnumYield a
x a
stride a
toMinus) =
        Step (EnumState a) a -> m (Step (EnumState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (EnumState a) a -> m (Step (EnumState a) a))
-> Step (EnumState a) a -> m (Step (EnumState a) a)
forall a b. (a -> b) -> a -> b
$
            if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
toMinus
            then a -> EnumState a -> Step (EnumState a) a
forall s a. a -> s -> Step s a
Yield a
x EnumState a
forall a. EnumState a
EnumStop
            else a -> EnumState a -> Step (EnumState a) a
forall s a. a -> s -> Step s a
Yield a
x (EnumState a -> Step (EnumState a) a)
-> EnumState a -> Step (EnumState a) a
forall a b. (a -> b) -> a -> b
$ a -> a -> a -> EnumState a
forall a. a -> a -> a -> EnumState a
EnumYield (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
stride) a
stride a
toMinus

    step p
_ EnumState a
EnumStop = Step (EnumState a) a -> m (Step (EnumState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (EnumState a) a
forall s a. Step s a
Stop

-- XXX This can perhaps be simplified and written in terms of
-- enumeratFromStepIntegral as we have done in unfolds. But anyway we should be
-- replacing the stream generation module with unfolds.
{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
enumerateFromThenToIntegral
    :: (Monad m, Integral a)
    => a -> a -> a -> Stream m a
enumerateFromThenToIntegral :: a -> a -> a -> Stream m a
enumerateFromThenToIntegral a
from a
next a
to
    | a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
from = a -> a -> a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp a
from a
next a
to
    | Bool
otherwise    = a -> a -> a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn a
from a
next a
to

{-# INLINE_NORMAL enumerateFromThenIntegral #-}
enumerateFromThenIntegral
    :: (Monad m, Integral a, Bounded a)
    => a -> a -> Stream m a
enumerateFromThenIntegral :: a -> a -> Stream m a
enumerateFromThenIntegral a
from a
next =
    if a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
from
    then a -> a -> a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp a
from a
next a
forall a. Bounded a => a
maxBound
    else a -> a -> a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn a
from a
next a
forall a. Bounded a => a
minBound

-- | Can be used to enumerate unbounded integrals. This does not check for
-- overflow or underflow for bounded integrals.
--
{-# INLINE_NORMAL enumerateFromStepIntegral #-}
enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
enumerateFromStepIntegral :: a -> a -> Stream m a
enumerateFromStepIntegral a
from a
stride =
    a
from a -> Stream m a -> Stream m a
`seq` a
stride a -> Stream m a -> Stream m a
`seq` (State Stream m a -> a -> m (Step a a)) -> a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> a -> m (Step a a)
forall (m :: * -> *) p. Monad m => p -> a -> m (Step a a)
step a
from
    where
        {-# INLINE_LATE step #-}
        step :: p -> a -> m (Step a a)
step p
_ !a
x = 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 -> Step a a) -> a -> Step a a
forall a b. (a -> b) -> a -> b
$! (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
stride)

-- | Enumerate upwards from @from@ to @to@. We are assuming that "to" is
-- constrained by the type to be within max/min bounds.
{-# INLINE enumerateFromToIntegral #-}
enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
enumerateFromToIntegral :: a -> a -> Stream m a
enumerateFromToIntegral a
from a
to =
    (a -> Bool) -> Stream m a -> Stream m a
forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to) (Stream m a -> Stream m a) -> Stream m a -> Stream m a
forall a b. (a -> b) -> a -> b
$ a -> a -> Stream m a
forall a (m :: * -> *).
(Integral a, Monad m) =>
a -> a -> Stream m a
enumerateFromStepIntegral a
from a
1

{-# INLINE enumerateFromIntegral #-}
enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
enumerateFromIntegral :: a -> Stream m a
enumerateFromIntegral a
from = a -> a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> Stream m a
enumerateFromToIntegral a
from a
forall a. Bounded a => a
maxBound

------------------------------------------------------------------------------
-- Enumeration of Fractionals
------------------------------------------------------------------------------

-- | We cannot write a general function for Num.  The only way to write code
-- portable between the two is to use a 'Real' constraint and convert between
-- Fractional and Integral using fromRational which is horribly slow.
{-# INLINE_NORMAL enumerateFromToFractional #-}
enumerateFromToFractional
    :: (Monad m, Fractional a, Ord a)
    => a -> a -> Stream m a
enumerateFromToFractional :: a -> a -> Stream m a
enumerateFromToFractional a
from a
to =
    (a -> Bool) -> Stream m a -> Stream m a
forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to a -> a -> a
forall a. Num a => a -> a -> a
+ a
1 a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2) (Stream m a -> Stream m a) -> Stream m a -> Stream m a
forall a b. (a -> b) -> a -> b
$ a -> a -> Stream m a
forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from a
1

{-# INLINE_NORMAL enumerateFromThenToFractional #-}
enumerateFromThenToFractional
    :: (Monad m, Fractional a, Ord a)
    => a -> a -> a -> Stream m a
enumerateFromThenToFractional :: a -> a -> a -> Stream m a
enumerateFromThenToFractional a
from a
next a
to =
    (a -> Bool) -> Stream m a -> Stream m a
forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile a -> Bool
predicate (Stream m a -> Stream m a) -> Stream m a -> Stream m a
forall a b. (a -> b) -> a -> b
$ a -> a -> Stream m a
forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromThenNum a
from a
next
    where
    mid :: a
mid = (a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2
    predicate :: a -> Bool
predicate | a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
from  = (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to a -> a -> a
forall a. Num a => a -> a -> a
+ a
mid)
              | Bool
otherwise     = (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
to a -> a -> a
forall a. Num a => a -> a -> a
+ a
mid)

------------------------------------------------------------------------------
-- Time Enumeration
------------------------------------------------------------------------------

{-# INLINE_NORMAL times #-}
times :: MonadIO m => Double -> Stream m (AbsTime, RelTime64)
times :: Double -> Stream m (AbsTime, RelTime64)
times Double
g = (State Stream m (AbsTime, RelTime64)
 -> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
 -> m (Step
         (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
         (AbsTime, RelTime64)))
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> Stream m (AbsTime, RelTime64)
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m (AbsTime, RelTime64)
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
forall (m :: * -> *) p.
MonadIO m =>
p
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
step Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
forall a. Maybe a
Nothing

    where

    {-# INLINE_LATE step #-}
    step :: p
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
step p
_ Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
Nothing = do
        (ThreadId, IORef MicroSecond64)
clock <- IO (ThreadId, IORef MicroSecond64)
-> m (ThreadId, IORef MicroSecond64)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (ThreadId, IORef MicroSecond64)
 -> m (ThreadId, IORef MicroSecond64))
-> IO (ThreadId, IORef MicroSecond64)
-> m (ThreadId, IORef MicroSecond64)
forall a b. (a -> b) -> a -> b
$ Clock -> Double -> IO (ThreadId, IORef MicroSecond64)
asyncClock Clock
Monotonic Double
g
        MicroSecond64
a <- IO MicroSecond64 -> m MicroSecond64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MicroSecond64 -> m MicroSecond64)
-> IO MicroSecond64 -> m MicroSecond64
forall a b. (a -> b) -> a -> b
$ (ThreadId, IORef MicroSecond64) -> IO MicroSecond64
readClock (ThreadId, IORef MicroSecond64)
clock
        Step
  (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
  (AbsTime, RelTime64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step
   (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
   (AbsTime, RelTime64)
 -> m (Step
         (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
         (AbsTime, RelTime64)))
-> Step
     (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
     (AbsTime, RelTime64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
forall a b. (a -> b) -> a -> b
$ Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> Step
     (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
     (AbsTime, RelTime64)
forall s a. s -> Step s a
Skip (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
 -> Step
      (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
      (AbsTime, RelTime64))
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> Step
     (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
     (AbsTime, RelTime64)
forall a b. (a -> b) -> a -> b
$ ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
forall a. a -> Maybe a
Just ((ThreadId, IORef MicroSecond64)
clock, MicroSecond64
a)

    step p
_ s :: Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
s@(Just ((ThreadId, IORef MicroSecond64)
clock, MicroSecond64
t0)) = do
        MicroSecond64
a <- IO MicroSecond64 -> m MicroSecond64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO MicroSecond64 -> m MicroSecond64)
-> IO MicroSecond64 -> m MicroSecond64
forall a b. (a -> b) -> a -> b
$ (ThreadId, IORef MicroSecond64) -> IO MicroSecond64
readClock (ThreadId, IORef MicroSecond64)
clock
        -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
        -- efficiency.  or maybe we can use a representation using Double for
        -- floating precision time
        Step
  (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
  (AbsTime, RelTime64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step
   (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
   (AbsTime, RelTime64)
 -> m (Step
         (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
         (AbsTime, RelTime64)))
-> Step
     (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
     (AbsTime, RelTime64)
-> m (Step
        (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
        (AbsTime, RelTime64))
forall a b. (a -> b) -> a -> b
$ (AbsTime, RelTime64)
-> Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
-> Step
     (Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64))
     (AbsTime, RelTime64)
forall s a. a -> s -> Step s a
Yield (MicroSecond64 -> AbsTime
forall a. TimeUnit a => a -> AbsTime
toAbsTime MicroSecond64
t0, MicroSecond64 -> RelTime64
forall a. TimeUnit64 a => a -> RelTime64
toRelTime64 (MicroSecond64
a MicroSecond64 -> MicroSecond64 -> MicroSecond64
forall a. Num a => a -> a -> a
- MicroSecond64
t0)) Maybe ((ThreadId, IORef MicroSecond64), MicroSecond64)
s

-------------------------------------------------------------------------------
-- From Generators
-------------------------------------------------------------------------------

{-# INLINE_NORMAL fromIndicesM #-}
fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
fromIndicesM :: (Int -> m a) -> Stream m a
fromIndicesM Int -> m a
gen = (State Stream m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Int -> m (Step Int a)
forall p. p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i = do
       a
x <- Int -> m a
gen Int
i
       Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> m (Step Int a)) -> Step Int a -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$ a -> Int -> Step Int a
forall s a. a -> s -> Step s a
Yield a
x (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE fromIndices #-}
fromIndices :: Monad m => (Int -> a) -> Stream m a
fromIndices :: (Int -> a) -> Stream m a
fromIndices Int -> a
gen = (Int -> m a) -> Stream m a
forall (m :: * -> *) a. Monad m => (Int -> m a) -> Stream m a
fromIndicesM (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> (Int -> a) -> Int -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
gen)

-- Adapted from the vector package
{-# INLINE_NORMAL generateM #-}
generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
generateM :: Int -> (Int -> m a) -> Stream m a
generateM Int
n Int -> m a
gen = Int
n Int -> Stream m a -> Stream m a
`seq` (State Stream m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Int -> m (Step Int a)
forall p. p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n     = do
                           a
x <- Int -> m a
gen Int
i
                           Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> m (Step Int a)) -> Step Int a -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$ a -> Int -> Step Int a
forall s a. a -> s -> Step s a
Yield a
x (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
             | Bool
otherwise = Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step Int a
forall s a. Step s a
Stop

{-# INLINE generate #-}
generate :: Monad m => Int -> (Int -> a) -> Stream m a
generate :: Int -> (Int -> a) -> Stream m a
generate Int
n Int -> a
gen = Int -> (Int -> m a) -> Stream m a
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> Stream m a
generateM Int
n (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> (Int -> a) -> Int -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
gen)

-------------------------------------------------------------------------------
-- Iteration
-------------------------------------------------------------------------------

{-# INLINE_NORMAL iterateM #-}
iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
iterateM :: (a -> m a) -> m a -> Stream m a
iterateM a -> m a
step = (State Stream m a -> m a -> m (Step (m a) a)) -> m a -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ m a
st -> m a
st m a -> (a -> m (Step (m a) a)) -> m (Step (m a) a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(!a
x) -> Step (m a) a -> m (Step (m a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (m a) a -> m (Step (m a) a))
-> Step (m a) a -> m (Step (m a) a)
forall a b. (a -> b) -> a -> b
$ a -> m a -> Step (m a) a
forall s a. a -> s -> Step s a
Yield a
x (a -> m a
step a
x))

{-# INLINE_NORMAL iterate #-}
iterate :: Monad m => (a -> a) -> a -> Stream m a
iterate :: (a -> a) -> a -> Stream m a
iterate a -> a
step a
st = (a -> m a) -> m a -> Stream m a
forall (m :: * -> *) a. Monad m => (a -> m a) -> m a -> Stream m a
iterateM (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> (a -> a) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
step) (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
st)

-------------------------------------------------------------------------------
-- From containers
-------------------------------------------------------------------------------

-- XXX we need the MonadAsync constraint because of a rewrite rule.
-- | Convert a list of monadic actions to a 'Stream'
{-# INLINE_LATE fromListM #-}
fromListM :: MonadAsync m => [m a] -> Stream m a
fromListM :: [m a] -> Stream m a
fromListM = (State Stream m a -> [m a] -> m (Step [m a] a))
-> [m a] -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> [m a] -> m (Step [m a] a)
forall (m :: * -> *) p a. Monad m => p -> [m a] -> m (Step [m a] a)
step
  where
    {-# INLINE_LATE step #-}
    step :: p -> [m a] -> m (Step [m a] a)
step p
_ (m a
m:[m a]
ms) = m a
m m a -> (a -> m (Step [m a] a)) -> m (Step [m a] a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> Step [m a] a -> m (Step [m a] a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step [m a] a -> m (Step [m a] a))
-> Step [m a] a -> m (Step [m a] a)
forall a b. (a -> b) -> a -> b
$ a -> [m a] -> Step [m a] a
forall s a. a -> s -> Step s a
Yield a
x [m a]
ms
    step p
_ []     = Step [m a] a -> m (Step [m a] a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step [m a] a
forall s a. Step s a
Stop