-- | Incremental parser interface
--
-- Intended for qualified import.
--
-- > import Network.GRPC.Spec.Util.Parser (Parser)
-- > import Network.GRPC.Spec.Util.Parser qualified as Parser
module Network.GRPC.Spec.Util.Parser (
    Parser -- opaque
    -- * Construction
  , consumeExactly
  , getExactly
    -- * Execution
  , IsFinal
  , Leftover
  , ProcessResult(..)
  , processAll
  ) where

import Control.Monad
import Data.Bifunctor
import Data.Binary (Get)
import Data.Binary.Get qualified as Binary
import Data.ByteString qualified as BS.Strict
import Data.ByteString qualified as Strict (ByteString)
import Data.ByteString.Lazy qualified as BS.Lazy
import Data.ByteString.Lazy qualified as Lazy (ByteString)
import Data.Int

{-------------------------------------------------------------------------------
  Definition
-------------------------------------------------------------------------------}

-- | Simple incremental parser
--
-- This is used to parse a stream of values, where we know ahead of time for
-- each value how much data to expect (perhaps based on the previous value).
-- Individual values are not parsed incrementally; see 'consumeExactly' or
-- 'getExactly'.
newtype Parser e a = Parser {
      forall e a. Parser e a -> Accumulator -> Result e a
runParser :: Accumulator -> Result e a
    }

data Result e a =
    -- | Parsing failed
    --
    -- This implies that we can stop parsing: getting more data won't fix the
    -- problem (see also 'NeedData')
    Failed e

    -- | We make some partial progress, but we need more data to continue
  | NeedData (Parser e a) Accumulator

    -- | Parsing succeeded
    --
    -- Also returns the left-over data
  | Done a Accumulator

instance Bifunctor Result where
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> Result a c -> Result b d
bimap a -> b
f c -> d
_ (Failed a
e)      = b -> Result b d
forall e a. e -> Result e a
Failed (a -> b
f a
e)
  bimap a -> b
_ c -> d
g (Done c
a Accumulator
bs)     = d -> Accumulator -> Result b d
forall e a. a -> Accumulator -> Result e a
Done (c -> d
g c
a) Accumulator
bs
  bimap a -> b
f c -> d
g (NeedData Parser a c
p Accumulator
bs) = Parser b d -> Accumulator -> Result b d
forall e a. Parser e a -> Accumulator -> Result e a
NeedData ((a -> b) -> (c -> d) -> Parser a c -> Parser b d
forall a b c d. (a -> b) -> (c -> d) -> Parser a c -> Parser b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f c -> d
g Parser a c
p) Accumulator
bs

instance Functor (Result e) where
  fmap :: forall a b. (a -> b) -> Result e a -> Result e b
fmap = (a -> b) -> Result e a -> Result e b
forall b c a. (b -> c) -> Result a b -> Result a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second

instance Bifunctor Parser where
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> Parser a c -> Parser b d
bimap a -> b
f c -> d
g (Parser Accumulator -> Result a c
p) = (Accumulator -> Result b d) -> Parser b d
forall e a. (Accumulator -> Result e a) -> Parser e a
Parser ((a -> b) -> (c -> d) -> Result a c -> Result b d
forall a b c d. (a -> b) -> (c -> d) -> Result a c -> Result b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f c -> d
g (Result a c -> Result b d)
-> (Accumulator -> Result a c) -> Accumulator -> Result b d
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Accumulator -> Result a c
p)

instance Functor (Parser e) where
  fmap :: forall a b. (a -> b) -> Parser e a -> Parser e b
fmap = (a -> b) -> Parser e a -> Parser e b
forall b c a. (b -> c) -> Parser a b -> Parser a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second

instance Applicative (Parser e) where
  pure :: forall a. a -> Parser e a
pure a
x = (Accumulator -> Result e a) -> Parser e a
forall e a. (Accumulator -> Result e a) -> Parser e a
Parser ((Accumulator -> Result e a) -> Parser e a)
-> (Accumulator -> Result e a) -> Parser e a
forall a b. (a -> b) -> a -> b
$ a -> Accumulator -> Result e a
forall e a. a -> Accumulator -> Result e a
Done a
x
  <*> :: forall a b. Parser e (a -> b) -> Parser e a -> Parser e b
(<*>)  = Parser e (a -> b) -> Parser e a -> Parser e b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad (Parser e) where
  Parser e a
f >>= :: forall a b. Parser e a -> (a -> Parser e b) -> Parser e b
>>= a -> Parser e b
g = (Accumulator -> Result e b) -> Parser e b
forall e a. (Accumulator -> Result e a) -> Parser e a
Parser ((Accumulator -> Result e b) -> Parser e b)
-> (Accumulator -> Result e b) -> Parser e b
forall a b. (a -> b) -> a -> b
$ \Accumulator
bs ->
      case Parser e a -> Accumulator -> Result e a
forall e a. Parser e a -> Accumulator -> Result e a
runParser Parser e a
f Accumulator
bs of
        Failed e
e        -> e -> Result e b
forall e a. e -> Result e a
Failed e
e
        NeedData Parser e a
f' Accumulator
bs' -> Parser e b -> Accumulator -> Result e b
forall e a. Parser e a -> Accumulator -> Result e a
NeedData (Parser e a
f' Parser e a -> (a -> Parser e b) -> Parser e b
forall a b. Parser e a -> (a -> Parser e b) -> Parser e b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Parser e b
g) Accumulator
bs'
        Done a
a Accumulator
bs'      -> Parser e b -> Accumulator -> Result e b
forall e a. Parser e a -> Accumulator -> Result e a
runParser (a -> Parser e b
g a
a) Accumulator
bs'

instance MonadFail (Parser String) where
  fail :: forall a. String -> Parser String a
fail String
err = (Accumulator -> Result String a) -> Parser String a
forall e a. (Accumulator -> Result e a) -> Parser e a
Parser ((Accumulator -> Result String a) -> Parser String a)
-> (Accumulator -> Result String a) -> Parser String a
forall a b. (a -> b) -> a -> b
$ \Accumulator
_ -> String -> Result String a
forall e a. e -> Result e a
Failed String
err

{-------------------------------------------------------------------------------
  Accumulated input

  This is an internal abstraction only; client code never sees this.
-------------------------------------------------------------------------------}

data Accumulator = Accumulator {
      -- | All the chunks received so far, in reverse order
      Accumulator -> [ByteString]
accumulatedChunks :: [Strict.ByteString]

      -- | Total accumulated length
    , Accumulator -> Int64
accumulatedLength :: !Int64
    }

nil :: Accumulator
nil :: Accumulator
nil = Accumulator {
      accumulatedChunks :: [ByteString]
accumulatedChunks = []
    , accumulatedLength :: Int64
accumulatedLength = Int64
0
    }

-- | Append chunks at the end of the accumulator
--
-- @O(1)@ (this is the raison d'être of this abstraction)
snoc :: Accumulator -> Strict.ByteString -> Accumulator
snoc :: Accumulator -> ByteString -> Accumulator
snoc Accumulator
acc ByteString
chunk = Accumulator {
      accumulatedChunks :: [ByteString]
accumulatedChunks =
        ByteString
chunk ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Accumulator -> [ByteString]
accumulatedChunks Accumulator
acc
    , accumulatedLength :: Int64
accumulatedLength =
        Accumulator -> Int64
accumulatedLength Accumulator
acc Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.Strict.length ByteString
chunk)
    }

toLazy :: Accumulator -> Lazy.ByteString
toLazy :: Accumulator -> ByteString
toLazy = [ByteString] -> ByteString
BS.Lazy.fromChunks ([ByteString] -> ByteString)
-> (Accumulator -> [ByteString]) -> Accumulator -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse ([ByteString] -> [ByteString])
-> (Accumulator -> [ByteString]) -> Accumulator -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Accumulator -> [ByteString]
accumulatedChunks

split :: Int64 -> Accumulator -> Maybe (Lazy.ByteString, Accumulator)
split :: Int64 -> Accumulator -> Maybe (ByteString, Accumulator)
split Int64
n Accumulator
acc
  | Int64
n Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Accumulator -> Int64
accumulatedLength Accumulator
acc
  = Maybe (ByteString, Accumulator)
forall a. Maybe a
Nothing

  | Bool
otherwise
  = let bs :: ByteString
bs            = Accumulator -> ByteString
toLazy Accumulator
acc
        (ByteString
front, ByteString
back) = Int64 -> ByteString -> (ByteString, ByteString)
BS.Lazy.splitAt Int64
n ByteString
bs
        remainder :: Accumulator
remainder     = Accumulator {
            accumulatedChunks :: [ByteString]
accumulatedChunks = [ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse (ByteString -> [ByteString]
BS.Lazy.toChunks ByteString
back)
          , accumulatedLength :: Int64
accumulatedLength = Accumulator -> Int64
accumulatedLength Accumulator
acc Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
n
          }
    in (ByteString, Accumulator) -> Maybe (ByteString, Accumulator)
forall a. a -> Maybe a
Just (ByteString
front, Accumulator
remainder)

{-------------------------------------------------------------------------------
  Construction
-------------------------------------------------------------------------------}

-- | Consume a specified number of bytes
--
-- In order to use the t'Parser' interface we must know for each value exactly
-- how big it will be ahead of time. Typically this will be done by first
-- calling 'consumeExactly' for some kind of fixed size header, indicating how
-- big the value actual value is, which will then inform the next call to
-- 'consumeExactly'.
consumeExactly :: forall e a.
     Int64                           -- ^ Length
  -> (Lazy.ByteString -> Either e a) -- ^ Parser
  -> Parser e a
consumeExactly :: forall e a. Int64 -> (ByteString -> Either e a) -> Parser e a
consumeExactly Int64
len ByteString -> Either e a
parse = Parser e a
go
  where
    go :: Parser e a
    go :: Parser e a
go = (Accumulator -> Result e a) -> Parser e a
forall e a. (Accumulator -> Result e a) -> Parser e a
Parser ((Accumulator -> Result e a) -> Parser e a)
-> (Accumulator -> Result e a) -> Parser e a
forall a b. (a -> b) -> a -> b
$ \Accumulator
acc ->
        case Int64 -> Accumulator -> Maybe (ByteString, Accumulator)
split Int64
len Accumulator
acc of
          Maybe (ByteString, Accumulator)
Nothing          -> Parser e a -> Accumulator -> Result e a
forall e a. Parser e a -> Accumulator -> Result e a
NeedData Parser e a
go Accumulator
acc
          Just (ByteString
raw, Accumulator
left) -> (e -> Result e a) -> (a -> Result e a) -> Either e a -> Result e a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either e -> Result e a
forall e a. e -> Result e a
Failed ((a -> Accumulator -> Result e a) -> Accumulator -> a -> Result e a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Accumulator -> Result e a
forall e a. a -> Accumulator -> Result e a
Done Accumulator
left) (Either e a -> Result e a) -> Either e a -> Result e a
forall a b. (a -> b) -> a -> b
$ ByteString -> Either e a
parse ByteString
raw

-- | Convenience wrapper around 'consumeExactly'
getExactly :: Int64 -> Get a -> Parser String a
getExactly :: forall a. Int64 -> Get a -> Parser String a
getExactly Int64
len Get a
get =
    Int64 -> (ByteString -> Either String a) -> Parser String a
forall e a. Int64 -> (ByteString -> Either e a) -> Parser e a
consumeExactly Int64
len ((ByteString -> Either String a) -> Parser String a)
-> (ByteString -> Either String a) -> Parser String a
forall a b. (a -> b) -> a -> b
$ ((ByteString, Int64, String) -> Either String a)
-> ((ByteString, Int64, a) -> Either String a)
-> Either (ByteString, Int64, String) (ByteString, Int64, a)
-> Either String a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (ByteString, Int64, String) -> Either String a
forall a. (ByteString, Int64, String) -> Either String a
failed (ByteString, Int64, a) -> Either String a
forall a. (ByteString, Int64, a) -> Either String a
ok (Either (ByteString, Int64, String) (ByteString, Int64, a)
 -> Either String a)
-> (ByteString
    -> Either (ByteString, Int64, String) (ByteString, Int64, a))
-> ByteString
-> Either String a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Get a
-> ByteString
-> Either (ByteString, Int64, String) (ByteString, Int64, a)
forall a.
Get a
-> ByteString
-> Either (ByteString, Int64, String) (ByteString, Int64, a)
Binary.runGetOrFail Get a
get
  where
    failed :: (Lazy.ByteString, Binary.ByteOffset, String) -> Either String a
    failed :: forall a. (ByteString, Int64, String) -> Either String a
failed (ByteString
_, Int64
_, String
err) = String -> Either String a
forall a b. a -> Either a b
Left String
err

    ok :: (Lazy.ByteString, Binary.ByteOffset, a) -> Either String a
    ok :: forall a. (ByteString, Int64, a) -> Either String a
ok (ByteString
unconsumed, Int64
lenConsumed, a
a)
      | Bool -> Bool
not (ByteString -> Bool
BS.Lazy.null ByteString
unconsumed) = String -> Either String a
forall a b. a -> Either a b
Left String
"Unconsumed data"
      | Int64
lenConsumed Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Int64
len            = String -> Either String a
forall a. HasCallStack => String -> a
error String
"impossible"
      | Bool
otherwise                     = a -> Either String a
forall a b. b -> Either a b
Right a
a

{-------------------------------------------------------------------------------
  Execution
-------------------------------------------------------------------------------}

-- | Is this the final chunk in the input?
type IsFinal = Bool

-- | Leftover data
type Leftover = Lazy.ByteString

-- | Result from processing all chunks in the input
--
-- See 'processAll'.
data ProcessResult e b =
    -- | Parse error during processing
    ProcessError e

    -- | Parsing succeeded (compare to 'ProcessedWithoutFinal')
  | ProcessedWithFinal b Leftover

    -- | Parsing succeeded, but we did not recognize the final message on time
    --
    -- There are two ways that parsing can terminate: the final few chunks may
    -- look like this:
    --
    -- > chunk1       -- not marked final
    -- > chunk2       -- not marked final
    -- > chunk3       -- marked final
    --
    -- or like this:
    --
    -- > chunk1       -- not marked final
    -- > chunk2       -- not marked final
    -- > chunk3       -- not marked final
    -- > empty chunk  -- marked final
    --
    -- In the former case, we know that we are processing the final message /as/
    -- we are processing it ('ProcessedWithFinal'); in the latter case, we
    -- realize this only after we receive the final empty chunk.
  | ProcessedWithoutFinal Leftover

-- | Process all incoming data
--
-- Returns any unprocessed data.
-- Also returns if we knew that the final result
-- was in fact the final result when we received it (this may or may not be the
-- case, depending on
processAll :: forall m e a b.
     Monad m
  => m (Strict.ByteString, IsFinal)  -- ^ Get next chunk
  -> (a -> m ())                     -- ^ Process value
  -> (a -> m b)                      -- ^ Process final value
  -> Parser e a                      -- ^ Parser
  -> m (ProcessResult e b)
processAll :: forall (m :: * -> *) e a b.
Monad m =>
m (ByteString, Bool)
-> (a -> m ()) -> (a -> m b) -> Parser e a -> m (ProcessResult e b)
processAll m (ByteString, Bool)
getChunk a -> m ()
processOne a -> m b
processFinal Parser e a
parser =
    Result e a -> m (ProcessResult e b)
go (Result e a -> m (ProcessResult e b))
-> Result e a -> m (ProcessResult e b)
forall a b. (a -> b) -> a -> b
$ Parser e a -> Accumulator -> Result e a
forall e a. Parser e a -> Accumulator -> Result e a
runParser Parser e a
parser Accumulator
nil
  where
    go :: Result e a -> m (ProcessResult e b)
    go :: Result e a -> m (ProcessResult e b)
go (Failed e
err)            = ProcessResult e b -> m (ProcessResult e b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (ProcessResult e b -> m (ProcessResult e b))
-> ProcessResult e b -> m (ProcessResult e b)
forall a b. (a -> b) -> a -> b
$ e -> ProcessResult e b
forall e b. e -> ProcessResult e b
ProcessError e
err
    go (Done a
a Accumulator
left)           = a -> m ()
processOne a
a m () -> m (ProcessResult e b) -> m (ProcessResult e b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Result e a -> m (ProcessResult e b)
go (Parser e a -> Accumulator -> Result e a
forall e a. Parser e a -> Accumulator -> Result e a
runParser Parser e a
parser Accumulator
left)
    go (NeedData Parser e a
parser' Accumulator
left) = do
        (bs, isFinal) <- m (ByteString, Bool)
getChunk
        if not isFinal
          then go         $ runParser parser' (left `snoc` bs)
          else goFinal [] $ runParser parser' (left `snoc` bs)

    -- We have received the final chunk; extract all messages until we are done
    goFinal :: [a] -> Result e a -> m (ProcessResult e b)
    goFinal :: [a] -> Result e a -> m (ProcessResult e b)
goFinal [a]
_   (Failed e
err)      = ProcessResult e b -> m (ProcessResult e b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (ProcessResult e b -> m (ProcessResult e b))
-> ProcessResult e b -> m (ProcessResult e b)
forall a b. (a -> b) -> a -> b
$ e -> ProcessResult e b
forall e b. e -> ProcessResult e b
ProcessError e
err
    goFinal [a]
acc (Done a
a Accumulator
left)     = [a] -> Result e a -> m (ProcessResult e b)
goFinal (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
acc) (Result e a -> m (ProcessResult e b))
-> Result e a -> m (ProcessResult e b)
forall a b. (a -> b) -> a -> b
$ Parser e a -> Accumulator -> Result e a
forall e a. Parser e a -> Accumulator -> Result e a
runParser Parser e a
parser Accumulator
left
    goFinal [a]
acc (NeedData Parser e a
_ Accumulator
left) = do
        mb <- [a] -> m (Maybe b)
processLastFew ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc)
        return $ case mb of
                   Just b
b  -> b -> ByteString -> ProcessResult e b
forall e b. b -> ByteString -> ProcessResult e b
ProcessedWithFinal b
b  (ByteString -> ProcessResult e b)
-> ByteString -> ProcessResult e b
forall a b. (a -> b) -> a -> b
$ Accumulator -> ByteString
toLazy Accumulator
left
                   Maybe b
Nothing -> ByteString -> ProcessResult e b
forall e b. ByteString -> ProcessResult e b
ProcessedWithoutFinal (ByteString -> ProcessResult e b)
-> ByteString -> ProcessResult e b
forall a b. (a -> b) -> a -> b
$ Accumulator -> ByteString
toLazy Accumulator
left

    processLastFew :: [a] -> m (Maybe b)
    processLastFew :: [a] -> m (Maybe b)
processLastFew []     = Maybe b -> m (Maybe b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe b
forall a. Maybe a
Nothing
    processLastFew [a
a]    = b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> m b -> m (Maybe b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m b
processFinal a
a
    processLastFew (a
a:[a]
as) = a -> m ()
processOne a
a m () -> m (Maybe b) -> m (Maybe b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [a] -> m (Maybe b)
processLastFew [a]
as