Minimal Hindley-Milner-Damas type inference implementation. Has recursive let.

SmallHindleyMilner.hs

{-# LANGUAGE LambdaCase, TupleSections, TemplateHaskell, DeriveFunctor, DeriveFoldable, DeriveTraversable #-}

import qualified Data.IntMap.Strict as IM
import qualified Data.IntSet as IS
import qualified Data.Foldable as F

import Control.Applicative
import Control.Lens
import Control.Monad.State.Strict
import Debug.Trace

data Term
    = App Term Term 
    | Lam !Int Term 
    | Let !Int Term Term 
    | Var !Int 
    | Unit
    deriving (Eq, Show)

data Type' a = TVar !a | Type' a :=> Type' a | TUnit 
    deriving (Eq, Show, Functor, F.Foldable, Traversable)

data Scheme = Forall [Int] Type 
    deriving (Eq, Show)

type Type = Type' Int
type TI = State Env 

data Env = Env {
    _subs     :: IM.IntMap Type, 
    _vars     :: IM.IntMap Scheme, 
    _counter  :: Int
    } deriving (Eq, Show)
makeLenses ''Env

intNub :: [Int] -> [Int]
intNub = IS.toList . IS.fromList

apply :: Type -> TI Type
apply (TVar v)  = maybe (pure $ TVar v) apply =<< use (subs . at v)
apply (a :=> b) = (:=>) <$> apply a <*> apply b 
apply other     = pure other

fresh :: TI Int
fresh = counter <<+= 1 

freshTVar :: TI Type
freshTVar = TVar <$> fresh

unify :: Type -> Type -> TI ()
unify (a :=> b) (c :=> d) = do
    unify a c
    (b, d) <- each apply (b, d) 
    unify b d
unify (TVar a) (TVar b) = case compare a b of 
    LT -> subs . at b ?= TVar a
    GT -> subs . at a ?= TVar b
    EQ -> pure ()
unify a (TVar b) = unify (TVar b) a
unify (TVar a) b | F.elem a b = error "occurs check"
unify (TVar a) b = subs . at a ?= b
unify a b | a == b = pure ()
unify a b = error "cannot unify"

assuming :: Int -> Scheme -> TI a -> TI a
assuming v t action = do
    save <- use vars
    vars . at v ?= t 
    action <* (vars .= save)

infer :: Term -> TI Type
infer Unit = pure TUnit
infer (Var v) = do
    Forall qs t <- use (vars . at v . non (error "var not in context"))
    subs <- IM.fromList <$> mapM (\v -> (v,) <$> fresh) qs
    pure $ t <&> \v -> subs^. at v . non v
infer (App f x) = do
    (tf, tx) <- each infer (f, x)
    tRes <- freshTVar
    unify tf (tx :=> tRes)
    apply tRes
infer (Lam x e) = do
    tx <- freshTVar
    te <- assuming x (Forall [] tx) (infer e)
    apply (tx :=> te)
infer (Let v e1 e2) = do
    tv@(TVar tvId) <- freshTVar
    assuming v (Forall [] tv) $ do
        te1 <- infer e1
        (te1, tv) <- each apply (te1, tv)
        unify tv te1
        te1 <- apply te1
        let qs = intNub . filter (>tvId) $ F.toList te1
        subs %= fst . IM.split tvId
        assuming v (Forall qs te1) (infer e2)

normalize :: Type -> Type
normalize t = t <&> \v -> sub^. at v . non v where
    sub = IM.fromList $ zip (intNub $ F.toList t) [0..]

inferTest :: Term -> Type
inferTest t = normalize $
    evalState (apply =<< infer t) (Env IM.empty IM.empty 0)

(!) = App 
infixl 9 !
[lx, ly, lz, lf, lg, lh] = map Lam [0..5]
[vx, vy, vz, vf, vg, vh] = map Var [0..5]
[x, y, z, f, g, h] = [0..5]
v = Var 

test1 = 
    Let f (lf $ lg $ lx $ vf ! (vg ! vx)) $
    Let g (lx $ vx) $
    vf ! vg ! vg

fix' = Let f (lg $ vg ! (vf ! vg)) vf

main = do
    print $ inferTest $ fix'
    print $ inferTest $ test1