-- Type Inference for the Simply-Typed Lambda Calculus w/ Let-Polymorphism
-- 20.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
type TyRange = (Int, Int)
data Type = TV VName | Arrow Type Type | Schema TyRange 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
    show (Schema ran t) = "∀ (" ++ show (fst ran) ++ ", " ++ show (snd ran) ++ "). " ++ show t

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

-- terms
data Term = V VName | Abs VName Term | App Term Term | Let VName Term Term | LetRec VName 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 ++ ")"
    show (Let x t1 t2) = "let " ++ showVar x ++ " = " ++ show t1 ++ " in " ++ show t2
    show (LetRec x t1 t2) = "letrec " ++ showVar x ++ " = " ++ show t1 ++ " in " ++ 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)
    | x == y = solve (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 t3 t4) : s', s) =
    solve ((t1, t3) : (t2, t4) : 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
------------------------------------------

fresh n = TV ("t", n)

freshs n k = map fresh [n..n+k]

-- instantiate type schema
instantiate :: Type -> Int -> (Int, Type)
instantiate (Schema ran ty) n =
    let
        (n', ty1) = instantiate ty n
        k = snd ran - fst ran
        vs = map (\ i -> ("t", i)) [fst ran..snd ran]
        tys = freshs n' k
        ty2 = lift (zip vs tys) ty1
    in (n' + k + 1, ty2)
instantiate t n = (n, t)

-- 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
        let
            ts = fromJust (get x env)
            (n', ty') = instantiate ts n
        in
            Just (n', (ty, ty') : cs)
    else Nothing
constraints (Abs x t) ty env (n, cs) =
    constraints t t2 ((x, t1) : env) (n + 2, (ty, t1 --> t2) : cs)
    where
        t1 = fresh n
        t2 = fresh (n + 1)
constraints (App t1 t2) ty env (n, cs) =
    case constraints t1 (ty1 --> ty) env (n + 1, cs) of
        Nothing -> Nothing
        Just (n', cs') ->
            constraints t2 ty1 env (n', cs')
    where
        ty1 = fresh n

constraints (Let x t1 t2) ty env (n, cs) =
    case constraints t1 ty1 env (n + 1, cs) of
        Nothing -> Nothing
        Just (n', cs') ->
            case solve (cs', []) of
                Nothing -> Nothing
                Just s ->
                    let
                        ty1' = fromJust (get ("t", n) s)
                        ran = (n, n' - 1)
                        ts = Schema ran ty1'
                    in
                        constraints t2 ty ((x, ts) : env) (n', cs')
    where
        ty1 = fresh n

constraints (LetRec x t1 t2) ty env (n, cs) =
    case constraints t1 ty1 ((x, ty1) : env) (n + 1, cs) of
        Nothing -> Nothing
        Just (n', cs') ->
            case solve (cs', []) of
                Nothing -> Nothing
                Just s ->
                    let
                        ty1' = fromJust (get ("t", n) s)
                        ran = (n, n' - 1)
                        ts = Schema ran ty1'
                    in
                        constraints t2 ty ((x, ts) : env) (n', cs')
    where
        ty1 = fresh n


-- 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)
ts1    = Schema (0,0) (tx --> tx)

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)
f = V ("f", 0)
tId = Abs ("x", 0) x
omega = Abs ("x", 0) (App x x)
tFst = abstr ["x", "y"] x
tSnd = abstr ["x", "y"] y
pair = abstr ["x", "y", "z"] (z $$ x $$ y)
comp = abstr ["x", "y", "z"] (x $$ (y $$ z))
tEx3b = abstr ["x", "y", "z"] (x $$ y $$ (y $$ z))

t1 = Let ("f", 0) tId (pair $$ (f $$ tFst) $$ (f $$ tId))

t2 = LetRec ("x", 0) (abstr ["y"] (x $$ (x $$ y))) (x $$ x)

t3 = LetRec ("x", 0) (abstr ["y"] (x $$ (x $$ y))) ((x $$ x) $$ (x $$ tId))