Fixed Parser issues and simplified AST

src/Windows12/Ast.hs

@@ -1,6 +1,9 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Windows12.Ast where
 
 import Data.Text (Text)
+import Prettyprinter
 
 data BinOp
   = Add
@@ -31,6 +34,20 @@ data UnOp
   | AddrOf
   deriving (Show, Eq)
 
+data AssignOp
+  = BaseAssign
+  | AddAssign
+  | SubAssign
+  | MulAssign
+  | DivAssign
+  | ModAssign
+  | BitAndAssign
+  | BitOrAssign
+  | BitXorAssign
+  | ShiftLAssign
+  | ShiftRAssign
+  deriving (Show, Eq)
+
 data Expr
   = Id Text
   | IntLit Int
@@ -41,7 +58,7 @@ data Expr
   | CharLit Char
   | BinOp BinOp Expr Expr
   | UnOp UnOp Expr
-  | Call Text [Expr]
+  | Call Expr [Expr]
   | Index Expr Expr
   | Member Expr Expr
   | Cast Type Expr
@@ -52,12 +69,11 @@ data Expr
 data Stmt
   = Expr Expr
   | Return Expr
-  | If Expr Stmt (Maybe Stmt)
-  | While Expr Stmt
-  | For (Maybe Stmt) (Maybe Expr) (Maybe Expr) Stmt
-  | Assign Expr Expr
+  | If Expr [Stmt] (Maybe [Stmt])
+  | While Expr [Stmt]
+  | Assign AssignOp Expr Expr
   | Block [Stmt]
-  | Var Bind Expr
+  | Var Text (Maybe Type) (Maybe Expr)
   deriving (Show, Eq)
 
 data Type
@@ -86,11 +102,112 @@ data TLEnum = Enum {enumName :: Text, enumFields :: [Text]}
 data TLFunc = Func {funcName :: Text, funcArgs :: [Bind], funcRetType :: Type, funcBody :: [Stmt]}
   deriving (Show, Eq)
 
-data TLMemberFunc = MemberFunc {memberFuncSelf :: Type, memberFuncName :: Text, memberFuncArgs :: [Bind], memberFuncRetType :: Type, memberFuncBody :: [Stmt]}
+data TL = TLStruct TLStruct | TLEnum TLEnum | TLFunc TLFunc
   deriving (Show, Eq)
 
-data TL = TLStruct TLStruct | TLEnum TLEnum | TLFunc TLFunc | TLMemberFunc TLMemberFunc
+data Program = Program [TLStruct] [TLEnum] [TLFunc]
   deriving (Show, Eq)
 
-data Program = Program [TLStruct] [TLEnum] [TLFunc] [TLMemberFunc]
-  deriving (Show, Eq)
+-- Pretty printing
+instance Pretty BinOp where
+  pretty Add = "+"
+  pretty Sub = "-"
+  pretty Mul = "*"
+  pretty Div = "/"
+  pretty Mod = "%"
+  pretty Eq = "=="
+  pretty Ne = "!="
+  pretty Lt = "<"
+  pretty Gt = ">"
+  pretty Le = "<="
+  pretty Ge = ">="
+  pretty And = "&&"
+  pretty Or = "||"
+  pretty BitAnd = "&"
+  pretty BitOr = "|"
+  pretty BitXor = "^"
+  pretty ShiftL = "<<"
+  pretty ShiftR = ">>"
+
+instance Pretty UnOp where
+  pretty Neg = "-"
+  pretty Not = "!"
+  pretty BitNot = "~"
+  pretty Deref = "*"
+  pretty AddrOf = "&"
+
+instance Pretty AssignOp where
+  pretty BaseAssign = "="
+  pretty AddAssign = "+="
+  pretty SubAssign = "-="
+  pretty MulAssign = "*="
+  pretty DivAssign = "/="
+  pretty ModAssign = "%="
+  pretty BitAndAssign = "&="
+  pretty BitOrAssign = "|="
+  pretty BitXorAssign = "^="
+  pretty ShiftLAssign = "<<="
+  pretty ShiftRAssign = ">>="
+
+instance Pretty Expr where
+  pretty (Id x) = pretty x
+  pretty (IntLit x) = pretty x
+  pretty (UIntLit x) = pretty x
+  pretty (FloatLit x) = pretty x
+  pretty (StrLit x) = dquotes (pretty x)
+  pretty (BoolLit x) = pretty x
+  pretty (CharLit x) = squotes (pretty x)
+  pretty (BinOp op l r) = parens (pretty l <+> pretty op <+> pretty r)
+  pretty (UnOp op e) = pretty op <> parens (pretty e)
+  pretty (Call f args) = parens (pretty f) <> parens (hsep (punctuate comma (map pretty args)))
+  pretty (Index arr idx) = parens (pretty arr) <> brackets (pretty idx)
+  pretty (Member e m) = pretty e <> "." <> pretty m
+  pretty (Cast t e) = parens (pretty t) <> parens (pretty e)
+  pretty (Sizeof t) = "sizeof" <> parens (pretty t)
+  pretty (StructInit s fields) = pretty s <+> lbrace <> line <> indent 4 (vsep (punctuate comma (map (\(n, e) -> pretty n <+> "=" <+> pretty e) fields))) <> line <> rbrace
+
+instance Pretty Stmt where
+  pretty (Expr e) = pretty e <> semi
+  pretty (Return e) = "return" <+> pretty e <> semi
+  pretty (If cond t f) = "if" <+> pretty cond <+> prettyBlock t <+> maybe "" (\f' -> "else" <+> prettyBlock f') f
+  pretty (While cond body) = "while" <+> pretty cond <+> prettyBlock body
+  pretty (Assign op l r) = pretty l <+> pretty op <+> pretty r <> semi
+  pretty (Block stmts) = braces (vsep (map pretty stmts))
+  pretty (Var n t e) = pretty n <+> maybe "" (\t' -> ":" <+> pretty t') t <+> maybe "" (\e' -> "=" <+> pretty e') e <> semi
+
+instance Pretty Type where
+  pretty IntType = "int"
+  pretty UIntType = "uint"
+  pretty FloatType = "float"
+  pretty StrType = "str"
+  pretty BoolType = "bool"
+  pretty CharType = "char"
+  pretty (PtrType t) = pretty t <> "*"
+  pretty (ArrayType t) = pretty t <> "[]"
+  pretty (StructType s) = pretty s
+  pretty (EnumType e) = pretty e
+  pretty VoidType = "void"
+
+instance Pretty Bind where
+  pretty (Bind n t) = pretty n <+> ":" <+> pretty t
+
+instance Pretty TLStruct where
+  pretty (Struct n fields) = "struct" <+> pretty n <+> prettyBlock fields
+
+instance Pretty TLEnum where
+  pretty (Enum n fields) = "enum" <+> pretty n <+> prettyBlock fields
+
+instance Pretty TLFunc where
+  pretty (Func n args ret body) =
+    pretty n <> parens (hsep (punctuate comma (map pretty args))) <+> "->" <+> pretty ret <+> prettyBlock body
+
+instance Pretty TL where
+  pretty (TLStruct s) = pretty s
+  pretty (TLEnum e) = pretty e
+  pretty (TLFunc f) = pretty f
+
+instance Pretty Program where
+  pretty (Program structs enums funcs) = vsep (map pretty structs) <> line <> vsep (map pretty enums) <> line <> vsep (map pretty funcs)
+
+prettyBlock :: (Pretty a) => [a] -> Doc ann
+prettyBlock stmts = lbrace <> line <> indent 4 (vsep (map pretty stmts)) <> line <> rbrace