-- | Useful combinators for 'ParserT' and 'Stream'.
module SimpleParser.Input
  ( peekToken
  , popToken
  , peekChunk
  , popChunk
  , dropChunk
  , isEnd
  , matchEnd
  , anyToken
  , anyChunk
  , satisfyToken
  , foldTokensWhile
  , takeTokensWhile
  , dropTokensWhile
  , matchToken
  , matchChunk
  ) where

import Control.Applicative (empty)
import Control.Monad.State (gets, state)
import Data.Bifunctor (first)
import Data.Maybe (isNothing)
import SimpleParser.Parser (ParserT)
import SimpleParser.Stream (Chunked (chunkLength), Stream (..))

-- | Return the next token, if any, but don't consume it.
peekToken :: (Stream s, Monad m) => ParserT e s m (Maybe (Token s))
peekToken :: ParserT e s m (Maybe (Token s))
peekToken = (s -> Maybe (Token s)) -> ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((Token s, s) -> Token s) -> Maybe (Token s, s) -> Maybe (Token s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Token s, s) -> Token s
forall a b. (a, b) -> a
fst (Maybe (Token s, s) -> Maybe (Token s))
-> (s -> Maybe (Token s, s)) -> s -> Maybe (Token s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Maybe (Token s, s)
forall s. Stream s => s -> Maybe (Token s, s)
streamTake1)

-- | Return the next token, if any, and consume it.
popToken :: (Stream s, Monad m) => ParserT e s m (Maybe (Token s))
popToken :: ParserT e s m (Maybe (Token s))
popToken = (s -> (Maybe (Token s), s)) -> ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\s
stream -> (Maybe (Token s), s)
-> ((Token s, s) -> (Maybe (Token s), s))
-> Maybe (Token s, s)
-> (Maybe (Token s), s)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe (Token s)
forall a. Maybe a
Nothing, s
stream) ((Token s -> Maybe (Token s))
-> (Token s, s) -> (Maybe (Token s), s)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Token s -> Maybe (Token s)
forall a. a -> Maybe a
Just) (s -> Maybe (Token s, s)
forall s. Stream s => s -> Maybe (Token s, s)
streamTake1 s
stream))

-- | Return the next chunk of the given size, if any, but don't consume it.
-- May return a smaller chunk at end of stream, but never returns an empty chunk.
peekChunk :: (Stream s, Monad m) => Int -> ParserT e s m (Maybe (Chunk s))
peekChunk :: Int -> ParserT e s m (Maybe (Chunk s))
peekChunk Int
n = (s -> Maybe (Chunk s)) -> ParserT e s m (Maybe (Chunk s))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((Chunk s, s) -> Chunk s) -> Maybe (Chunk s, s) -> Maybe (Chunk s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Chunk s, s) -> Chunk s
forall a b. (a, b) -> a
fst (Maybe (Chunk s, s) -> Maybe (Chunk s))
-> (s -> Maybe (Chunk s, s)) -> s -> Maybe (Chunk s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> s -> Maybe (Chunk s, s)
forall s. Stream s => Int -> s -> Maybe (Chunk s, s)
streamTakeN Int
n)

-- | Return the next chunk of the given size, if any, and consume it.
-- May return a smaller chunk at end of stream, but never returns an empty chunk.
popChunk :: (Stream s, Monad m) => Int -> ParserT e s m (Maybe (Chunk s))
popChunk :: Int -> ParserT e s m (Maybe (Chunk s))
popChunk Int
n = (s -> (Maybe (Chunk s), s)) -> ParserT e s m (Maybe (Chunk s))
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\s
stream -> (Maybe (Chunk s), s)
-> ((Chunk s, s) -> (Maybe (Chunk s), s))
-> Maybe (Chunk s, s)
-> (Maybe (Chunk s), s)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe (Chunk s)
forall a. Maybe a
Nothing, s
stream) ((Chunk s -> Maybe (Chunk s))
-> (Chunk s, s) -> (Maybe (Chunk s), s)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Chunk s -> Maybe (Chunk s)
forall a. a -> Maybe a
Just) (Int -> s -> Maybe (Chunk s, s)
forall s. Stream s => Int -> s -> Maybe (Chunk s, s)
streamTakeN Int
n s
stream))

-- | Drop the next chunk of the given size, if any, and consume it.
-- May return a smaller size at end of stream, but never returns size 0.
dropChunk :: (Stream s, Monad m) => Int -> ParserT e s m (Maybe Int)
dropChunk :: Int -> ParserT e s m (Maybe Int)
dropChunk Int
n = (s -> (Maybe Int, s)) -> ParserT e s m (Maybe Int)
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\s
stream -> (Maybe Int, s)
-> ((Int, s) -> (Maybe Int, s)) -> Maybe (Int, s) -> (Maybe Int, s)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe Int
forall a. Maybe a
Nothing, s
stream) ((Int -> Maybe Int) -> (Int, s) -> (Maybe Int, s)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Int -> Maybe Int
forall a. a -> Maybe a
Just) (Int -> s -> Maybe (Int, s)
forall s. Stream s => Int -> s -> Maybe (Int, s)
streamDropN Int
n s
stream))

-- | Is this the end of the stream?
isEnd :: (Stream s, Monad m) => ParserT e s m Bool
isEnd :: ParserT e s m Bool
isEnd = Maybe (Token s) -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe (Token s) -> Bool)
-> ParserT e s m (Maybe (Token s)) -> ParserT e s m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
ParserT e s m (Maybe (Token s))
peekToken

-- | Match the end of the stream or terminate the parser.
matchEnd :: (Stream s, Monad m) => ParserT e s m ()
matchEnd :: ParserT e s m ()
matchEnd = ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
ParserT e s m (Maybe (Token s))
peekToken ParserT e s m (Maybe (Token s))
-> (Maybe (Token s) -> ParserT e s m ()) -> ParserT e s m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ParserT e s m ()
-> (Token s -> ParserT e s m ())
-> Maybe (Token s)
-> ParserT e s m ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> ParserT e s m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (ParserT e s m () -> Token s -> ParserT e s m ()
forall a b. a -> b -> a
const ParserT e s m ()
forall (f :: * -> *) a. Alternative f => f a
empty)

-- | Return the next token or terminate the parser at end of stream.
anyToken :: (Stream s, Monad m) => ParserT e s m (Token s)
anyToken :: ParserT e s m (Token s)
anyToken = ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
ParserT e s m (Maybe (Token s))
popToken ParserT e s m (Maybe (Token s))
-> (Maybe (Token s) -> ParserT e s m (Token s))
-> ParserT e s m (Token s)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ParserT e s m (Token s)
-> (Token s -> ParserT e s m (Token s))
-> Maybe (Token s)
-> ParserT e s m (Token s)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ParserT e s m (Token s)
forall (f :: * -> *) a. Alternative f => f a
empty Token s -> ParserT e s m (Token s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Return the next chunk of the given size or terminate the parser at end of stream.
-- May return a smaller chunk at end of stream, but never returns an empty chunk.
anyChunk :: (Stream s, Monad m) => Int -> ParserT e s m (Chunk s)
anyChunk :: Int -> ParserT e s m (Chunk s)
anyChunk Int
n = Int -> ParserT e s m (Maybe (Chunk s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
Int -> ParserT e s m (Maybe (Chunk s))
popChunk Int
n ParserT e s m (Maybe (Chunk s))
-> (Maybe (Chunk s) -> ParserT e s m (Chunk s))
-> ParserT e s m (Chunk s)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ParserT e s m (Chunk s)
-> (Chunk s -> ParserT e s m (Chunk s))
-> Maybe (Chunk s)
-> ParserT e s m (Chunk s)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ParserT e s m (Chunk s)
forall (f :: * -> *) a. Alternative f => f a
empty Chunk s -> ParserT e s m (Chunk s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Match the next token with the given predicate or terminate the parser at predicate false or end of stream.
satisfyToken :: (Stream s, Monad m) => (Token s -> Bool) -> ParserT e s m (Token s)
satisfyToken :: (Token s -> Bool) -> ParserT e s m (Token s)
satisfyToken Token s -> Bool
p = do
  Maybe (Token s)
m <- ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
ParserT e s m (Maybe (Token s))
popToken
  case Maybe (Token s)
m of
    Just Token s
c | Token s -> Bool
p Token s
c -> Token s -> ParserT e s m (Token s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Token s
c
    Maybe (Token s)
_ -> ParserT e s m (Token s)
forall (f :: * -> *) a. Alternative f => f a
empty

-- | Folds over a stream of tokens while the boolean value is true.
-- Always succeeds, even at end of stream.
foldTokensWhile :: (Stream s, Monad m) => (Token s -> x -> (Bool, x)) -> (x -> x) -> x -> ParserT e s m x
foldTokensWhile :: (Token s -> x -> (Bool, x)) -> (x -> x) -> x -> ParserT e s m x
foldTokensWhile Token s -> x -> (Bool, x)
f x -> x
g = x -> ParserT e s m x
go where
  go :: x -> ParserT e s m x
go !x
x = do
    Maybe (Token s)
m <- ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
ParserT e s m (Maybe (Token s))
peekToken
    case Maybe (Token s)
m of
      Maybe (Token s)
Nothing -> x -> ParserT e s m x
forall (f :: * -> *) a. Applicative f => a -> f a
pure (x -> x
g x
x)
      Just Token s
c ->
        let (Bool
ok, x
newX) = Token s -> x -> (Bool, x)
f Token s
c x
x
        in if Bool
ok
          then ParserT e s m (Maybe (Token s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
ParserT e s m (Maybe (Token s))
popToken ParserT e s m (Maybe (Token s))
-> ParserT e s m x -> ParserT e s m x
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> x -> ParserT e s m x
go x
newX
          else x -> ParserT e s m x
forall (f :: * -> *) a. Applicative f => a -> f a
pure x
x

-- | Take tokens into a chunk while they satisfy the given predicate.
-- Always succeeds, even at end of stream. May return an empty chunk.
takeTokensWhile :: (Stream s, Monad m) => (Token s -> Bool) -> ParserT e s m (Chunk s)
takeTokensWhile :: (Token s -> Bool) -> ParserT e s m (Chunk s)
takeTokensWhile = (s -> (Chunk s, s)) -> ParserT e s m (Chunk s)
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((s -> (Chunk s, s)) -> ParserT e s m (Chunk s))
-> ((Token s -> Bool) -> s -> (Chunk s, s))
-> (Token s -> Bool)
-> ParserT e s m (Chunk s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Token s -> Bool) -> s -> (Chunk s, s)
forall s. Stream s => (Token s -> Bool) -> s -> (Chunk s, s)
streamTakeWhile

-- | Drop tokens and return chunk size while they satisfy the given predicate.
-- Always succeeds, even at end of stream. May return empty chunk size 0.
dropTokensWhile :: (Stream s, Monad m) => (Token s -> Bool) -> ParserT e s m Int
dropTokensWhile :: (Token s -> Bool) -> ParserT e s m Int
dropTokensWhile = (s -> (Int, s)) -> ParserT e s m Int
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((s -> (Int, s)) -> ParserT e s m Int)
-> ((Token s -> Bool) -> s -> (Int, s))
-> (Token s -> Bool)
-> ParserT e s m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Token s -> Bool) -> s -> (Int, s)
forall s. Stream s => (Token s -> Bool) -> s -> (Int, s)
streamDropWhile

-- | Match token with equality or terminate the parser at inequality or end of stream.
matchToken :: (Stream s, Monad m, Eq (Token s)) => Token s -> ParserT e s m (Token s)
matchToken :: Token s -> ParserT e s m (Token s)
matchToken = (Token s -> Bool) -> ParserT e s m (Token s)
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
(Token s -> Bool) -> ParserT e s m (Token s)
satisfyToken ((Token s -> Bool) -> ParserT e s m (Token s))
-> (Token s -> Token s -> Bool)
-> Token s
-> ParserT e s m (Token s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Token s -> Token s -> Bool
forall a. Eq a => a -> a -> Bool
(==)

-- | Match chunk with equality or terminate the parser at inequality or end of stream.
matchChunk :: (Stream s, Monad m, Eq (Chunk s)) => Chunk s -> ParserT e s m (Chunk s)
matchChunk :: Chunk s -> ParserT e s m (Chunk s)
matchChunk Chunk s
k = Int -> ParserT e s m (Maybe (Chunk s))
forall s (m :: * -> *) e.
(Stream s, Monad m) =>
Int -> ParserT e s m (Maybe (Chunk s))
popChunk (Chunk s -> Int
forall chunk token. Chunked chunk token => chunk -> Int
chunkLength Chunk s
k) ParserT e s m (Maybe (Chunk s))
-> (Maybe (Chunk s) -> ParserT e s m (Chunk s))
-> ParserT e s m (Chunk s)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ParserT e s m (Chunk s)
-> (Chunk s -> ParserT e s m (Chunk s))
-> Maybe (Chunk s)
-> ParserT e s m (Chunk s)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ParserT e s m (Chunk s)
forall (f :: * -> *) a. Alternative f => f a
empty (\Chunk s
j -> if Chunk s
k Chunk s -> Chunk s -> Bool
forall a. Eq a => a -> a -> Bool
== Chunk s
j then Chunk s -> ParserT e s m (Chunk s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Chunk s
j else ParserT e s m (Chunk s)
forall (f :: * -> *) a. Alternative f => f a
empty)