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

import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import qualified Data.Set as Set
import Data.Set (Set)
import Data.List

import CC.Pretty
import CC.Types


data Program = Program [Def] (Map Name TypeDef)
  deriving (Show, Read)

data Def = Def Name Expr
  deriving (Show, Read)

-- Named type with type parameters
data TypeDef = TypeDef Name [Int] Type
  deriving (Show, Read)

data Type = TFun Type Type
          | TInt
          | TTup [Type]
          | TNamed Name [Type]  -- named type with type arguments
          | TUnion (Set Type)
          | TyVar Rigidity Int
  deriving (Eq, Ord, Show, Read)

data Rigidity = Rigid | Instantiable
  deriving (Eq, Ord, Show, Read)

data TypeScheme = TypeScheme [Int] Type
  deriving (Show, Read)

data Expr = Lam Type Occ Expr
          | Let Occ Expr Expr
          | Call Type Expr Expr
          | Int Int
          | Tup [Expr]
          | Var Occ
          | Constr Type Name  -- Type is 'argument -> TNamed'
  deriving (Show, Read)

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

exprType :: Expr -> Type
exprType (Lam typ _ _) = typ
exprType (Let _ _ body) = exprType body
exprType (Call typ _ _) = typ
exprType (Int _) = TInt
exprType (Tup es) = TTup (map exprType es)
exprType (Var (Occ _ typ)) = typ
exprType (Constr typ _) = typ

instance Pretty Type where
    prettyPrec _ TInt = "Int"
    prettyPrec p (TFun a b) =
        precParens p 2 (prettyPrec 3 a ++ " -> " ++ prettyPrec 2 b)
    prettyPrec _ (TTup ts) =
        "(" ++ intercalate ", " (map pretty ts) ++ ")"
    prettyPrec _ (TNamed n ts) =
        n ++ "[" ++ intercalate ", " (map pretty ts) ++ "]"
    prettyPrec _ (TUnion ts) =
        "{ " ++ intercalate " | " (map pretty (Set.toList ts)) ++ " }"
    prettyPrec _ (TyVar Rigid i) = 't' : show i ++ "R"
    prettyPrec _ (TyVar Instantiable i) = 't' : show i

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

instance Pretty Expr where
    prettyPrec p (Lam ty (Occ n t) e) =
        precParens p 2 $
            "(\\(" ++ n ++ " :: " ++ pretty t ++ ") -> "
                ++ prettyPrec 2 e ++ ") :: " ++ pretty ty
    prettyPrec p (Let (Occ n t) rhs e) =
        precParens p 2 $
            "let (" ++ n ++ " :: " ++ pretty t ++ ") = " ++ pretty rhs ++ " " ++
                (case e of
                   Let _ _ _ -> pretty e
                   _ -> "in " ++ pretty e)
    prettyPrec p (Call ty e1 e2) =
        precParens p 2 $
            prettyPrec 3 e1 ++ " " ++ prettyPrec 3 e2 ++ " :: " ++ pretty ty
    prettyPrec _ (Int i) = show i
    prettyPrec _ (Tup es) = "(" ++ intercalate ", " (map pretty es) ++ ")"
    prettyPrec p (Var (Occ n t)) =
        precParens p 2 $
            n ++ " :: " ++ pretty t
    prettyPrec p (Constr t n) =
        precParens p 2 $
            n ++ " :: " ++ pretty t

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

instance Pretty TypeDef where
    pretty (TypeDef n [] t) =
        "type " ++ n ++ " = " ++ pretty t
    pretty (TypeDef n vs t) =
        "type " ++ n ++ " " ++ intercalate " " (map (pretty . TyVar Instantiable) vs) ++
            " = " ++ pretty t

instance Pretty Program where
    pretty (Program defs tdefs) =
        concatMap (++ "\n") $
            map pretty (Map.elems tdefs) ++ map pretty defs