From bff68fbafa3711186b19a1c5dcc16becb49f162a Mon Sep 17 00:00:00 2001
From: =?utf8?q?Tom=C3=A1=C5=A1=20Musil?= <tomik.musil@gmail.com>
Date: Tue, 25 Aug 2015 16:28:31 +0200
Subject: [PATCH] working on alg W

---
 src/HM.hs      | 116 ++++++++++++++++++++++++++++++++++++++++++---------------
 src/HM/Term.hs |   4 +-
 2 files changed, 87 insertions(+), 33 deletions(-)

diff --git a/src/HM.hs b/src/HM.hs
index 065ccbb..18c1f30 100644
--- a/src/HM.hs
+++ b/src/HM.hs
@@ -6,7 +6,7 @@
 -- Maintainer  :  tomik.musil@gmail.com
 -- Stability   :  experimental
 --
--- This is a toy implementation of λ-calculus with Hindley-Milner type system.
+-- This is a toy implementation of \-calculus with Hindley-Milner type system.
 
 module HM
   ( -- * Types
@@ -16,42 +16,96 @@ module HM
   , TypedTerm(..)
     -- * Type inference
   , algW
+  , runTI
   ) where
 
+import Control.Monad.Except
 import Control.Monad.State
+import qualified Data.Set as Set
+import qualified Data.Map as Map
 
 import HM.Term
 import HM.Parser
 
-type Substitution = TypeScheme -> TypeScheme
+type Substitution = Map.Map TypeVarName Type
 type VarS a = State Int a
+type TypeEnv = Map.Map VarName TypeScheme
+data TIState = TIState {tiSupply :: Int} deriving (Show)
+type TI a = ExceptT String (State TIState) a
 
-fresh :: TypeVarName
-fresh = undefined
---fresh = do
---  n <- get
---  "a" ++ show n
-  
-
-substitute :: TypeScheme -> TypeVarName -> TypeScheme -> TypeScheme
-substitute = undefined
-
-unify :: TypeScheme -> TypeScheme -> Either String Substitution
-unify (TScheme (Primitive a)) (TScheme (Primitive b)) | a == b = Right id
-unify (TScheme (TypeVar a)) (TScheme (TypeVar b)) | a == b = Right id
-unify a b = Left $ "cannot unify " ++ show a ++ " with " ++ show b
-
-algW :: TypedTerm -> Either String TypeScheme
-algW (TTerm (Var _) t) = Right t
---algW (TTerm (Lam x t) (TScheme p)) = do
---  let v = TScheme (TypeVar fresh)
---      np = substitute v x t
---  unify p np
---algW (TTerm (App u v) t) = do
---  tu <- algW u
---  tv <- algW v
---  case tu of
---    (TScheme (TypeFunction a b)) -> do
---      unify a tv
---      return b
---    _ -> Left $ "cannot apply " ++ show tu ++ " to " ++ show tv
+runTI :: TI a -> (Either String a, TIState)
+runTI t = runState (runExceptT t) $ TIState 0
+
+newVar :: TI Type
+newVar = do
+  s <- get
+  put s {tiSupply = tiSupply s + 1}
+  return (TypeVar $ "a" ++ show (tiSupply s))
+
+freeVarsT :: Type -> Set.Set TypeVarName
+freeVarsT (Primitive _) = Set.empty
+freeVarsT (TypeVar t) = Set.singleton t
+freeVarsT (TypeFunction a b) = freeVarsT a `Set.union` freeVarsT b
+
+freeVarsS :: TypeScheme -> Set.Set TypeVarName
+freeVarsS (TScheme t) = freeVarsT t
+freeVarsS (TSForAll v s) = v `Set.delete` freeVarsS s
+
+substituteT :: Substitution -> Type -> Type
+substituteT _ t@(Primitive _) = t
+substituteT s t@(TypeVar v) = Map.findWithDefault t v s
+substituteT s (TypeFunction a b) = TypeFunction (substituteT s a) (substituteT s b)
+
+substituteS :: Substitution -> TypeScheme -> TypeScheme
+substituteS s (TScheme t) = TScheme $ substituteT s t
+substituteS s (TSForAll v t) = TSForAll v $ substituteS (v `Map.delete` s) t
+
+idSub :: Substitution
+idSub = Map.empty
+
+composeSub :: Substitution -> Substitution -> Substitution
+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
+            | otherwise = return $ Map.singleton v t
+
+instantiate :: TypeScheme -> TI Type
+instantiate (TScheme t) = return t
+instantiate (TSForAll v t) = do 
+  nv <- newVar
+  instantiate $ substituteS (Map.singleton v nv) t
+
+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'
+  return $ s1 `composeSub` s2
+unify (Primitive a) (Primitive b) | a == b = return idSub
+unify a b = fail $ "cannot unify " ++ show a ++ " with " ++ show b
+
+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 e (NTTerm (Var v)) = case Map.lookup v e of
+  Nothing -> fail $ "unbound variable: " ++ v
+  Just sigma -> do
+    t <- instantiate sigma
+    return (idSub, t)
+ti e (NTTerm (Lam x y)) = do
+  tv <- newVar
+  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
+  (s, u) <- ti Map.empty t
+  return $ substituteT s u
diff --git a/src/HM/Term.hs b/src/HM/Term.hs
index 0a686f9..cf4cf36 100644
--- a/src/HM/Term.hs
+++ b/src/HM/Term.hs
@@ -23,8 +23,8 @@ type VarName = String
 type TypeVarName = String
 type TypeName = String
 
-data Type = Primitive TypeName | TypeVar TypeVarName | TypeFunction Type Type 
-data TypeScheme = TScheme Type | TSForAll TypeVarName TypeScheme
+data Type = Primitive TypeName | TypeVar TypeVarName | TypeFunction Type Type deriving (Eq)
+data TypeScheme = TScheme Type | TSForAll TypeVarName TypeScheme deriving (Eq)
 
 data Term = Var VarName | Lam VarName TypedTerm | App TypedTerm TypedTerm | Let VarName TypedTerm TypedTerm
 data TypedTerm = NTTerm Term | TTerm Term TypeScheme
-- 
2.4.2