SMT: Support writing to register references

Add a new AE_write_ref constructor in the ANF representation to
make writes to register references explicit in Jib_compile

6 files changed, 50 insertions, 4 deletions

diff --git a/src/jib/anf.ml b/src/jib/anf.ml
index 77f220dd..787d5b17 100644
--- a/src/jib/anf.ml
+++ b/src/jib/anf.ml
@@ -68,6 +68,7 @@ and 'a aexp_aux =
   | AE_app of id * ('a aval) list * 'a
   | AE_cast of 'a aexp * 'a
   | AE_assign of id * 'a * 'a aexp
+  | AE_write_ref of id * 'a * 'a aexp
   | AE_let of mut * id * 'a * 'a aexp * 'a aexp * 'a
   | AE_block of ('a aexp) list * 'a aexp * 'a
   | AE_return of 'a aval * 'a
@@ -169,6 +170,7 @@ let aexp_typ (AE_aux (aux, _, _)) =
   | AE_app (_, _, typ) -> typ
   | AE_cast (_, typ) -> typ
   | AE_assign _ -> unit_typ
+  | AE_write_ref _ -> unit_typ
   | AE_let (_, _, _, _, _, typ) -> typ
   | AE_block (_, _, typ) -> typ
   | AE_return (_, typ) -> typ
@@ -202,6 +204,8 @@ let rec aexp_rename from_id to_id (AE_aux (aexp, env, l)) =
     | AE_cast (aexp, typ) -> AE_cast (recur aexp, typ)
     | AE_assign (id, typ, aexp) when Id.compare from_id id = 0 -> AE_assign (to_id, typ, aexp_rename from_id to_id aexp)
     | AE_assign (id, typ, aexp) -> AE_assign (id, typ, aexp_rename from_id to_id aexp)
+    | AE_write_ref (id, typ, aexp) when Id.compare from_id id = 0 -> AE_write_ref (to_id, typ, aexp_rename from_id to_id aexp)
+    | AE_write_ref (id, typ, aexp) -> AE_write_ref (id, typ, aexp_rename from_id to_id aexp)
     | AE_let (mut, id, typ1, aexp1, aexp2, typ2) when Id.compare from_id id = 0 -> AE_let (mut, id, typ1, recur aexp1, aexp2, typ2)
     | AE_let (mut, id, typ1, aexp1, aexp2, typ2) -> AE_let (mut, id, typ1, recur aexp1, recur aexp2, typ2)
     | AE_block (aexps, aexp, typ) -> AE_block (List.map recur aexps, recur aexp, typ)
@@ -238,6 +242,7 @@ let rec no_shadow ids (AE_aux (aexp, env, l)) =
     | AE_app (id, avals, typ) -> AE_app (id, avals, typ)
     | AE_cast (aexp, typ) -> AE_cast (no_shadow ids aexp, typ)
     | AE_assign (id, typ, aexp) -> AE_assign (id, typ, no_shadow ids aexp)
+    | AE_write_ref (id, typ, aexp) -> AE_write_ref (id, typ, no_shadow ids aexp)
     | AE_let (mut, id, typ1, aexp1, aexp2, typ2) when IdSet.mem id ids ->
        let shadow_id = new_shadow id in
        let aexp1 = no_shadow ids aexp1 in
@@ -282,6 +287,7 @@ let rec map_aval f (AE_aux (aexp, env, l)) =
     | AE_val v -> AE_val (f env l v)
     | AE_cast (aexp, typ) -> AE_cast (map_aval f aexp, typ)
     | AE_assign (id, typ, aexp) -> AE_assign (id, typ, map_aval f aexp)
+    | AE_write_ref (id, typ, aexp) -> AE_write_ref (id, typ, map_aval f aexp)
     | AE_app (id, vs, typ) -> AE_app (id, List.map (f env l) vs, typ)
     | AE_let (mut, id, typ1, aexp1, aexp2, typ2) ->
        AE_let (mut, id, typ1, map_aval f aexp1, map_aval f aexp2, typ2)
@@ -311,6 +317,7 @@ let rec map_functions f (AE_aux (aexp, env, l)) =
     | AE_app (id, vs, typ) -> f env l id vs typ
     | AE_cast (aexp, typ) -> AE_cast (map_functions f aexp, typ)
     | AE_assign (id, typ, aexp) -> AE_assign (id, typ, map_functions f aexp)
+    | AE_write_ref (id, typ, aexp) -> AE_write_ref (id, typ, map_functions f aexp)
     | AE_short_circuit (op, aval, aexp) -> AE_short_circuit (op, aval, map_functions f aexp)
     | AE_let (mut, id, typ1, aexp1, aexp2, typ2) -> AE_let (mut, id, typ1, map_functions f aexp1, map_functions f aexp2, typ2)
     | AE_block (aexps, aexp, typ) -> AE_block (List.map (map_functions f) aexps, map_functions f aexp, typ)
@@ -332,6 +339,7 @@ let rec fold_aexp f (AE_aux (aexp, env, l)) =
     | AE_app (id, vs, typ) -> AE_app (id, vs, typ)
     | AE_cast (aexp, typ) -> AE_cast (fold_aexp f aexp, typ)
     | AE_assign (id, typ, aexp) -> AE_assign (id, typ, fold_aexp f aexp)
+    | AE_write_ref (id, typ, aexp) -> AE_write_ref (id, typ, fold_aexp f aexp)
     | AE_short_circuit (op, aval, aexp) -> AE_short_circuit (op, aval, fold_aexp f aexp)
     | AE_let (mut, id, typ1, aexp1, aexp2, typ2) -> AE_let (mut, id, typ1, fold_aexp f aexp1, fold_aexp f aexp2, typ2)
     | AE_block (aexps, aexp, typ) -> AE_block (List.map (fold_aexp f) aexps, fold_aexp f aexp, typ)
@@ -377,6 +385,8 @@ let rec pp_aexp (AE_aux (aexp, _, _)) =
      pp_annot typ (string "$" ^^ pp_aexp aexp)
   | AE_assign (id, typ, aexp) ->
      pp_annot typ (pp_id id) ^^ string " := " ^^ pp_aexp aexp
+  | AE_write_ref (id, typ, aexp) ->
+     string "*" ^^ parens (pp_annot typ (pp_id id)) ^^ string " := " ^^ pp_aexp aexp
   | AE_app (id, args, typ) ->
      pp_annot typ (pp_id id ^^ parens (separate_map (comma ^^ space) pp_aval args))
   | AE_short_circuit (SC_or, aval, aexp) ->
@@ -538,7 +548,7 @@ let rec anf (E_aux (e_aux, ((l, _) as exp_annot)) as exp) =
       | AE_block (_, _, typ) ->
        let id = gensym () in
        (AV_id (id, Local (Immutable, typ)), fun x -> mk_aexp (AE_let (Immutable, id, typ, aexp, x, typ_of exp)))
-    | AE_assign _ | AE_for _ | AE_loop _ ->
+    | AE_assign _ | AE_for _ | AE_loop _ | AE_write_ref _ ->
        let id = gensym () in
        (AV_id (id, Local (Immutable, unit_typ)), fun x -> mk_aexp (AE_let (Immutable, id, unit_typ, aexp, x, typ_of exp)))
   in
@@ -557,7 +567,7 @@ let rec anf (E_aux (e_aux, ((l, _) as exp_annot)) as exp) =
 
   | E_assign (LEXP_aux (LEXP_deref dexp, _), exp) ->
      let gs = gensym () in
-     mk_aexp (AE_let (Mutable, gs, typ_of dexp, anf dexp, mk_aexp (AE_assign (gs, typ_of dexp, anf exp)), unit_typ))
+     mk_aexp (AE_let (Mutable, gs, typ_of dexp, anf dexp, mk_aexp (AE_write_ref (gs, typ_of dexp, anf exp)), unit_typ))
 
   | E_assign (LEXP_aux (LEXP_id id, _), exp)
     | E_assign (LEXP_aux (LEXP_cast (_, id), _), exp) ->
diff --git a/src/jib/anf.mli b/src/jib/anf.mli
index 571546cb..3ce0b286 100644
--- a/src/jib/anf.mli
+++ b/src/jib/anf.mli
@@ -88,6 +88,7 @@ and 'a aexp_aux =
   | AE_app of id * ('a aval) list * 'a
   | AE_cast of 'a aexp * 'a
   | AE_assign of id * 'a * 'a aexp
+  | AE_write_ref of id * 'a * 'a aexp
   | AE_let of mut * id * 'a * 'a aexp * 'a aexp * 'a
   | AE_block of ('a aexp) list * 'a aexp * 'a
   | AE_return of 'a aval * 'a
diff --git a/src/jib/c_backend.ml b/src/jib/c_backend.ml
index 41684b73..1b9fa1dd 100644
--- a/src/jib/c_backend.ml
+++ b/src/jib/c_backend.ml
@@ -351,6 +351,8 @@ let rec analyze_functions ctx f (AE_aux (aexp, env, l)) =
 
     | AE_assign (id, typ, aexp) -> AE_assign (id, typ, analyze_functions ctx f aexp)
 
+    | AE_write_ref (id, typ, aexp) -> AE_write_ref (id, typ, analyze_functions ctx f aexp)
+
     | AE_short_circuit (op, aval, aexp) -> AE_short_circuit (op, aval, analyze_functions ctx f aexp)
 
     | AE_let (mut, id, typ1, aexp1, (AE_aux (_, env2, _) as aexp2), typ2) ->
diff --git a/src/jib/jib_compile.ml b/src/jib/jib_compile.ml
index a8a0c640..fbc97f85 100644
--- a/src/jib/jib_compile.ml
+++ b/src/jib/jib_compile.ml
@@ -733,6 +733,15 @@ let rec compile_aexp ctx (AE_aux (aexp_aux, env, l)) =
      let setup, call, cleanup = compile_aexp ctx aexp in
      setup @ [call (CL_id (name id, assign_ctyp))], (fun clexp -> icopy l clexp unit_cval), cleanup
 
+  | AE_write_ref (id, assign_typ, aexp) ->
+     let assign_ctyp =
+       match Bindings.find_opt id ctx.locals with
+       | Some (_, ctyp) -> ctyp
+       | None -> ctyp_of_typ ctx assign_typ
+     in
+     let setup, call, cleanup = compile_aexp ctx aexp in
+     setup @ [call (CL_addr (CL_id (name id, assign_ctyp)))], (fun clexp -> icopy l clexp unit_cval), cleanup
+
   | AE_block (aexps, aexp, _) ->
      let block = compile_block ctx aexps in
      let setup, call, cleanup = compile_aexp ctx aexp in
diff --git a/src/jib/jib_smt.ml b/src/jib/jib_smt.ml
index eec8ca1e..8e9d0105 100644
--- a/src/jib/jib_smt.ml
+++ b/src/jib/jib_smt.ml
@@ -1262,6 +1262,9 @@ let smt_instr ctx =
        let smt_args = List.map (smt_cval ctx) args in
        [define_const ctx id ret_ctyp (Fn (zencode_id function_id, smt_args))]
 
+  | I_aux (I_copy (CL_addr (CL_id (_, _)), _), (_, l)) ->
+     Reporting.unreachable l __POS__ "Register reference write should be re-written by now"
+
   | I_aux (I_init (ctyp, id, cval), _) | I_aux (I_copy (CL_id (id, ctyp), cval), _) ->
      [define_const ctx id ctyp
         (smt_conversion ctx (cval_ctyp cval) ctyp (smt_cval ctx cval))]
@@ -1404,7 +1407,8 @@ let smt_header ctx stack cdefs =
 (* For generating SMT when we have a reg_deref(r : register(t))
    function, we have to expand it into a if-then-else cascade that
    checks if r is any one of the registers with type t, and reads that
-   register if it is. *)
+   register if it is. We also do a similar thing for *r = x
+*)
 let expand_reg_deref ctx = function
   | I_aux (I_funcall (clexp, false, function_id, [reg_ref]), (_, l)) as instr ->
      let open Type_check in
@@ -1432,6 +1436,26 @@ let expand_reg_deref ctx = function
         end
      | _ -> instr
      end
+  | I_aux (I_copy (CL_addr (CL_id (id, ctyp)), cval), (_, l)) ->
+     begin match ctyp with
+     | CT_ref reg_ctyp ->
+        begin match CTMap.find_opt reg_ctyp ctx.register_map with
+        | Some regs ->
+           let end_label = label "end_reg_write_" in
+           let try_reg r =
+             let next_label = label "next_reg_write_" in
+             [ijump (V_op (V_ref (name r, reg_ctyp), "!=", V_id (id, ctyp))) next_label;
+              icopy l (CL_id (name r, reg_ctyp)) cval;
+              igoto end_label;
+              ilabel next_label]
+           in
+           iblock (List.concat (List.map try_reg regs) @ [ilabel end_label])
+        | None ->
+           raise (Reporting.err_general l ("Could not find any registers with type " ^ string_of_ctyp reg_ctyp))
+        end
+     | _ ->
+        raise (Reporting.err_general l "Register reference assignment must take a register reference as an argument")
+     end
   | instr -> instr
 
 let smt_cdef props lets name_file ctx all_cdefs = function
diff --git a/src/jib/jib_ssa.ml b/src/jib/jib_ssa.ml
index 92cfe03b..1e008362 100644
--- a/src/jib/jib_ssa.ml
+++ b/src/jib/jib_ssa.ml
@@ -527,7 +527,7 @@ let rename_variables graph root children =
        CL_id (ssa_name i id, ctyp)
     | CL_rmw _ -> assert false
     | CL_field (clexp, field) -> CL_field (fold_clexp true clexp, field)
-    | CL_addr clexp -> CL_addr (fold_clexp true clexp)
+    | CL_addr clexp -> CL_addr (fold_clexp false clexp)
     | CL_tuple (clexp, n) -> CL_tuple (fold_clexp true clexp, n)
     | CL_void -> CL_void
   in