Algorithm W

W.hs

module W where

import qualified Data.Map.Strict as M
import Control.Monad.State.Lazy

data LExpr = Var String | App LExpr LExpr | Abs String LExpr | Let String LExpr LExpr

instance Show LExpr where
    show (Var s) = s
    show (App l r) = "(" ++ show l ++ ")(" ++ show r ++ show ")"
    show (Abs v e) = "\\" ++ v ++ "." ++ show e
    show (Let v e1 e2) = "let " ++ v ++ " = " ++ show e1 ++ " in " ++ show e2

data Type = TVar Int | Func Type Type

instance Show Type where
    show (TVar i) = "a" ++ show i
    show (Func l r) = "(" ++ show l ++ " -> " ++ show r ++ ")"

type W a = State (M.Map String Type, [(Int,Type)]) a

freeInType :: Type -> Int
freeInType (TVar i) = i + 1
freeInType (Func l r) = max (freeInType l) (freeInType r)

free :: W Int
free = do
    (tV,tI) <- get
    return $ max (M.foldr max 0 (M.map freeInType tV)) (foldr max 0 (map (freeInType.snd) tI))

regType :: String -> Type -> W ()
regType s t = do
    (tV,tI) <- get
    put (M.insert s t tV,tI)

typeOf :: String -> W Type
typeOf s = do
    (tV,tI) <- get
    if s `M.member` tV then
        return $ tV M.! s
    else do 
        f <- free
        regType s $ TVar f
        return $ TVar f

unify :: Type -> Type -> [(Int,Type)]
unify (TVar i) (TVar j) | i == j = [] | otherwise = [(i,TVar j)]
unify (TVar i) f = [(i,f)]
unify f (TVar j) = [(j,f)]
unify (Func da ia) (Func db ib) = 
    let dm = unify da db in
    let im = unify ia ib in
    dm ++ im

w :: LExpr -> W Type
w (Var s) = typeOf s
w (App l r) = do
    tl <- w l
    tr <- w r
    e <- free
    let m = unify tl (Func tr (TVar e))
    (tv,iv) <- get
    put (tv,iv ++ m)
    return (TVar e)
w (Abs v e) = do
    te <- w e
    tv <- typeOf v
    return $ Func tv te