Writing Hindley-Milner from memory for the first time

Hemem.hs

{-# OPTIONS_GHC -Wno-tabs #-}
{-# LANGUAGE ImportQualifiedPost, LambdaCase, BlockArguments #-}

module Hemem where

-- roughly Hindley-Milner type inference
-- my first implementation
-- in 1½ hours with no reference
-- just things I pulled from memory (of reading articles about HM)
-- missing an occurs check

import Data.Map qualified as M
import Data.Set qualified as S
import Data.Maybe
import Data.Foldable
import Control.Monad.Trans.State
import Control.Monad.Trans.Class
import Debug.Trace

data TV = TV {
	tvName :: String,
	tvIndex :: Int
} deriving (Show, Eq, Ord)
data Type
	= TyVar TV
	| TyArr Type Type
	| TyCon String [Type]
	deriving (Show)
data Typing = Typing (S.Set TV) Type deriving (Show)
data Term
	= TeVar String
	| TeLam String Term
	| TeApp Term Term
	| TeLet String Term Term
	deriving (Show)

subst :: M.Map TV Type -> Type -> Type
subst s (TyVar v) = fromMaybe (TyVar v) (M.lookup v s)
subst s (TyArr a r) = TyArr (subst s a) (subst s r)
subst s (TyCon nm ts) = TyCon nm $ fmap (subst s) ts

unify' :: Type -> Type -> Either String (M.Map TV Type)
unify' (TyVar a) b = pure $ M.singleton a b
unify' a (TyVar b) = pure $ M.singleton b a
unify' (TyArr aa ar) (TyArr ba br) = do
	sa <- unify' aa ba
	sr <- unify' (subst sa ar) (subst sa br)
	pure $ M.unionWithKey (\k a b -> error $ show (k, a, b)) sa sr
unify' (TyCon an ats) (TyCon bn bts) | an == bn && length ats == length bts = do
	foldlM
		(\s (at, bt) -> do
			s' <- unify' (subst s at) (subst s bt)
			pure $ M.unionWithKey (\k a b -> error $ show (k, a, b)) s s')
		M.empty (zip ats bts)
unify' a b = Left $ "failed to unify " <> show a <> " with " <> show b

type M = StateT (M.Map TV (Maybe Type)) (Either String)

latest :: Type -> M Type
latest t = do
	e <- get
	pure $ subst (M.mapMaybe id e) t

newTyVar :: String -> M Type
newTyVar n = state $ \e ->
	let
		i = let
				go i | M.member (TV n i) e = go (i + 1)
				go i = i
			in go 0
		v = TV n i
	in (TyVar v, M.insert v Nothing e)

unify :: Type -> Type -> M ()
unify a b = do
	a <- latest a
	b <- latest b
	StateT $ \e -> do
		s <- unify' a b
		pure ((), M.unionWithKey
			(\k -> curry \case
				(Just a, Just b) -> error $ show (k, a, b)
				(Nothing, b) -> b
				(a, Nothing) -> a)
			e (fmap Just s))

inst :: Typing -> M Type
inst (Typing vs t) = do
	s <- sequence $ M.fromSet (newTyVar . ("i:" <>) . tvName) vs
	pure $ subst s t

gen :: M a -> M (S.Set TV, a)
gen b = do
	e <- get
	(r, e') <- lift $ runStateT b e
	let new = M.difference e' e
	put $ M.difference e' new
	pure (M.keysSet $ M.filter isNothing new, r)

genTyping :: M Type -> M Typing
genTyping = fmap (uncurry Typing) . gen . (>>= latest)

infer :: M.Map String Typing -> Term -> M Type
infer e (TeVar v) = case M.lookup v e of
	Just t -> inst t
	Nothing -> lift $ Left $ "no variable " <> show v
infer e (TeLam n b) = do
	at <- newTyVar $ "l:" <> n
	rt <- infer (M.insert n (Typing S.empty at) e) b
	pure $ TyArr at rt
infer e (TeApp f a) = do
	ft <- infer e f >>= latest
	at <- infer e a >>= latest
	rt <- newTyVar "ar"
	unify ft (TyArr at rt)
	pure rt
infer e (TeLet n v b) = do
	(vs, vt) <- gen $ infer e v >>= latest
	infer (M.insert n (Typing vs vt) e) b

defaultEnv :: M.Map String Typing
defaultEnv = M.fromList [
		("Unit1", Typing S.empty $ TyCon "Unit1" []),
		("Unit2", Typing S.empty $ TyCon "Unit2" [])
	]