Make `if cond { ... return() };` assert cond in the type environment

Avoids generating an assert expression with an escape effect.
Mirrors existing case for `if cond { throw(...); };`.
No longer requires `-non_lexical_flow`.

3 files changed, 12 insertions, 20 deletions

diff --git a/src/nl_flow.ml b/src/nl_flow.ml
index bf99f869..6196f23b 100644
--- a/src/nl_flow.ml
+++ b/src/nl_flow.ml
@@ -60,18 +60,6 @@ let rec escapes (E_aux (aux, _)) =
   | E_block exps -> escapes (List.hd (List.rev exps))
   | _ -> false
 
-let rec returns (E_aux (aux, _)) =
-  match aux with
-  | E_return _ -> true
-  | E_block [] -> false
-  | E_block exps -> returns (List.hd (List.rev exps))
-  | _ -> false
-
-let is_assert (E_aux (aux, _)) =
-  match aux with
-  | E_assert _ -> true
-  | _ -> false
-
 let is_bitvector_literal (L_aux (aux, _)) =
   match aux with
   | L_bin _ | L_hex _ -> true
@@ -124,12 +112,6 @@ let analyze' exps =
         List.map (modify_unsigned id value) exps
      | _ -> exps
      end
-  | (E_aux (E_if (cond, then_exp, _), (l, _)) as exp) :: next :: rest
-    when returns then_exp && not (is_assert next) ->
-     let msg = mk_lit_exp (L_string "") in
-     let not_cond = locate (fun _ -> gen_loc l) (mk_exp (E_app (mk_id "not_bool", [cond]))) in
-     let assertion = locate (fun _ -> gen_loc l) (mk_exp (E_assert (not_cond, msg))) in
-     exp :: assertion :: next :: rest
   | _ -> exps
 
 let analyze exps =
diff --git a/src/type_check.ml b/src/type_check.ml
index 0de68a61..b03ef3c4 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -2845,6 +2845,13 @@ let fresh_var =
             let () = counter := n+1 in
             mk_id ("v#" ^ string_of_int n)
 
+let rec exp_unconditionally_returns (E_aux (aux, _)) =
+  match aux with
+  | E_return _ -> true
+  | E_block [] -> false
+  | E_block exps -> exp_unconditionally_returns (List.hd (List.rev exps))
+  | _ -> false
+
 let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ_aux, _) as typ) : tannot exp =
   let annot_exp_effect exp typ' eff = E_aux (exp, (l, mk_expected_tannot env typ' eff (Some typ))) in
   let add_effect exp eff = match exp with
@@ -3119,6 +3126,11 @@ and check_block l env exps ret_typ =
      let cond' = crule check_exp env cond (mk_typ (Typ_id (mk_id "bool"))) in
      let env = add_opt_constraint (option_map nc_not (assert_constraint env false cond')) env in
      texp :: check_block l env exps ret_typ
+  | ((E_aux (E_if (cond, true_exp, _), _) as exp) :: exps) when exp_unconditionally_returns true_exp ->
+     let texp = crule check_exp env exp (mk_typ (Typ_id (mk_id "unit"))) in
+     let cond' = crule check_exp env cond (mk_typ (Typ_id (mk_id "bool"))) in
+     let env = add_opt_constraint (option_map nc_not (assert_constraint env false cond')) env in
+     texp :: check_block l env exps ret_typ
   | (exp :: exps) ->
      let texp = crule check_exp env exp (mk_typ (Typ_id (mk_id "unit"))) in
      texp :: check_block l env exps ret_typ
diff --git a/test/typecheck/pass/if_return.sail b/test/typecheck/pass/if_return.sail
index 9b373c2d..0331bbab 100644
--- a/test/typecheck/pass/if_return.sail
+++ b/test/typecheck/pass/if_return.sail
@@ -1,5 +1,3 @@
-$option -non_lexical_flow
-
 default Order dec
 
 $include <prelude.sail>