Lift TypeScheme to AST.Typed

2 files changed, 22 insertions, 15 deletions

diff --git a/ast/CC/AST/Typed.hs b/ast/CC/AST/Typed.hs
index f4c4d56..537c7a4 100644
--- a/ast/CC/AST/Typed.hs
+++ b/ast/CC/AST/Typed.hs
@@ -21,6 +21,9 @@ data TypeT = TFunT TypeT 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
@@ -46,6 +49,12 @@ instance Pretty TypeT where
         "(" ++ 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 $
diff --git a/typecheck/CC/Typecheck.hs b/typecheck/CC/Typecheck.hs
index 195fc19..0c926f3 100644
--- a/typecheck/CC/Typecheck.hs
+++ b/typecheck/CC/Typecheck.hs
@@ -42,9 +42,7 @@ runTM :: TM a -> Either TCError a
 runTM m = evalState (runExceptT m) 1
 
 
-data Scheme = Scheme [Int] TypeT
-
-newtype Env = Env (Map Name Scheme)
+newtype Env = Env (Map Name TypeSchemeT)
 
 newtype Subst = Subst (Map Int TypeT)
 
@@ -58,8 +56,8 @@ instance FreeTypeVars TypeT where
     freeTypeVars (TTupT ts) = Set.unions (map freeTypeVars ts)
     freeTypeVars (TyVar var) = Set.singleton var
 
-instance FreeTypeVars Scheme where
-    freeTypeVars (Scheme bnds ty) = foldr Set.delete (freeTypeVars ty) bnds
+instance FreeTypeVars TypeSchemeT where
+    freeTypeVars (TypeSchemeT bnds ty) = foldr Set.delete (freeTypeVars ty) bnds
 
 instance FreeTypeVars Env where
     freeTypeVars (Env mp) = foldMap freeTypeVars (Map.elems mp)
@@ -76,9 +74,9 @@ instance Substitute TypeT where
         TTupT ts -> TTupT (map (theta >>!) ts)
         TyVar i -> fromMaybe ty (Map.lookup i mp)
 
-instance Substitute Scheme where
-    Subst mp >>! Scheme bnds ty =
-        Scheme bnds (Subst (foldr Map.delete mp bnds) >>! ty)
+instance Substitute TypeSchemeT where
+    Subst mp >>! TypeSchemeT bnds ty =
+        TypeSchemeT bnds (Subst (foldr Map.delete mp bnds) >>! ty)
 
 instance Substitute Env where
     theta >>! Env mp = Env (Map.map (theta >>!) mp)
@@ -103,21 +101,21 @@ instance Monoid Subst where
 emptyEnv :: Env
 emptyEnv = Env mempty
 
-envAdd :: Name -> Scheme -> Env -> Env
+envAdd :: Name -> TypeSchemeT -> Env -> Env
 envAdd name sty (Env mp) = Env (Map.insert name sty mp)
 
-envFind :: Name -> Env -> Maybe Scheme
+envFind :: Name -> Env -> Maybe TypeSchemeT
 envFind name (Env mp) = Map.lookup name mp
 
 substVar :: Int -> TypeT -> Subst
 substVar var ty = Subst (Map.singleton var ty)
 
-generalise :: Env -> TypeT -> Scheme
+generalise :: Env -> TypeT -> TypeSchemeT
 generalise env ty =
-    Scheme (Set.toList (freeTypeVars ty Set.\\ freeTypeVars env)) ty
+    TypeSchemeT (Set.toList (freeTypeVars ty Set.\\ freeTypeVars env)) ty
 
-instantiate :: Scheme -> TM TypeT
-instantiate (Scheme bnds ty) = do
+instantiate :: TypeSchemeT -> TM TypeT
+instantiate (TypeSchemeT bnds ty) = do
     vars <- traverse (const genTyVar) bnds
     let theta = Subst (Map.fromList (zip bnds vars))
     return (theta >>! ty)
@@ -147,7 +145,7 @@ infer env expr = case expr of
     Lam sr args@(_:_:_) body -> infer env (foldr (Lam sr . pure) body args)
     Lam _ [(arg, _)] body -> do
         argVar <- genTyVar
-        let augEnv = envAdd arg (Scheme [] argVar) env
+        let augEnv = envAdd arg (TypeSchemeT [] argVar) env
         (theta, body') <- infer augEnv body
         let argType = theta >>! argVar
         return (theta, LamT (TFunT argType (exprType body')) (Occ arg argType) body')