-- |
-- Module      : Streamly.Internal.Data.Array.Stream.Fold.Foreign
-- Copyright   : (c) 2021 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- 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.Foreign as ArrayStream
-- >>> import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ArrayFold
-- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream (arraysOf)
-- >>> import qualified Streamly.Prelude as Stream
--
-- >>> ArrayStream.foldArr (ArrayFold.fromFold (Fold.take 7 Fold.toList)) $ Stream.arraysOf 5 $ Stream.fromList "hello world"
-- "hello w"
--
module Streamly.Internal.Data.Array.Stream.Fold.Foreign
    (
      Fold (..)

    -- * Construction
    , fromFold
    , fromParser
    , fromArrayFold

    -- * Mapping
    , rmapM

    -- * Applicative
    , fromPure
    , fromEffect
    , serialWith

    -- * Monad
    , concatMap

    -- * Transformation
    , take
    )
where

import Control.Applicative (liftA2)
import Control.Exception (assert)
import Control.Monad.Catch (MonadThrow)
import Control.Monad.IO.Class (MonadIO(..))
import Foreign.Ptr (minusPtr, plusPtr)
import Foreign.Storable (Storable(..))
import GHC.Types (SPEC(..))
import Streamly.Internal.Data.Array.Foreign.Mut.Type (touch)
import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
import Streamly.Internal.Data.Parser.ParserD (Initial(..), Step(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

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

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 Fold m a b = Fold (ParserD.Parser m (Array a) 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, Storable a) =>
    Fold.Fold m a b -> Fold m a b
fromFold :: Fold m a b -> Fold m a b
fromFold (Fold.Fold s -> a -> m (Step s b)
fstep m (Step s b)
finitial s -> m b
fextract) =
    Parser m (Array a) b -> Fold m a b
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold ((s -> Array a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m (Array a) b
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m a b
ParserD.Parser s -> Array a -> m (Step s b)
step m (Initial s b)
initial s -> m b
fextract)

    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 ArrayContents
contents Ptr a
start Ptr a
end) = do
        SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC Ptr a
start s
s

        where

        goArray :: SPEC -> Ptr a -> s -> m (Step s b)
goArray !SPEC
_ !Ptr a
cur !s
fs | Ptr a
cur Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
end = do
            Bool -> m () -> m ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
cur Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
end) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
            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
_ !Ptr a
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
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
cur
            Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
            let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. (?callStack::CallStack) => a
undefined :: a)
                next :: Ptr b
next = Ptr a
cur Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize
            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 ((Ptr a
end Ptr a -> Ptr Any -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Any
forall b. Ptr b
next) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
elemSize) b
b
                Fold.Partial s
fs1 ->
                    SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC Ptr a
forall b. Ptr b
next s
fs1

-- | Convert an element '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, Storable a) =>
    ParserD.Parser m a b -> Fold m a b
fromParser :: Parser m a b -> Fold m a b
fromParser (ParserD.Parser s -> a -> m (Step s b)
step1 m (Initial s b)
initial1 s -> m b
extract1) =
    Parser m (Array a) b -> Fold m a b
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold ((s -> Array a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m (Array a) b
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m a b
ParserD.Parser s -> Array a -> m (Step s b)
step m (Initial s b)
initial1 s -> m b
extract1)

    where

    step :: s -> Array a -> m (Step s b)
step s
s (Array ArrayContents
contents Ptr a
start Ptr a
end) = do
        if Ptr a
start Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
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 -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC Ptr a
start s
s

        where

        {-# INLINE partial #-}
        partial :: Int
-> Ptr a
-> Ptr a
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Ptr a
cur Ptr a
next Int
elemSize Int -> s -> Step s b
st Int
n s
fs1 = do
            let next1 :: Ptr b
next1 = Ptr a
next Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int
forall a. Num a => a -> a
negate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
elemSize)
            if Ptr a
forall b. Ptr b
next1 Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
start Bool -> Bool -> Bool
&& Ptr a
cur Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
< Ptr a
end
            then SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC Ptr a
forall b. Ptr b
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 -> Ptr a -> s -> m (Step s b)
goArray !SPEC
_ !Ptr a
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
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
cur
            IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
            Step s b
res <- s -> a -> m (Step s b)
step1 s
fs a
x
            let elemSize :: Int
elemSize = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. (?callStack::CallStack) => a
undefined :: a)
                next :: Ptr b
next = Ptr a
cur Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
elemSize
                arrRem :: Int
arrRem = (Ptr a
end Ptr a -> Ptr Any -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Any
forall b. Ptr b
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
-> Ptr a
-> Ptr a
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Ptr a
cur Ptr a
forall b. Ptr b
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
-> Ptr a
-> Ptr a
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Ptr a
cur Ptr a
forall b. Ptr b
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

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

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

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

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

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

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

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

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

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

    {-# INLINE (<*>) #-}
    <*> :: Fold m a (a -> b) -> Fold m a a -> Fold m a b
(<*>) = ((a -> b) -> a -> b)
-> Fold m a (a -> b) -> Fold m a a -> Fold m a b
forall (m :: * -> *) a b c x.
MonadThrow m =>
(a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
serialWith (a -> b) -> a -> b
forall a. a -> a
id

    {-# INLINE (*>) #-}
    *> :: Fold m a a -> Fold m a b -> Fold m a b
(*>) = Fold m a a -> Fold m a b -> Fold m a b
forall (m :: * -> *) x a b.
MonadThrow m =>
Fold m x a -> Fold m x b -> Fold m x b
serial_

#if MIN_VERSION_base(4,10,0)
    {-# INLINE liftA2 #-}
    liftA2 :: (a -> b -> c) -> Fold m a a -> Fold m a b -> Fold m a c
liftA2 a -> b -> c
f Fold m a a
x = Fold m a (b -> c) -> Fold m a b -> Fold m a c
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) ((a -> b -> c) -> Fold m a a -> Fold m a (b -> c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b -> c
f Fold m a a
x)
#endif

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

-- | 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 :: MonadThrow m =>
    (b -> Fold m a c) -> Fold m a b -> Fold m a c
concatMap :: (b -> Fold m a c) -> Fold m a b -> Fold m a c
concatMap b -> Fold m a c
func (Fold Parser m (Array a) b
p) =
    Parser m (Array a) c -> Fold m a c
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold (Parser m (Array a) c -> Fold m a c)
-> Parser m (Array a) c -> Fold m a c
forall a b. (a -> b) -> a -> b
$ (b -> Parser m (Array a) c)
-> Parser m (Array a) b -> Parser m (Array a) c
forall (m :: * -> *) b a c.
MonadThrow m =>
(b -> Parser m a c) -> Parser m a b -> Parser m a c
ParserD.noErrorUnsafeConcatMap (\b
x -> let Fold Parser m (Array a) c
y = b -> Fold m a c
func b
x in Parser m (Array a) c
y) Parser m (Array a) b
p

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

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

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

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

{-# INLINE take #-}
take :: forall m a b. (Monad m, Storable a) => Int -> Fold m a b -> Fold m a b
take :: Int -> Fold m a b -> Fold m a b
take Int
n (Fold (ParserD.Parser s -> Array a -> m (Step s b)
step1 m (Initial s b)
initial1 s -> m b
extract1)) =
    Parser m (Array a) b -> Fold m a b
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold (Parser m (Array a) b -> Fold m a b)
-> Parser m (Array a) b -> Fold 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 b)
-> Parser m (Array a) b
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m a 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 b
forall a. Tuple' a s -> m b
extract

    where

    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 b -> Initial (Tuple' Int s) b
forall s b. b -> Initial s b
IDone (b -> Initial (Tuple' Int s) b)
-> m b -> m (Initial (Tuple' Int s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
extract1 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 :: a -> (a -> Tuple' a s -> Step s b) -> a -> s -> m (Step s b)
partial a
i1 a -> Tuple' a s -> Step s b
st a
j s
s =
        let i2 :: a
i2 = a
i1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
j
         in if a
i2 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0
            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
$ a -> Tuple' a s -> Step s b
st a
j (a -> s -> Tuple' a s
forall a b. a -> b -> Tuple' a b
Tuple' a
i2 s
s)
            else Int -> b -> Step s b
forall s b. Int -> b -> Step s b
Done Int
0 (b -> Step s b) -> m b -> m (Step s b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
extract1 s
s -- i2 == i1 == j == 0

    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. Storable 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 a s.
(Ord a, Num a) =>
a -> (a -> Tuple' a s -> Step s b) -> a -> s -> m (Step 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 a s.
(Ord a, Num a) =>
a -> (a -> Tuple' a s -> Step s b) -> a -> s -> m (Step 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 ArrayContents
contents Ptr a
start Ptr a
_) = Array a
arr
                sz :: Int
sz = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. (?callStack::CallStack) => a
undefined :: a)
                end :: Ptr b
end = Ptr a
start Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
sz)
                arr1 :: Array a
arr1 = ArrayContents -> Ptr a -> Ptr a -> Array a
forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
start Ptr a
forall b. Ptr b
end
                remaining :: Int
remaining = Int -> Int
forall a. Num a => a -> a
negate Int
i1
            Step s b
res <- s -> Array a -> m (Step s b)
step1 s
r Array a
arr1
            case Step s b
res of
                Partial Int
0 s
s -> Int -> b -> Step (Tuple' Int s) b
forall s b. Int -> b -> Step s b
Done Int
remaining (b -> Step (Tuple' Int s) b) -> m b -> m (Step (Tuple' Int s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m 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 -> b -> Step (Tuple' Int s) b
forall s b. Int -> b -> Step s b
Done Int
remaining (b -> Step (Tuple' Int s) b) -> m b -> m (Step (Tuple' Int s) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m 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 b
extract (Tuple' a
_ s
r) = s -> m b
extract1 s
r