-- Type Inference for the Simply-Typed Lambda Calculus
-- 09.12.2017 Simon Wimmer

import Data.Maybe

-- variable names 
type VName = (String, Int)

showVar (s, n) = s ++ if n == 0 then "" else show n

-- types
data Type = TV VName | Arrow Type Type

instance Show Type where
    show (TV x) = showVar x
    show (Arrow (TV x) b) = showVar x ++ " --> " ++ show b
    show (Arrow a b) = "(" ++ show a ++ ") --> " ++ show b

infixr 7 -->
(-->) = Arrow

-- terms
data Term = V VName | Abs VName Term | App Term Term

instance Show Term where
    show (V x) = showVar x
    show (Abs x t) = "\\" ++ showVar x ++ ". " ++ show t
    show (App t1 (V x)) = show t1 ++ " " ++ showVar x
    show (App t1 t2) = show t1 ++ " (" ++ show t2 ++ ")"

infixl 7 $$
($$) = App

-- define == on Type
instance Eq Type where (==) = curry isequal

-- equality for terms
isequal :: (Type,Type) -> Bool
isequal (TV x, TV y)         = x == y
isequal (Arrow t1 t2, Arrow u1 u2) = isequal (t1, u1) && isequal (t2, u2)

-- substitutions
type Subst = [(VName,Type)]

-- type environments
type Env = [(VName, Type)]

-- get the binding of x
get :: VName -> Env -> Maybe Type
get x [] = Nothing
get x ((y,ty):env) = if x == y then Just ty else get x env

-- is variable x in the domain of substitution s?
indom :: VName -> Subst -> Bool
indom x s = isJust (get x s)

-- applies substitution to variable
app :: Subst -> VName -> Type
app ((y,t):s) x = if x == y then t else app s x

-- applies substition to Type
lift :: Subst -> Type -> Type
lift s (TV x)      = if indom x s then app s x else TV x
lift s (Arrow t1 t2) = Arrow (lift s t1) (lift s t2)

-- does variable x occur in a Type?
occurs :: VName -> Type -> Bool
occurs x (TV y)     = x == y
occurs x (Arrow t1 t2) = occurs x t1 || occurs x t2



------------------------------------------
-- unification
------------------------------------------

type Unifier = Maybe Subst

-- given (list of disagreement pairs, current substitution)
-- produces a unifier
solve :: ([(Type,Type)], Subst) -> Unifier
solve ([], s)                  = Just s
solve ((TV x, TV y) : s', s)   = if x == y then solve(s',s) else elim(x,TV y,s',s)
solve ((TV x, t) : s', s)      = elim(x,t,s',s)
solve ((t, TV x) : s', s)      = elim(x,t,s',s)
solve ((Arrow t1 t2, Arrow u1 u2):s', s) = solve((t1, u1) : (t2, u2) : s', s)

-- given (x, t, list of disagreement pairs, current substitution)
-- propagates the constraint x = t and solves the remaining constraints
elim :: (VName, Type, [(Type,Type)], Subst) -> Unifier
elim (x,t,s',s) = if occurs x t then 
                    Nothing 
                  else
                    solve( (map (\(t1,t2) -> (xt t1, xt t2)) s'), 
                           ((x,t):(map (\(y,u) -> (y, xt u)) s)) )
                    where xt = lift [(x,t)]


-- embedding of solve
unify :: (Type, Type) -> Unifier
unify (t1,t2) = solve([(t1,t2)], [])



------------------------------------------
-- type inference
------------------------------------------

-- accumulate type constraints
constraints :: Term -> Type -> Env -> (Int, [(Type, Type)]) -> Maybe (Int, [(Type, Type)])
constraints (V x) ty env (n, cs) = if indom x env then Just (n, (ty, fromJust (get x env)) : cs) else Nothing
constraints (Abs x t) ty env (n, cs) = constraints t tau2 ((x, tau1) : env) (n + 2, (ty, Arrow tau1 tau2) : cs)
    where
        tau1 = TV ("t", n)
        tau2 = TV ("t", n + 1)
constraints (App t1 t2) ty env (n, cs) =
    case cs' of
        Nothing -> Nothing
        Just (n', cs') -> constraints t2 tau env (n', cs')
    where
        tau = TV ("t", n)
        cs' = constraints t1 (Arrow tau ty) env (n + 1, cs)


-- type inference procedure
infer :: Term -> Maybe Type
infer t = case cs of
    Nothing -> Nothing
    Just (_, cs) ->
        case solve (cs, []) of
            Just s -> get tau s
            Nothing -> Nothing
    where
        tau = ("t", 0)
        cs = constraints t (TV tau) [] (1, [])

--
-- terms & types for testing:

-- building blocks
tx     = TV("x",0)
ty     = TV("y",0)
tz     = TV("z",0)
f a b = a --> b
g a b = a --> a --> b

abstr xs t = foldr (\ x -> Abs (x, 0)) t xs

apply [t] = t
apply (t : ts) = t $$ apply ts

x = V ("x", 0)
y = V ("y", 0)
z = V ("z", 0)
tId = Abs ("x", 0) x
omega = Abs ("x", 0) (App x x)
tFst = abstr ["x", "y"] x
tSnd = abstr ["x", "y"] y
comp = abstr ["x", "y", "z"] (x $$ (y $$ z))
tEx3b = abstr ["x", "y", "z"] (x $$ y $$ (y $$ z))