Add some optional experimental rewrites to help with monomorphisation

5 files changed, 228 insertions, 4 deletions

diff --git a/lib/mono_rewrites.sail b/lib/mono_rewrites.sail
new file mode 100644
index 00000000..3190003f
--- /dev/null
+++ b/lib/mono_rewrites.sail
@@ -0,0 +1,69 @@
+(* Definitions for use with the -mono_rewrites option *)
+
+(* External definitions not in the usual asl prelude *)
+
+val "shiftleft" : forall 'n 'o ('ord : Order).
+    (vector('o, 'n, 'ord, bit), int) -> vector('o, 'n, 'ord, bit) effect pure
+
+overload operator << = {shiftleft}
+
+val "shiftright" : forall 'n 'o ('ord : Order).
+    (vector('o, 'n, 'ord, bit), int) -> vector('o, 'n, 'ord, bit) effect pure
+
+overload operator >> = {shiftright}
+
+val "extz" : forall 'n 'm 'o 'p. (atom('p),atom('m),vector('o, 'n, dec, bit)) -> vector('p, 'm, dec, bit) effect pure
+val extzv : forall 'n 'm 'o 'p. vector('o, 'n, dec, bit) -> vector('p, 'm, dec, bit) effect pure
+function extzv(v) = extz(sizeof('p), sizeof('m),v)
+
+(* Definitions for the rewrites *)
+
+val slice_mask : forall 'n, 'n >= 0. (int, int) -> bits('n) effect pure
+function slice_mask(i,l) =
+  let one : bits('n) = extzv(0b1) in
+  ((one << l) - one) << i
+
+val is_zero_subrange : forall 'n, 'n >= 0.
+  (bits('n), int, int) -> bool effect pure
+
+function is_zero_subrange (xs, i, j) = {
+  (xs & slice_mask(j, i-j)) == extzv(0b0)
+}
+
+val is_ones_subrange : forall 'n, 'n >= 0.
+  (bits('n), int, int) -> bool effect pure
+
+function is_ones_subrange (xs, i, j) = {
+  let m : bits('n) = slice_mask(j,j-i) in
+  (xs & m) == m
+}
+
+val slice_slice_concat : forall 'n 'm 'p 'q 'r, 'r = 'p + 'q & 'n >= 0 & 'm >= 0 & 'p >= 0 & 'q >= 0.
+  (bits('n), int, atom('p), bits('m), int, atom('q)) -> bits('r) effect pure
+
+function slice_slice_concat (xs, i, l, ys, i', l') = {
+  let xs = (xs & slice_mask(i,l)) >> i in
+  let ys = (ys & slice_mask(i',l')) >> i' in
+  extzv(xs) << l' | extzv(ys)
+}
+
+(* Assumes initial vectors are of equal size *)
+
+val subrange_subrange_eq : forall 'n, 'n >= 0.
+  (bits('n), int, int, bits('n), int, int) -> bool effect pure
+
+function subrange_subrange_eq (xs, i, j, ys, i', j') = {
+  let xs = (xs & slice_mask(j,i-j)) >> j in
+  let ys = (ys & slice_mask(j',i'-j')) >> j' in
+  xs == ys
+}
+
+val subrange_subrange_concat : forall 'n 'o 'p 'm 'q 'r 's, 's = 'o - ('p - 1) + 'q - ('r - 1) & 'n >= 0 & 'm >= 0.
+  (bits('n), atom('o), atom('p), bits('m), atom('q), atom('r)) -> bits('s) effect pure
+
+function subrange_subrange_concat (xs, i, j, ys, i', j') = {
+  let xs = (xs & slice_mask(j,i-j)) >> j in
+  let ys = (ys & slice_mask(j',i'-j')) >> j' in
+  extzv(xs) << i' - (j' - 1) | extzv(ys)
+}
+
diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index 0cbeda49..eef94982 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -2460,11 +2460,156 @@ let argset_to_list splits =
   List.map argelt l
 end
 
+module MonoRewrites =
+struct
+
+let is_constant_range = function
+  | E_aux (E_lit _,_), E_aux (E_lit _,_) -> true
+  | _ -> false
+
+let is_constant = function
+  | E_aux (E_lit _,_) -> true
+  | _ -> false
+
+let is_constant_vec_typ env typ =
+  let typ = Env.base_typ_of env typ in
+  match destruct_vector env typ with
+  | Some (_,size,_,_) ->
+     (match nexp_simp size with
+     | Nexp_aux (Nexp_constant _,_) -> true
+     | _ -> false)
+  | _ -> false
+
+let is_id env id =
+  let ids = Env.get_overloads (Id_aux (id,Parse_ast.Unknown)) env in
+  let ids = id :: List.map (fun (Id_aux (id,_)) -> id) ids in
+  fun (Id_aux (x,_)) -> List.mem x ids
+
+(* We have to add casts in here with appropriate length information so that the
+   type checker knows the expected return types. *)
 
-let monomorphise mwords auto debug_analysis splits env defs =
+let rewrite_app env typ (id,args) =
+  let is_append = is_id env (Id "append") in
+  if is_append id then
+    let is_subrange = is_id env (Id "vector_subrange") in
+    let is_slice = is_id env (Id "slice") in
+    match args with
+      (* (known-size-vector @ variable-vector) @ variable-vector *)
+    | [E_aux (E_app (append,
+              [e1;
+               E_aux (E_app (subrange1,
+                             [vector1; start1; end1]),_)]),_);
+       E_aux (E_app (subrange2,
+                     [vector2; start2; end2]),_)]
+        when is_append append && is_subrange subrange1 && is_subrange subrange2 &&
+          is_constant_vec_typ env (typ_of e1) &&
+          not (is_constant_range (start1, end1) || is_constant_range (start2, end2)) ->
+       let (start,size,order,bittyp) = vector_typ_args_of (Env.base_typ_of env typ) in
+       let (_,size1,_,_) = vector_typ_args_of (Env.base_typ_of env (typ_of e1)) in
+       let midsize = nminus size size1 in
+       let midstart =
+         if is_order_inc order
+         then nconstant zero_big_int
+         else nminus midsize (nconstant unit_big_int) in
+       let midtyp = vector_typ midstart midsize order bittyp in
+       E_app (append,
+              [e1;
+               E_aux (E_cast (midtyp,
+                              E_aux (E_app (mk_id "subrange_subrange_concat",
+                                            [vector1; start1; end1; vector2; start2; end2]),
+                                     (Unknown,None))),(Unknown,None))])
+
+    (* variable-range @ variable-range *)
+    | [E_aux (E_app (subrange1,
+                     [vector1; start1; end1]),_);
+       E_aux (E_app (subrange2,
+                     [vector2; start2; end2]),_)]
+        when is_subrange subrange1 && is_subrange subrange2 &&
+          not (is_constant_range (start1, end1) || is_constant_range (start2, end2)) ->
+       E_cast (typ,
+               E_aux (E_app (mk_id "subrange_subrange_concat",
+                             [vector1; start1; end1; vector2; start2; end2]),
+                      (Unknown,None)))
+
+    (* variable-slice @ variable-slice *)
+    | [E_aux (E_app (slice1,
+                     [vector1; start1; length1]),_);
+       E_aux (E_app (slice2,
+                     [vector2; start2; length2]),_)]
+        when is_slice slice1 && is_slice slice2 &&
+          not (is_constant length1 || is_constant length2) ->
+       E_cast (typ,
+               E_aux (E_app (mk_id "slice_slice_concat",
+                             [vector1; start1; length1; vector2; start2; length2]),(Unknown,None)))
+
+    | _ -> E_app (id,args)
+
+  else if is_id env (Id "eq_vec") id then
+    (* variable-range == variable_range *)
+    let is_subrange = is_id env (Id "vector_subrange") in
+    match args with
+    | [E_aux (E_app (subrange1,
+                     [vector1; start1; end1]),_);
+       E_aux (E_app (subrange2,
+                     [vector2; start2; end2]),_)]
+        when is_subrange subrange1 && is_subrange subrange2 &&
+          not (is_constant_range (start1, end1) || is_constant_range (start2, end2)) ->
+       E_app (mk_id "subrange_subrange_eq",
+              [vector1; start1; end1; vector2; start2; end2])
+    | _ -> E_app (id,args)
+
+  else if is_id env (Id "IsZero") id then
+    match args with
+    | [E_aux (E_app (subrange1, [vector1; start1; end1]),_)]
+        when is_id env (Id "vector_subrange") subrange1 &&
+          not (is_constant_range (start1,end1)) ->
+       E_app (mk_id "is_zero_subrange",
+              [vector1; start1; end1])
+    | _ -> E_app (id,args)
+
+  else if is_id env (Id "IsOnes") id then
+    match args with
+    | [E_aux (E_app (subrange1, [vector1; start1; end1]),_)]
+        when is_id env (Id "vector_subrange") subrange1 &&
+          not (is_constant_range (start1,end1)) ->
+       E_app (mk_id "is_ones_subrange",
+              [vector1; start1; end1])
+    | _ -> E_app (id,args)
+
+  else E_app (id,args)
+
+let rewrite_aux = function
+  | (E_app (id,args),((_,Some (env,ty,_)) as annot)) ->
+     E_aux (rewrite_app env ty (id,args),annot)
+  | exp,annot -> E_aux (exp,annot)
+
+let mono_rewrite defs =
+  let open Rewriter in
+  rewrite_defs_base
+    { rewriters_base with
+      rewrite_exp = fun _ -> fold_exp { id_exp_alg with e_aux = rewrite_aux } }
+    defs
+end
+
+(* TODO: put in a proper mli for this stuff *)
+type options = {
+  auto : bool;
+  debug_analysis : int;
+  rewrites : bool
+}
+
+let monomorphise mwords opts splits env defs =
+  let (defs,env) =
+    if opts.rewrites then
+      let defs = MonoRewrites.mono_rewrite defs in
+    (* TODO: is this necessary? *)
+      Type_check.check (Type_check.Env.no_casts Type_check.initial_env) defs
+    else (defs,env)
+  in
+(*let _ = Pretty_print.pp_defs stdout defs in*)
   let new_splits =
-    if auto
-    then Analysis.argset_to_list (Analysis.analyse_defs debug_analysis env defs)
+    if opts.auto
+    then Analysis.argset_to_list (Analysis.analyse_defs opts.debug_analysis env defs)
     else [] in
   let defs = split_defs (new_splits@splits) defs in
   (* TODO: currently doing this because constant propagation leaves numeric literals as
diff --git a/src/process_file.ml b/src/process_file.ml
index 83cfd8a3..127a5aee 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -126,9 +126,15 @@ let check_ast (defs : unit Ast.defs) : Type_check.tannot Ast.defs * Type_check.E
 let opt_ddump_raw_mono_ast = ref false
 let opt_dmono_analysis = ref 0
 let opt_auto_mono = ref false
+let opt_mono_rewrites = ref false
 
 let monomorphise_ast locs type_env ast =
-  let ast = Monomorphise.monomorphise (!Pretty_print_lem.opt_mwords) (!opt_auto_mono) (!opt_dmono_analysis)
+  let open Monomorphise in
+  let opts = {
+    auto = !opt_auto_mono;
+    debug_analysis = !opt_dmono_analysis;
+    rewrites = !opt_mono_rewrites } in
+  let ast = monomorphise (!Pretty_print_lem.opt_mwords) opts
     locs type_env ast in
   let () = if !opt_ddump_raw_mono_ast then Pretty_print.pp_defs stdout ast else () in
   let ienv = Type_check.Env.no_casts Type_check.initial_env in
diff --git a/src/process_file.mli b/src/process_file.mli
index 88c9b9fb..10510e5c 100644
--- a/src/process_file.mli
+++ b/src/process_file.mli
@@ -70,6 +70,7 @@ val opt_dno_cast : bool ref
 val opt_ddump_raw_mono_ast : bool ref
 val opt_dmono_analysis : int ref
 val opt_auto_mono : bool ref
+val opt_mono_rewrites : bool ref
 
 type out_type =
   | Lem_ast_out
diff --git a/src/sail.ml b/src/sail.ml
index 9bb7cfcb..1a0f6282 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -132,6 +132,9 @@ let options = Arg.align ([
   ( "-auto_mono",
     Arg.Set opt_auto_mono,
     " automatically infer how to monomorphise code");
+  ( "-mono_rewrites",
+    Arg.Set Process_file.opt_mono_rewrites,
+    " turn on rewrites for combining bitvector operations");
   ( "-verbose",
     Arg.Set opt_print_verbose,
     " (debug) pretty-print the input to standard output");