Avoid adding explicit type annotations with generated type variables

3 files changed, 22 insertions, 8 deletions

diff --git a/src/rewriter.ml b/src/rewriter.ml
index 8cf8d87c..203e8a58 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -92,19 +92,31 @@ let lookup_generated_kid env kid =
   in
   List.fold_left match_kid_nc kid (Env.get_constraints env)
 
+let generated_kids typ = KidSet.filter is_kid_generated (tyvars_of_typ typ)
+
+let resolve_generated_kids env typ =
+  let subst_kid kid typ = typ_subst_kid kid (lookup_generated_kid env kid) typ in
+  KidSet.fold subst_kid (generated_kids typ) typ
+
 let rec remove_p_typ = function
   | P_aux (P_typ (typ, pat), _) -> remove_p_typ pat
   | pat -> pat
 
 let add_p_typ typ (P_aux (paux, annot) as pat) =
-  let env = pat_env_of pat in
-  let generated_kids typ = KidSet.filter is_kid_generated (tyvars_of_typ typ) in
-  let subst_kid kid typ = typ_subst_kid kid (lookup_generated_kid env kid) typ in
-  let typ' = KidSet.fold subst_kid (generated_kids typ) typ in
+  let typ' = resolve_generated_kids (pat_env_of pat) typ in
   if KidSet.is_empty (generated_kids typ') then
     P_aux (P_typ (typ', remove_p_typ pat), annot)
   else pat
 
+let rec remove_e_cast = function
+  | E_aux (E_cast (_, exp), _) -> remove_e_cast exp
+  | exp -> exp
+
+let add_e_cast typ (E_aux (eaux, annot) as exp) =
+  let typ' = resolve_generated_kids (env_of exp) typ in
+  if KidSet.is_empty (generated_kids typ') then
+    E_aux (E_cast (typ', remove_e_cast exp), annot)
+  else exp
 
 let rec small (E_aux (exp,_)) = match exp with
   | E_id _
diff --git a/src/rewriter.mli b/src/rewriter.mli
index 83434764..f8982d69 100644
--- a/src/rewriter.mli
+++ b/src/rewriter.mli
@@ -206,6 +206,8 @@ val simple_annot : Parse_ast.l -> typ -> Parse_ast.l * tannot
 
 val add_p_typ : typ -> tannot pat -> tannot pat
 
+val add_e_cast : typ -> tannot exp -> tannot exp
+
 val effect_of_pexp : tannot pexp -> effect
 
 val union_eff_exps : (tannot exp) list -> effect
diff --git a/src/rewrites.ml b/src/rewrites.ml
index 265e0551..4a213970 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -1814,7 +1814,7 @@ let rewrite_defs_early_return (Defs defs) =
       (* Add escape effect annotation, since we use the exception mechanism
          of the state monad to implement early return in the Lem backend *)
       let annot' = Some (env, typ, union_effects eff (mk_effect [BE_escape])) in
-      let exp' = annot_exp (E_cast (typ_of exp, exp)) l env (typ_of exp) in
+      let exp' = add_e_cast (typ_of exp) exp in
       E_aux (E_app (mk_id "early_return", [exp']), (l, annot'))
     | _ -> full_exp in
 
@@ -2532,7 +2532,7 @@ let rewrite_defs_letbind_effects =
     let exp =
       if newreturn then
         (* let typ = try typ_of exp with _ -> unit_typ in *)
-        let exp = annot_exp (E_cast (typ_of exp, exp)) l (env_of exp) (typ_of exp) in
+        let exp = add_e_cast (typ_of exp) exp in
         annot_exp (E_internal_return exp) l (env_of exp) (typ_of exp)
       else
         exp in
@@ -2738,7 +2738,7 @@ let rewrite_defs_internal_lets =
        let (lhs, rhs) = rewrite_lexp_to_rhs le in
        let (LEXP_aux (_, lannot)) = lhs in
        let ltyp = typ_of_annot lannot in
-       let rhs = annot_exp (E_cast (ltyp, rhs exp)) l (env_of_annot lannot) ltyp in
+       let rhs = add_e_cast ltyp (rhs exp) in
        E_let (LB_aux (LB_val (pat_of_local_lexp lhs, rhs), annot), body)
     | LB_aux (LB_val (pat,exp'),annot') ->
        if effectful exp'
@@ -3102,7 +3102,7 @@ let rewrite_defs_remove_superfluous_returns =
 
   let add_opt_cast typopt1 typopt2 annot exp =
     match typopt1, typopt2 with
-    | Some typ, _ | _, Some typ -> E_aux (E_cast (typ, exp), annot)
+    | Some typ, _ | _, Some typ -> add_e_cast typ exp
     | None, None -> exp
   in