fix more bugs

[fp.git] / src / HM.hs

diff --git a/src/HM.hs b/src/HM.hs
index 18c1f30..3cbf5c9 100644 (file)
--- a/src/HM.hs
+++ b/src/HM.hs
@@ -16,7 +16,6 @@ module HM
   , TypedTerm(..)
     -- * Type inference
   , algW
-  , runTI
   ) where
 
 import Control.Monad.Except
@@ -68,7 +67,7 @@ composeSub s1 s2 = Map.map (substituteT s1) s2 `Map.union` s1
 
 varBind :: TypeVarName -> Type -> TI Substitution
 varBind v t | t == TypeVar v = return idSub
-            | v `Set.member` freeVarsT t = fail $ "occur check failed: " ++ v ++ " in " ++ show t
+            | v `Set.member` freeVarsT t = fail $ "occur check failed: " ++ v ++ " ~ " ++ show t
             | otherwise = return $ Map.singleton v t
 
 instantiate :: TypeScheme -> TI Type
@@ -77,23 +76,35 @@ instantiate (TSForAll v t) = do
   nv <- newVar
   instantiate $ substituteS (Map.singleton v nv) t
 
+generalize :: TypeEnv -> Type -> TypeScheme
+generalize e t = foldr TSForAll (TScheme t) vars
+  where
+    vars = Set.toList $ freeVarsT t Set.\\ Set.unions (map freeVarsS $ Map.elems e)
+
 unify :: Type -> Type -> TI Substitution
 unify (TypeVar a) t = varBind a t
 unify t (TypeVar a) = varBind a t
 unify (TypeFunction a b) (TypeFunction a' b') = do
   s1 <- unify a a'
-  s2 <- unify b b'
+  s2 <- unify (substituteT s1 b) (substituteT s1 b')
   return $ s1 `composeSub` s2
 unify (Primitive a) (Primitive b) | a == b = return idSub
 unify a b = fail $ "cannot unify " ++ show a ++ " with " ++ show b
 
+tiLit :: Literal -> TI (Substitution, Type)
+tiLit (LBool _) = return (idSub, Primitive "Bool")
+tiLit (LInt _) = return (idSub, Primitive "Integer")
+tiLit (LFunc If) = do
+  a <- newVar
+  return (idSub, Primitive "Bool" `TypeFunction` (a `TypeFunction` (a `TypeFunction` a)))
+
 ti :: TypeEnv -> TypedTerm -> TI (Substitution, Type)
---ti _ (TTerm (Var v) (TScheme t@(Primitive _))) = return (idSub, t)
 ti e (TTerm tr sch) = do
   (s, t) <- ti e (NTTerm tr)
   sch' <- instantiate sch
   s' <- unify t sch'
   return (s', substituteT s' sch')
+ti _ (NTTerm (Lit l)) = tiLit l
 ti e (NTTerm (Var v)) = case Map.lookup v e of
   Nothing -> fail $ "unbound variable: " ++ v
   Just sigma -> do
@@ -104,8 +115,20 @@ ti e (NTTerm (Lam x y)) = do
   let e' = Map.insert x (TScheme tv) e
   (s, t) <-  ti e' y
   return (s, TypeFunction (substituteT s tv) t)
-
-algW :: TypedTerm -> TI Type
-algW t = do
+ti e (NTTerm (App a b)) = do
+  tv <- newVar 
+  (s1, t1) <- ti e a
+  (s2, t2) <- ti (Map.map (substituteS s1) e) b
+  s3 <- unify (substituteT s2 t1) (TypeFunction t2 tv)
+  return (s3 `composeSub` s2 `composeSub` s1, substituteT s3 tv)
+ti e (NTTerm (Let x a b)) = do
+  (s1, t1) <- ti e a
+  let t' = generalize (Map.map (substituteS s1) e) t1
+      e' = Map.insert x t' e
+  (s2, t2) <- ti (Map.map (substituteS s1) e') b
+  return (s1 `composeSub` s2, t2)
+ 
+algW :: TypedTerm -> Either String Type
+algW t = fst . runTI $ do
   (s, u) <- ti Map.empty t
   return $ substituteT s u