-- |
-- Module      : Streamly.Internal.Data.Fold.Chunked
-- Copyright   : (c) 2021 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Use "Streamly.Data.Parser" instead.
--
-- Fold a stream of foreign arrays.  @Fold m a b@ in this module works
-- on a stream of "Array a" and produces an output of type @b@.
--
-- Though @Fold m a b@ in this module works on a stream of @Array a@ it is
-- different from @Data.Fold m (Array a) b@.  While the latter works on arrays
-- as a whole treating them as atomic elements, the folds in this module can
-- work on the stream of arrays as if it is an element stream with all the
-- arrays coalesced together. This module allows adapting the element stream
-- folds in Data.Fold to correctly work on an array stream as if it is an
-- element stream. For example:
--
-- >> import qualified Streamly.Data.Fold as Fold
-- >> import qualified Streamly.Internal.Data.Array.Stream as ArrayStream
-- >> import qualified Streamly.Internal.Data.Fold.Chunked as ChunkFold
-- >> import qualified Streamly.Data.Stream as Stream
-- >> import qualified Streamly.Data.StreamK as StreamK
--
-- >> f = ChunkFold.fromFold (Fold.take 7 Fold.toList)
-- >> s = Stream.chunksOf 5 $ Stream.fromList "hello world"
-- >> ArrayStream.runArrayFold f (StreamK.fromStream s)
-- Right "hello w"
--
module Streamly.Internal.Data.Fold.Chunked
    {-# DEPRECATED "Please use Streamly.Data.Parser instead." #-}
    (
      ChunkFold (..)

    -- * Construction
    , fromFold
    , adaptFold
    , fromParser
    , fromParserD

    -- * Mapping
    , rmapM

    -- * Applicative
    , fromPure
    , fromEffect
    , splitWith

    -- * Monad
    , concatMap

    -- * Combinators
    , take
    )
where

#include "ArrayMacros.h"

#if !MIN_VERSION_base(4,18,0)
import Control.Applicative (liftA2)
#endif
import Control.Exception (assert)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Bifunctor (first)
import Data.Proxy (Proxy(..))
import Streamly.Internal.Data.Unbox (Unbox(..))
import GHC.Types (SPEC(..))
import Streamly.Internal.Data.Array (Array(..))
import Streamly.Internal.Data.Parser (Initial(..), Step(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

import qualified Streamly.Internal.Data.Array.Type as Array
import qualified Streamly.Internal.Data.Fold as Fold
import qualified Streamly.Internal.Data.Parser as ParserD
import qualified Streamly.Internal.Data.Parser as Parser

import Prelude hiding (concatMap, take)

-- | Array stream fold.
--
-- An array stream fold is basically an array stream "Parser" that does not
-- fail.  In case of array stream folds the count in 'Partial', 'Continue' and
-- 'Done' is a count of elements that includes the leftover element count in
-- the array that is currently being processed by the parser. If none of the
-- elements is consumed by the parser the count is at least the whole array
-- length. If the whole array is consumed by the parser then the count will be
-- 0.
--
-- /Pre-release/
--
newtype ChunkFold m a b = ChunkFold (ParserD.Parser (Array a) m b)

-------------------------------------------------------------------------------
-- Constructing array stream folds from element folds and parsers
-------------------------------------------------------------------------------

-- | Convert an element 'Fold' into an array stream fold.
--
-- /Pre-release/
{-# INLINE fromFold #-}
fromFold :: forall m a b. (MonadIO m, Unbox a) =>
    Fold.Fold m a b -> ChunkFold m a b
fromFold :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Fold m a b -> ChunkFold m a b
fromFold (Fold.Fold s -> a -> m (Step s b)
fstep m (Step s b)
finitial s -> m b
_ s -> m b
ffinal) =
    Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold ((s -> Array a -> m (Step s b))
-> m (Initial s b) -> (s -> m (Step s b)) -> Parser (Array a) m b
forall a (m :: * -> *) b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m (Step s b)) -> Parser a m b
ParserD.Parser s -> Array a -> m (Step s b)
forall {a}. s -> Array a -> m (Step s b)
step m (Initial s b)
initial s -> m (Step s b)
forall {s}. s -> m (Step s b)
extract)

    where

    initial :: m (Initial s b)
initial = do
        Step s b
res <- m (Step s b)
finitial
        Initial s b -> m (Initial s b)
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Initial s b -> m (Initial s b)) -> Initial s b -> m (Initial s b)
forall a b. (a -> b) -> a -> b
$ case Step s b
res of
                  Fold.Partial s
s1 -> s -> Initial s b
forall s b. s -> Initial s b
IPartial s
s1
                  Fold.Done b
b -> b -> Initial s b
forall s b. b -> Initial s b
IDone b
b

    step :: s -> Array a -> m (Step s b)
step s
s (Array MutByteArray
contents Int
start Int
end) = do
        SPEC -> Int -> s -> m (Step s b)
goArray SPEC
SPEC Int
start s
s

        where

        goArray :: SPEC -> Int -> s -> m (Step s b)
goArray !SPEC
_ !Int
cur !s
fs | Int
cur Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end = do
            Bool -> m () -> m ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
cur Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            Step s b -> m (Step s b)
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
$ Int -> s -> Step s b
forall s b. Int -> s -> Step s b
Partial Int
0 s
fs
        goArray !SPEC
_ !Int
cur !s
fs = do
            a
x <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Int -> MutByteArray -> IO a
forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
cur MutByteArray
contents
            Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
            let elemSize :: Int
elemSize = SIZE_OF(a)
                next :: Int
next = INDEX_NEXT(cur,a)
            case Step s b
res of
                Fold.Done b
b ->
                    Step s b -> m (Step s b)
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
$ Int -> b -> Step s b
forall s b. Int -> b -> Step s b
Done ((Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
next) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
elemSize) b
b
                Fold.Partial s
fs1 ->
                    SPEC -> Int -> s -> m (Step s b)
goArray SPEC
SPEC Int
next s
fs1

    extract :: s -> m (Step s b)
extract = (b -> Step s b) -> m b -> m (Step s b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> b -> Step s b
forall s b. Int -> b -> Step s b
Done Int
0) (m b -> m (Step s b)) -> (s -> m b) -> s -> m (Step s b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> m b
ffinal

-- | Convert an element 'ParserD.Parser' into an array stream fold. If the
-- parser fails the fold would throw an exception.
--
-- /Pre-release/
{-# INLINE fromParserD #-}
fromParserD :: forall m a b. (MonadIO m, Unbox a) =>
    ParserD.Parser a m b -> ChunkFold m a b
fromParserD :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Parser a m b -> ChunkFold m a b
fromParserD (ParserD.Parser s -> a -> m (Step s b)
step1 m (Initial s b)
initial1 s -> m (Step s b)
extract1) =
    Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold ((s -> Array a -> m (Step s b))
-> m (Initial s b) -> (s -> m (Step s b)) -> Parser (Array a) m b
forall a (m :: * -> *) b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m (Step s b)) -> Parser a m b
ParserD.Parser s -> Array a -> m (Step s b)
forall {a}. s -> Array a -> m (Step s b)
step m (Initial s b)
initial1 s -> m (Step s b)
extract1)

    where

    step :: s -> Array a -> m (Step s b)
step s
s (Array MutByteArray
contents Int
start Int
end) = do
        if Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
end
        then Step s b -> m (Step s b)
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
$ Int -> s -> Step s b
forall s b. Int -> s -> Step s b
Continue Int
0 s
s
        else SPEC -> Int -> s -> m (Step s b)
goArray SPEC
SPEC Int
start s
s

        where

        {-# INLINE partial #-}
        partial :: Int
-> Int
-> Int
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Int
cur Int
next Int
elemSize Int -> s -> Step s b
st Int
n s
fs1 = do
            let next1 :: Int
next1 = Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
elemSize)
            if Int
next1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
start Bool -> Bool -> Bool
&& Int
cur Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end
            then SPEC -> Int -> s -> m (Step s b)
goArray SPEC
SPEC Int
next1 s
fs1
            else Step s b -> m (Step s b)
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
$ Int -> s -> Step s b
st (Int
arrRem Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) s
fs1

        goArray :: SPEC -> Int -> s -> m (Step s b)
goArray !SPEC
_ !Int
cur !s
fs = do
            a
x <- IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Int -> MutByteArray -> IO a
forall a. Unbox a => Int -> MutByteArray -> IO a
peekAt Int
cur MutByteArray
contents
            Step s b
res <- s -> a -> m (Step s b)
step1 s
fs a
x
            let elemSize :: Int
elemSize = SIZE_OF(a)
                next :: Int
next = INDEX_NEXT(cur,a)
                arrRem :: Int
arrRem = (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
next) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
elemSize
            case Step s b
res of
                ParserD.Done Int
n b
b -> do
                    Step s b -> m (Step s b)
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
$ Int -> b -> Step s b
forall s b. Int -> b -> Step s b
Done (Int
arrRem Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) b
b
                ParserD.Partial Int
n s
fs1 ->
                    Int
-> Int
-> Int
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Int
cur Int
next Int
elemSize Int -> s -> Step s b
forall s b. Int -> s -> Step s b
Partial Int
n s
fs1
                ParserD.Continue Int
n s
fs1 -> do
                    Int
-> Int
-> Int
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Int
cur Int
next Int
elemSize Int -> s -> Step s b
forall s b. Int -> s -> Step s b
Continue Int
n s
fs1
                Error String
err -> Step s b -> m (Step s b)
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
$ String -> Step s b
forall s b. String -> Step s b
Error String
err

-- | Convert an element 'Parser.Parser' into an array stream fold. If the
-- parser fails the fold would throw an exception.
--
-- /Pre-release/
{-# INLINE fromParser #-}
fromParser :: forall m a b. (MonadIO m, Unbox a) =>
    Parser.Parser a m b -> ChunkFold m a b
fromParser :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Parser a m b -> ChunkFold m a b
fromParser = Parser a m b -> ChunkFold m a b
forall (m :: * -> *) a b.
(MonadIO m, Unbox a) =>
Parser a m b -> ChunkFold m a b
fromParserD

-- | Adapt an array stream fold.
--
-- /Pre-release/
{-# INLINE adaptFold #-}
adaptFold :: forall m a b. (MonadIO m) =>
    Fold.Fold m (Array a) b -> ChunkFold m a b
adaptFold :: forall (m :: * -> *) a b.
MonadIO m =>
Fold m (Array a) b -> ChunkFold m a b
adaptFold Fold m (Array a) b
f = Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m b -> ChunkFold m a b)
-> Parser (Array a) m b -> ChunkFold m a b
forall a b. (a -> b) -> a -> b
$ Fold m (Array a) b -> Parser (Array a) m b
forall (m :: * -> *) a b. Monad m => Fold m a b -> Parser a m b
ParserD.fromFold Fold m (Array a) b
f

-------------------------------------------------------------------------------
-- Functor
-------------------------------------------------------------------------------

-- | Maps a function over the result of fold.
--
-- /Pre-release/
instance Functor m => Functor (ChunkFold m a) where
    {-# INLINE fmap #-}
    fmap :: forall a b. (a -> b) -> ChunkFold m a a -> ChunkFold m a b
fmap a -> b
f (ChunkFold Parser (Array a) m a
p) = Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m b -> ChunkFold m a b)
-> Parser (Array a) m b -> ChunkFold m a b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Parser (Array a) m a -> Parser (Array a) m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Parser (Array a) m a
p

-- | Map a monadic function on the output of a fold.
--
-- /Pre-release/
{-# INLINE rmapM #-}
rmapM :: Monad m => (b -> m c) -> ChunkFold m a b -> ChunkFold m a c
rmapM :: forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> ChunkFold m a b -> ChunkFold m a c
rmapM b -> m c
f (ChunkFold Parser (Array a) m b
p) = Parser (Array a) m c -> ChunkFold m a c
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m c -> ChunkFold m a c)
-> Parser (Array a) m c -> ChunkFold m a c
forall a b. (a -> b) -> a -> b
$ (b -> m c) -> Parser (Array a) m b -> Parser (Array a) m c
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Parser a m b -> Parser a m c
ParserD.rmapM b -> m c
f Parser (Array a) m b
p

-------------------------------------------------------------------------------
-- Sequential applicative
-------------------------------------------------------------------------------

-- | A fold that always yields a pure value without consuming any input.
--
-- /Pre-release/
--
{-# INLINE fromPure #-}
fromPure :: Monad m => b -> ChunkFold m a b
fromPure :: forall (m :: * -> *) b a. Monad m => b -> ChunkFold m a b
fromPure = Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m b -> ChunkFold m a b)
-> (b -> Parser (Array a) m b) -> b -> ChunkFold m a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Parser (Array a) m b
forall (m :: * -> *) b a. Monad m => b -> Parser a m b
ParserD.fromPure

-- | A fold that always yields the result of an effectful action without
-- consuming any input.
--
-- /Pre-release/
--
{-# INLINE fromEffect #-}
fromEffect :: Monad m => m b -> ChunkFold m a b
fromEffect :: forall (m :: * -> *) b a. Monad m => m b -> ChunkFold m a b
fromEffect = Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m b -> ChunkFold m a b)
-> (m b -> Parser (Array a) m b) -> m b -> ChunkFold m a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m b -> Parser (Array a) m b
forall (m :: * -> *) b a. Monad m => m b -> Parser a m b
ParserD.fromEffect

-- | Applies two folds sequentially on the input stream and combines their
-- results using the supplied function.
--
-- /Pre-release/
{-# INLINE split_ #-}
split_ :: Monad m =>
    ChunkFold m x a -> ChunkFold m x b -> ChunkFold m x b
split_ :: forall (m :: * -> *) x a b.
Monad m =>
ChunkFold m x a -> ChunkFold m x b -> ChunkFold m x b
split_ (ChunkFold Parser (Array x) m a
p1) (ChunkFold Parser (Array x) m b
p2) =
    Parser (Array x) m b -> ChunkFold m x b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array x) m b -> ChunkFold m x b)
-> Parser (Array x) m b -> ChunkFold m x b
forall a b. (a -> b) -> a -> b
$ Parser (Array x) m a
-> Parser (Array x) m b -> Parser (Array x) m b
forall (m :: * -> *) x a b.
Monad m =>
Parser x m a -> Parser x m b -> Parser x m b
ParserD.noErrorUnsafeSplit_ Parser (Array x) m a
p1 Parser (Array x) m b
p2

-- | Applies two folds sequentially on the input stream and combines their
-- results using the supplied function.
--
-- /Pre-release/
{-# INLINE splitWith #-}
splitWith :: Monad m
    => (a -> b -> c) -> ChunkFold m x a -> ChunkFold m x b -> ChunkFold m x c
splitWith :: forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c)
-> ChunkFold m x a -> ChunkFold m x b -> ChunkFold m x c
splitWith a -> b -> c
f (ChunkFold Parser (Array x) m a
p1) (ChunkFold Parser (Array x) m b
p2) =
    Parser (Array x) m c -> ChunkFold m x c
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array x) m c -> ChunkFold m x c)
-> Parser (Array x) m c -> ChunkFold m x c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c)
-> Parser (Array x) m a
-> Parser (Array x) m b
-> Parser (Array x) m c
forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Parser x m a -> Parser x m b -> Parser x m c
ParserD.noErrorUnsafeSplitWith a -> b -> c
f Parser (Array x) m a
p1 Parser (Array x) m b
p2

-- | 'Applicative' form of 'splitWith'.
-- > (<*>) = splitWith id
instance Monad m => Applicative (ChunkFold m a) where
    {-# INLINE pure #-}
    pure :: forall a. a -> ChunkFold m a a
pure = a -> ChunkFold m a a
forall (m :: * -> *) b a. Monad m => b -> ChunkFold m a b
fromPure

    {-# INLINE (<*>) #-}
    <*> :: forall a b.
ChunkFold m a (a -> b) -> ChunkFold m a a -> ChunkFold m a b
(<*>) = ((a -> b) -> a -> b)
-> ChunkFold m a (a -> b) -> ChunkFold m a a -> ChunkFold m a b
forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c)
-> ChunkFold m x a -> ChunkFold m x b -> ChunkFold m x c
splitWith (a -> b) -> a -> b
forall a. a -> a
id

    {-# INLINE (*>) #-}
    *> :: forall a b. ChunkFold m a a -> ChunkFold m a b -> ChunkFold m a b
(*>) = ChunkFold m a a -> ChunkFold m a b -> ChunkFold m a b
forall (m :: * -> *) x a b.
Monad m =>
ChunkFold m x a -> ChunkFold m x b -> ChunkFold m x b
split_

    {-# INLINE liftA2 #-}
    liftA2 :: forall a b c.
(a -> b -> c)
-> ChunkFold m a a -> ChunkFold m a b -> ChunkFold m a c
liftA2 a -> b -> c
f ChunkFold m a a
x = ChunkFold m a (b -> c) -> ChunkFold m a b -> ChunkFold m a c
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) ((a -> b -> c) -> ChunkFold m a a -> ChunkFold m a (b -> c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b -> c
f ChunkFold m a a
x)

-------------------------------------------------------------------------------
-- Monad
-------------------------------------------------------------------------------

-- XXX This should be implemented using CPS
--
-- | Applies a fold on the input stream, generates the next fold from the
-- output of the previously applied fold and then applies that fold.
--
-- /Pre-release/
--
{-# INLINE concatMap #-}
concatMap :: Monad m =>
    (b -> ChunkFold m a c) -> ChunkFold m a b -> ChunkFold m a c
concatMap :: forall (m :: * -> *) b a c.
Monad m =>
(b -> ChunkFold m a c) -> ChunkFold m a b -> ChunkFold m a c
concatMap b -> ChunkFold m a c
func (ChunkFold Parser (Array a) m b
p) =
    let f :: b -> Parser (Array a) m c
f b
x = let ChunkFold Parser (Array a) m c
y = b -> ChunkFold m a c
func b
x in Parser (Array a) m c
y
     in Parser (Array a) m c -> ChunkFold m a c
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m c -> ChunkFold m a c)
-> Parser (Array a) m c -> ChunkFold m a c
forall a b. (a -> b) -> a -> b
$ (b -> Parser (Array a) m c)
-> Parser (Array a) m b -> Parser (Array a) m c
forall (m :: * -> *) b a c.
Monad m =>
(b -> Parser a m c) -> Parser a m b -> Parser a m c
ParserD.noErrorUnsafeConcatMap b -> Parser (Array a) m c
f Parser (Array a) m b
p

-- | Monad instance applies folds sequentially. Next fold can depend on the
-- output of the previous fold. See 'concatMap'.
--
-- > (>>=) = flip concatMap
instance Monad m => Monad (ChunkFold m a) where
    {-# INLINE return #-}
    return :: forall a. a -> ChunkFold m a a
return = a -> ChunkFold m a a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    {-# INLINE (>>=) #-}
    >>= :: forall a b.
ChunkFold m a a -> (a -> ChunkFold m a b) -> ChunkFold m a b
(>>=) = ((a -> ChunkFold m a b) -> ChunkFold m a a -> ChunkFold m a b)
-> ChunkFold m a a -> (a -> ChunkFold m a b) -> ChunkFold m a b
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> ChunkFold m a b) -> ChunkFold m a a -> ChunkFold m a b
forall (m :: * -> *) b a c.
Monad m =>
(b -> ChunkFold m a c) -> ChunkFold m a b -> ChunkFold m a c
concatMap

    {-# INLINE (>>) #-}
    >> :: forall a b. ChunkFold m a a -> ChunkFold m a b -> ChunkFold m a b
(>>) = ChunkFold m a a -> ChunkFold m a b -> ChunkFold m a b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

-------------------------------------------------------------------------------
-- Array to Array folds
-------------------------------------------------------------------------------

-- | Take @n@ array elements (@a@) from a stream of arrays (@Array a@).
{-# INLINE take #-}
take :: forall m a b. (Monad m, Unbox a) =>
    Int -> ChunkFold m a b -> ChunkFold m a b
take :: forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Int -> ChunkFold m a b -> ChunkFold m a b
take Int
n (ChunkFold (ParserD.Parser s -> Array a -> m (Step s b)
step1 m (Initial s b)
initial1 s -> m (Step s b)
extract1)) =
    Parser (Array a) m b -> ChunkFold m a b
forall (m :: * -> *) a b. Parser (Array a) m b -> ChunkFold m a b
ChunkFold (Parser (Array a) m b -> ChunkFold m a b)
-> Parser (Array a) m b -> ChunkFold m a b
forall a b. (a -> b) -> a -> b
$ (Tuple' Int s -> Array a -> m (Step (Tuple' Int s) b))
-> m (Initial (Tuple' Int s) b)
-> (Tuple' Int s -> m (Step (Tuple' Int s) b))
-> Parser (Array a) m b
forall a (m :: * -> *) b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m (Step s b)) -> Parser a m b
ParserD.Parser Tuple' Int s -> Array a -> m (Step (Tuple' Int s) b)
step m (Initial (Tuple' Int s) b)
initial Tuple' Int s -> m (Step (Tuple' Int s) b)
forall {a}. Tuple' a s -> m (Step (Tuple' a s) b)
extract

    where

    -- XXX Need to make the Initial type Step to remove this
    iextract :: s -> m (Initial s b)
iextract s
s = do
        Step s b
r <- s -> m (Step s b)
extract1 s
s
        Initial s b -> m (Initial s b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Initial s b -> m (Initial s b)) -> Initial s b -> m (Initial s b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Done Int
_ b
b -> b -> Initial s b
forall s b. b -> Initial s b
IDone b
b
            Error String
err -> String -> Initial s b
forall s b. String -> Initial s b
IError String
err
            Step s b
_ -> String -> Initial s b
forall a. (?callStack::CallStack) => String -> a
error String
"Bug: ChunkFold take invalid state in initial"

    initial :: m (Initial (Tuple' Int s) b)
initial = do
        Initial s b
res <- m (Initial s b)
initial1
        case Initial s b
res of
            IPartial s
s ->
                if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                then Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b))
-> Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ Tuple' Int s -> Initial (Tuple' Int s) b
forall s b. s -> Initial s b
IPartial (Tuple' Int s -> Initial (Tuple' Int s) b)
-> Tuple' Int s -> Initial (Tuple' Int s) b
forall a b. (a -> b) -> a -> b
$ Int -> s -> Tuple' Int s
forall a b. a -> b -> Tuple' a b
Tuple' Int
n s
s
                else s -> m (Initial (Tuple' Int s) b)
forall {s}. s -> m (Initial s b)
iextract s
s
            IDone b
b -> Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b))
-> Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ b -> Initial (Tuple' Int s) b
forall s b. b -> Initial s b
IDone b
b
            IError String
err -> Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b))
-> Initial (Tuple' Int s) b -> m (Initial (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ String -> Initial (Tuple' Int s) b
forall s b. String -> Initial s b
IError String
err

    {-# INLINE partial #-}
    partial :: t
-> (t -> Tuple' t s -> Step (Tuple' t s) b)
-> t
-> s
-> m (Step (Tuple' t s) b)
partial t
i1 t -> Tuple' t s -> Step (Tuple' t s) b
st t
j s
s =
        let i2 :: t
i2 = t
i1 t -> t -> t
forall a. Num a => a -> a -> a
+ t
j
         in if t
i2 t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
0
            then Step (Tuple' t s) b -> m (Step (Tuple' t s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' t s) b -> m (Step (Tuple' t s) b))
-> Step (Tuple' t s) b -> m (Step (Tuple' t s) b)
forall a b. (a -> b) -> a -> b
$ t -> Tuple' t s -> Step (Tuple' t s) b
st t
j (t -> s -> Tuple' t s
forall a b. a -> b -> Tuple' a b
Tuple' t
i2 s
s)
            else do
                -- i2 == i1 == j == 0
                Step s b
r <- s -> m (Step s b)
extract1 s
s
                Step (Tuple' t s) b -> m (Step (Tuple' t s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' t s) b -> m (Step (Tuple' t s) b))
-> Step (Tuple' t s) b -> m (Step (Tuple' t s) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r of
                    Error String
err -> String -> Step (Tuple' t s) b
forall s b. String -> Step s b
Error String
err
                    Done Int
n1 b
b -> Int -> b -> Step (Tuple' t s) b
forall s b. Int -> b -> Step s b
Done Int
n1 b
b
                    Continue Int
n1 s
s1 -> Int -> Tuple' t s -> Step (Tuple' t s) b
forall s b. Int -> s -> Step s b
Continue Int
n1 (t -> s -> Tuple' t s
forall a b. a -> b -> Tuple' a b
Tuple' t
i2 s
s1)
                    Partial Int
_ s
_ -> String -> Step (Tuple' t s) b
forall a. (?callStack::CallStack) => String -> a
error String
"Partial in extract"

    -- Tuple' (how many more items to take) (fold state)
    step :: Tuple' Int s -> Array a -> m (Step (Tuple' Int s) b)
step (Tuple' Int
i s
r) Array a
arr = do
        let len :: Int
len = Array a -> Int
forall a. Unbox a => Array a -> Int
Array.length Array a
arr
            i1 :: Int
i1 = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len
        if Int
i1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
        then do
            Step s b
res <- s -> Array a -> m (Step s b)
step1 s
r Array a
arr
            case Step s b
res of
                Partial Int
j s
s -> Int
-> (Int -> Tuple' Int s -> Step (Tuple' Int s) b)
-> Int
-> s
-> m (Step (Tuple' Int s) b)
forall {t}.
(Ord t, Num t) =>
t
-> (t -> Tuple' t s -> Step (Tuple' t s) b)
-> t
-> s
-> m (Step (Tuple' t s) b)
partial Int
i1 Int -> Tuple' Int s -> Step (Tuple' Int s) b
forall s b. Int -> s -> Step s b
Partial Int
j s
s
                Continue Int
j s
s -> Int
-> (Int -> Tuple' Int s -> Step (Tuple' Int s) b)
-> Int
-> s
-> m (Step (Tuple' Int s) b)
forall {t}.
(Ord t, Num t) =>
t
-> (t -> Tuple' t s -> Step (Tuple' t s) b)
-> t
-> s
-> m (Step (Tuple' t s) b)
partial Int
i1 Int -> Tuple' Int s -> Step (Tuple' Int s) b
forall s b. Int -> s -> Step s b
Continue Int
j s
s
                Done Int
j b
b -> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b))
-> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ Int -> b -> Step (Tuple' Int s) b
forall s b. Int -> b -> Step s b
Done Int
j b
b
                Error String
err -> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b))
-> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ String -> Step (Tuple' Int s) b
forall s b. String -> Step s b
Error String
err
        else do
            let !(Array MutByteArray
contents Int
start Int
_) = Array a
arr
                end :: Int
end = INDEX_OF(start,i,a)
                -- Supply only the required slice of array
                arr1 :: Array a
arr1 = MutByteArray -> Int -> Int -> Array a
forall a. MutByteArray -> Int -> Int -> Array a
Array MutByteArray
contents Int
start Int
end
                remaining :: Int
remaining = Int -> Int
forall a. Num a => a -> a
negate Int
i1 -- i1 is negative here
            Step s b
res <- s -> Array a -> m (Step s b)
step1 s
r Array a
forall {a}. Array a
arr1
            case Step s b
res of
                Partial Int
0 s
s ->
                    Int
-> (s -> Tuple' Int s)
-> (b -> b)
-> Step s b
-> Step (Tuple' Int s) b
forall s s1 b b1.
Int -> (s -> s1) -> (b -> b1) -> Step s b -> Step s1 b1
ParserD.bimapOverrideCount
                        Int
remaining (Int -> s -> Tuple' Int s
forall a b. a -> b -> Tuple' a b
Tuple' Int
0) b -> b
forall a. a -> a
id (Step s b -> Step (Tuple' Int s) b)
-> m (Step s b) -> m (Step (Tuple' Int s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
extract1 s
s
                Partial Int
j s
s -> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b))
-> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ Int -> Tuple' Int s -> Step (Tuple' Int s) b
forall s b. Int -> s -> Step s b
Partial (Int
remaining Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (Int -> s -> Tuple' Int s
forall a b. a -> b -> Tuple' a b
Tuple' Int
j s
s)
                Continue Int
0 s
s ->
                    Int
-> (s -> Tuple' Int s)
-> (b -> b)
-> Step s b
-> Step (Tuple' Int s) b
forall s s1 b b1.
Int -> (s -> s1) -> (b -> b1) -> Step s b -> Step s1 b1
ParserD.bimapOverrideCount
                        Int
remaining (Int -> s -> Tuple' Int s
forall a b. a -> b -> Tuple' a b
Tuple' Int
0) b -> b
forall a. a -> a
id (Step s b -> Step (Tuple' Int s) b)
-> m (Step s b) -> m (Step (Tuple' Int s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
extract1 s
s
                Continue Int
j s
s -> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b))
-> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ Int -> Tuple' Int s -> Step (Tuple' Int s) b
forall s b. Int -> s -> Step s b
Continue (Int
remaining Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (Int -> s -> Tuple' Int s
forall a b. a -> b -> Tuple' a b
Tuple' Int
j s
s)
                Done Int
j b
b -> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b))
-> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ Int -> b -> Step (Tuple' Int s) b
forall s b. Int -> b -> Step s b
Done (Int
remaining Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) b
b
                Error String
err -> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b))
-> Step (Tuple' Int s) b -> m (Step (Tuple' Int s) b)
forall a b. (a -> b) -> a -> b
$ String -> Step (Tuple' Int s) b
forall s b. String -> Step s b
Error String
err

    extract :: Tuple' a s -> m (Step (Tuple' a s) b)
extract (Tuple' a
i s
r) = (s -> Tuple' a s) -> Step s b -> Step (Tuple' a s) b
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a -> s -> Tuple' a s
forall a b. a -> b -> Tuple' a b
Tuple' a
i) (Step s b -> Step (Tuple' a s) b)
-> m (Step s b) -> m (Step (Tuple' a s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
extract1 s
r