Handle more uses of mutable variables during monomorphisation cast insertion

In particular, bitvector subrange updates work with this version.

4 files changed, 115 insertions, 33 deletions

diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index ac68ba60..1c06e2d9 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -3227,6 +3227,7 @@ let make_bitvector_cast_fns cast_name top_env env quant_kids src_typ target_typ
     check_for_spec env cast_name;
     let src_ann = mk_tannot env src_typ no_effect in
     let tar_ann = mk_tannot env target_typ no_effect in
+    let asg_ann = mk_tannot env unit_typ no_effect in
     match src_typ' with
       (* Simple case with just the bitvector; don't need to pull apart value *)
     | Typ_aux (Typ_app _,_) ->
@@ -3236,9 +3237,15 @@ let make_bitvector_cast_fns cast_name top_env env quant_kids src_typ target_typ
                                        E_aux (E_app (Id_aux (Id cast_name,genunk),
                                                      [E_aux (E_id var,(genunk,src_ann))]),(genunk,tar_ann))),(genunk,tar_ann)),
                        exp),(genunk,exp_ann))),
+       (fun var ->
+         [E_aux (E_assign (LEXP_aux (LEXP_cast (one_target_typ, var),(genunk,tar_ann)),
+                          E_aux (E_app (Id_aux (Id cast_name,genunk),
+                                        [E_aux (E_id var,(genunk,src_ann))]),(genunk,tar_ann)
+                  )),(genunk,asg_ann))]),
       (fun (E_aux (_,(exp_l,exp_ann)) as exp) ->
         E_aux (E_cast (one_target_typ,
                        E_aux (E_app (Id_aux (Id cast_name, genunk), [exp]), (Generated exp_l,tar_ann))),
+
                (Generated exp_l,tar_ann)))
     | _ ->
        (fun var exp ->
@@ -3246,17 +3253,58 @@ let make_bitvector_cast_fns cast_name top_env env quant_kids src_typ target_typ
          E_aux (E_let (LB_aux (LB_val (pat, E_aux (E_id var,(genunk,src_ann))),(genunk,src_ann)),
                        E_aux (E_let (LB_aux (LB_val (P_aux (P_id var,(genunk,tar_ann)),e'),(genunk,tar_ann)),
                                      exp),(genunk,exp_ann))),(genunk,exp_ann))),
+       (fun var ->
+         [E_aux (E_let (LB_aux (LB_val (pat, E_aux (E_id var,(genunk,src_ann))),(genunk,src_ann)),
+                       E_aux (E_assign (LEXP_aux (LEXP_cast (one_target_typ, var),(genunk,tar_ann)),
+                                        e'),(genunk,asg_ann))),(genunk,asg_ann))]),
       (fun (E_aux (_,(exp_l,exp_ann)) as exp) ->
         E_aux (E_let (LB_aux (LB_val (pat, exp),(Generated exp_l,exp_ann)), e'),(Generated exp_l,tar_ann)))
   end
-  | None -> (fun _ e -> e),(fun e -> e)
+  | None -> (fun _ e -> e),(fun _ -> []),(fun e -> e)
+let make_bitvector_cast_let cast_name top_env env quant_kids src_typ target_typ =
+  let f,_,_ = make_bitvector_cast_fns cast_name top_env env quant_kids src_typ target_typ
+  in f
+let make_bitvector_cast_assign cast_name top_env env quant_kids src_typ target_typ =
+  let _,f,_ = make_bitvector_cast_fns cast_name top_env env quant_kids src_typ target_typ
+  in f
+let make_bitvector_cast_cast cast_name top_env env quant_kids src_typ target_typ =
+  let _,_,f = make_bitvector_cast_fns cast_name top_env env quant_kids src_typ target_typ
+  in f
+
+let ids_in_exp exp =
+  let open Rewriter in
+  fold_exp {
+      (pure_exp_alg IdSet.empty IdSet.union) with
+      e_id = IdSet.singleton;
+      lEXP_id = IdSet.singleton;
+      lEXP_memory = (fun (id,s) -> List.fold_left IdSet.union (IdSet.singleton id) s);
+      lEXP_cast = (fun (_,id) -> IdSet.singleton id)
+    } exp
 
-(* TODO: bound vars *)
 let make_bitvector_env_casts env quant_kids (kid,i) exp =
-  let mk_cast var typ exp = (fst (make_bitvector_cast_fns "bitvector_cast_in" env env quant_kids typ (subst_kids_typ (KBindings.singleton kid (nconstant i)) typ))) var exp in
+  let mk_cast var typ exp = (make_bitvector_cast_let "bitvector_cast_in" env env quant_kids typ (subst_kids_typ (KBindings.singleton kid (nconstant i)) typ)) var exp in
+  let mk_assign_in var typ =
+    make_bitvector_cast_assign "bitvector_cast_in" env env quant_kids typ
+      (subst_kids_typ (KBindings.singleton kid (nconstant i)) typ) var
+  in
+  let mk_assign_out var typ =
+    make_bitvector_cast_assign "bitvector_cast_out" env env quant_kids
+      (subst_kids_typ (KBindings.singleton kid (nconstant i)) typ) typ var
+  in
   let locals = Env.get_locals env in
+  let used_ids = ids_in_exp exp in
+  let locals = Bindings.filter (fun id _ -> IdSet.mem id used_ids) locals in
+  let immutables,mutables = Bindings.partition (fun _ (mut,_) -> mut = Immutable) locals in
+  let assigns_in = Bindings.fold (fun var (_,typ) acc -> mk_assign_in var typ @ acc) mutables [] in
+  let assigns_out = Bindings.fold (fun var (_,typ) acc -> mk_assign_out var typ @ acc) mutables [] in
+  let exp = match assigns_in, exp with
+    | [], _ -> exp
+    | _::_, E_aux (E_block es,ann) -> E_aux (E_block (assigns_in @ es @ assigns_out),ann)
+    | _::_, E_aux (_,(l,ann)) ->
+       E_aux (E_block (assigns_in @ [exp] @ assigns_out), (Generated l,ann))
+  in
   Bindings.fold (fun var (mut,typ) exp ->
-    if mut = Immutable then mk_cast var typ exp else exp) locals exp
+    if mut = Immutable then mk_cast var typ exp else exp) immutables exp
 
 let make_bitvector_cast_exp cast_name cast_env quant_kids typ target_typ exp =
   if alpha_equivalent cast_env typ target_typ then exp else
@@ -3300,7 +3348,7 @@ let make_bitvector_cast_exp cast_name cast_env quant_kids typ target_typ exp =
        let tgt_arg_typ = infer_arg_typ (env_of exp) f l target_typ in
        E_aux (E_app (f,[aux arg (src_arg_typ, tgt_arg_typ)]),(l,ann))
     | _ ->
-       (snd (make_bitvector_cast_fns cast_name cast_env (env_of exp) quant_kids typ target_typ)) exp
+       (make_bitvector_cast_cast cast_name cast_env (env_of exp) quant_kids typ target_typ) exp
   in
   aux exp (typ, target_typ)
 
@@ -3419,26 +3467,6 @@ let add_bitvector_casts (Defs defs) =
          E_aux (E_if (e1,e2',e3'), ann)
       | E_return e' ->
          E_aux (E_return (make_bitvector_cast_exp "bitvector_cast_out" top_env quant_kids (fill_in_type (env_of e') (typ_of e')) ret_typ e'),ann)
-      | E_assign (LEXP_aux (_,lexp_annot) as lexp,e') -> begin
-         (* The type in the lexp_annot might come from e' rather than being the
-            type of the storage, so ask the type checker what it really is. *)
-         match infer_lexp (env_of_annot lexp_annot) (strip_lexp lexp) with
-         | LEXP_aux (_,lexp_annot') ->
-            E_aux (E_assign (lexp,
-                             make_bitvector_cast_exp "bitvector_cast_out" top_env quant_kids (fill_in_type (env_of e') (typ_of e'))
-                               (typ_of_annot lexp_annot') e'),ann)
-         | exception _ -> E_aux (e,ann)
-        end
-      | E_id id -> begin
-        let env = env_of_annot ann in
-        match Env.lookup_id id env with
-        | Local (Mutable, vtyp) ->
-           make_bitvector_cast_exp "bitvector_cast_in" top_env quant_kids
-             (fill_in_type (env_of_annot ann) (typ_of_annot ann))
-             vtyp
-             (E_aux (e,ann))
-        | _ -> E_aux (e,ann)
-        end
       | E_block es ->
          let env = env_of_annot ann in
          let result_typ = Env.base_typ_of env (typ_of_annot ann) in
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index cf083eb5..4890b8d1 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -552,6 +552,28 @@ let contains_early_return exp =
   { (Rewriter.compute_exp_alg false (||))
     with e_return = (fun (_, r) -> (true, E_return r)); e_app = e_app } exp)
 
+(* Does the expression have the form of a bitvector cast from the monomorphiser? *)
+type is_bitvector_cast = BVC_yes | BVC_allowed | BVC_not
+let is_bitvector_cast_out exp =
+  let merge x y = match x,y with
+    | BVC_allowed, _ -> y
+    | _, BVC_allowed -> x
+    | BVC_not, _ -> BVC_not
+    | _, BVC_not -> BVC_not
+    | _ -> BVC_yes
+  in
+  let rec aux (E_aux (e,_)) =
+    match e with
+    | E_tuple es -> List.fold_left merge BVC_allowed (List.map aux es)
+    | E_cast (_,e) -> aux e
+    | E_app (Id_aux (Id "bitvector_cast_out",_),_) -> BVC_yes
+    | E_id _ -> BVC_allowed
+    | _ -> BVC_not
+  in aux exp = BVC_yes
+
+let replace_env_for_cast_out new_env pat =
+  map_pat_annot (fun (l,a) -> (l,replace_env new_env a)) pat
+
 let find_e_ids exp =
   let e_id id = IdSet.singleton id, E_id id in
   fst (fold_exp
@@ -978,6 +1000,7 @@ let doc_exp_lem, doc_let_lem =
     else_pp
   and let_exp ctxt (LB_aux(lb,_)) = match lb with
     | LB_val(pat,e) ->
+       let pat = if is_bitvector_cast_out e then replace_env_for_cast_out ctxt.top_env pat else pat in
        prefix 2 1
               (separate space [string "let"; doc_pat_lem ctxt true pat; equals])
               (top_exp ctxt false e)
diff --git a/test/mono/castreq.sail b/test/mono/castreq.sail
index e8fbabb0..ce7080c4 100644
--- a/test/mono/castreq.sail
+++ b/test/mono/castreq.sail
@@ -21,6 +21,13 @@ val bitvector_cast_out = "zeroExtend" : forall 'n. bits('n) -> bits('n) effect p
 val bitvector_length = "length" : forall 'n. bits('n) -> atom('n)
 overload length = {bitvector_length}
 overload __size = {length}
+val add_bits = {ocaml: "add_vec", lem: "add_vec"} : forall 'n. (bits('n), bits('n)) -> bits('n)
+overload operator + = {add_bits}
+val vector_update_subrange = {
+  ocaml: "update_subrange",
+  lem: "update_subrange_vec_dec"
+} : forall 'n 'm 'o, 0 <= 'o <= 'm < 'n. (bits('n), atom('m), atom('o), bits('m - ('o - 1))) -> bits('n)
+
 
 /* Test generation of casts across case splits (e.g., going from bits('m) to bits(32)) */
 
@@ -69,11 +76,9 @@ function ret(n, x) =
     64 => let z = y@y@y@y in { dfsf = 4; return z; undefined }
   }
 
-/* TODO: Assignments need more plumbing
-
-val assign : forall 'm 'n, 'm in {8,16} & 'n in {32,64}. bits('m) -> bits('n) effect {undef}
+val assign : forall 'm 'n, 'm in {8,16} & 'n in {32,64}. (implicit('n), bits('m)) -> bits('n) effect {undef}
 
-function assign(x) = {
+function assign(n, x) = {
   let y : bits(16) = extzv(x);
   r : bits('n) = undefined;
   match 'n {
@@ -82,7 +87,29 @@ function assign(x) = {
   };
   r
 }
-*/
+
+val assign2 : forall 'm, 'm in {8,16}. bits('m) -> bits(32)
+
+function assign2(x) = {
+  y : bits('m) = x;
+  r : bits(32) = 0x00000000;
+  match 'm {
+     8 => { y = y + 0x01; r = extzv(y) },
+    16 => r = extzv(y)
+  };
+  r
+}
+
+val assign3 : forall 'm, 'm in {8,16}. bits('m) -> bits('m)
+
+function assign3(x) = {
+  y : bits('m) = x;
+  match 'm {
+    8 => y = y + 0x01,
+   16 => y[7..0] = 0x89
+  };
+  y
+}
 
 /* Adding casts for top-level pattern matches */
 
@@ -123,10 +150,14 @@ function run () = {
     assert((ret(0x34) : bits(64)) == 0x0034003400340034);
     assert((ret(0x3456) : bits(32)) == 0x34563456);
     assert((ret(0x3456) : bits(64)) == 0x3456345634563456);
-/*    assert((assign(0x12) : bits(32)) == 0x00120012);
+    assert((assign(0x12) : bits(32)) == 0x00120012);
     assert((assign(0x1234) : bits(32)) == 0x12341234);
     assert((assign(0x12) : bits(64)) == 0x0012001200120012);
-    assert((assign(0x1234) : bits(64)) == 0x1234123412341234);*/
+    assert((assign(0x1234) : bits(64)) == 0x1234123412341234);
+    assert(assign2(0x12) == 0x00000013);
+    assert(assign2(0x1234) == 0x00001234);
+    assert(assign3(0x12) == 0x13);
+    assert(assign3(0x1234) == 0x1289);
     assert(foo2(32,0x12) == 0x00120012);
     assert(foo2(64,0x12) == 0x0012001200120012);
     assert(foo3(4,0x12) == 0x00120012);
diff --git a/test/mono/run_tests.sh b/test/mono/run_tests.sh
index 08926aaa..d2023229 100755
--- a/test/mono/run_tests.sh
+++ b/test/mono/run_tests.sh
@@ -2,7 +2,7 @@
 set -e
 
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
-SAILDIR="$DIR/../.."
+SAILDIR=${SAIL_DIR:-"$DIR/../.."}
 TESTSDIR="$DIR"
 OUTPUTDIR="$DIR/test-output"