Remove unnecessary size_itself_int uses in guards (for Lem)

Doesn't remove them from function bodies because that can produce
more work for the sizeof rewriting.

3 files changed, 97 insertions, 9 deletions

diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index 8ec07f7f..4d87070c 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -2121,13 +2121,25 @@ let change_parameter_pat = function
   | P_aux (_,(l,_)) -> raise (Reporting_basic.err_unreachable l
                                 "Expected variable pattern")
 
+(* TODO: make more precise, preferably with a proper free variables function
+   which deals with shadowing *)
+let var_maybe_used_in_exp exp var =
+  let open Rewriter in
+  fst (fold_exp {
+    (compute_exp_alg false (||)) with
+      e_id = fun id -> (Id.compare id var == 0, E_id id) } exp)
+
 (* We add code to change the itself('n) parameter into the corresponding
-   integer. *)
-let add_var_rebind exp var =
-  let l = Generated Unknown in
-  let annot = (l,None) in
-  E_aux (E_let (LB_aux (LB_val (P_aux (P_id var,annot),
-                                E_aux (E_app (mk_id "size_itself_int",[E_aux (E_id var,annot)]),annot)),annot),exp),annot)
+   integer.  We always do this for the function body (otherwise we'd have to do
+   something clever with E_sizeof to avoid making things more complex), but
+   only for guards when they actually use the variable. *)
+let add_var_rebind unconditional exp var =
+  if unconditional || var_maybe_used_in_exp exp var then
+    let l = Generated Unknown in
+    let annot = (l,None) in
+    E_aux (E_let (LB_aux (LB_val (P_aux (P_id var,annot),
+                                  E_aux (E_app (mk_id "size_itself_int",[E_aux (E_id var,annot)]),annot)),annot),exp),annot)
+  else exp
 
 (* atom('n) arguments to function calls need to be rewritten *)
 let replace_with_the_value bound_nexps (E_aux (_,(l,_)) as exp) =
@@ -2266,11 +2278,10 @@ let rewrite_size_parameters env (Defs defs) =
                    pat, [var]
               end
          in
-         (* TODO: only add bindings that are necessary (esp for guards) *)
-         let body = List.fold_left add_var_rebind body vars in
+         let body = List.fold_left (add_var_rebind true) body vars in
          let guard = match guard with
            | None -> None
-           | Some exp -> Some (List.fold_left add_var_rebind exp vars)
+           | Some exp -> Some (List.fold_left (add_var_rebind false) exp vars)
          in
          pat,guard,body,nexps
       | exception Not_found -> pat,guard,body,NexpSet.empty
diff --git a/test/mono/itself_rewriting.sail b/test/mono/itself_rewriting.sail
new file mode 100644
index 00000000..f3c5ed62
--- /dev/null
+++ b/test/mono/itself_rewriting.sail
@@ -0,0 +1,76 @@
+$include <smt.sail>
+$include <flow.sail>
+default Order dec
+type bits ('n : Int) = vector('n, dec, bit)
+val operator & = "and_bool" : (bool, bool) -> bool
+val eq_vec = {ocaml: "eq_list", lem: "eq_vec"} : forall 'n. (bits('n), bits('n)) -> bool
+overload operator == = {eq_int, eq_vec}
+val neq_vec = {lem: "neq"} : forall 'n. (bits('n), bits('n)) -> bool
+function neq_vec (x, y) = not_bool(eq_vec(x, y))
+overload operator != = {neq_atom, neq_vec}
+val vector_subrange = {ocaml: "subrange", lem: "subrange_vec_dec"} : forall ('n : Int) ('m : Int) ('o : Int), 'o <= 'm <= 'n.
+  (bits('n), atom('m), atom('o)) -> bits('m - ('o - 1))
+val mult_int = {ocaml: "mult", lem: "integerMult"} : (int, int) -> int
+overload operator * = {mult_range, mult_int, mult_real}
+val UInt = {
+  ocaml: "uint",
+  lem: "uint",
+  interpreter: "uint",
+  c: "sail_uint"
+} : forall 'n. bits('n) -> range(0, 2 ^ 'n - 1)
+val bitvector_cast = "zeroExtend" : forall 'n. bits('n) -> bits('n) effect pure
+val slice = "slice" : forall ('n : Int) ('m : Int), 'm >= 0 & 'n >= 0.
+  (bits('m), int, atom('n)) -> bits('n)
+val replicate_bits = "replicate_bits" : forall 'n 'm. (bits('n), atom('m)) -> bits('n * 'm)
+overload operator > = {gt_int}
+
+/* This was written to manually inspect that "let n = size_itself_int n in" gets
+   added in the right places, but it's also worth running in case that gets
+   broken. */
+
+val needs_size_in_guard : forall 'n. atom('n) -> unit
+
+function needs_size_in_guard(n if n > 8) = {
+  let x : bits('n) = replicate_bits(0b0,n);
+  ()
+}
+and needs_size_in_guard(n) = {
+  let x : bits('n) = replicate_bits(0b1,n);
+  ()
+}
+
+val no_size_in_guard : forall 'n. (atom('n), int) -> unit
+
+function no_size_in_guard((n,m) if m > 8) = {
+  let x : bits('n) = replicate_bits(0b0,n);
+  ()
+}
+and no_size_in_guard(n,m) = {
+  let x : bits('n) = replicate_bits(0b1,n);
+  ()
+}
+
+val shadowed : forall 'n. atom('n) -> unit
+
+function shadowed(n) = {
+  let n = 8;
+  let x : bits(8) = replicate_bits(0b0,n);
+  ()
+}
+
+val willsplit : bool -> unit
+
+function willsplit(x) = {
+  let 'n : int = if x then 8 else 16;
+  needs_size_in_guard(n);
+  no_size_in_guard(n,n);
+  shadowed(n);
+}
+
+val run : unit -> unit effect {escape}
+
+function run () = {
+  assert(true); /* To force us into the monad */
+  willsplit(true);
+  willsplit(false);
+}
\ No newline at end of file
diff --git a/test/mono/pass/itself_rewriting b/test/mono/pass/itself_rewriting
new file mode 100644
index 00000000..f500694d
--- /dev/null
+++ b/test/mono/pass/itself_rewriting
@@ -0,0 +1 @@
+itself_rewriting.sail -auto_mono