-- Type Inference for the Simply-Typed Lambda Calculus
import Data.Maybe (fromJust, isJust)

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

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

infixr 7 :-->

-- types
data Type = TV VName | Type :--> Type
  deriving Eq

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

-- terms
infixl 7 :$:

data Term = V VName | Abs VName Term | Term :$: Term

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

-- 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 substition to Type
appSubst :: Subst -> Type -> Type
appSubst s (TV x) = if inDom x s then app s x else TV x
  where
    app :: Subst -> VName -> Type
    app ((y,t) : s) x = if x == y then t else app s x
appSubst s (t1 :--> t2) = appSubst s t1 :--> appSubst s t2

-- does variable x occur in a Type?
occurs :: VName -> Type -> Bool
occurs = undefined

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

type Unifier = Maybe Subst

-- given list of disagreement pairs, current substitution
-- produces a unifier
solve :: [(Type,Type)] -> Subst -> Unifier
solve das s = undefined

-- 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 das s = undefined

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

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

-- accumulate type constraints
constraints :: Term -> Type -> Env -> (Int, [(Type, Type)]) -> Maybe (Int, [(Type, Type)])
constraints _ ty env (n, cs) = undefined

-- type inference procedure
infer :: Term -> Maybe Type
infer t = undefined

--
-- 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) (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))