{-# LANGUAGE OverloadedStrings #-}

module Futhark.IR.Primitive.Parse
  ( pPrimValue,
    pPrimType,
    pFloatType,
    pIntType,

    -- * Building blocks
    constituent,
    lexeme,
    keyword,
    whitespace,
  )
where

import Data.Char (isAlphaNum)
import Data.Functor
import qualified Data.Text as T
import Data.Void
import Futhark.IR.Primitive
import Futhark.Util.Pretty hiding (empty)
import Text.Megaparsec
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L

type Parser = Parsec Void T.Text

constituent :: Char -> Bool
constituent :: Char -> Bool
constituent Char
c = Char -> Bool
isAlphaNum Char
c Bool -> Bool -> Bool
|| (Char
c Char -> [Char] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ([Char]
"_/'+-=!&^.<>*|" :: String))

whitespace :: Parser ()
whitespace :: Parser ()
whitespace = Parser () -> Parser () -> Parser () -> Parser ()
forall e s (m :: * -> *).
MonadParsec e s m =>
m () -> m () -> m () -> m ()
L.space Parser ()
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m ()
space1 (Tokens Text -> Parser ()
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
Tokens s -> m ()
L.skipLineComment Tokens Text
"--") Parser ()
forall (f :: * -> *) a. Alternative f => f a
empty

lexeme :: Parser a -> Parser a
lexeme :: forall a. Parser a -> Parser a
lexeme = ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (ParsecT Void Text Identity a -> ParsecT Void Text Identity a)
-> (ParsecT Void Text Identity a -> ParsecT Void Text Identity a)
-> ParsecT Void Text Identity a
-> ParsecT Void Text Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser ()
-> ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m () -> m a -> m a
L.lexeme Parser ()
whitespace

keyword :: T.Text -> Parser ()
keyword :: Text -> Parser ()
keyword Text
s = Parser () -> Parser ()
forall a. Parser a -> Parser a
lexeme (Parser () -> Parser ()) -> Parser () -> Parser ()
forall a b. (a -> b) -> a -> b
$ Tokens Text -> ParsecT Void Text Identity (Tokens Text)
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
chunk Text
Tokens Text
s ParsecT Void Text Identity Text -> Parser () -> Parser ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> ParsecT Void Text Identity Char -> Parser ()
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m ()
notFollowedBy ((Token Text -> Bool) -> ParsecT Void Text Identity (Token Text)
forall e s (m :: * -> *).
MonadParsec e s m =>
(Token s -> Bool) -> m (Token s)
satisfy Char -> Bool
Token Text -> Bool
constituent)

pIntValue :: Parser IntValue
pIntValue :: Parser IntValue
pIntValue = Parser IntValue -> Parser IntValue
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Parser IntValue -> Parser IntValue)
-> Parser IntValue -> Parser IntValue
forall a b. (a -> b) -> a -> b
$ do
  Integer
x <- Parser ()
-> ParsecT Void Text Identity Integer
-> ParsecT Void Text Identity Integer
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m () -> m a -> m a
L.signed (() -> Parser ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) ParsecT Void Text Identity Integer
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
L.decimal
  IntType
t <- Parser IntType
pIntType
  IntValue -> Parser IntValue
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntValue -> Parser IntValue) -> IntValue -> Parser IntValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t (Integer
x :: Integer)

pFloatValue :: Parser FloatValue
pFloatValue :: Parser FloatValue
pFloatValue =
  [Parser FloatValue] -> Parser FloatValue
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice
    [ Parser FloatValue
pNum,
      Text -> Parser ()
keyword Text
"f32.nan" Parser () -> FloatValue -> Parser FloatValue
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Float -> FloatValue
Float32Value (Float
0 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
0),
      Text -> Parser ()
keyword Text
"f32.inf" Parser () -> FloatValue -> Parser FloatValue
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Float -> FloatValue
Float32Value (Float
1 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
0),
      Text -> Parser ()
keyword Text
"-f32.inf" Parser () -> FloatValue -> Parser FloatValue
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Float -> FloatValue
Float32Value (-Float
1 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
0),
      Text -> Parser ()
keyword Text
"f64.nan" Parser () -> FloatValue -> Parser FloatValue
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Double -> FloatValue
Float64Value (Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0),
      Text -> Parser ()
keyword Text
"f64.inf" Parser () -> FloatValue -> Parser FloatValue
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Double -> FloatValue
Float64Value (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0),
      Text -> Parser ()
keyword Text
"-f64.inf" Parser () -> FloatValue -> Parser FloatValue
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Double -> FloatValue
Float64Value (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    ]
  where
    pNum :: Parser FloatValue
pNum = Parser FloatValue -> Parser FloatValue
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Parser FloatValue -> Parser FloatValue)
-> Parser FloatValue -> Parser FloatValue
forall a b. (a -> b) -> a -> b
$ do
      Double
x <- Parser ()
-> ParsecT Void Text Identity Double
-> ParsecT Void Text Identity Double
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m () -> m a -> m a
L.signed (() -> Parser ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) ParsecT Void Text Identity Double
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, RealFloat a) =>
m a
L.float
      FloatType
t <- Parser FloatType
pFloatType
      FloatValue -> Parser FloatValue
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FloatValue -> Parser FloatValue)
-> FloatValue -> Parser FloatValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
t (Double
x :: Double)

pBoolValue :: Parser Bool
pBoolValue :: Parser Bool
pBoolValue =
  [Parser Bool] -> Parser Bool
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice
    [ Text -> Parser ()
keyword Text
"true" Parser () -> Bool -> Parser Bool
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool
True,
      Text -> Parser ()
keyword Text
"false" Parser () -> Bool -> Parser Bool
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool
False
    ]

-- | Defined in this module for convenience.
pPrimValue :: Parser PrimValue
pPrimValue :: Parser PrimValue
pPrimValue =
  [Parser PrimValue] -> Parser PrimValue
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice
    [ FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue) -> Parser FloatValue -> Parser PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser FloatValue
pFloatValue,
      IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> Parser IntValue -> Parser PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser IntValue
pIntValue,
      Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> Parser Bool -> Parser PrimValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Bool
pBoolValue,
      PrimValue
UnitValue PrimValue -> ParsecT Void Text Identity Text -> Parser PrimValue
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ParsecT Void Text Identity Text
"()"
    ]
    Parser PrimValue -> [Char] -> Parser PrimValue
forall e s (m :: * -> *) a.
MonadParsec e s m =>
m a -> [Char] -> m a
<?> [Char]
"primitive value"

pFloatType :: Parser FloatType
pFloatType :: Parser FloatType
pFloatType = [Parser FloatType] -> Parser FloatType
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice ([Parser FloatType] -> Parser FloatType)
-> [Parser FloatType] -> Parser FloatType
forall a b. (a -> b) -> a -> b
$ (FloatType -> Parser FloatType)
-> [FloatType] -> [Parser FloatType]
forall a b. (a -> b) -> [a] -> [b]
map FloatType -> Parser FloatType
forall {b}. Pretty b => b -> ParsecT Void Text Identity b
p [FloatType]
allFloatTypes
  where
    p :: b -> ParsecT Void Text Identity b
p b
t = Text -> Parser ()
keyword (b -> Text
forall a. Pretty a => a -> Text
prettyText b
t) Parser () -> b -> ParsecT Void Text Identity b
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> b
t

pIntType :: Parser IntType
pIntType :: Parser IntType
pIntType = [Parser IntType] -> Parser IntType
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice ([Parser IntType] -> Parser IntType)
-> [Parser IntType] -> Parser IntType
forall a b. (a -> b) -> a -> b
$ (IntType -> Parser IntType) -> [IntType] -> [Parser IntType]
forall a b. (a -> b) -> [a] -> [b]
map IntType -> Parser IntType
forall {b}. Pretty b => b -> ParsecT Void Text Identity b
p [IntType]
allIntTypes
  where
    p :: b -> ParsecT Void Text Identity b
p b
t = Text -> Parser ()
keyword (b -> Text
forall a. Pretty a => a -> Text
prettyText b
t) Parser () -> b -> ParsecT Void Text Identity b
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> b
t

pPrimType :: Parser PrimType
pPrimType :: Parser PrimType
pPrimType =
  [Parser PrimType] -> Parser PrimType
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice [PrimType -> Parser PrimType
forall {b}. Pretty b => b -> ParsecT Void Text Identity b
p PrimType
Bool, PrimType -> Parser PrimType
forall {b}. Pretty b => b -> ParsecT Void Text Identity b
p PrimType
Unit, FloatType -> PrimType
FloatType (FloatType -> PrimType) -> Parser FloatType -> Parser PrimType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser FloatType
pFloatType, IntType -> PrimType
IntType (IntType -> PrimType) -> Parser IntType -> Parser PrimType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser IntType
pIntType]
  where
    p :: b -> ParsecT Void Text Identity b
p b
t = Text -> Parser ()
keyword (b -> Text
forall a. Pretty a => a -> Text
prettyText b
t) Parser () -> b -> ParsecT Void Text Identity b
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> b
t