module CC.AST.Typed(
    module CC.AST.Typed,
    module CC.Types
) where

import Data.List

import CC.Pretty
import CC.Types


data ProgramT = ProgramT [DefT]
  deriving (Show, Read)

data DefT = DefT Name ExprT
  deriving (Show, Read)

data TypeT = TFunT TypeT TypeT
           | TIntT
           | TTupT [TypeT]
           | TyVar Int
  deriving (Show, Read)

data TypeSchemeT = TypeSchemeT [Int] TypeT
  deriving (Show, Read)

data ExprT = LamT TypeT Occ ExprT
           | CallT TypeT ExprT ExprT
           | IntT Int
           | TupT [ExprT]
           | VarT Occ
  deriving (Show, Read)

data Occ = Occ Name TypeT
  deriving (Show, Read)

exprType :: ExprT -> TypeT
exprType (LamT typ _ _) = typ
exprType (CallT typ _ _) = typ
exprType (IntT _) = TIntT
exprType (TupT es) = TTupT (map exprType es)
exprType (VarT (Occ _ typ)) = typ

instance Pretty TypeT where
    prettyPrec _ TIntT = "Int"
    prettyPrec p (TFunT a b) =
        precParens p 2 (prettyPrec 3 a ++ " -> " ++ prettyPrec 2 b)
    prettyPrec _ (TTupT ts) =
        "(" ++ intercalate ", " (map pretty ts) ++ ")"
    prettyPrec _ (TyVar i) = 't' : show i

instance Pretty TypeSchemeT where
    prettyPrec p (TypeSchemeT bnds ty) =
        precParens p 2
            ("forall " ++ intercalate " " (map (pretty . TyVar) bnds) ++ ". " ++
                prettyPrec 2 ty)

instance Pretty ExprT where
    prettyPrec p (LamT ty (Occ n t) e) =
        precParens p 2 $
            "(\\(" ++ n ++ " :: " ++ pretty t ++ ") -> "
                ++ prettyPrec 2 e ++ ") :: " ++ pretty ty
    prettyPrec p (CallT ty e1 e2) =
        precParens p 2 $
            prettyPrec 3 e1 ++ " " ++ prettyPrec 3 e2 ++ " :: " ++ pretty ty
    prettyPrec _ (IntT i) = show i
    prettyPrec _ (TupT es) = "(" ++ intercalate ", " (map pretty es) ++ ")"
    prettyPrec p (VarT (Occ n t)) =
        precParens p 2 $
            show n ++ " :: " ++ pretty t

instance Pretty DefT where
    pretty (DefT n e) = n ++ " = " ++ pretty e

instance Pretty ProgramT where
    pretty (ProgramT defs) = concatMap ((++ "\n") . pretty) defs