-- |
-- Module      : Streamly.Internal.Data.Producer.Source
-- Copyright   : (c) 2021 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- A 'Source' is a seed that can be unfolded to a stream with a buffer.  Allows
-- to 'unread' data i.e.  push unused data back to the source buffer. This is
-- useful in parsing applications with backtracking.
--

module Streamly.Internal.Data.Producer.Source
    ( Source

    -- * Creation
    , source

    -- * Transformation
    , unread

    -- * Consumption
    , isEmpty
    , producer

    -- * Parsing
    , parse
    , parseMany
    , parseManyD
    )
where

#include "inline.hs"

import Control.Exception (assert)
import Control.Monad.Catch (MonadThrow, throwM)
import GHC.Exts (SpecConstrAnnotation(..))
import GHC.Types (SPEC(..))
import Streamly.Internal.Data.Parser.ParserD (ParseError(..), Step(..))
import Streamly.Internal.Data.Producer.Type (Producer(..))
import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))

import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
import qualified Streamly.Internal.Data.Parser.ParserK.Type as ParserK

import Prelude hiding (read)

-- | A seed with a buffer. It allows us to 'unread' or return some data
-- after reading it. Useful in backtracked parsing.
--
data Source a b = Source [b] (Maybe a)

-- | Make a source from a seed value. The buffer would start as empty.
--
-- /Pre-release/
source :: Maybe a -> Source a b
source :: forall a b. Maybe a -> Source a b
source = forall a b. [b] -> Maybe a -> Source a b
Source []

-- | Return some unused data back to the source. The data is prepended (or
-- consed) to the source.
--
-- /Pre-release/
unread :: [b] -> Source a b -> Source a b
unread :: forall b a. [b] -> Source a b -> Source a b
unread [b]
xs (Source [b]
ys Maybe a
seed) = forall a b. [b] -> Maybe a -> Source a b
Source ([b]
xs forall a. [a] -> [a] -> [a]
++ [b]
ys) Maybe a
seed

-- | Determine if the source is empty.
isEmpty :: Source a b -> Bool
isEmpty :: forall a b. Source a b -> Bool
isEmpty (Source [] Maybe a
Nothing) = Bool
True
isEmpty Source a b
_ = Bool
False

-- | Convert a producer to a producer from a buffered source. Any buffered data
-- is read first and then the seed is unfolded.
--
-- /Pre-release/
{-# INLINE_NORMAL producer #-}
producer :: Monad m => Producer m a b -> Producer m (Source a b) b
producer :: forall (m :: * -> *) a b.
Monad m =>
Producer m a b -> Producer m (Source a b) b
producer (Producer s -> m (Step s b)
step1 a -> m s
inject1 s -> m a
extract1) = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer Either s ([b], Maybe a) -> m (Step (Either s ([b], Maybe a)) b)
step forall {b}. Source a b -> m (Either s ([b], Maybe a))
inject forall {b}. Either s ([b], Maybe a) -> m (Source a b)
extract

    where

    inject :: Source a b -> m (Either s ([b], Maybe a))
inject (Source [] (Just a
a)) = do
        s
s <- a -> m s
inject1 a
a
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left s
s
    inject (Source [b]
xs Maybe a
a) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right ([b]
xs, Maybe a
a)

    {-# INLINE_LATE step #-}
    step :: Either s ([b], Maybe a) -> m (Step (Either s ([b], Maybe a)) b)
step (Left s
s) = do
        Step s b
r <- s -> m (Step s b)
step1 s
s
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s1 -> forall s a. a -> s -> Step s a
Yield b
x (forall a b. a -> Either a b
Left s
s1)
            Skip s
s1 -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
s1)
            Step s b
Stop -> forall s a. Step s a
Stop
    step (Right ([], Maybe a
Nothing)) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step (Right ([], Just a
_)) = forall a. HasCallStack => [Char] -> a
error [Char]
"Bug: unreachable"
    step (Right (b
x:[], Just a
a)) = do
        s
s <- a -> m s
inject1 a
a
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
x (forall a b. a -> Either a b
Left s
s)
    step (Right (b
x:[b]
xs, Maybe a
a)) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
x (forall a b. b -> Either a b
Right ([b]
xs, Maybe a
a))

    extract :: Either s ([b], Maybe a) -> m (Source a b)
extract (Left s
s) = forall a b. [b] -> Maybe a -> Source a b
Source [] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m a
extract1 s
s
    extract (Right ([b]
xs, Maybe a
a)) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. [b] -> Maybe a -> Source a b
Source [b]
xs Maybe a
a

-------------------------------------------------------------------------------
-- Parsing
-------------------------------------------------------------------------------

-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
-- to make a newtype.
{-# ANN type List NoSpecConstr #-}
newtype List a = List {forall a. List a -> [a]
getList :: [a]}

{-# INLINE_NORMAL parseD #-}
parseD
    :: MonadThrow m
    => ParserD.Parser m a b
    -> Producer m (Source s a) a
    -> Source s a
    -> m (b, Source s a)
parseD :: forall (m :: * -> *) a b s.
MonadThrow m =>
Parser m a b
-> Producer m (Source s a) a -> Source s a -> m (b, Source s a)
parseD
    (ParserD.Parser s -> a -> m (Step s b)
pstep m (Initial s b)
initial s -> m b
extract)
    (Producer s -> m (Step s a)
ustep Source s a -> m s
uinject s -> m (Source s a)
uextract)
    Source s a
seed = do

    Initial s b
res <- m (Initial s b)
initial
    case Initial s b
res of
        ParserD.IPartial s
s -> do
            s
state <- Source s a -> m s
uinject Source s a
seed
            SPEC -> s -> List a -> s -> m (b, Source s a)
go SPEC
SPEC s
state (forall a. [a] -> List a
List []) s
s
        ParserD.IDone b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (b
b, Source s a
seed)
        ParserD.IError [Char]
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ [Char] -> ParseError
ParseError [Char]
err

    where

    -- XXX currently we are using a dumb list based approach for backtracking
    -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
    -- That will allow us more efficient random back and forth movement.
    go :: SPEC -> s -> List a -> s -> m (b, Source s a)
go !SPEC
_ s
st List a
buf !s
pst = do
        Step s a
r <- s -> m (Step s a)
ustep s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Step s b
pRes <- s -> a -> m (Step s b)
pstep s
pst a
x
                case Step s b
pRes of
                    Partial Int
0 s
pst1 -> SPEC -> s -> List a -> s -> m (b, Source s a)
go SPEC
SPEC s
s (forall a. [a] -> List a
List []) s
pst1
                    Partial Int
n s
pst1 -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let src0 :: [a]
src0 = forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)
                            src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0
                        SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf SPEC
SPEC s
s (forall a. [a] -> List a
List []) (forall a. [a] -> List a
List [a]
src) s
pst1
                    Continue Int
0 s
pst1 -> SPEC -> s -> List a -> s -> m (b, Source s a)
go SPEC
SPEC s
s (forall a. [a] -> List a
List (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) s
pst1
                    Continue Int
n s
pst1 -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let ([a]
src0, [a]
buf1) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
n (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)
                            src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0
                        SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf SPEC
SPEC s
s (forall a. [a] -> List a
List [a]
buf1) (forall a. [a] -> List a
List [a]
src) s
pst1
                    Done Int
n b
b -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let src0 :: [a]
src0 = forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)
                            src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0
                        Source s a
s1 <- s -> m (Source s a)
uextract s
s
                        forall (m :: * -> *) a. Monad m => a -> m a
return (b
b, forall b a. [b] -> Source a b -> Source a b
unread [a]
src Source s a
s1)
                    Error [Char]
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ [Char] -> ParseError
ParseError [Char]
err
            Skip s
s -> SPEC -> s -> List a -> s -> m (b, Source s a)
go SPEC
SPEC s
s List a
buf s
pst
            Step s a
Stop   -> do
                b
b <- s -> m b
extract s
pst
                forall (m :: * -> *) a. Monad m => a -> m a
return (b
b, forall b a. [b] -> Source a b -> Source a b
unread (forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a. List a -> [a]
getList List a
buf) (forall a b. Maybe a -> Source a b
source forall a. Maybe a
Nothing))

    gobuf :: SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf !SPEC
_ s
s List a
buf (List []) !s
pst = SPEC -> s -> List a -> s -> m (b, Source s a)
go SPEC
SPEC s
s List a
buf s
pst
    gobuf !SPEC
_ s
s List a
buf (List (a
x:[a]
xs)) !s
pst = do
        Step s b
pRes <- s -> a -> m (Step s b)
pstep s
pst a
x
        case Step s b
pRes of
            Partial Int
0 s
pst1 ->
                SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf SPEC
SPEC s
s (forall a. [a] -> List a
List []) (forall a. [a] -> List a
List [a]
xs) s
pst1
            Partial Int
n s
pst1 -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let src0 :: [a]
src0 = forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)
                    src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0 forall a. [a] -> [a] -> [a]
++ [a]
xs
                SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf SPEC
SPEC s
s (forall a. [a] -> List a
List []) (forall a. [a] -> List a
List [a]
src) s
pst1
            Continue Int
0 s
pst1 ->
                SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf SPEC
SPEC s
s (forall a. [a] -> List a
List (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall a. [a] -> List a
List [a]
xs) s
pst1
            Continue Int
n s
pst1 -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let ([a]
src0, [a]
buf1) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
n (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)
                    src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0 forall a. [a] -> [a] -> [a]
++ [a]
xs
                SPEC -> s -> List a -> List a -> s -> m (b, Source s a)
gobuf SPEC
SPEC s
s (forall a. [a] -> List a
List [a]
buf1) (forall a. [a] -> List a
List [a]
src) s
pst1
            Done Int
n b
b -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let src0 :: [a]
src0 = forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:forall a. List a -> [a]
getList List a
buf)
                    src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0
                Source s a
s1 <- s -> m (Source s a)
uextract s
s
                forall (m :: * -> *) a. Monad m => a -> m a
return (b
b, forall b a. [b] -> Source a b -> Source a b
unread [a]
src Source s a
s1)
            Error [Char]
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ [Char] -> ParseError
ParseError [Char]
err

-- | Parse a buffered source using a parser, returning the parsed value and the
-- remaining source.
--
-- /Pre-release/
{-# INLINE [3] parse #-}
parse
    :: MonadThrow m
    => ParserK.Parser m a b
    -> Producer m (Source s a) a
    -> Source s a
    -> m (b, Source s a)
parse :: forall (m :: * -> *) a b s.
MonadThrow m =>
Parser m a b
-> Producer m (Source s a) a -> Source s a -> m (b, Source s a)
parse = forall (m :: * -> *) a b s.
MonadThrow m =>
Parser m a b
-> Producer m (Source s a) a -> Source s a -> m (b, Source s a)
parseD forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
MonadThrow m =>
Parser m a b -> Parser m a b
ParserK.fromParserK

-------------------------------------------------------------------------------
-- Nested parsing
-------------------------------------------------------------------------------

{-# INLINE parseManyD #-}
parseManyD :: MonadThrow m =>
       ParserD.Parser m a b
    -> Producer m (Source x a) a
    -> Producer m (Source x a) b
parseManyD :: forall (m :: * -> *) a b x.
MonadThrow m =>
Parser m a b
-> Producer m (Source x a) a -> Producer m (Source x a) b
parseManyD Parser m a b
parser Producer m (Source x a) a
reader = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> (s -> m a) -> Producer m a b
Producer Source x a -> m (Step (Source x a) b)
step forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    {-# INLINE_LATE step #-}
    step :: Source x a -> m (Step (Source x a) b)
step Source x a
src = do
        if forall a b. Source a b -> Bool
isEmpty Source x a
src
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
        else do
            (b
b, Source x a
s1) <- forall (m :: * -> *) a b s.
MonadThrow m =>
Parser m a b
-> Producer m (Source s a) a -> Source s a -> m (b, Source s a)
parseD Parser m a b
parser Producer m (Source x a) a
reader Source x a
src
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b Source x a
s1

-- | Apply a parser repeatedly on a buffered source producer to generate a
-- producer of parsed values.
--
-- /Pre-release/
{-# INLINE parseMany #-}
parseMany :: MonadThrow m =>
       ParserK.Parser m a b
    -> Producer m (Source x a) a
    -> Producer m (Source x a) b
parseMany :: forall (m :: * -> *) a b x.
MonadThrow m =>
Parser m a b
-> Producer m (Source x a) a -> Producer m (Source x a) b
parseMany Parser m a b
parser = forall (m :: * -> *) a b x.
MonadThrow m =>
Parser m a b
-> Producer m (Source x a) a -> Producer m (Source x a) b
parseManyD (forall (m :: * -> *) a b.
MonadThrow m =>
Parser m a b -> Parser m a b
ParserK.fromParserK Parser m a b
parser)