44 lines
1.5 KiB
Haskell
44 lines
1.5 KiB
Haskell
|
module TypeInference where
|
||
|
|
|
||
|
|
import Control.Monad.State.Strict
|
||
|
|
( gets, modify, evalState, State )
|
||
|
|
import qualified Data.Map.Strict as Map
|
||
|
|
|
||
|
|
import Terms ( Term(..) )
|
||
|
|
import Types ( TermInContext(..), Context, Type(TypeVar, Arrow) )
|
||
|
|
import Unification ( Unifier, unify, applyMgu )
|
||
|
|
|
||
|
|
-- state needed for algorithm W
|
||
|
|
newtype PTState = PTState {
|
||
|
|
varCount :: Int
|
||
|
|
}
|
||
|
|
type PTMonad = State PTState
|
||
|
|
|
||
|
|
-- returns a fresh type var in algorithm W
|
||
|
|
getNewTypeVar :: PTMonad String
|
||
|
|
getNewTypeVar = do
|
||
|
|
count <- gets varCount
|
||
|
|
modify (\s -> s {varCount = count + 1})
|
||
|
|
return $ 'a' : show count
|
||
|
|
|
||
|
|
-- implementation of algorithm W, this builds a list of terms to be unified in the next step
|
||
|
|
pt :: Context -> Term -> Type -> PTMonad Unifier
|
||
|
|
pt gamma (Var x) alpha = return [(alpha, gamma Map.! x)]
|
||
|
|
pt gamma (Abs x t) alpha = do
|
||
|
|
a <- TypeVar <$> getNewTypeVar
|
||
|
|
b <- TypeVar <$> getNewTypeVar
|
||
|
|
result <- pt (Map.insert x a gamma) t b
|
||
|
|
return $ result ++ [(Arrow a b, alpha)]
|
||
|
|
pt gamma (App t s) alpha = do
|
||
|
|
a <- TypeVar <$> getNewTypeVar
|
||
|
|
result1 <- pt gamma t (Arrow a alpha)
|
||
|
|
result2 <- pt gamma s a
|
||
|
|
return $ result1 ++ result2
|
||
|
|
|
||
|
|
-- infer the type of a term in context (if any)
|
||
|
|
inferType :: Context -> Term -> Maybe TermInContext
|
||
|
|
inferType ctx term = do
|
||
|
|
let unificationProblem = evalState (pt ctx term (TypeVar "a0")) (PTState 1)
|
||
|
|
mgu <- unify unificationProblem
|
||
|
|
let principalType = applyMgu (TypeVar "a0") mgu
|
||
|
|
return $ TermInContext ctx term principalType
|