content: {-# OPTIONS_GHC -Wno-deprecations #-}
  {-# OPTIONS_GHC -Wno-incomplete-patterns #-}
  
  -- |
  -- 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 = Array.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(..), Final(..))
  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 (Fold.Fold fstep finitial _ ffinal) =
      ChunkFold (ParserD.Parser step initial extract)
  
      where
  
      initial = do
          res <- finitial
          return
              $ case res of
                    Fold.Partial s1 -> IPartial s1
                    Fold.Done b -> IDone b
  
      step s (Array contents start end) = do
          goArray SPEC start s
  
          where
  
          goArray !_ !cur !fs | cur >= end = do
              assert (cur == end) (return ())
              return $ Partial 0 fs
          goArray !_ !cur !fs = do
              x <- liftIO $ peekAt cur contents
              res <- fstep fs x
              let elemSize = SIZE_OF(a)
                  next = INDEX_NEXT(cur,a)
              case res of
                  Fold.Done b ->
                      return $ Done ((end - next) `div` elemSize) b
                  Fold.Partial fs1 ->
                      goArray SPEC next fs1
  
      extract = fmap (FDone 0) . 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 (ParserD.Parser step1 initial1 extract1) =
      ChunkFold (ParserD.Parser step initial1 extract1)
  
      where
  
      step s (Array contents start end) = do
          if start >= end
          then return $ Continue 0 s
          else goArray SPEC start s
  
          where
  
          {-# INLINE partial #-}
          partial arrRem cur next elemSize st n fs1 = do
              let next1 = next - (n * elemSize)
              if next1 >= start && cur < end
              then goArray SPEC next1 fs1
              else return $ st (arrRem + n) fs1
  
          goArray !_ !cur !fs = do
              x <- liftIO $ peekAt cur contents
              res <- step1 fs x
              let elemSize = SIZE_OF(a)
                  next = INDEX_NEXT(cur,a)
                  arrRem = (end - next) `div` elemSize
              case res of
                  ParserD.Done n b -> do
                      return $ Done (arrRem + n) b
                  ParserD.Partial n fs1 ->
                      partial arrRem cur next elemSize Partial n fs1
                  ParserD.Continue n fs1 -> do
                      partial arrRem cur next elemSize Continue n fs1
                  SError err -> return $ SError 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 = 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 f = ChunkFold $ ParserD.fromFold f
  
  -------------------------------------------------------------------------------
  -- Functor
  -------------------------------------------------------------------------------
  
  -- | Maps a function over the result of fold.
  --
  -- /Pre-release/
  instance Functor m => Functor (ChunkFold m a) where
      {-# INLINE fmap #-}
      fmap f (ChunkFold p) = ChunkFold $ fmap f 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 f (ChunkFold p) = ChunkFold $ ParserD.rmapM f 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 = ChunkFold . 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 = ChunkFold . 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_ (ChunkFold p1) (ChunkFold p2) =
      ChunkFold $ ParserD.noErrorUnsafeSplit_ p1 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 f (ChunkFold p1) (ChunkFold p2) =
      ChunkFold $ ParserD.noErrorUnsafeSplitWith f p1 p2
  
  -- | 'Applicative' form of 'splitWith'.
  -- > (<*>) = splitWith id
  instance Monad m => Applicative (ChunkFold m a) where
      {-# INLINE pure #-}
      pure = fromPure
  
      {-# INLINE (<*>) #-}
      (<*>) = splitWith id
  
      {-# INLINE (*>) #-}
      (*>) = split_
  
      {-# INLINE liftA2 #-}
      liftA2 f x = (<*>) (fmap f 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 func (ChunkFold p) =
      let f x = let ChunkFold y = func x in y
       in ChunkFold $ ParserD.noErrorUnsafeConcatMap f 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 = pure
  
      {-# INLINE (>>=) #-}
      (>>=) = flip concatMap
  
      {-# INLINE (>>) #-}
      (>>) = (*>)
  
  -------------------------------------------------------------------------------
  -- 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 n (ChunkFold (ParserD.Parser step1 initial1 extract1)) =
      ChunkFold $ ParserD.Parser step initial extract
  
      where
  
      -- XXX Need to make the Initial type Step to remove this
      iextract s = do
          r <- extract1 s
          return $ case r of
              FDone _ b -> IDone b
              FError err -> IError err
              _ -> error "Bug: ChunkFold take invalid state in initial"
  
      initial = do
          res <- initial1
          case res of
              IPartial s ->
                  if n > 0
                  then return $ IPartial $ Tuple' n s
                  else iextract s
              IDone b -> return $ IDone b
              IError err -> return $ IError err
  
      {-# INLINE partial #-}
      partial i1 st j s =
          let i2 = i1 + j
           in if i2 > 0
              then return $ st j (Tuple' i2 s)
              else do
                  -- i2 == i1 == j == 0
                  r <- extract1 s
                  return $ case r of
                      FError err -> SError err
                      FDone n1 b -> Done n1 b
                      FContinue n1 s1 -> Continue n1 (Tuple' i2 s1)
  
      -- Tuple' (how many more items to take) (fold state)
      step (Tuple' i r) arr = do
          let len = Array.length arr
              i1 = i - len
          if i1 >= 0
          then do
              res <- step1 r arr
              case res of
                  Partial j s -> partial i1 Partial j s
                  Continue j s -> partial i1 Continue j s
                  Done j b -> return $ Done j b
                  SError err -> return $ SError err
          else do
              let !(Array contents start _) = arr
                  end = INDEX_OF(start,i,a)
                  -- Supply only the required slice of array
                  arr1 = Array contents start end
                  remaining = negate i1 -- i1 is negative here
              res <- step1 r arr1
              case res of
                  Partial 0 s ->
                      ParserD.bimapMorphOverrideCount
                          remaining (Tuple' 0) id <$> extract1 s
                  Partial j s -> return $ Partial (remaining + j) (Tuple' j s)
                  Continue 0 s ->
                      ParserD.bimapMorphOverrideCount
                          remaining (Tuple' 0) id <$> extract1 s
                  Continue j s -> return $ Continue (remaining + j) (Tuple' j s)
                  Done j b -> return $ Done (remaining + j) b
                  SError err -> return $ SError err
  
      extract (Tuple' i r) = first (Tuple' i) <$> extract1 r
  