stack alloc: relax the checker

Before this patch, an assignment a = b, where a and b are two arrays, was allowed only if b was fully valid, and a was marked as fully valid too. After this patch, b does not need to be fully valid, and a is marked as having the same validity as b.
jasmin-lang · Jun 19, 2024 · cc0fdc2 · cc0fdc2
1 parent ee7b3bd
commit cc0fdc2
Show file tree

Hide file tree

Showing 6 changed files with 345 additions and 196 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,19 @@
 
 ## New features
 
+- The stack allocation checker accepts more programs. This checker is run
+  during the stack allocation pass to verify that the transformation done
+  by the pass (mostly, turning arrays accesses into memory accesses)
+  is correct. It ensures that there is no aliasing problem when two arrays
+  are put in the same place in memory. Before, an assignment `a1 = a2`, where
+  `a1` and `a2` are two arrays, was accepted only if there was no aliasing
+  issue on `a2`, and `a1` was marked as having no aliasing issue.
+  Now, there is no such requirement on `a2`, and `a1` is marked as having
+  the same aliasing issues as `a2`.
+  This gives in particular more freedom for spilling and unspilling reg ptr,
+  see `compiler/tests/success/subarrays/x86-64/spill_partial.jazz`.
+  ([PR #841](https://github.com/jasmin-lang/jasmin/pull/841)).
+
 - ARM now compiles `x = imm;` smartly: for small immediates, a single `MOV`; for
   immediates whose negation is small, a single `MVN`; and for large immediates
   a pair of `MOV` and `MOVT`.

diff --git a/compiler/tests/success/subarrays/x86-64/copy_partial.jazz b/compiler/tests/success/subarrays/x86-64/copy_partial.jazz
@@ -0,0 +1,22 @@
+export fn test () -> reg u64 {
+  stack u64[2] s;
+  reg ptr u64[2] r1 r2;
+
+  s[0] = 0;
+  s[1] = 1;
+
+  r1 = s;
+  r1[0] = 2;
+
+  r2 = s; // we copy s while it is not fully valid
+  r2[1] = 3;
+
+  s[0:1] = r1[0:1];
+  s[1:1] = r2[1:1];
+
+  reg u64 res;
+  res = s[0];
+  res += s[1];
+
+  return res;
+}
diff --git a/compiler/tests/success/subarrays/x86-64/spill_partial.jazz b/compiler/tests/success/subarrays/x86-64/spill_partial.jazz
@@ -0,0 +1,19 @@
+fn f (reg ptr u64[2] r) -> reg ptr u64[2] {
+  r[0] = 1;
+  return r;
+}
+
+export fn test (reg ptr u64[4] r) -> reg ptr u64[4], reg u64 {
+  reg ptr u64[2] r2;
+
+  r2 = r[1:2];
+  () = #spill(r);
+  r2 = f(r2);
+  () = #unspill(r); // we unspill r, while it is not fully valid
+  r[1:2] = r2;
+
+  reg u64 res;
+  res = r[1];
+
+  return r, res;
+}
diff --git a/compiler/tests/success/subarrays/x86-64/swap_partial.jazz b/compiler/tests/success/subarrays/x86-64/swap_partial.jazz
@@ -0,0 +1,25 @@
+export fn test () -> reg u64 {
+  stack u64[2] s;
+  reg ptr u64[2] r1 r2;
+
+  s[0] = 0;
+  s[1] = 1;
+
+  r1 = s;
+  r1[0] = 2;
+
+  r2 = s; // we copy s while it is not fully valid
+  r2[1] = 3;
+
+  // we swap two arrays that are both not full valid
+  (r1, r2) = #swap(r1, r2);
+
+  s[0:1] = r2[0:1];
+  s[1:1] = r1[1:1];
+
+  reg u64 res;
+  res = 0;
+  res += s[1];
+
+  return res;
+}
diff --git a/proofs/compiler/stack_alloc.v b/proofs/compiler/stack_alloc.v
@@ -236,16 +236,20 @@ Definition sub_region_at_ofs sr ofs len :=
      sr_zone   := sub_zone_at_ofs sr.(sr_zone) ofs len
   |}.
 
-Definition check_valid (rmap:region_map) (x:var_i) ofs len :=
+Definition get_sub_region_bytes (rmap:region_map) (x:var_i) ofs len :=
   (* we get the bytes associated to variable [x] *)
   Let sr := get_sub_region rmap x in
   let bytes := get_var_bytes rmap sr.(sr_region) x in
   let sr' := sub_region_at_ofs sr ofs len in
   let isub_ofs := interval_of_zone sr'.(sr_zone) in
-  (* we check if [isub_ofs] is a subset of one of the intervals of [bytes] *)
-  Let _   := assert (ByteSet.mem bytes isub_ofs)
-                    (stk_error x (pp_box [:: pp_s "the region associated to variable"; pp_var x; pp_s "is partial"])) in
-  ok (sr, sr').
+  ok (sr, sr', ByteSet.inter bytes (ByteSet.full isub_ofs)).
+
+Definition check_valid x sr bytes :=
+  Let _ :=
+    assert (ByteSet.mem bytes (interval_of_zone sr.(sr_zone)))
+           (stk_error x (pp_box [:: pp_s "the region associated to variable"; pp_var x; pp_s "is partial"]))
+  in
+  ok tt.
 
 Definition clear_bytes i bytes := ByteSet.remove bytes i.
 (* TODO: check optim
@@ -329,18 +333,20 @@ Definition set_arr_word (rmap:region_map) al (x:var_i) ofs ws :=
 
 Definition set_arr_call rmap x sr := set_sub_region rmap x sr (Some 0)%Z (size_slot x).
 
-Definition set_move_bytes rv x sr :=
+Definition set_move_bytes rv x sr bytesy :=
   let bm := get_bytes_map sr.(sr_region) rv in
   let bytes := get_bytes x bm in
-  let bm := Mvar.set bm x (ByteSet.add (interval_of_zone sr.(sr_zone)) bytes) in
+  let bytes := ByteSet.remove bytes (interval_of_zone sr.(sr_zone)) in
+  let bytes := ByteSet.union bytes bytesy in
+  let bm := Mvar.set bm x bytes in
   Mr.set rv sr.(sr_region) bm.
 
-Definition set_move_sub (rmap:region_map) x sr :=
-  let rv := set_move_bytes rmap x sr in
+Definition set_move_sub (rmap:region_map) x sr bytes :=
+  let rv := set_move_bytes rmap x sr bytes in
   {| var_region := rmap.(var_region);
      region_var := rv |}.
 
-Definition set_arr_sub (rmap:region_map) (x:var_i) ofs len sr_from :=
+Definition set_arr_sub (rmap:region_map) (x:var_i) ofs len sr_from bytesy :=
   Let sr := get_sub_region rmap x in
   let sr' := sub_region_at_ofs sr (Some ofs) len in
   Let _ := assert (sr' == sr_from)
@@ -349,13 +355,13 @@ Definition set_arr_sub (rmap:region_map) (x:var_i) ofs len sr_from :=
                       pp_s "the assignment to sub-array"; pp_var x;
                       pp_s "cannot be turned into a nop: source and destination regions are not equal"]))
   in
-  ok (set_move_sub rmap x sr').
+  ok (set_move_sub rmap x sr' bytesy).
 
 (* identical to [set_sub_region], except clearing
    TODO: fusion with set_arr_sub ? not sure its worth
 *)
-Definition set_move (rmap:region_map) (x:var) sr :=
-  let rv := set_move_bytes rmap x sr in
+Definition set_move (rmap:region_map) (x:var) sr bytesy :=
+  let rv := set_move_bytes rmap x sr bytesy in
   {| var_region := Mvar.set rmap.(var_region) x sr;
      region_var := rv |}.
 
@@ -565,26 +571,20 @@ Definition check_vpk rmap (x:var_i) vpk ofs len :=
   match vpk with
   | VKglob (_, ws) =>
     let sr := sub_region_glob x ws in
-    ok (sr, sub_region_at_ofs sr ofs len)
+    let sr' := sub_region_at_ofs sr ofs len in
+    (* useless inter, but allows to state a uniform lemma *)
+    let bytes := ByteSet.inter (ByteSet.full (interval_of_zone sr.(sr_zone)))
+                               (ByteSet.full (interval_of_zone sr'.(sr_zone)))
+    in
+    ok (sr, sr', bytes)
   | VKptr _pk =>
-    check_valid rmap x ofs len
+    get_sub_region_bytes rmap x ofs len
   end.
 
-(* We could write [check_vpk] as follows.
-  Definition check_vpk' rmap (x : gvar) ofs len :=
-    let (sr, bytes) := check_gvalid rmap x in
-    let sr' := sub_region_at_ofs sr.(sr_zone) ofs len in
-    let isub_ofs := interval_of_zone sr'.(sr_zone) in
-    (* we check if [isub_ofs] is a subset of one of the intervals of [bytes] *)
-    (* useless test when [x] is glob, but factorizes call to [sub_region_at_ofs] *)
-    Let _   := assert (ByteSet.mem bytes isub_ofs)
-                      (Cerr_stk_alloc "check_valid: the region is partial") in
-    ok sr'.
-*)
-
 Definition check_vpk_word rmap al x vpk ofs ws :=
-  Let srs := check_vpk rmap x vpk ofs (wsize_size ws) in
-  check_align al x srs.1 ws.
+  Let: (sr, sr', bytes) := check_vpk rmap x vpk ofs (wsize_size ws) in
+  Let _ := check_valid x sr' bytes in
+  check_align al x sr ws.
 
 Fixpoint alloc_e (e:pexpr) :=
   match e with
@@ -715,11 +715,11 @@ Definition is_nop is_spilling rmap (x:var) (sry:sub_region) : bool :=
   else false.
 
 (* TODO: better error message *)
-Definition get_addr is_spilling rmap x dx tag sry vpk y ofs :=
+Definition get_addr is_spilling rmap x dx tag sry bytesy vpk y ofs :=
   let ir := if is_nop is_spilling rmap x sry
             then Some nop
             else sap_mov_ofs saparams dx tag vpk y ofs in
-  let rmap := Region.set_move rmap x sry in
+  let rmap := Region.set_move rmap x sry bytesy in
   (rmap, ir).
 
 Definition get_ofs_sub aa ws x e1 :=
@@ -791,13 +791,12 @@ Definition alloc_array_move rmap r tag e :=
     match vk with
     | None => Error (stk_ierror_basic vy "register array remains")
     | Some vpk =>
-      Let srs := check_vpk rmap vy vpk (Some ofs) len in
-      let sry := srs.2 in
+      Let: (_, sry, bytesy) := check_vpk rmap vy vpk (Some ofs) len in
       Let eofs := mk_addr_pexpr rmap vy vpk in
-      ok (sry, vpk, eofs.1, (eofs.2 + ofs)%Z)
+      ok (sry, vpk, bytesy, eofs.1, (eofs.2 + ofs)%Z)
     end
   in
-  let '(sry, vpk, ey, ofs) := sryl in
+  let '(sry, vpk, bytesy, ey, ofs) := sryl in
   match subx with
   | None =>
     match get_local (v_var x) with
@@ -813,11 +812,11 @@ Definition alloc_array_move rmap r tag e :=
                       pp_s "the assignment to array"; pp_var x;
                       pp_s "cannot be turned into a nop: source and destination regions are not equal"]))
         in
-        let rmap := Region.set_move rmap x sry in
+        let rmap := Region.set_move rmap x sry bytesy in
         ok (rmap, nop)
       | Pregptr p =>
         let (rmap, oir) :=
-            get_addr None rmap x (Lvar (with_var x p)) tag sry vpk ey ofs in
+            get_addr None rmap x (Lvar (with_var x p)) tag sry bytesy vpk ey ofs in
         match oir with
         | None =>
           let err_pp := pp_box [:: pp_s "cannot compute address"; pp_var x] in
@@ -829,7 +828,7 @@ Definition alloc_array_move rmap r tag e :=
         let is_spilling := Some (slot, ws, z, x') in
         let dx_ofs := cast_const (ofsx + z.(z_ofs)) in
         let dx := Lmem Aligned Uptr (with_var x pmap.(vrsp)) dx_ofs in
-        let (rmap, oir) := get_addr is_spilling rmap x dx tag sry vpk ey ofs in
+        let (rmap, oir) := get_addr is_spilling rmap x dx tag sry bytesy vpk ey ofs in
         match oir with
         | None =>
           let err_pp := pp_box [:: pp_s "cannot compute address"; pp_var x] in
@@ -843,7 +842,7 @@ Definition alloc_array_move rmap r tag e :=
     match get_local (v_var x) with
     | None   => Error (stk_ierror_basic x "register array remains")
     | Some _ =>
-      Let rmap := Region.set_arr_sub rmap x ofs len sry in
+      Let rmap := Region.set_arr_sub rmap x ofs len sry bytesy in
       ok (rmap, nop)
     end
   end.
@@ -880,13 +879,12 @@ Definition alloc_protect_ptr rmap ii r t e msf :=
                       (stk_error_no_var "argument of protect_ptr should be a reg ptr") in
       Let _ := assert (if r is Lvar _ then true else false)
                       (stk_error_no_var "destination of protect_ptr should be a reg ptr") in
-      Let srs := check_vpk rmap vy vpk (Some ofs) len in
-      let sry := srs.2 in
+      Let: (_, sry, bytesy) := check_vpk rmap vy vpk (Some ofs) len in
       Let: (e, _ofs) := mk_addr_pexpr rmap vy vpk in (* ofs is ensured to be 0 *)
-      ok (sry, vpk, e)
+      ok (sry, bytesy, vpk, e)
     end
   in
-  let '(sry, vpk, ey) := sryl in
+  let '(sry, bytesy, vpk, ey) := sryl in
   match subx with
   | None =>
     match get_local (v_var x) with
@@ -897,7 +895,7 @@ Definition alloc_protect_ptr rmap ii r t e msf :=
         let dx := Lvar (with_var x px) in
         Let msf := add_iinfo ii (alloc_e rmap msf) in
         Let ir := lower_protect_ptr_fail ii [::dx] t [:: ey; msf] in
-        let rmap := Region.set_move rmap x sry in
+        let rmap := Region.set_move rmap x sry bytesy in
         ok (rmap, ir)
       | _ => Error (stk_error_no_var "only reg ptr can receive the result of protect_ptr")
       end
@@ -932,7 +930,8 @@ Definition alloc_array_move_init rmap r tag e :=
         end
       end in
     let sr := sub_region_at_ofs sr (Some ofs) len in
-    let rmap := Region.set_move_sub rmap x sr in
+    let bytesy := ByteSet.full (interval_of_zone sr.(sr_zone)) in
+    let rmap := Region.set_move_sub rmap x sr bytesy in
     ok (rmap, nop )
   else alloc_array_move rmap r tag e.
 
@@ -1043,10 +1042,14 @@ Definition alloc_call_arg_aux rmap0 rmap (sao_param: option param_info) (e:pexpr
     ok (rmap, (None, Pvar x))
   | None, Some _ => Error (stk_ierror_basic xv "argument not a reg")
   | Some pi, Some (Pregptr p) =>
-    Let srs := Region.check_valid rmap0 xv (Some 0%Z) (size_slot xv) in
-    let sr := srs.1 in
-    Let rmap := if pi.(pp_writable) then set_clear rmap xv sr (Some 0%Z) (size_slot xv) else ok rmap in
+    Let: (sr, sr', bytes) := Region.get_sub_region_bytes rmap0 xv (Some 0%Z) (size_slot xv) in
+    Let _ := Region.check_valid xv sr' bytes in
     Let _  := check_align Aligned xv sr pi.(pp_align) in
+    Let rmap :=
+      if pi.(pp_writable) then
+        set_clear rmap xv sr (Some 0%Z) (size_slot xv)
+      else ok rmap
+    in
     ok (rmap, (Some (pi.(pp_writable),sr), Pvar (mk_lvar (with_var xv p))))
   | Some _, _ => Error (stk_ierror_basic xv "the argument should be a reg ptr")
   end.
@@ -1186,12 +1189,12 @@ Definition alloc_array_swap rmap rs t es :=
   | [:: Lvar x; Lvar y], [::Pvar z'; Pvar w'] =>
     let z := z'.(gv) in
     Let pz  := get_regptr z in
-    Let: (_, srz) := check_valid rmap z (Some 0%Z) (size_of z.(vtype)) in
+    Let: (_, srz, bytesz) := get_sub_region_bytes rmap z (Some 0%Z) (size_of z.(vtype)) in
     let w := w'.(gv) in
     Let pw := get_regptr w in
-    Let: (_, srw) := check_valid rmap w (Some 0%Z) (size_of w.(vtype)) in
-    let rmap := Region.set_move rmap x srw in
-    let rmap := Region.set_move rmap y srz in
+    Let: (_, srw, bytesw) := get_sub_region_bytes rmap w (Some 0%Z) (size_of w.(vtype)) in
+    let rmap := Region.set_move rmap x srw bytesw in
+    let rmap := Region.set_move rmap y srz bytesz in
     Let px := get_regptr x in
     Let py := get_regptr y in
     Let _ := assert ((is_lvar z') && (is_lvar w'))
@@ -1301,7 +1304,8 @@ Definition add_alloc globals stack (xpk:var * ptr_kind_init) (lrx: Mvar.t ptr_ki
             let rmap :=
               if sc is Slocal then
                 let sr := sub_region_stack x' ws' z in
-                Region.set_arr_init rmap x sr
+                let bytes := ByteSet.full (interval_of_zone sr.(sr_zone)) in
+                Region.set_arr_init rmap x sr bytes
               else
                 rmap
             in
@@ -1371,12 +1375,12 @@ Definition check_result pmap rmap paramsi params oi (x:var_i) :=
   match oi with
   | Some i =>
     match nth None paramsi i with
-    | Some sr =>
+    | Some sr_param =>
       Let _ := assert (x.(vtype) == (nth x params i).(vtype))
                       (stk_ierror_no_var "reg ptr in result not corresponding to a parameter") in
-      Let srs := check_valid rmap x (Some 0%Z) (size_slot x) in
-      let sr' := srs.1 in
-      Let _  := assert (sr == sr') (stk_ierror_no_var "invalid reg ptr in result") in
+      Let: (sr, sr', bytes) := get_sub_region_bytes rmap x (Some 0%Z) (size_slot x) in
+      Let _ := check_valid x sr' bytes in
+      Let _  := assert (sr_param == sr) (stk_ierror_no_var "invalid reg ptr in result") in
       Let p  := get_regptr pmap x in
       ok p
     | None => Error (stk_ierror_no_var "invalid function info")
@@ -1422,9 +1426,10 @@ Definition init_param (mglob stack : Mvar.t (Z * wsize)) accu pi (x:var_i) :=
       {| r_slot := x;
          r_align := pi.(pp_align); r_writable := pi.(pp_writable) |} in
     let sr := sub_region_full x r in
+    let bytes := ByteSet.full (interval_of_zone sr.(sr_zone)) in
     ok (Sv.add pi.(pp_ptr) disj,
         Mvar.set lmap x (Pregptr pi.(pp_ptr)),
-        set_move rmap x sr,
+        set_move rmap x sr bytes,
         (Some sr, with_var x pi.(pp_ptr)))
   end.