-- |
-- Module      : Streamly.Internal.Data.Unfold.Enumeration
-- Copyright   : (c) 2019, 2021 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- The functions defined in this module should be rarely needed for direct use,
-- try to use the operations from the 'Enumerable' type class
-- instances instead.
--
-- This module provides an 'Enumerable' type class to enumerate 'Enum' types
-- into a stream. The operations in this type class correspond to similar
-- operations in the 'Enum' type class, the only difference is that they produce
-- a stream instead of a list. These operations cannot be defined generically
-- based on the 'Enum' type class. We provide instances for commonly used
-- types. If instances for other types are needed convenience functions defined
-- in this module can be used to define them. Alternatively, these functions
-- can be used directly.
--
module Streamly.Internal.Data.Unfold.Enumeration
    (
      Enumerable (..)

    -- ** Enumerating 'Num' Types
    , enumerateFromStepNum
    , enumerateFromNum
    , enumerateFromThenNum

    -- ** Enumerating unbounded 'Integral' Types
    , enumerateFromStepIntegral
    , enumerateFromIntegral
    , enumerateFromThenIntegral
    , enumerateFromToIntegral
    , enumerateFromThenToIntegral

    -- ** Enumerating 'Bounded' 'Integral' Types
    , enumerateFromIntegralBounded
    , enumerateFromThenIntegralBounded
    , enumerateFromToIntegralBounded
    , enumerateFromThenToIntegralBounded

    -- ** Enumerating small 'Integral' Types
    -- | Small types are always bounded.
    , enumerateFromSmallBounded
    , enumerateFromThenSmallBounded
    , enumerateFromToSmall
    , enumerateFromThenToSmall

    -- ** Enumerating 'Fractional' Types
    -- | Enumeration of 'Num' specialized to 'Fractional' types.
    , enumerateFromFractional
    , enumerateFromThenFractional
    , enumerateFromToFractional
    , enumerateFromThenToFractional
    )
where

#include "inline.hs"
import Data.Fixed
import Data.Bifunctor (bimap)
import Data.Int
import Data.Ratio
import Data.Word
import Numeric.Natural
import Data.Functor.Identity (Identity(..))
import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
import Streamly.Internal.Data.Unfold.Type
import Prelude
       hiding (map, mapM, takeWhile, take, filter, const, zipWith
              , drop, dropWhile)

-- $setup
-- >>> :m
-- >>> import qualified Streamly.Prelude as Stream
-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
-- >>> import Streamly.Internal.Data.Unfold.Type
-- >>> import Data.Word

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

-- | Unfolds @(from, stride)@ generating an infinite stream starting from
-- @from@ and incrementing every time by @stride@.  For 'Bounded' types, after
-- the value overflows it keeps enumerating in a cycle:
--
-- @
-- >>> Stream.toList $ Stream.take 10 $ Stream.unfold Unfold.enumerateFromStepNum (255::Word8,1)
-- [255,0,1,2,3,4,5,6,7,8]
--
-- @
--
-- The implementation is numerically stable for floating point values.
--
-- Note 'enumerateFromStepIntegral' is faster for integrals.
--
-- /Internal/
--
{-# INLINE enumerateFromStepNum #-}
enumerateFromStepNum :: (Monad m, Num a) => Unfold m (a, a) a
enumerateFromStepNum :: Unfold m (a, a) a
enumerateFromStepNum = ((a, a, a) -> m (Step (a, a, a) a))
-> ((a, a) -> m (a, a, a)) -> Unfold m (a, a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (a, a, a) -> m (Step (a, a, a) a)
forall (m :: * -> *) c.
(Monad m, Num c) =>
(c, c, c) -> m (Step (c, c, c) c)
step (a, a) -> m (a, a, a)
forall (m :: * -> *) c a b.
(Monad m, Num c) =>
(a, b) -> m (a, b, c)
inject

    where

    inject :: (a, b) -> m (a, b, c)
inject (!a
from, !b
stride) = (a, b, c) -> m (a, b, c)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
from, b
stride, c
0)

    -- Note that the counter "i" is the same type as the type being enumerated.
    -- It may overflow, for example, if we are enumerating Word8, after 255 the
    -- counter will become 0, but the overflow does not affect the enumeration
    -- behavior.
    {-# INLINE_LATE step #-}
    step :: (c, c, c) -> m (Step (c, c, c) c)
step (c
from, c
stride, c
i) =
        Step (c, c, c) c -> m (Step (c, c, c) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (c, c, c) c -> m (Step (c, c, c) c))
-> Step (c, c, c) c -> m (Step (c, c, c) c)
forall a b. (a -> b) -> a -> b
$
            (c -> (c, c, c) -> Step (c, c, c) c
forall s a. a -> s -> Step s a
Yield (c -> (c, c, c) -> Step (c, c, c) c)
-> c -> (c, c, c) -> Step (c, c, c) c
forall a b. (a -> b) -> a -> b
$! (c
from c -> c -> c
forall a. Num a => a -> a -> a
+ c
i c -> c -> c
forall a. Num a => a -> a -> a
* c
stride)) ((c, c, c) -> Step (c, c, c) c) -> (c, c, c) -> Step (c, c, c) c
forall a b. (a -> b) -> a -> b
$! (c
from, c
stride, c
i c -> c -> c
forall a. Num a => a -> a -> a
+ c
1)

-- | Same as 'enumerateFromStepNum (from, next)' using a stride of @next - from@:
--
-- @
-- >>> enumerateFromThenNum = lmap (\(from, next) -> (from, next - from)) Unfold.enumerateFromStepNum
--
-- @
--
-- Example:
-- @
-- >>> Stream.toList $ Stream.take 10 $ Stream.unfold enumerateFromThenNum (255::Word8,0)
-- [255,0,1,2,3,4,5,6,7,8]
--
-- @
--
-- The implementation is numerically stable for floating point values.
--
-- Note that 'enumerateFromThenIntegral' is faster for integrals.
--
-- Note that in the strange world of floating point numbers, using
-- @enumerateFromThenNum (from, from + 1)@ is almost exactly the same as
-- @enumerateFromStepNum (from, 1) but not precisely the same. Because @(from +
-- 1) - from@ is not exactly 1, it may lose some precision, the loss may also
-- be aggregated in each step, if you want that precision then use
-- 'enumerateFromStepNum' instead.
--
-- /Internal/
--
{-# INLINE enumerateFromThenNum #-}
enumerateFromThenNum :: (Monad m, Num a) => Unfold m (a, a) a
enumerateFromThenNum :: Unfold m (a, a) a
enumerateFromThenNum =
    ((a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (\(a
from, a
next) -> (a
from, a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from)) Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Num a) => Unfold m (a, a) a
enumerateFromStepNum

-- | Same as 'enumerateFromStepNum' using a stride of 1:
--
-- @
-- >>> enumerateFromNum = lmap (\from -> (from, 1)) Unfold.enumerateFromStepNum
-- >>> Stream.toList $ Stream.take 6 $ Stream.unfold enumerateFromNum (0.9)
-- [0.9,1.9,2.9,3.9,4.9,5.9]
--
-- @
--
-- Also, same as 'enumerateFromThenNum' using a stride of 1 but see the note in
-- 'enumerateFromThenNum' about the loss of precision:
--
-- @
-- >>> enumerateFromNum = lmap (\from -> (from, from + 1)) Unfold.enumerateFromThenNum
-- >>> Stream.toList $ Stream.take 6 $ Stream.unfold enumerateFromNum (0.9)
-- [0.9,1.9,2.9,3.8999999999999995,4.8999999999999995,5.8999999999999995]
--
-- @
--
-- /Internal/
--
{-# INLINE enumerateFromNum #-}
enumerateFromNum :: (Monad m, Num a) => Unfold m a a
enumerateFromNum :: Unfold m a a
enumerateFromNum = (a -> (a, a)) -> Unfold m (a, a) a -> Unfold m a a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (\a
from -> (a
from, a
1)) Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Num a) => Unfold m (a, a) a
enumerateFromStepNum

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

-- | Can be used to enumerate unbounded integrals. This does not check for
-- overflow or underflow for bounded integrals.
--
-- /Internal/
{-# INLINE_NORMAL enumerateFromStepIntegral #-}
enumerateFromStepIntegral :: (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromStepIntegral :: Unfold m (a, a) a
enumerateFromStepIntegral = ((a, a) -> m (Step (a, a) a))
-> ((a, a) -> m (a, a)) -> Unfold m (a, a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (a, a) -> m (Step (a, a) a)
forall (m :: * -> *) b.
(Monad m, Num b) =>
(b, b) -> m (Step (b, b) b)
step (a, a) -> m (a, a)
forall (m :: * -> *) a b. Monad m => (a, b) -> m (a, b)
inject

    where

    inject :: (a, b) -> m (a, b)
inject (a
from, b
stride) = a
from a -> m (a, b) -> m (a, b)
`seq` b
stride b -> m (a, b) -> m (a, b)
`seq` (a, b) -> m (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
from, b
stride)

    {-# INLINE_LATE step #-}
    step :: (b, b) -> m (Step (b, b) b)
step (b
x, b
stride) = Step (b, b) b -> m (Step (b, b) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (b, b) b -> m (Step (b, b) b))
-> Step (b, b) b -> m (Step (b, b) b)
forall a b. (a -> b) -> a -> b
$ b -> (b, b) -> Step (b, b) b
forall s a. a -> s -> Step s a
Yield b
x ((b, b) -> Step (b, b) b) -> (b, b) -> Step (b, b) b
forall a b. (a -> b) -> a -> b
$! (b
x b -> b -> b
forall a. Num a => a -> a -> a
+ b
stride, b
stride)

-- Enumerate Unbounded Integrals ----------------------------------------------
{-# INLINE enumerateFromIntegral #-}
enumerateFromIntegral :: (Monad m, Integral a) => Unfold m a a
enumerateFromIntegral :: Unfold m a a
enumerateFromIntegral = (a -> (a, a)) -> Unfold m (a, a) a -> Unfold m a a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (\a
from -> (a
from, a
1)) Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromStepIntegral

{-# INLINE enumerateFromThenIntegral #-}
enumerateFromThenIntegral :: (Monad m, Integral a ) => Unfold m (a, a) a
enumerateFromThenIntegral :: Unfold m (a, a) a
enumerateFromThenIntegral =
    ((a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (\(a
from, a
next) -> (a
from, a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from)) Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromStepIntegral

{-# INLINE enumerateFromToIntegral #-}
enumerateFromToIntegral :: (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromToIntegral :: Unfold m (a, a) a
enumerateFromToIntegral =
    ((a, a) -> a -> m Bool) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (\(a
_, a
to) a
b -> Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to)
        (Unfold m (a, a) a -> Unfold m (a, a) a)
-> Unfold m (a, a) a -> Unfold m (a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (\(a
from, a
_) -> (a
from, a
1)) Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromStepIntegral

{-# INLINE enumerateFromThenToIntegral #-}
enumerateFromThenToIntegral :: (Monad m, Integral a) => Unfold m (a, a, a) a
enumerateFromThenToIntegral :: Unfold m (a, a, a) a
enumerateFromThenToIntegral =
    ((a, a, a) -> a -> m Bool)
-> Unfold m (a, a, a) a -> Unfold m (a, a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (a, a, a) -> a -> m Bool
forall (m :: * -> *) a a.
(Monad m, Ord a, Ord a) =>
(a, a, a) -> a -> m Bool
cond (Unfold m (a, a, a) a -> Unfold m (a, a, a) a)
-> Unfold m (a, a, a) a -> Unfold m (a, a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a, a) -> (a, a)
forall b c. Num b => (b, b, c) -> (b, b)
toFromStep Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromStepIntegral

    where

    toFromStep :: (b, b, c) -> (b, b)
toFromStep (b
from, b
next, c
_) = (b
from, b
next b -> b -> b
forall a. Num a => a -> a -> a
- b
from)

    cond :: (a, a, a) -> a -> m Bool
cond (a
from, a
next, a
to) a
b =
        Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ if a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
from
              then a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to
              else a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
to

-- Enumerate Bounded Integrals ------------------------------------------------
{-# INLINE enumerateFromIntegralBounded #-}
enumerateFromIntegralBounded :: (Monad m, Integral a, Bounded a) =>
    Unfold m a a
enumerateFromIntegralBounded :: Unfold m a a
enumerateFromIntegralBounded = a -> Unfold m (a, a) a -> Unfold m a a
forall b (m :: * -> *) a c. b -> Unfold m (a, b) c -> Unfold m a c
supplySecond a
forall a. Bounded a => a
maxBound Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromToIntegral

{-# INLINE enumerateFromThenIntegralBounded #-}
enumerateFromThenIntegralBounded :: (Monad m, Integral a, Bounded a ) =>
    Unfold m (a, a) a
enumerateFromThenIntegralBounded :: Unfold m (a, a) a
enumerateFromThenIntegralBounded =
    ((a, a) -> a -> m Bool) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (a, a) -> a -> m Bool
forall (m :: * -> *) a a.
(Monad m, Ord a, Ord a, Bounded a) =>
(a, a) -> a -> m Bool
cond (Unfold m (a, a) a -> Unfold m (a, a) a)
-> Unfold m (a, a) a -> Unfold m (a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a) -> (a, a)
forall b. Num b => (b, b) -> (b, b)
toFromStep Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromStepIntegral

    where

    toFromStep :: (b, b) -> (b, b)
toFromStep (b
from, b
next) = (b
from, b
next b -> b -> b
forall a. Num a => a -> a -> a
- b
from)

    cond :: (a, a) -> a -> m Bool
cond (a
from, a
next) a
b =
        Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ if a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
from
              then a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
forall a. Bounded a => a
maxBound
              else a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
forall a. Bounded a => a
minBound

{-# INLINE enumerateFromToIntegralBounded #-}
enumerateFromToIntegralBounded :: (Monad m, Integral a, Bounded a) =>
    Unfold m (a, a) a
enumerateFromToIntegralBounded :: Unfold m (a, a) a
enumerateFromToIntegralBounded =
    ((a, a) -> a -> m Bool) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (\(a
_, a
to) a
b -> Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to)
        (Unfold m (a, a) a -> Unfold m (a, a) a)
-> Unfold m (a, a) a -> Unfold m (a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a) -> a) -> Unfold m a a -> Unfold m (a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a) -> a
forall a b. (a, b) -> a
fst Unfold m a a
forall (m :: * -> *) a.
(Monad m, Integral a, Bounded a) =>
Unfold m a a
enumerateFromIntegralBounded

{-# INLINE enumerateFromThenToIntegralBounded #-}
enumerateFromThenToIntegralBounded :: (Monad m, Integral a, Bounded a) =>
    Unfold m (a, a, a) a
enumerateFromThenToIntegralBounded :: Unfold m (a, a, a) a
enumerateFromThenToIntegralBounded =
    ((a, a, a) -> a -> m Bool)
-> Unfold m (a, a, a) a -> Unfold m (a, a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (a, a, a) -> a -> m Bool
forall (m :: * -> *) a a.
(Monad m, Ord a, Ord a) =>
(a, a, a) -> a -> m Bool
cond (Unfold m (a, a, a) a -> Unfold m (a, a, a) a)
-> Unfold m (a, a, a) a -> Unfold m (a, a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a, a) -> (a, a)
forall a b c. (a, b, c) -> (a, b)
toFromThen Unfold m (a, a) a
forall (m :: * -> *) a.
(Monad m, Integral a, Bounded a) =>
Unfold m (a, a) a
enumerateFromThenIntegralBounded

    where

    toFromThen :: (a, b, c) -> (a, b)
toFromThen (a
from, b
next, c
_) = (a
from, b
next)

    cond :: (a, a, a) -> a -> m Bool
cond (a
from, a
next, a
to) a
b =
        Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ if a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
from
              then a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to
              else a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
to

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

{-# INLINE_NORMAL enumerateFromFractional #-}
enumerateFromFractional :: (Monad m, Fractional a) => Unfold m a a
enumerateFromFractional :: Unfold m a a
enumerateFromFractional = Unfold m a a
forall (m :: * -> *) a. (Monad m, Num a) => Unfold m a a
enumerateFromNum

{-# INLINE_NORMAL enumerateFromThenFractional #-}
enumerateFromThenFractional :: (Monad m, Fractional a) => Unfold m (a, a) a
enumerateFromThenFractional :: Unfold m (a, a) a
enumerateFromThenFractional = Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Num a) => Unfold m (a, a) a
enumerateFromThenNum

-- | Same as 'enumerateFromStepNum' with a step of 1 and enumerating up to the
-- specified upper limit rounded to the nearest integral value:
--
-- @
-- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromToFractional (0.1, 6.3)
-- [0.1,1.1,2.1,3.1,4.1,5.1,6.1]
--
-- @
--
-- /Internal/
--
{-# INLINE_NORMAL enumerateFromToFractional #-}
enumerateFromToFractional :: (Monad m, Fractional a, Ord a) =>
    Unfold m (a, a) a
enumerateFromToFractional :: Unfold m (a, a) a
enumerateFromToFractional =
    ((a, a) -> a -> m Bool) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (\(a
_, a
to) a
b -> Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ a
b 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)
        (Unfold m (a, a) a -> Unfold m (a, a) a)
-> Unfold m (a, a) a -> Unfold m (a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (\(a
from, a
_) -> (a
from, a
1)) Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Num a) => Unfold m (a, a) a
enumerateFromStepNum

{-# INLINE enumerateFromThenToFractional #-}
enumerateFromThenToFractional :: (Monad m, Fractional a, Ord a) =>
    Unfold m (a, a, a) a
enumerateFromThenToFractional :: Unfold m (a, a, a) a
enumerateFromThenToFractional =
    ((a, a, a) -> a -> m Bool)
-> Unfold m (a, a, a) a -> Unfold m (a, a, a) a
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput (a, a, a) -> a -> m Bool
forall (m :: * -> *) a.
(Monad m, Ord a, Fractional a) =>
(a, a, a) -> a -> m Bool
cond (Unfold m (a, a, a) a -> Unfold m (a, a, a) a)
-> Unfold m (a, a, a) a -> Unfold m (a, a, a) a
forall a b. (a -> b) -> a -> b
$ ((a, a, a) -> (a, a)) -> Unfold m (a, a) a -> Unfold m (a, a, a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a, a) -> (a, a)
forall b c. Num b => (b, b, c) -> (b, b)
toFromStep Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Num a) => Unfold m (a, a) a
enumerateFromStepNum

    where

    toFromStep :: (b, b, c) -> (b, b)
toFromStep (b
from, b
next, c
_) = (b
from, b
next b -> b -> b
forall a. Num a => a -> a -> a
- b
from)

    cond :: (a, a, a) -> a -> m Bool
cond (a
from, a
next, a
to) a
b =
        let stride :: a
stride = a
next a -> a -> a
forall a. Num a => a -> a -> a
- a
from
         in Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return
                (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ if a
next a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
from
                  then a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
to a -> a -> a
forall a. Num a => a -> a -> a
+ a
stride a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2
                  else a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
to a -> a -> a
forall a. Num a => a -> a -> a
+ a
stride a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2

-------------------------------------------------------------------------------
-- Enumeration of Enum types not larger than Int
-------------------------------------------------------------------------------

-- | Enumerate from given starting Enum value 'from' and to Enum value 'to'
-- with stride of 1 till to value.
--
-- /Internal/
--
{-# INLINE enumerateFromToSmall #-}
enumerateFromToSmall :: (Monad m, Enum a) => Unfold m (a, a) a
enumerateFromToSmall :: Unfold m (a, a) a
enumerateFromToSmall =
    (Int -> a) -> Unfold m (a, a) Int -> Unfold m (a, a) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> a
forall a. Enum a => Int -> a
toEnum (((a, a) -> (Int, Int))
-> Unfold m (Int, Int) Int -> Unfold m (a, a) Int
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap ((a -> Int) -> (a -> Int) -> (a, a) -> (Int, Int)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> Int
forall a. Enum a => a -> Int
fromEnum a -> Int
forall a. Enum a => a -> Int
fromEnum) Unfold m (Int, Int) Int
forall (m :: * -> *) a. (Monad m, Integral a) => Unfold m (a, a) a
enumerateFromToIntegral)

-- | Enumerate from given starting Enum value 'from' and then Enum value 'next'
-- and to Enum value 'to' with stride of (fromEnum next - fromEnum from)
-- till to value.
--
-- /Internal/
--
{-# INLINE enumerateFromThenToSmall #-}
enumerateFromThenToSmall :: (Monad m, Enum a) => Unfold m (a, a, a) a
enumerateFromThenToSmall :: Unfold m (a, a, a) a
enumerateFromThenToSmall =
    let toInts :: (a, a, a) -> (Int, Int, Int)
toInts (a
x, a
y, a
z) = (a -> Int
forall a. Enum a => a -> Int
fromEnum a
x, a -> Int
forall a. Enum a => a -> Int
fromEnum a
y, a -> Int
forall a. Enum a => a -> Int
fromEnum a
z)
     in (Int -> a) -> Unfold m (a, a, a) Int -> Unfold m (a, a, a) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> a
forall a. Enum a => Int -> a
toEnum (((a, a, a) -> (Int, Int, Int))
-> Unfold m (Int, Int, Int) Int -> Unfold m (a, a, a) Int
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a, a) -> (Int, Int, Int)
forall a a a.
(Enum a, Enum a, Enum a) =>
(a, a, a) -> (Int, Int, Int)
toInts Unfold m (Int, Int, Int) Int
forall (m :: * -> *) a.
(Monad m, Integral a) =>
Unfold m (a, a, a) a
enumerateFromThenToIntegral)

-------------------------------------------------------------------------------
-- Bounded Enumeration of Enum types not larger than Int
-------------------------------------------------------------------------------

-- | Enumerate from given starting Enum value 'from' with stride of 1 till
-- maxBound
--
-- /Internal/
--
{-# INLINE enumerateFromSmallBounded #-}
enumerateFromSmallBounded :: (Monad m, Enum a, Bounded a) => Unfold m a a
enumerateFromSmallBounded :: Unfold m a a
enumerateFromSmallBounded = a -> Unfold m (a, a) a -> Unfold m a a
forall b (m :: * -> *) a c. b -> Unfold m (a, b) c -> Unfold m a c
supplySecond a
forall a. Bounded a => a
maxBound Unfold m (a, a) a
forall (m :: * -> *) a. (Monad m, Enum a) => Unfold m (a, a) a
enumerateFromToSmall

-- | Enumerate from given starting Enum value 'from' and next Enum value 'next'
-- with stride of (fromEnum next - fromEnum from) till maxBound.
--
-- /Internal/
--
{-# INLINE enumerateFromThenSmallBounded #-}
enumerateFromThenSmallBounded :: forall m a. (Monad m, Enum a, Bounded a) =>
    Unfold m (a, a) a
enumerateFromThenSmallBounded :: Unfold m (a, a) a
enumerateFromThenSmallBounded =
    let adapt :: (a, b) -> (Int, Int, Int)
adapt (a
from, b
next) =
            let frm :: Int
frm = a -> Int
forall a. Enum a => a -> Int
fromEnum a
from
                nxt :: Int
nxt = b -> Int
forall a. Enum a => a -> Int
fromEnum b
next
                stride :: Int
stride = Int
nxt Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
frm
                to :: Int
to = if Int
stride Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
                     then a -> Int
forall a. Enum a => a -> Int
fromEnum (a
forall a. Bounded a => a
maxBound :: a)
                     else a -> Int
forall a. Enum a => a -> Int
fromEnum (a
forall a. Bounded a => a
minBound :: a)
             in (Int
frm, Int
nxt, Int
to)
     in (Int -> a) -> Unfold m (a, a) Int -> Unfold m (a, a) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> a
forall a. Enum a => Int -> a
toEnum (((a, a) -> (Int, Int, Int))
-> Unfold m (Int, Int, Int) Int -> Unfold m (a, a) Int
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, a) -> (Int, Int, Int)
forall b a. (Enum b, Enum a) => (a, b) -> (Int, Int, Int)
adapt Unfold m (Int, Int, Int) Int
forall (m :: * -> *) a.
(Monad m, Integral a) =>
Unfold m (a, a, a) a
enumerateFromThenToIntegral)

-------------------------------------------------------------------------------
-- Enumerable type class
-------------------------------------------------------------------------------

-- | Types that can be enumerated as a stream. The operations in this type
-- class are equivalent to those in the 'Enum' type class, except that these
-- generate a stream instead of a list. Use the functions in
-- "Streamly.Internal.Data.Unfold.Enumeration" module to define new instances.
--
-- /Pre-release/
class Enum a => Enumerable a where

    -- | Unfolds @from@ generating a stream starting with the element
    -- @from@, enumerating up to 'maxBound' when the type is 'Bounded' or
    -- generating an infinite stream when the type is not 'Bounded'.
    --
    -- >>> import qualified Streamly.Prelude as Stream
    -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
    --
    -- @
    -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFrom (0 :: Int)
    -- [0,1,2,3]
    --
    -- @
    --
    -- For 'Fractional' types, enumeration is numerically stable. However, no
    -- overflow or underflow checks are performed.
    --
    -- @
    -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFrom 1.1
    -- [1.1,2.1,3.1,4.1]
    --
    -- @
    --
    -- /Pre-release/
    --
    enumerateFrom :: Monad m => Unfold m a a

    -- | Unfolds @(from, to)@ generating a finite stream starting with the element
    -- @from@, enumerating the type up to the value @to@. If @to@ is smaller than
    -- @from@ then an empty stream is returned.
    --
    -- >>> import qualified Streamly.Prelude as Stream
    -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
    --
    -- @
    -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromTo (0, 4)
    -- [0,1,2,3,4]
    --
    -- @
    --
    -- For 'Fractional' types, the last element is equal to the specified @to@
    -- value after rounding to the nearest integral value.
    --
    -- @
    -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromTo (1.1, 4)
    -- [1.1,2.1,3.1,4.1]
    --
    -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromTo (1.1, 4.6)
    -- [1.1,2.1,3.1,4.1,5.1]
    --
    -- @
    --
    -- /Pre-release/
    enumerateFromTo :: Monad m => Unfold m (a, a) a

    -- | Unfolds @(from, then)@ generating a stream whose first element is
    -- @from@ and the successive elements are in increments of @then@.  Enumeration
    -- can occur downwards or upwards depending on whether @then@ comes before or
    -- after @from@. For 'Bounded' types the stream ends when 'maxBound' is
    -- reached, for unbounded types it keeps enumerating infinitely.
    --
    -- >>> import qualified Streamly.Prelude as Stream
    -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
    --
    -- @
    -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFromThen (0, 2)
    -- [0,2,4,6]
    --
    -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFromThen (0,(-2))
    -- [0,-2,-4,-6]
    --
    -- @
    --
    -- /Pre-release/
    enumerateFromThen :: Monad m => Unfold m (a, a) a

    -- | Unfolds @(from, then, to)@ generating a finite stream whose first element
    -- is @from@ and the successive elements are in increments of @then@ up to
    -- @to@. Enumeration can occur downwards or upwards depending on whether @then@
    -- comes before or after @from@.
    --
    -- >>> import qualified Streamly.Prelude as Stream
    -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
    --
    -- @
    -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromThenTo (0, 2, 6)
    -- [0,2,4,6]
    --
    -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromThenTo (0, (-2), (-6))
    -- [0,-2,-4,-6]
    --
    -- @
    --
    -- /Pre-release/
    enumerateFromThenTo :: Monad m => Unfold m (a, a, a) a

-------------------------------------------------------------------------------
-- Enumerable Instances
-------------------------------------------------------------------------------
--
-- For Enum types smaller than or equal to Int size.
#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE)           \
instance Enumerable SMALL_TYPE where {                 \
    {-# INLINE enumerateFrom #-};                      \
    enumerateFrom = enumerateFromSmallBounded;         \
    {-# INLINE enumerateFromThen #-};                  \
    enumerateFromThen = enumerateFromThenSmallBounded; \
    {-# INLINE enumerateFromTo #-};                    \
    enumerateFromTo = enumerateFromToSmall;            \
    {-# INLINE enumerateFromThenTo #-};                \
    enumerateFromThenTo = enumerateFromThenToSmall }

ENUMERABLE_BOUNDED_SMALL(())
ENUMERABLE_BOUNDED_SMALL(Bool)
ENUMERABLE_BOUNDED_SMALL(Ordering)
ENUMERABLE_BOUNDED_SMALL(Char)

-- For bounded Integral Enum types, may be larger than Int.
#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE)          \
instance Enumerable INTEGRAL_TYPE where {                   \
    {-# INLINE enumerateFrom #-};                           \
    enumerateFrom = enumerateFromIntegralBounded;           \
    {-# INLINE enumerateFromThen #-};                       \
    enumerateFromThen = enumerateFromThenIntegralBounded;   \
    {-# INLINE enumerateFromTo #-};                         \
    enumerateFromTo = enumerateFromToIntegralBounded;       \
    {-# INLINE enumerateFromThenTo #-};                     \
    enumerateFromThenTo = enumerateFromThenToIntegralBounded }

ENUMERABLE_BOUNDED_INTEGRAL(Int)
ENUMERABLE_BOUNDED_INTEGRAL(Int8)
ENUMERABLE_BOUNDED_INTEGRAL(Int16)
ENUMERABLE_BOUNDED_INTEGRAL(Int32)
ENUMERABLE_BOUNDED_INTEGRAL(Int64)
ENUMERABLE_BOUNDED_INTEGRAL(Word)
ENUMERABLE_BOUNDED_INTEGRAL(Word8)
ENUMERABLE_BOUNDED_INTEGRAL(Word16)
ENUMERABLE_BOUNDED_INTEGRAL(Word32)
ENUMERABLE_BOUNDED_INTEGRAL(Word64)

-- For unbounded Integral Enum types.
#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE)              \
instance Enumerable INTEGRAL_TYPE where {                         \
    {-# INLINE enumerateFrom #-};                                 \
    enumerateFrom = enumerateFromIntegral;                        \
    {-# INLINE enumerateFromThen #-};                             \
    enumerateFromThen = enumerateFromThenIntegral;                \
    {-# INLINE enumerateFromTo #-};                               \
    enumerateFromTo = enumerateFromToIntegral;                    \
    {-# INLINE enumerateFromThenTo #-};                           \
    enumerateFromThenTo = enumerateFromThenToIntegral }

ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)
ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT)         \
instance (CONSTRAINT) => Enumerable FRACTIONAL_TYPE where {       \
    {-# INLINE enumerateFrom #-};                                 \
    enumerateFrom = enumerateFromFractional;                      \
    {-# INLINE enumerateFromThen #-};                             \
    enumerateFromThen = enumerateFromThenFractional;              \
    {-# INLINE enumerateFromTo #-};                               \
    enumerateFromTo = enumerateFromToFractional;                  \
    {-# INLINE enumerateFromThenTo #-};                           \
    enumerateFromThenTo = enumerateFromThenToFractional }

ENUMERABLE_FRACTIONAL(Float,)
ENUMERABLE_FRACTIONAL(Double,)
ENUMERABLE_FRACTIONAL((Fixed a),HasResolution a)
ENUMERABLE_FRACTIONAL((Ratio a),Integral a)

instance Enumerable a => Enumerable (Identity a) where
    {-# INLINE enumerateFrom #-}
    enumerateFrom :: Unfold m (Identity a) (Identity a)
enumerateFrom =
        (a -> Identity a)
-> Unfold m (Identity a) a -> Unfold m (Identity a) (Identity a)
forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Unfold m a b -> Unfold m a c
map a -> Identity a
forall a. a -> Identity a
Identity (Unfold m (Identity a) a -> Unfold m (Identity a) (Identity a))
-> Unfold m (Identity a) a -> Unfold m (Identity a) (Identity a)
forall a b. (a -> b) -> a -> b
$ (Identity a -> a) -> Unfold m a a -> Unfold m (Identity a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap Identity a -> a
forall a. Identity a -> a
runIdentity Unfold m a a
forall a (m :: * -> *). (Enumerable a, Monad m) => Unfold m a a
enumerateFrom
    {-# INLINE enumerateFromThen #-}
    enumerateFromThen :: Unfold m (Identity a, Identity a) (Identity a)
enumerateFromThen =
        (a -> Identity a)
-> Unfold m (Identity a, Identity a) a
-> Unfold m (Identity a, Identity a) (Identity a)
forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Unfold m a b -> Unfold m a c
map a -> Identity a
forall a. a -> Identity a
Identity (Unfold m (Identity a, Identity a) a
 -> Unfold m (Identity a, Identity a) (Identity a))
-> Unfold m (Identity a, Identity a) a
-> Unfold m (Identity a, Identity a) (Identity a)
forall a b. (a -> b) -> a -> b
$ ((Identity a, Identity a) -> (a, a))
-> Unfold m (a, a) a -> Unfold m (Identity a, Identity a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap ((Identity a -> a)
-> (Identity a -> a) -> (Identity a, Identity a) -> (a, a)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Identity a -> a
forall a. Identity a -> a
runIdentity Identity a -> a
forall a. Identity a -> a
runIdentity) Unfold m (a, a) a
forall a (m :: * -> *).
(Enumerable a, Monad m) =>
Unfold m (a, a) a
enumerateFromThen
    {-# INLINE enumerateFromTo #-}
    enumerateFromTo :: Unfold m (Identity a, Identity a) (Identity a)
enumerateFromTo  =
        (a -> Identity a)
-> Unfold m (Identity a, Identity a) a
-> Unfold m (Identity a, Identity a) (Identity a)
forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Unfold m a b -> Unfold m a c
map a -> Identity a
forall a. a -> Identity a
Identity (Unfold m (Identity a, Identity a) a
 -> Unfold m (Identity a, Identity a) (Identity a))
-> Unfold m (Identity a, Identity a) a
-> Unfold m (Identity a, Identity a) (Identity a)
forall a b. (a -> b) -> a -> b
$ ((Identity a, Identity a) -> (a, a))
-> Unfold m (a, a) a -> Unfold m (Identity a, Identity a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap ((Identity a -> a)
-> (Identity a -> a) -> (Identity a, Identity a) -> (a, a)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Identity a -> a
forall a. Identity a -> a
runIdentity Identity a -> a
forall a. Identity a -> a
runIdentity) Unfold m (a, a) a
forall a (m :: * -> *).
(Enumerable a, Monad m) =>
Unfold m (a, a) a
enumerateFromThen
    {-# INLINE enumerateFromThenTo #-}
    enumerateFromThenTo :: Unfold m (Identity a, Identity a, Identity a) (Identity a)
enumerateFromThenTo  =
        (a -> Identity a)
-> Unfold m (Identity a, Identity a, Identity a) a
-> Unfold m (Identity a, Identity a, Identity a) (Identity a)
forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Unfold m a b -> Unfold m a c
map a -> Identity a
forall a. a -> Identity a
Identity (Unfold m (Identity a, Identity a, Identity a) a
 -> Unfold m (Identity a, Identity a, Identity a) (Identity a))
-> Unfold m (Identity a, Identity a, Identity a) a
-> Unfold m (Identity a, Identity a, Identity a) (Identity a)
forall a b. (a -> b) -> a -> b
$
            ((Identity a, Identity a, Identity a) -> (a, a, a))
-> Unfold m (a, a, a) a
-> Unfold m (Identity a, Identity a, Identity a) a
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap
            (\(Identity a
from, Identity a
next, Identity a
to) ->
                 (Identity a -> a
forall a. Identity a -> a
runIdentity Identity a
from, Identity a -> a
forall a. Identity a -> a
runIdentity Identity a
next, Identity a -> a
forall a. Identity a -> a
runIdentity Identity a
to))
            Unfold m (a, a, a) a
forall a (m :: * -> *).
(Enumerable a, Monad m) =>
Unfold m (a, a, a) a
enumerateFromThenTo