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 _ -> None
  | Term.IntConst _ -> Some Type.Int
  | Term.Binop (t1, _, t2) ->
    (match typeof t1, typeof t2 with
     | ty1, ty2 when ty1 = ty2 -> Some Type.Int
     | _ -> None)
  | Term.Pair (t1, t2) ->
    (* Pair give Products *)
    (match typeof t1, typeof t2 with
     | Some ty1, Some ty2 -> Some (Type.Product (ty1, ty2))
     | _ -> None)
  | Term.Proj (proj, t) ->
    (* Projections returns type of product based on the projection type *)
    (match proj, typeof t with
     | Term.First, Some (Type.Product (ty, _)) | Term.Second, Some (Type.Product (_, ty))
       -> Some ty
     | _, _ -> None)
  | Term.Fun (id, t) ->
    (match typeof t with
     | Some body_type -> Some (Type.Arrow (Type.Int, body_type))
     | None -> Some (Type.Var id))
  | Term.App (t1, t2) ->
    (match typeof t1, typeof t2 with
     | Some (Type.Arrow (ty_arg, ty_res)), Some ty_arg' ->
       (* Unification for application *)
       (match Unification.unify ty_arg ty_arg' with
        | Some subst -> Some (TypeSubstitution.apply subst ty_res)
        | None -> None)
     | _ -> None)
;;