unification-pfa/lib/inference.ml

(** Infer the type of a given term and, if exists, returns the type of the term *)
let rec typeof = function
  | Term.Var id ->
    (match Unification.unify (Type.Var id) Type.Int with
     | Some _ -> Some Type.Int
     | None -> Some (Type.Var id))
  | Term.IntConst _ -> Some Type.Int
  | Term.Binop (t1, _, t2) ->
    (* Both operands must have type Int *)
    (match typeof t1, typeof t2 with
     | Some ty1, Some ty2 ->
       (match Unification.unify ty1 Type.Int, Unification.unify ty2 Type.Int with
        | Some _, Some _ -> Some Type.Int
        | _ -> None)
     | _, _ -> None)
  | Term.Pair (t1, t2) ->
    (match typeof t1, typeof t2 with
     | Some ty1, Some ty2 -> Some (Type.Product (ty1, ty2))
     | _, _ -> None)
  | Term.Proj (proj, t) ->
    (* Check if the term is a pair *)
    (match typeof t with
     | Some (Type.Product (ty1, ty2)) ->
       (match proj with
        | Term.First -> Some ty1
        | Term.Second -> Some ty2)
     | _ -> None)
  | Term.Fun (id, body) ->
    (match typeof body with
     | Some ty_body ->
       (match typeof (Term.Var id) with
        | Some tt -> Some (Type.Arrow (tt, ty_body))
        | None -> Some (Type.Arrow (Type.Var id, ty_body)))
     | _ -> None)
  | Term.App (t1, t2) ->
    (* Check if the function type matches the arguments *)
    (match typeof t1, typeof t2 with
     | Some (Type.Arrow (ty_param, ty_fn)), Some ty_args ->
       (match Unification.unify ty_param ty_args with
        | Some _ -> Some ty_fn
        | None -> None)
     | _, _ -> None)
;;