external val print : string -> unit = "print"

type var <- string

type expr =
  | Ref of var
  | App of expr * expr
  | Lam of var  * expr
  | Lit of int
  | Pair of expr * expr

let show_expr = function
  | Ref x -> x
  | App (f, x) -> "(" ^ show_expr f ^ ") (" ^ show_expr x ^ ")"
  | Lam (v, x) -> "fun " ^ v ^ ". " ^ show_expr x
  | Lit x -> tostring x
  | Pair (a, b) -> "(" ^ show_expr a ^ ", " ^ show_expr b ^ ")"

let f :$  x = App (f, x)
let v :-> x = Lam (v, x)

type value =
  | Closure of value -> value
  | Free_var of var
  | Int of int
  | Tup of value * value

let eval context = function
  | Ref v ->
    match find (fun (x, _) -> x == v) context with
    | Some (_, vl) -> vl
    | None -> Free_var v
  | Lam (name, body) ->
    Closure @@ fun argument ->
      eval ((name, argument) :: context) body
  | App (f, x) ->
    match eval context f with
    | Closure lambda -> lambda (eval context x)
  | Lit x ->
    Int x
  | Pair (x, y) ->
    Tup (eval context x, eval context y)

type ty =
  | Tv of var * ref (option ty)
  | Fn of ty * ty
  | Int_t
  | Tup_t of ty * ty

type context <- list (var * ty)
type problem <- list (ty * ty)

let fresh_type =
  let counter = ref 0
  fun () ->
    counter := !counter + 1
    Tv ("t" ^ tostring !counter, ref None)

let infer (context : context) expression : either string (ty * problem) =
  let open Either
  match expression with
  | Ref v ->
    match find (fun (x, _) -> x == v) context with
    | Some (_, t) -> Right (t, [])
    | None        -> Left ("variable not in scope: " ^ v)
  | App (fn, arg) -> begin
      with (fn_t, fn_q)   <- infer context fn
      with (arg_t, arg_q) <- infer context arg
      let res_t = fresh_type ()
      Right (res_t, (fn_t, Fn (arg_t, res_t)) :: fn_q ++ arg_q)
    end
  | Lam (v, body) -> begin
      let arg_t = fresh_type ()
      with (t, q) <- infer ((v, arg_t) :: context) body
      Right (Fn (arg_t, t), q)
    end
  | Lit _ ->
    Right (Int_t, [])
  | Pair (x, y) -> begin
      with (x_t, _) <- infer context x
      with (y_t, _) <- infer context y
      Right (Tup_t (x_t, y_t), [])
    end

let ftv = function
  | Tv (v, _) -> [v]
  | Fn (x, y) -> ftv x ++ ftv y
  | Int_t -> []
  | Tup_t (x, y) -> ftv x ++ ftv y

let unify a b =
  match a, b with
  | Int_t, Int_t -> Right ()
  | Int_t, Fn _ -> Left "can't match integer with function type"
  | Fn _, Int_t -> Left "can't match integer with function type"
  | Tv (a, contents), ty ->
    match !contents with
    | Some tau -> unify tau ty
  | Fn (a, b), Fn (c, d) -> begin
      match unify a c with
      | Right _ ->
        match unify b d with
        | Right _ -> Right ()
        | _ -> Left "can't match functions"
      | _ -> Left "can't match functions"
    end
  | Tup_t (a, b), Tup_t (c, d) -> begin
      match unify a c with
      | Right _ ->
        match unify b d with
        | Right _ -> Right ()
        | _ -> Left "can't match tuples"
      | _ -> Left "can't match tuples"
    end
and solve = function
  | [] -> Right ()
  | Cons ((a, b), xs) -> Either.( unify a b >>= fun _ -> solve xs )

let free_in v t =
  match find (== v) (ftv t) with
  | Some _ -> true
  | None -> false

let norm_ty = function
  | Tv (v, contents) -> Option.fold norm_ty (Tv (v, ref None)) !contents
  | Fn (a, b) -> Fn (norm_ty a, norm_ty b)
  | Int_t -> Int_t
  | Tup_t (a, b) -> Tup_t (norm_ty a, norm_ty b)

let display_ty = function
  | Fn (x, b) -> paren_fn x ^ " -> " ^ display_ty b
  | Tup_t (a, b) -> "(" ^ display_ty a ^ ", " ^ display_ty b ^ ")"
  | Int_t -> "int"
  | Tv (v, _) -> v
and paren_fn = function
  | Fn _ as t -> "(" ^ display_ty t ^ ")"
  | x -> display_ty x

let type_check expr =
  let open Either in
  fold (fun x -> x) (fun x -> x) @@
    begin
      with (ty, qs) <- infer [] expr
      with _ <- solve qs
      Right (display_ty (norm_ty ty))
    end

let () =
  let id = "x" :-> Ref "x"
  let x = Ref "x"
  print @@ show_expr id