-- | Abstract syntax tree for the untyped lambda calculus, plus some helpers.
module AST (
  MonoType (..),
  PolyType (..),
  Pattern (..),
  Term (..),
  GADT (..),
  Program (..),
  Lit (..),
  Binop (..),
  intType,
  boolType,
) where

newtype Program = Program ([GADT], Term)

newtype GADT = GADT (String, Int, [(String, MonoType)]) deriving (Eq, Ord, Show)

data MonoType =
  MTVar String
  | MTFunc MonoType MonoType
  | MTConstr String [MonoType]
  deriving (Eq, Ord)

intType :: MonoType
intType = MTConstr "Int" []
boolType :: MonoType
boolType = MTConstr "Bool" []

data PolyType =
  PTForall [String] MonoType
  deriving (Eq, Ord, Show)

data Pattern =
  PVar String
  | PConstr String [Pattern]
  | PAnn Pattern MonoType
  deriving (Eq, Ord, Show)

data Term =
  TVar String
  | TConstr String
  | TLit Lit
  | TBinop Binop Term Term
  | TIf Term Term Term
  | TLam String Term
  | TApp Term Term
  | TLet String Term Term
  | TLetAnn String PolyType Term Term
  | TCase Term [(Pattern, Term)]
  deriving (Eq, Ord)

data Lit =
  LInt Integer
  | LBool Bool
  deriving (Eq, Ord)

data Binop = Add | Sub | Mul | Equal
  deriving (Eq, Ord, Show)

instance Show MonoType where
  show (MTVar x)       = x
  show (MTFunc a b)    = "(" ++ show a ++ " -> " ++ show b ++ ")"
  show (MTConstr c []) = c
  show (MTConstr c as) = "(" ++ c ++ " " ++ unwords (map show as) ++ ")"

instance Show Program where
  show (Program (gadts, t)) = concatMap (\g -> show g ++ "\n") gadts ++ show t

instance Show Term where
  show = aux 0
    where
      aux n t@(TLet _ e1 e2) =
        replicate (2*n) ' ' ++ auxName t ++ "\n"
        ++ aux (n+1) e1 ++ replicate (2*n) ' '++ "in\n" ++ aux (n+1) e2
      aux n t@(TLetAnn _ _ e1 e2) =
        replicate (2*n) ' ' ++ auxName t ++ "\n"
        ++ aux (n+1) e1 ++ replicate (2*n) ' ' ++ "in\n" ++ aux (n+1) e2
      aux n t@(TCase e cs) =
        replicate (2*n) ' ' ++ auxName t ++ "\n"
        ++ aux (n+1) e ++ replicate (2*n) ' ' ++ "of\n"
        ++ concatMap (\(p,ec) -> replicate (2*n+2) ' ' ++ show p ++ " ->\n" ++ aux (n+2) ec) cs
      aux n t@(TIf e1 e2 e3) =
        replicate (2*n) ' ' ++ auxName t ++ "\n"
        ++ aux (n+1) e1 ++ replicate (2*n) ' ' ++ "then\n"
        ++ aux (n+1) e2 ++ replicate (2*n) ' ' ++ "else\n"
        ++ aux (n+1) e3
      aux n t =
        replicate (2*n) ' ' ++ auxName t ++ "\n"
        ++ concatMap (aux (n+1)) (children t)
      auxName (TVar i)          = "TVar " ++ i
      auxName (TConstr i)       = "TConstr " ++ i
      auxName (TLit (LInt i))   = "IntLit " ++ show i
      auxName (TLit (LBool b))  = "BoolLit " ++ show b
      auxName (TBinop b _ _)    = show b
      auxName (TLam x _)        = "TLam " ++ x ++ " ."
      auxName (TApp {})         = "TApp"
      auxName (TLet x _ _)      = "TLet " ++ x ++ " ="
      auxName (TLetAnn x a _ _) = "TLetAnn " ++ x ++ "::" ++ show a ++ " ="
      auxName (TCase _ _)       = "TCase "
      auxName (TIf {})          = "TIf"
      children (TVar _)            = []
      children (TConstr _)         = []
      children (TLit _)            = []
      children (TBinop _ e1 e2)    = [e1, e2]
      children (TLam _ e)          = [e]
      children (TApp e1 e2)        = [e1, e2]
      children (TLet _ e1 e2)      = [e1, e2]
      children (TLetAnn _ _ e1 e2) = [e1, e2]
      children (TCase e cs)        = e : map snd cs
      children (TIf e1 e2 e3)      = [e1, e2, e3]
