Parser.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Megaparsec-based parser for 'Value's in the textual value format.
-- The difference between this and the reader defined in
-- "Futhark.Data.Reader" is that we don't try to handle both the
-- textual and binary format - only the former.  On the other hand,
-- this parser has (much) better error messages and can be easily used
-- by other parsers (like the ones for FutharkScript or test blocks).
module Futhark.Data.Parser
  ( parsePrimType,
    parseType,
    parsePrimValue,
    parseValue,
  )
where

import Control.Monad (unless)
import Data.Char (digitToInt, isDigit, isHexDigit)
import Data.Functor
import qualified Data.Scientific as Sci
import qualified Data.Set as S
import qualified Data.Text as T
import qualified Data.Vector.Storable as SVec
import Data.Void
import Futhark.Data
import Text.Megaparsec
import Text.Megaparsec.Char (char)
import Text.Megaparsec.Char.Lexer (charLiteral, signed)
import Prelude hiding (exponent)

-- | Parse the name of a primitive type.  Does *not* consume any
-- trailing whitespace, nor does it permit any internal whitespace.
parsePrimType :: Parsec Void T.Text PrimType
parsePrimType =
  choice
    [ "i8" $> I8,
      "i16" $> I16,
      "i32" $> I32,
      "i64" $> I64,
      "u8" $> U8,
      "u16" $> U16,
      "u32" $> U32,
      "u64" $> U64,
      "f16" $> F16,
      "f32" $> F32,
      "f64" $> F64,
      "bool" $> Bool
    ]

allowUnderscores :: String -> (Char -> Bool) -> Parsec Void T.Text T.Text
allowUnderscores desc p =
  T.filter (/= '_')
    <$> ( (<>)
            <$> takeWhile1P (Just desc) p
            <*> takeWhileP (Just descOrUnderscore) pOrUnderscore
        )
  where
    descOrUnderscore = desc <> " or underscore"
    pOrUnderscore c = p c || c == '_'

-- Adapted from megaparsec.
decimal :: (Num a) => Parsec Void T.Text a
decimal =
  mkNum <$> allowUnderscores "digit" isDigit
  where
    mkNum = T.foldl' step 0
    step a c = a * 10 + fromIntegral (digitToInt c)

-- Adapted from megaparsec.
binary :: (Num a) => Parsec Void T.Text a
binary =
  mkNum <$> allowUnderscores "binary digit" isBinDigit
  where
    mkNum = T.foldl' step 0
    step a c = a * 2 + fromIntegral (digitToInt c)
    isBinDigit x = x == '0' || x == '1'

-- Adapted from megaparsec.
hexadecimal :: (Num a) => Parsec Void T.Text a
hexadecimal =
  mkNum <$> allowUnderscores "hexadecimal digit" isHexDigit
  where
    mkNum = T.foldl' step 0
    step a c = a * 16 + fromIntegral (digitToInt c)

parseInteger :: Parsec Void T.Text Integer
parseInteger =
  signed (pure ()) $
    choice
      [ "0b" *> binary,
        "0x" *> hexadecimal,
        decimal
      ]

scalar :: (SVec.Storable a) => (Vector Int -> Vector a -> Value) -> a -> Value
scalar f x = f mempty (SVec.singleton x)

parseIntConst :: (Integer -> Value) -> Parsec Void T.Text Value
parseIntConst def = do
  x <- parseInteger
  notFollowedBy $ choice ["f16", "f32", "f64", ".", "e"]
  choice
    [ intV I8Value x "i8",
      intV I16Value x "i16",
      intV I32Value x "i32",
      intV I64Value x "i64",
      intV U8Value x "u8",
      intV U16Value x "u16",
      intV U32Value x "u32",
      intV U64Value x "u64",
      pure $ def x
    ]
  where
    intV mk x suffix =
      suffix $> scalar mk (fromInteger x)

-- Adapted from megaparsec.
float :: (RealFloat a) => Parsec Void T.Text a
float = do
  c' <- decimal
  Sci.toRealFloat
    <$> ( ( do
              (c, e') <- dotDecimal c'
              e <- option e' $ try $ exponent e'
              pure $ Sci.scientific c e
          )
            <|> (Sci.scientific c' <$> exponent 0)
        )
  where
    exponent e' = do
      void $ choice ["e", "E"]
      (+ e') <$> signed (pure ()) decimal
    dotDecimal c' = do
      void "."
      mkNum <$> allowUnderscores "digit" isDigit
      where
        mkNum = T.foldl' step (c', 0)
        step (a, e') c =
          (a * 10 + fromIntegral (digitToInt c), e' - 1)

parseFloatConst :: Parsec Void T.Text Value
parseFloatConst =
  choice
    [ "f16.nan" $> scalar F16Value (0 / 0),
      "f32.nan" $> scalar F32Value (0 / 0),
      "f64.nan" $> scalar F64Value (0 / 0),
      --
      "f16.inf" $> scalar F16Value (1 / 0),
      "f32.inf" $> scalar F32Value (1 / 0),
      "f64.inf" $> scalar F64Value (1 / 0),
      --
      "-f16.inf" $> scalar F16Value (-1 / 0),
      "-f32.inf" $> scalar F32Value (-1 / 0),
      "-f64.inf" $> scalar F64Value (-1 / 0),
      numeric
    ]
  where
    numeric = do
      x <-
        signed (pure ()) $ choice [try float, fromInteger <$> decimal]
      choice
        [ floatV F16Value x "f16",
          floatV F32Value x "f32",
          floatV F64Value x "f64",
          floatV F64Value x ""
        ]

    floatV mk x suffix =
      suffix $> scalar mk (realToFrac (x :: Double))

pBool :: Parsec Void T.Text Bool
pBool = choice ["true" $> True, "false" $> False]

-- | Parse a primitive value.  Does *not* consume any trailing
-- whitespace, nor does it permit any internal whitespace.
parsePrimValue :: Parsec Void T.Text Value
parsePrimValue =
  choice
    [ try (parseIntConst $ I32Value mempty . SVec.singleton . fromInteger),
      parseFloatConst,
      scalar BoolValue <$> pBool
    ]

parseStringConst :: Parsec Void T.Text Value
parseStringConst =
  char '"' *> (putValue1 . T.pack <$> manyTill charLiteral (char '"'))

-- | Parse float that is a specific type.
pFloat ::
  (SVec.Storable a, Fractional a) =>
  (Vector Int -> Vector a -> Value) ->
  T.Text ->
  Parsec Void T.Text Value
pFloat mk suffix =
  choice
    [ chunk nan $> scalar mk (0 / 0),
      chunk inf $> scalar mk (1 / 0),
      chunk neginf $> scalar mk (-1 / 0),
      numeric
    ]
  where
    nan = suffix <> ".nan"
    inf = suffix <> ".inf"
    neginf = "-" <> inf
    numeric = do
      x <- signed (pure ()) $ choice [try float, fromInteger <$> decimal]
      void $ optional $ chunk suffix
      pure $ scalar mk $ realToFrac (x :: Double)

pPrimOfType :: PrimType -> Parsec Void T.Text Value
pPrimOfType Bool = scalar BoolValue <$> pBool
pPrimOfType F16 = pFloat F16Value "f16"
pPrimOfType F32 = pFloat F32Value "f32"
pPrimOfType F64 = pFloat F64Value "f64"
pPrimOfType I8 = scalar I8Value . fromInteger <$> parseInteger <* optional "i8"
pPrimOfType I16 = scalar I16Value . fromInteger <$> parseInteger <* optional "i16"
pPrimOfType I32 = scalar I32Value . fromInteger <$> parseInteger <* optional "i32"
pPrimOfType I64 = scalar I64Value . fromInteger <$> parseInteger <* optional "i64"
pPrimOfType U8 = scalar U8Value . fromInteger <$> parseInteger <* optional "u8"
pPrimOfType U16 = scalar U16Value . fromInteger <$> parseInteger <* optional "u16"
pPrimOfType U32 = scalar U32Value . fromInteger <$> parseInteger <* optional "u32"
pPrimOfType U64 = scalar U64Value . fromInteger <$> parseInteger <* optional "u64"

lexeme :: Parsec Void T.Text () -> Parsec Void T.Text a -> Parsec Void T.Text a
lexeme sep p = p <* sep

inBrackets :: Parsec Void T.Text () -> Parsec Void T.Text a -> Parsec Void T.Text a
inBrackets sep = between (lexeme sep "[") (lexeme sep "]")

-- | Parse a type.  Does *not* consume any trailing whitespace, nor
-- does it permit any internal whitespace.
parseType :: Parsec Void T.Text ValueType
parseType = ValueType <$> many parseDim <*> parsePrimType
  where
    parseDim = fromInteger <$> ("[" *> parseInteger <* "]")

parseEmpty :: Parsec Void T.Text Value
parseEmpty = do
  ValueType dims t <- parseType
  unless (product dims == 0) $ fail "Expected at least one empty dimension"
  pure $ case t of
    I8 -> I8Value (SVec.fromList dims) mempty
    I16 -> I16Value (SVec.fromList dims) mempty
    I32 -> I32Value (SVec.fromList dims) mempty
    I64 -> I64Value (SVec.fromList dims) mempty
    U8 -> U8Value (SVec.fromList dims) mempty
    U16 -> U16Value (SVec.fromList dims) mempty
    U32 -> U32Value (SVec.fromList dims) mempty
    U64 -> U64Value (SVec.fromList dims) mempty
    F16 -> F16Value (SVec.fromList dims) mempty
    F32 -> F32Value (SVec.fromList dims) mempty
    F64 -> F64Value (SVec.fromList dims) mempty
    Bool -> BoolValue (SVec.fromList dims) mempty

-- | Parse a value, given a post-lexeme parser for whitespace and a parser for primitive values..
parseValueWith :: Parsec Void T.Text () -> Parsec Void T.Text Value -> Parsec Void T.Text Value
parseValueWith sep pPrim =
  choice
    [ lexeme sep pPrim,
      lexeme sep parseStringConst,
      putValue' . inBrackets sep $ do
        v <- parseValue sep
        (v :)
          <$> choice
            [ pComma *> parseValueWith sep (pPrimOfType (valueElemType v)) `sepEndBy` pComma,
              pure []
            ],
      lexeme sep $ "empty(" *> parseEmpty <* ")"
    ]
  where
    pComma = lexeme sep ","

    putValue' :: (PutValue v) => Parsec Void T.Text v -> Parsec Void T.Text Value
    putValue' p = do
      o <- getOffset
      x <- p
      case putValue x of
        Nothing ->
          parseError . FancyError o . S.singleton $
            ErrorFail "array is irregular or has elements of multiple types."
        Just v ->
          pure v

-- | Parse a value, given a post-lexeme parser for whitespace.
parseValue :: Parsec Void T.Text () -> Parsec Void T.Text Value
parseValue sep = parseValueWith sep parsePrimValue