smallc parser in haskell using parser combinators (w/ Megaparsec)

smallc.hs

module Main where

{-
  Libraries:
  - Megaparsec
  - pretty-simple
-}

import           Data.Functor                     (($>))
import           Data.Void                        (Void)
import           Control.Applicative              ((<|>))
import qualified Text.Megaparsec            as P
import qualified Text.Megaparsec.Char       as P
import qualified Text.Megaparsec.Char.Lexer as L
import qualified Text.Pretty.Simple         as PS

type Parser = P.Parsec Void String

data Expr =
    EID           String
  | EInt          Int
  | EBool         Bool
  | EAdd          Expr Expr
  | ESub          Expr Expr
  | EMult         Expr Expr
  | EDiv          Expr Expr
  | EPow          Expr Expr
  | EGreater      Expr Expr
  | ELess         Expr Expr
  | EGreaterEqual Expr Expr
  | ELessEqual    Expr Expr
  | EEqual        Expr Expr
  | ENotEqual     Expr Expr
  | EOr           Expr Expr
  | EAnd          Expr Expr
  | ENot          Expr
  deriving (Show)

data DataType =
    IntT
  | BoolT
  deriving (Show)

data Stmt =
    SNoOp
  | SSeq      Stmt Stmt
  | SDeclare  DataType String
  | SAssign   String Expr
  | SIf       Expr Stmt Stmt
  | SFor      String Expr Expr Stmt
  | SWhile    Expr Stmt
  | SPrint    Expr
  deriving (Show)

-- | Parse a SmallC program.

main :: IO ()
main =
  do input <- getContents
     case P.parse parser "" input of
       Left  x -> putStrLn "Error!"
                  >> PS.pPrint x
       Right x -> PS.pPrint x

{- Helper parsers. -}

lexeme :: Parser a -> Parser a
lexeme = L.lexeme P.space

tok :: String -> Parser String
tok x = lexeme $ P.chunk x

{- Main parser. -}

parser :: Parser Stmt
parser =
  tok "int"
  *> tok "main"
  *> tok "("
  *> tok ")"
  *> tok "{"
  *> stmt
  <* tok "}"
  <* P.eof

{- Expression parser. -}

expr :: Parser Expr
expr = orExpr

orExpr :: Parser Expr
orExpr =
  P.choice
    [ P.try $ EOr <$> (andExpr <* tok "||") <*> orExpr
    , andExpr
    ]

andExpr :: Parser Expr
andExpr =
  P.choice
    [ P.try $ EAnd <$> (equalityExpr <* tok "&&") <*> andExpr
    , equalityExpr
    ]

equalityExpr :: Parser Expr
equalityExpr =
  P.choice
    [ P.try $ EEqual    <$> (relationalExpr <* tok "==") <*> equalityExpr
    , P.try $ ENotEqual <$> (relationalExpr <* tok "!=") <*> equalityExpr
    , relationalExpr
    ]

relationalExpr :: Parser Expr
relationalExpr =
  P.choice
    [ P.try $ ELess         <$> (additiveExpr <* tok "<")  <*> relationalExpr
    , P.try $ ELessEqual    <$> (additiveExpr <* tok "<=") <*> relationalExpr
    , P.try $ EGreater      <$> (additiveExpr <* tok ">")  <*> relationalExpr
    , P.try $ EGreaterEqual <$> (additiveExpr <* tok ">=") <*> relationalExpr
    , additiveExpr
    ]

additiveExpr :: Parser Expr
additiveExpr =
  P.choice
    [ P.try $ EAdd <$> (multiplicativeExpr <* tok "+") <*> additiveExpr
    , P.try $ ESub <$> (multiplicativeExpr <* tok "-") <*> additiveExpr
    , multiplicativeExpr
    ]

multiplicativeExpr :: Parser Expr
multiplicativeExpr =
  P.choice
    [ P.try $ EMult <$> (powerExpr <* tok "*") <*> multiplicativeExpr
    , P.try $ EDiv  <$> (powerExpr <* tok "/") <*> multiplicativeExpr
    , powerExpr
    ]

powerExpr :: Parser Expr
powerExpr =
  P.choice
    [ P.try $ EPow <$> (unaryExpr <* tok "^") <*> powerExpr
    , unaryExpr
    ]

unaryExpr :: Parser Expr
unaryExpr =
  P.choice
    [ P.try $ tok "!" *> unaryExpr
    , primaryExpr
    ]

primaryExpr :: Parser Expr
primaryExpr =
  P.choice
    [ P.try $ EInt  <$> L.signed (pure ()) (lexeme L.decimal)
    , P.try $ EBool <$> (tok "true"  $> True)
    , P.try $ EBool <$> (tok "false" $> False)
    , P.try $ EID   <$> idString
    , tok "(" *> expr <* tok ")"
    ]

idString :: Parser String
idString = lexeme ((:) <$> P.letterChar <*> P.many P.alphaNumChar)

{- Statement parser. -}

stmt :: Parser Stmt
stmt =
  P.choice
    [ P.try $ SSeq <$> stmtOptions <*> stmt
    , pure SNoOp
    ]

stmtOptions :: Parser Stmt
stmtOptions =
  P.choice
    [ P.try declareStmt
    , P.try printStmt
    , P.try ifStmt
    , P.try forStmt
    , P.try whileStmt
    , assignStmt
    ]

declareStmt :: Parser Stmt
declareStmt =
  P.choice
    [ SDeclare <$> (tok "int"  $> IntT)  <*> idString <* tok ";"
    , SDeclare <$> (tok "bool" $> BoolT) <*> idString <* tok ";"
    ]

printStmt :: Parser Stmt
printStmt = SPrint <$> (tok "printf" *> tok "(" *> expr <* tok ")" <* tok ";")

ifStmt :: Parser Stmt
ifStmt =
  SIf
  <$> (tok "if" *> tok "(" *> expr <* tok ")")
  <*> (tok "{" *> stmt <* tok "}")
  <*> elseBranch

elseBranch :: Parser Stmt
elseBranch =
  P.choice
    [ tok "else" *> tok "{" *> stmt <* tok "}"
    , pure SNoOp
    ]

forStmt :: Parser Stmt
forStmt =
  SFor
  <$> (tok "for" *> tok "(" *> idString)
  <*> (tok "from" *> expr <* tok "to")
  <*> (expr <* tok ")")
  <*> (tok "{" *> stmt <* tok "}")

whileStmt :: Parser Stmt
whileStmt =
  SWhile
  <$> (tok "while" *> tok "(" *> expr <* tok ")")
  <*> (tok "{" *> stmt <* tok "}")

assignStmt :: Parser Stmt
assignStmt = SAssign <$> idString <*> (tok "=" *> expr <* tok ";")