unification-pfa/lib/inference.ml

(** Infer the type of a given term and, if exists, returns the type of the term *)
let typeof t =
  let rec infer env = function
    | Term.Var id -> TypeSubstitution.find id env, env
    | Term.IntConst _ -> Some Type.Int, env
    | Term.Binop (t1, _, t2) ->
      (* Both operands must have type Int *)
      (match infer env t1, infer env t2 with
       | (Some ty1, _), (Some ty2, _) ->
         (match Unification.unify ty1 Type.Int, Unification.unify ty2 Type.Int with
          | Some env1, Some env2 -> Some Type.Int, TypeSubstitution.compose env1 env2
          | _ -> None, env)
       | _, _ -> None, env)
    | Term.Pair (t1, t2) ->
      (match infer env t1, infer env t2 with
       | (Some ty1, _), (Some ty2, _) -> Some (Type.Product (ty1, ty2)), env
       | _, _ -> None, env)
    | Term.Proj (proj, t) ->
      (* Check if the term is a pair *)
      (match infer env t with
       | Some (Type.Product (ty1, ty2)), _ ->
         (match proj with
          | Term.First -> Some ty1, env
          | Term.Second -> Some ty2, env)
       | Some (Type.Var id), env' ->
         (* Variable is a product *)
         (match TypeSubstitution.find id env' with
          | Some ty -> Some ty, env'
          | None -> None, env')
       (* Bad typed *)
       | _ -> None, env)
    | Term.Fun (id, body) ->
      let arg = TypeSubstitution.singleton id (Type.Var id) in
      (match infer (TypeSubstitution.compose arg env) body with
       | Some ty_body, env' ->
         (match infer env' (Term.Var id) with
          | Some ty, _ -> Some (Type.Arrow (ty, ty_body)), env'
          | None, _ -> Some (Type.Arrow (Type.Var id, ty_body)), env')
       | _ -> None, env)
    | Term.App (t1, t2) ->
      (* Infer the type of t1 and t2 *)
      (match infer env t1, infer env t2 with
       | (Some (Type.Arrow (ty_param, ty_fn)), _), (Some ty_args, _) ->
         (* Check if the argument type matches the parameter type *)
         (match Unification.unify ty_param ty_args with
          | Some env' ->
            let rec nested_infer oldenv currenv final =
              if TypeSubstitution.equal oldenv currenv
              then Some final, currenv
              else (
                match infer currenv t1 with
                | Some (Type.Arrow (_, ty_fn)), newenv ->
                  nested_infer currenv newenv ty_fn
                | _ -> None, currenv)
            in
            (* Nested *)
            nested_infer env env' ty_fn
          | None -> None, env)
       | _ -> None, env)
  in
  fst (infer TypeSubstitution.empty t)
;;