{-# LANGUAGE OverloadedStrings #-}

module Attoparsec.Extra (
    module Attoparsec,
    char,
    endOfInputEx,
    isSuccessful,
    label,
    parseOnlyL,
    takeL,
    definiteDouble,
    withInputSize,
    (??),
    (<+>),
) where

import           RON.Prelude

import           Data.Attoparsec.ByteString.Char8 (anyChar, decimal, isDigit_w8,
                                                   signed)
import           Data.Attoparsec.ByteString.Lazy as Attoparsec
import qualified Data.Attoparsec.Internal.Types as Internal
import qualified Data.ByteString as BS
import           Data.ByteString.Lazy (fromStrict)
import qualified Data.Scientific as Sci
import           GHC.Real (toInteger)

-- | TODO(2020-06-17, cblp) make parser lazy/incremental
-- TODO(2021-04-25, cblp) remove in favor of attoparsec-0.14 lazy version
parseOnlyL :: Parser a -> ByteStringL -> Either String a
parseOnlyL :: Parser a -> ByteStringL -> Either String a
parseOnlyL Parser a
p = Result a -> Either String a
forall r. Result r -> Either String r
eitherResult (Result a -> Either String a)
-> (ByteStringL -> Result a) -> ByteStringL -> Either String a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser a -> ByteStringL -> Result a
forall a. Parser a -> ByteStringL -> Result a
parse Parser a
p

-- | 'Attoparsec.take' adapter to 'ByteStringL'
takeL :: Int -> Parser ByteStringL
takeL :: Int -> Parser ByteStringL
takeL = (ByteString -> ByteStringL)
-> Parser ByteString ByteString -> Parser ByteStringL
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteStringL
fromStrict (Parser ByteString ByteString -> Parser ByteStringL)
-> (Int -> Parser ByteString ByteString)
-> Int
-> Parser ByteStringL
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Parser ByteString ByteString
Attoparsec.take

getPos :: Parser Int
getPos :: Parser Int
getPos =
    (forall r.
 State ByteString
 -> Pos
 -> More
 -> Failure ByteString (State ByteString) r
 -> Success ByteString (State ByteString) Int r
 -> IResult ByteString r)
-> Parser Int
forall i a.
(forall r.
 State i
 -> Pos
 -> More
 -> Failure i (State i) r
 -> Success i (State i) a r
 -> IResult i r)
-> Parser i a
Internal.Parser ((forall r.
  State ByteString
  -> Pos
  -> More
  -> Failure ByteString (State ByteString) r
  -> Success ByteString (State ByteString) Int r
  -> IResult ByteString r)
 -> Parser Int)
-> (forall r.
    State ByteString
    -> Pos
    -> More
    -> Failure ByteString (State ByteString) r
    -> Success ByteString (State ByteString) Int r
    -> IResult ByteString r)
-> Parser Int
forall a b. (a -> b) -> a -> b
$ \State ByteString
t Pos
pos More
more Failure ByteString (State ByteString) r
_ Success ByteString (State ByteString) Int r
suc -> Success ByteString (State ByteString) Int r
suc State ByteString
t Pos
pos More
more (Int -> IResult ByteString r) -> Int -> IResult ByteString r
forall a b. (a -> b) -> a -> b
$ Pos -> Int
Internal.fromPos Pos
pos

withInputSize :: Parser a -> Parser (Int, a)
withInputSize :: Parser a -> Parser (Int, a)
withInputSize Parser a
p = do
    Int
posBefore <- Parser Int
getPos
    a
r <- Parser a
p
    Int
posAfter <- Parser Int
getPos
    (Int, a) -> Parser (Int, a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
posAfter Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
posBefore, a
r)

label :: String -> Parser a -> Parser a
label :: String -> Parser a -> Parser a
label = (Parser a -> String -> Parser a) -> String -> Parser a -> Parser a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Parser a -> String -> Parser a
forall i a. Parser i a -> String -> Parser i a
(<?>)

-- label' :: String -> Parser a -> Parser a
-- label' name p = do
--     pos <- getPos
--     label (name ++ ':' : show pos) p

-- | Variant of 'endOfInput' with a more debuggable message.
endOfInputEx :: Parser ()
endOfInputEx :: Parser ()
endOfInputEx = do
    Bool
weAreAtEnd <- Parser ByteString Bool
forall t. Chunk t => Parser t Bool
atEnd
    Bool -> Parser () -> Parser ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
weAreAtEnd (Parser () -> Parser ()) -> Parser () -> Parser ()
forall a b. (a -> b) -> a -> b
$ do
        Int
pos <- Parser Int
getPos
        ByteString
rest <- Int -> Parser ByteString ByteString
takeAtMost Int
11
        let cite :: ByteString
cite
                | ByteString -> Int
BS.length ByteString
rest Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
11 = ByteString
rest
                | Bool
otherwise           = Int -> ByteString -> ByteString
BS.take Int
10 ByteString
rest ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"..."
        String -> Parser ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Parser ()) -> String -> Parser ()
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a s. (Show a, IsString s) => a -> s
show Int
pos String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": extra input: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a s. (Show a, IsString s) => a -> s
show ByteString
cite

takeAtMost :: Int -> Parser ByteString
takeAtMost :: Int -> Parser ByteString ByteString
takeAtMost Int
limit = do
    Int
pos0 <- Parser Int
getPos
    [Word8] -> ByteString
BS.pack ([Word8] -> ByteString)
-> Parser ByteString [Word8] -> Parser ByteString ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString Word8 -> Parser () -> Parser ByteString [Word8]
forall (f :: * -> *) a b. Alternative f => f a -> f b -> f [a]
manyTill Parser ByteString Word8
anyWord8 (Int -> Parser ()
checkLimit (Int -> Parser ()) -> Int -> Parser ()
forall a b. (a -> b) -> a -> b
$ Int
pos0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
limit)
  where
    checkLimit :: Int -> Parser ()
checkLimit Int
maxPos = do
        Int
pos <- Parser Int
getPos
        Bool -> Parser ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxPos) Parser () -> Parser () -> Parser ()
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser ()
forall t. Chunk t => Parser t ()
endOfInput

(??) :: Applicative f => Maybe a -> f a -> f a
Maybe a
a ?? :: Maybe a -> f a -> f a
?? f a
alt = f a -> (a -> f a) -> Maybe a -> f a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe f a
alt a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
a

-- | Apply parser and check it is applied successfully.
-- Kinda opposite to 'guard'.
isSuccessful :: Alternative f => f a -> f Bool
isSuccessful :: f a -> f Bool
isSuccessful f a
p = f a
p f a -> Bool -> f Bool
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool
True f Bool -> f Bool -> f Bool
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Bool -> f Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

char :: Char -> Parser Char
char :: Char -> Parser Char
char Char
c = do
    Char
c' <- Parser Char
anyChar
    if Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
c' then
        Char -> Parser Char
forall (f :: * -> *) a. Applicative f => a -> f a
pure Char
c
    else
        String -> Parser Char
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Parser Char) -> String -> Parser Char
forall a b. (a -> b) -> a -> b
$ String
"Expected " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char -> String
forall a s. (Show a, IsString s) => a -> s
show Char
c String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", got " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char -> String
forall a s. (Show a, IsString s) => a -> s
show Char
c'

-- | Parses a definite double, i.e. it is not an integer. For this, the double has either a '.', and 'e'/'E' part or both.
{-# INLINE definiteDouble #-}
definiteDouble :: Parser Double
definiteDouble :: Parser Double
definiteDouble = do
    let parseIntegerPart :: Parser Integer
parseIntegerPart = Parser Integer -> Parser Integer
forall a. Num a => Parser a -> Parser a
signed Parser Integer
forall a. Integral a => Parser a
decimal
    let parseDot :: Parser Char
parseDot = Char -> Parser Char
char Char
'.'
    let parseFractionalPartWithLength :: Parser ByteString (Integer, Int)
parseFractionalPartWithLength =
            ((Integer, Int) -> Word8 -> (Integer, Int))
-> (Integer, Int) -> ByteString -> (Integer, Int)
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' (Integer, Int) -> Word8 -> (Integer, Int)
forall a a b. (Integral a, Num a, Num b) => (a, b) -> a -> (a, b)
step (Integer
0, Int
0) (ByteString -> (Integer, Int))
-> Parser ByteString ByteString -> Parser ByteString (Integer, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` (Word8 -> Bool) -> Parser ByteString ByteString
Attoparsec.takeWhile1 Word8 -> Bool
isDigit_w8
                where step :: (a, b) -> a -> (a, b)
step (a
a, b
l) a
w = (a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
10 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
48), b
l b -> b -> b
forall a. Num a => a -> a -> a
+ b
1)
    let parseExponent :: Parser Int
parseExponent = (Char -> Parser Char
char Char
'e' Parser Char -> Parser Char -> Parser Char
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Char -> Parser Char
char Char
'E') Parser Char -> Parser Int -> Parser Int
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser Int -> Parser Int
forall a. Num a => Parser a -> Parser a
signed Parser Int
forall a. Integral a => Parser a
decimal

    let withDot :: Parser Double
withDot = do
          Maybe Integer
i <- Parser Integer -> Parser ByteString (Maybe Integer)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
optional Parser Integer
parseIntegerPart
          Char
_ <- Parser Char
parseDot
          (Integer
f, Int
l) <- Parser ByteString (Integer, Int)
parseFractionalPartWithLength
          Maybe Int
e <- Parser Int -> Parser ByteString (Maybe Int)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
optional Parser Int
parseExponent
          Double -> Parser Double
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Parser Double) -> Double -> Parser Double
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Int -> Int -> Double
buildDouble (Integer -> Maybe Integer -> Integer
forall a. a -> Maybe a -> a
fromMaybe Integer
0 Maybe Integer
i) Integer
f Int
l (Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 Maybe Int
e)

    let withE :: Parser Double
withE = do
          Maybe Integer
i <- Parser Integer -> Parser ByteString (Maybe Integer)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
optional Parser Integer
parseIntegerPart
          Integer -> Integer -> Int -> Int -> Double
buildDouble (Integer -> Maybe Integer -> Integer
forall a. a -> Maybe a -> a
fromMaybe Integer
0 Maybe Integer
i) Integer
0 Int
0 (Int -> Double) -> Parser Int -> Parser Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Int
parseExponent

    Parser Double
withDot Parser Double -> Parser Double -> Parser Double
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser Double
withE

buildDouble :: Integer -> Integer -> Int -> Int -> Double
buildDouble :: Integer -> Integer -> Int -> Int -> Double
buildDouble Integer
integerPart Integer
fractionalPart Int
fractionalPartLength Int
exponentPart =
    let addOrSubFractionalPart :: Integer
addOrSubFractionalPart = if Integer
integerPart Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then -Integer
fractionalPart else Integer
fractionalPart
        coeff :: Integer
coeff = Integer
integerPart Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
10 Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
fractionalPartLength Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer
forall a. Integral a => a -> Integer
toInteger Integer
addOrSubFractionalPart
        exponent :: Int
exponent = Int
exponentPart Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
fractionalPartLength
     in Scientific -> Double
forall a. RealFloat a => Scientific -> a
Sci.toRealFloat (Scientific -> Double) -> Scientific -> Double
forall a b. (a -> b) -> a -> b
$ Integer -> Int -> Scientific
Sci.scientific Integer
coeff Int
exponent

(<+>) :: Parser a -> Parser a -> Parser a
<+> :: Parser a -> Parser a -> Parser a
(<+>) Parser a
p1 Parser a
p2 = (forall r.
 State ByteString
 -> Pos
 -> More
 -> Failure ByteString (State ByteString) r
 -> Success ByteString (State ByteString) a r
 -> IResult ByteString r)
-> Parser a
forall i a.
(forall r.
 State i
 -> Pos
 -> More
 -> Failure i (State i) r
 -> Success i (State i) a r
 -> IResult i r)
-> Parser i a
Internal.Parser ((forall r.
  State ByteString
  -> Pos
  -> More
  -> Failure ByteString (State ByteString) r
  -> Success ByteString (State ByteString) a r
  -> IResult ByteString r)
 -> Parser a)
-> (forall r.
    State ByteString
    -> Pos
    -> More
    -> Failure ByteString (State ByteString) r
    -> Success ByteString (State ByteString) a r
    -> IResult ByteString r)
-> Parser a
forall a b. (a -> b) -> a -> b
$ \State ByteString
t Pos
pos More
more Failure ByteString (State ByteString) r
lose Success ByteString (State ByteString) a r
suc -> let
    lose1 :: Buffer -> p -> More -> [String] -> String -> IResult ByteString r
lose1 Buffer
t' p
_pos More
more1 [String]
ctx1 String
msg1 = Parser a
-> State ByteString
-> Pos
-> More
-> Failure ByteString (State ByteString) r
-> Success ByteString (State ByteString) a r
-> IResult ByteString r
forall i a.
Parser i a
-> forall r.
   State i
   -> Pos
   -> More
   -> Failure i (State i) r
   -> Success i (State i) a r
   -> IResult i r
Internal.runParser Parser a
p2 Buffer
State ByteString
t' Pos
pos More
more1 Failure ByteString (State ByteString) r
lose2 Success ByteString (State ByteString) a r
suc
      where
        lose2 :: Failure ByteString (State ByteString) r
lose2 State ByteString
t2 Pos
pos2 More
more2 [String]
ctx2 String
msg2 = Failure ByteString (State ByteString) r
lose State ByteString
t2 Pos
pos2 More
more2 [] (String -> IResult ByteString r) -> String -> IResult ByteString r
forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords
            [ String
"Many fails:\n"
            , String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
" > " [String]
ctx1, String
":", String
msg1, String
"|\n"
            , String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
" > " [String]
ctx2, String
":", String
msg2
            ]
    in Parser a
-> State ByteString
-> Pos
-> More
-> Failure ByteString (State ByteString) r
-> Success ByteString (State ByteString) a r
-> IResult ByteString r
forall i a.
Parser i a
-> forall r.
   State i
   -> Pos
   -> More
   -> Failure i (State i) r
   -> Success i (State i) a r
   -> IResult i r
Internal.runParser Parser a
p1 State ByteString
t Pos
pos More
more Failure ByteString (State ByteString) r
forall p.
Buffer -> p -> More -> [String] -> String -> IResult ByteString r
lose1 Success ByteString (State ByteString) a r
suc
infixl 3 <+>