-- |
-- Module      : Streamly.Internal.Data.Refold.Type
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- The 'Fold' type embeds a default initial value, therefore, it is like a
-- 'Monoid' whereas the 'Refold' type has to be supplied with an initial
-- value, therefore, it is more like a 'Semigroup' operation.
--
-- Refolds can be appended to each other or to a fold to build the fold
-- incrementally. This is useful in incremental builder like use cases.
--
-- See the file splitting example in the @streamly-examples@ repository for an
-- application of the 'Refold' type. The 'Fold' type does not perform as well
-- in this situation.
--
-- 'Refold' type is to 'Fold' as 'Unfold' type is to 'Stream'. 'Unfold'
-- provides better optimizaiton than stream in nested operations, similarly,
-- 'Refold' provides better optimization than 'Fold'.
--
module Streamly.Internal.Data.Refold.Type
    (
    -- * Types
      Refold (..)

    -- * Constructors
    , foldl'

    -- * Refolds
    -- ** Accumulators
    , sconcat
    , drainBy
    , iterate

    -- * Combinators
    , lmapM
    , rmapM
    , append
    , take
    )
where

import Control.Monad ((>=>))
import Fusion.Plugin.Types (Fuse(..))
import Streamly.Internal.Data.Fold.Step (Step(..), mapMStep)

import Prelude hiding (Foldable(..), take, iterate)

-- $setup
-- >>> :m
-- >>> import qualified Streamly.Internal.Data.Refold.Type as Refold
-- >>> import qualified Streamly.Internal.Data.Fold as Fold
-- >>> import qualified Streamly.Internal.Data.Stream as Stream

-- All folds in the Fold module should be implemented using Refolds.
--
-- | Like 'Fold' except that the initial state of the accmulator can be
-- generated using a dynamically supplied input. This affords better stream
-- fusion optimization in nested fold operations where the initial fold state
-- is determined based on a dynamic value.
--
-- /Internal/
data Refold m c a b =
  -- | @Fold @ @ step @ @ inject @ @ extract@
  forall s. Refold (s -> a -> m (Step s b)) (c -> m (Step s b)) (s -> m b)

------------------------------------------------------------------------------
-- Left fold constructors
------------------------------------------------------------------------------

-- | Make a consumer from a left fold style pure step function.
--
-- If your 'Fold' returns only 'Partial' (i.e. never returns a 'Done') then you
-- can use @foldl'*@ constructors.
--
-- See also: @Streamly.Data.Fold.foldl'@
--
-- /Internal/
--
{-# INLINE foldl' #-}
foldl' :: Monad m => (b -> a -> b) -> Refold m b a b
foldl' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> Refold m b a b
foldl' b -> a -> b
step =
    (b -> a -> m (Step b b))
-> (b -> m (Step b b)) -> (b -> m b) -> Refold m b a b
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold
        (\b
s a
a -> Step b b -> m (Step b b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step b b -> m (Step b b)) -> Step b b -> m (Step b b)
forall a b. (a -> b) -> a -> b
$ b -> Step b b
forall s b. s -> Step s b
Partial (b -> Step b b) -> b -> Step b b
forall a b. (a -> b) -> a -> b
$ b -> a -> b
step b
s a
a)
        (Step b b -> m (Step b b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step b b -> m (Step b b)) -> (b -> Step b b) -> b -> m (Step b b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Step b b
forall s b. s -> Step s b
Partial)
        b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return

------------------------------------------------------------------------------
-- Mapping on input
------------------------------------------------------------------------------

-- | @lmapM f fold@ maps the monadic function @f@ on the input of the fold.
--
-- /Internal/
{-# INLINE lmapM #-}
lmapM :: Monad m => (a -> m b) -> Refold m c b r -> Refold m c a r
lmapM :: forall (m :: * -> *) a b c r.
Monad m =>
(a -> m b) -> Refold m c b r -> Refold m c a r
lmapM a -> m b
f (Refold s -> b -> m (Step s r)
step c -> m (Step s r)
inject s -> m r
extract) = (s -> a -> m (Step s r))
-> (c -> m (Step s r)) -> (s -> m r) -> Refold m c a r
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold s -> a -> m (Step s r)
step1 c -> m (Step s r)
inject s -> m r
extract

    where

    step1 :: s -> a -> m (Step s r)
step1 s
x a
a = a -> m b
f a
a m b -> (b -> m (Step s r)) -> m (Step s r)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= s -> b -> m (Step s r)
step s
x

------------------------------------------------------------------------------
-- Mapping on the output
------------------------------------------------------------------------------

-- | Map a monadic function on the output of a fold.
--
-- /Internal/
{-# INLINE rmapM #-}
rmapM :: Monad m => (b -> m c) -> Refold m x a b -> Refold m x a c
rmapM :: forall (m :: * -> *) b c x a.
Monad m =>
(b -> m c) -> Refold m x a b -> Refold m x a c
rmapM b -> m c
f (Refold s -> a -> m (Step s b)
step x -> m (Step s b)
inject s -> m b
extract) = (s -> a -> m (Step s c))
-> (x -> m (Step s c)) -> (s -> m c) -> Refold m x a c
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold s -> a -> m (Step s c)
step1 x -> m (Step s c)
inject1 (s -> m b
extract (s -> m b) -> (b -> m c) -> s -> m c
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> b -> m c
f)

    where

    inject1 :: x -> m (Step s c)
inject1 x
x = x -> m (Step s b)
inject x
x m (Step s b) -> (Step s b -> m (Step s c)) -> m (Step s c)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (b -> m c) -> Step s b -> m (Step s c)
forall (m :: * -> *) a b s.
Applicative m =>
(a -> m b) -> Step s a -> m (Step s b)
mapMStep b -> m c
f
    step1 :: s -> a -> m (Step s c)
step1 s
s a
a = s -> a -> m (Step s b)
step s
s a
a m (Step s b) -> (Step s b -> m (Step s c)) -> m (Step s c)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (b -> m c) -> Step s b -> m (Step s c)
forall (m :: * -> *) a b s.
Applicative m =>
(a -> m b) -> Step s a -> m (Step s b)
mapMStep b -> m c
f

------------------------------------------------------------------------------
-- Refolds
------------------------------------------------------------------------------

-- |
--
-- /Internal/
{-# INLINE drainBy #-}
drainBy ::  Monad m => (c -> a -> m b) -> Refold m c a ()
drainBy :: forall (m :: * -> *) c a b.
Monad m =>
(c -> a -> m b) -> Refold m c a ()
drainBy c -> a -> m b
f = (c -> a -> m (Step c ()))
-> (c -> m (Step c ())) -> (c -> m ()) -> Refold m c a ()
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold c -> a -> m (Step c ())
forall {b}. c -> a -> m (Step c b)
step c -> m (Step c ())
forall {a} {b}. a -> m (Step a b)
inject c -> m ()
forall {m :: * -> *} {p}. Monad m => p -> m ()
extract

    where

    inject :: a -> m (Step a b)
inject = Step a b -> m (Step a b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step a b -> m (Step a b)) -> (a -> Step a b) -> a -> m (Step a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Step a b
forall s b. s -> Step s b
Partial

    step :: c -> a -> m (Step c b)
step c
c a
a = c -> a -> m b
f c
c a
a m b -> m (Step c b) -> m (Step c b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step c b -> m (Step c b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (c -> Step c b
forall s b. s -> Step s b
Partial c
c)

    extract :: p -> m ()
extract p
_ = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

------------------------------------------------------------------------------
-- Semigroup
------------------------------------------------------------------------------

-- | Append the elements of an input stream to a provided starting value.
--
-- >>> stream = fmap Data.Monoid.Sum $ Stream.enumerateFromTo 1 10
-- >>> Stream.fold (Fold.fromRefold Refold.sconcat 10) stream
-- Sum {getSum = 65}
--
-- >>> sconcat = Refold.foldl' (<>)
--
-- /Internal/
{-# INLINE sconcat #-}
sconcat :: (Monad m, Semigroup a) => Refold m a a a
sconcat :: forall (m :: * -> *) a. (Monad m, Semigroup a) => Refold m a a a
sconcat = (a -> a -> a) -> Refold m a a a
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> Refold m b a b
foldl' a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

------------------------------------------------------------------------------
-- append
------------------------------------------------------------------------------

-- | Supply the output of the first consumer as input to the second consumer.
--
-- /Internal/
{-# INLINE append #-}
append :: Monad m => Refold m x a b -> Refold m b a b -> Refold m x a b
append :: forall (m :: * -> *) x a b.
Monad m =>
Refold m x a b -> Refold m b a b -> Refold m x a b
append (Refold s -> a -> m (Step s b)
step1 x -> m (Step s b)
inject1 s -> m b
extract1) (Refold s -> a -> m (Step s b)
step2 b -> m (Step s b)
inject2 s -> m b
extract2) =
    (Either s s -> a -> m (Step (Either s s) b))
-> (x -> m (Step (Either s s) b))
-> (Either s s -> m b)
-> Refold m x a b
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold Either s s -> a -> m (Step (Either s s) b)
step x -> m (Step (Either s s) b)
inject Either s s -> m b
extract

    where

    goLeft :: Step a b -> m (Step (Either a s) b)
goLeft Step a b
r = do
        case Step a b
r of
            Partial a
s -> Step (Either a s) b -> m (Step (Either a s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either a s) b -> m (Step (Either a s) b))
-> Step (Either a s) b -> m (Step (Either a s) b)
forall a b. (a -> b) -> a -> b
$ Either a s -> Step (Either a s) b
forall s b. s -> Step s b
Partial (Either a s -> Step (Either a s) b)
-> Either a s -> Step (Either a s) b
forall a b. (a -> b) -> a -> b
$ a -> Either a s
forall a b. a -> Either a b
Left a
s
            Done b
b -> do
                Step s b
r1 <- b -> m (Step s b)
inject2 b
b
                Step (Either a s) b -> m (Step (Either a s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either a s) b -> m (Step (Either a s) b))
-> Step (Either a s) b -> m (Step (Either a s) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r1 of
                    Partial s
s -> Either a s -> Step (Either a s) b
forall s b. s -> Step s b
Partial (Either a s -> Step (Either a s) b)
-> Either a s -> Step (Either a s) b
forall a b. (a -> b) -> a -> b
$ s -> Either a s
forall a b. b -> Either a b
Right s
s
                    Done b
b1 -> b -> Step (Either a s) b
forall s b. b -> Step s b
Done b
b1

    inject :: x -> m (Step (Either s s) b)
inject x
x = x -> m (Step s b)
inject1 x
x m (Step s b)
-> (Step s b -> m (Step (Either s s) b)) -> m (Step (Either s s) b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Step s b -> m (Step (Either s s) b)
forall {a}. Step a b -> m (Step (Either a s) b)
goLeft

    step :: Either s s -> a -> m (Step (Either s s) b)
step (Left s
s) a
a = s -> a -> m (Step s b)
step1 s
s a
a m (Step s b)
-> (Step s b -> m (Step (Either s s) b)) -> m (Step (Either s s) b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Step s b -> m (Step (Either s s) b)
forall {a}. Step a b -> m (Step (Either a s) b)
goLeft

    step (Right s
s) a
a = do
        Step s b
r <- s -> a -> m (Step s b)
step2 s
s a
a
        case Step s b
r of
            Partial s
s1 -> Step (Either s s) b -> m (Step (Either s s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s s) b -> m (Step (Either s s) b))
-> Step (Either s s) b -> m (Step (Either s s) b)
forall a b. (a -> b) -> a -> b
$ Either s s -> Step (Either s s) b
forall s b. s -> Step s b
Partial (s -> Either s s
forall a b. b -> Either a b
Right s
s1)
            Done b
b -> Step (Either s s) b -> m (Step (Either s s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s s) b -> m (Step (Either s s) b))
-> Step (Either s s) b -> m (Step (Either s s) b)
forall a b. (a -> b) -> a -> b
$ b -> Step (Either s s) b
forall s b. b -> Step s b
Done b
b

    extract :: Either s s -> m b
extract (Left s
s) = s -> m b
extract1 s
s
    extract (Right s
s) = s -> m b
extract2 s
s

-- | Keep running the same consumer over and over again on the input, feeding
-- the output of the previous run to the next.
--
-- /Internal/
iterate :: Monad m => Refold m b a b -> Refold m b a b
iterate :: forall (m :: * -> *) b a.
Monad m =>
Refold m b a b -> Refold m b a b
iterate (Refold s -> a -> m (Step s b)
step1 b -> m (Step s b)
inject1 s -> m b
extract1) =
    (s -> a -> m (Step s b))
-> (b -> m (Step s b)) -> (s -> m b) -> Refold m b a b
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold s -> a -> m (Step s b)
forall {b}. s -> a -> m (Step s b)
step b -> m (Step s b)
forall {b}. b -> m (Step s b)
inject s -> m b
extract1

    where

    go :: Step s b -> m (Step s b)
go Step s b
r =
        case Step s b
r of
            Partial s
s -> Step s b -> m (Step s b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step s b -> m (Step s b)) -> Step s b -> m (Step s b)
forall a b. (a -> b) -> a -> b
$ s -> Step s b
forall s b. s -> Step s b
Partial s
s
            Done b
b -> b -> m (Step s b)
inject b
b

    inject :: b -> m (Step s b)
inject b
x = b -> m (Step s b)
inject1 b
x m (Step s b) -> (Step s b -> m (Step s b)) -> m (Step s b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Step s b -> m (Step s b)
go

    step :: s -> a -> m (Step s b)
step s
s a
a = s -> a -> m (Step s b)
step1 s
s a
a m (Step s b) -> (Step s b -> m (Step s b)) -> m (Step s b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Step s b -> m (Step s b)
forall {b}. Step s b -> m (Step s b)
go

------------------------------------------------------------------------------
-- Transformation
------------------------------------------------------------------------------

-- Required to fuse "take" with "many" in "chunksOf", for ghc-9.x
{-# ANN type Tuple'Fused Fuse #-}
data Tuple'Fused a b = Tuple'Fused !a !b deriving Int -> Tuple'Fused a b -> ShowS
[Tuple'Fused a b] -> ShowS
Tuple'Fused a b -> String
(Int -> Tuple'Fused a b -> ShowS)
-> (Tuple'Fused a b -> String)
-> ([Tuple'Fused a b] -> ShowS)
-> Show (Tuple'Fused a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> Tuple'Fused a b -> ShowS
forall a b. (Show a, Show b) => [Tuple'Fused a b] -> ShowS
forall a b. (Show a, Show b) => Tuple'Fused a b -> String
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> Tuple'Fused a b -> ShowS
showsPrec :: Int -> Tuple'Fused a b -> ShowS
$cshow :: forall a b. (Show a, Show b) => Tuple'Fused a b -> String
show :: Tuple'Fused a b -> String
$cshowList :: forall a b. (Show a, Show b) => [Tuple'Fused a b] -> ShowS
showList :: [Tuple'Fused a b] -> ShowS
Show

-- | Take at most @n@ input elements and fold them using the supplied fold. A
-- negative count is treated as 0.
--
-- /Internal/
{-# INLINE take #-}
take :: Monad m => Int -> Refold m x a b -> Refold m x a b
take :: forall (m :: * -> *) x a b.
Monad m =>
Int -> Refold m x a b -> Refold m x a b
take Int
n (Refold s -> a -> m (Step s b)
fstep x -> m (Step s b)
finject s -> m b
fextract) = (Tuple'Fused Int s -> a -> m (Step (Tuple'Fused Int s) b))
-> (x -> m (Step (Tuple'Fused Int s) b))
-> (Tuple'Fused Int s -> m b)
-> Refold m x a b
forall (m :: * -> *) c a b s.
(s -> a -> m (Step s b))
-> (c -> m (Step s b)) -> (s -> m b) -> Refold m c a b
Refold Tuple'Fused Int s -> a -> m (Step (Tuple'Fused Int s) b)
step x -> m (Step (Tuple'Fused Int s) b)
forall {a}. Num a => x -> m (Step (Tuple'Fused a s) b)
inject Tuple'Fused Int s -> m b
forall {a}. Tuple'Fused a s -> m b
extract

    where

    inject :: x -> m (Step (Tuple'Fused a s) b)
inject x
x = do
        Step s b
res <- x -> m (Step s b)
finject x
x
        case Step s b
res of
            Partial s
s ->
                if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                then Step (Tuple'Fused a s) b -> m (Step (Tuple'Fused a s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple'Fused a s) b -> m (Step (Tuple'Fused a s) b))
-> Step (Tuple'Fused a s) b -> m (Step (Tuple'Fused a s) b)
forall a b. (a -> b) -> a -> b
$ Tuple'Fused a s -> Step (Tuple'Fused a s) b
forall s b. s -> Step s b
Partial (Tuple'Fused a s -> Step (Tuple'Fused a s) b)
-> Tuple'Fused a s -> Step (Tuple'Fused a s) b
forall a b. (a -> b) -> a -> b
$ a -> s -> Tuple'Fused a s
forall a b. a -> b -> Tuple'Fused a b
Tuple'Fused a
0 s
s
                else b -> Step (Tuple'Fused a s) b
forall s b. b -> Step s b
Done (b -> Step (Tuple'Fused a s) b)
-> m b -> m (Step (Tuple'Fused a s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
fextract s
s
            Done b
b -> Step (Tuple'Fused a s) b -> m (Step (Tuple'Fused a s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple'Fused a s) b -> m (Step (Tuple'Fused a s) b))
-> Step (Tuple'Fused a s) b -> m (Step (Tuple'Fused a s) b)
forall a b. (a -> b) -> a -> b
$ b -> Step (Tuple'Fused a s) b
forall s b. b -> Step s b
Done b
b

    step :: Tuple'Fused Int s -> a -> m (Step (Tuple'Fused Int s) b)
step (Tuple'Fused Int
i s
r) a
a = do
        Step s b
res <- s -> a -> m (Step s b)
fstep s
r a
a
        case Step s b
res of
            Partial s
sres -> do
                let i1 :: Int
i1 = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                    s1 :: Tuple'Fused Int s
s1 = Int -> s -> Tuple'Fused Int s
forall a b. a -> b -> Tuple'Fused a b
Tuple'Fused Int
i1 s
sres
                if Int
i1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
                then Step (Tuple'Fused Int s) b -> m (Step (Tuple'Fused Int s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple'Fused Int s) b -> m (Step (Tuple'Fused Int s) b))
-> Step (Tuple'Fused Int s) b -> m (Step (Tuple'Fused Int s) b)
forall a b. (a -> b) -> a -> b
$ Tuple'Fused Int s -> Step (Tuple'Fused Int s) b
forall s b. s -> Step s b
Partial Tuple'Fused Int s
s1
                else b -> Step (Tuple'Fused Int s) b
forall s b. b -> Step s b
Done (b -> Step (Tuple'Fused Int s) b)
-> m b -> m (Step (Tuple'Fused Int s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
fextract s
sres
            Done b
bres -> Step (Tuple'Fused Int s) b -> m (Step (Tuple'Fused Int s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple'Fused Int s) b -> m (Step (Tuple'Fused Int s) b))
-> Step (Tuple'Fused Int s) b -> m (Step (Tuple'Fused Int s) b)
forall a b. (a -> b) -> a -> b
$ b -> Step (Tuple'Fused Int s) b
forall s b. b -> Step s b
Done b
bres

    extract :: Tuple'Fused a s -> m b
extract (Tuple'Fused a
_ s
r) = s -> m b
fextract s
r