Merge pull request #573 from hacspec/un-monomorphize-updateat

fix(engine): monomorphize update_at only for vecs, slices and arrays
cryspen · Mar 14, 2024 · d668de4 · d668de4
2 parents e4f3580 + 4751624
commit d668de4
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 18 deletions.
diff --git a/engine/lib/phases/phase_trivialize_assign_lhs.ml b/engine/lib/phases/phase_trivialize_assign_lhs.ml
@@ -53,25 +53,33 @@ module%inlined_contents Make (F : Features.T) = struct
       | LhsArrayAccessor { e; typ = _; index; _ } ->
           let lhs = UA.expr_of_lhs span e |> dexpr in
           let update_at : Concrete_ident.name =
-            let index_typ =
-              match index.typ with TRef { typ; _ } -> typ | _ -> index.typ
+            let is_array_slice_or_vec =
+              match lhs.typ with
+              | TSlice _ | TArray _ -> true
+              | TApp { ident; _ } -> Global_ident.eq_name Alloc__vec__Vec ident
+              | _ -> false
             in
-            match index_typ with
-            | TInt { size = SSize; signedness = Unsigned } ->
-                Rust_primitives__hax__monomorphized_update_at__update_at_usize
-            | TApp { ident; _ }
-              when Global_ident.eq_name Core__ops__range__Range ident ->
-                Rust_primitives__hax__monomorphized_update_at__update_at_range
-            | TApp { ident; _ }
-              when Global_ident.eq_name Core__ops__range__RangeFrom ident ->
-                Rust_primitives__hax__monomorphized_update_at__update_at_range_from
-            | TApp { ident; _ }
-              when Global_ident.eq_name Core__ops__range__RangeTo ident ->
-                Rust_primitives__hax__monomorphized_update_at__update_at_range_to
-            | TApp { ident; _ }
-              when Global_ident.eq_name Core__ops__range__RangeFull ident ->
-                Rust_primitives__hax__monomorphized_update_at__update_at_range_full
-            | _ -> Rust_primitives__hax__update_at
+            if is_array_slice_or_vec then
+              let index_typ =
+                match index.typ with TRef { typ; _ } -> typ | _ -> index.typ
+              in
+              match index_typ with
+              | TInt { size = SSize; signedness = Unsigned } ->
+                  Rust_primitives__hax__monomorphized_update_at__update_at_usize
+              | TApp { ident; _ }
+                when Global_ident.eq_name Core__ops__range__Range ident ->
+                  Rust_primitives__hax__monomorphized_update_at__update_at_range
+              | TApp { ident; _ }
+                when Global_ident.eq_name Core__ops__range__RangeFrom ident ->
+                  Rust_primitives__hax__monomorphized_update_at__update_at_range_from
+              | TApp { ident; _ }
+                when Global_ident.eq_name Core__ops__range__RangeTo ident ->
+                  Rust_primitives__hax__monomorphized_update_at__update_at_range_to
+              | TApp { ident; _ }
+                when Global_ident.eq_name Core__ops__range__RangeFull ident ->
+                  Rust_primitives__hax__monomorphized_update_at__update_at_range_full
+              | _ -> Rust_primitives__hax__update_at
+            else Rust_primitives__hax__update_at
           in
           let rhs = UB.call update_at [ lhs; dexpr index; rhs ] span lhs.typ in
           updater_of_lhs e rhs span

diff --git a/test-harness/src/snapshots/toolchain__attributes into-fstar.snap b/test-harness/src/snapshots/toolchain__attributes into-fstar.snap
@@ -90,6 +90,24 @@ let v_MAX: usize = sz 10
 
 type t_MyArray = | MyArray : t_Array u8 (sz 10) -> t_MyArray
 
+let mutation_example
+      (use_generic_update_at: t_MyArray)
+      (use_specialized_update_at: t_Slice u8)
+      (specialized_as_well: Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global)
+    : (t_MyArray & t_Slice u8 & Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global) =
+  let use_generic_update_at:t_MyArray =
+    Rust_primitives.Hax.update_at use_generic_update_at (sz 2) 0uy
+  in
+  let use_specialized_update_at:t_Slice u8 =
+    Rust_primitives.Hax.Monomorphized_update_at.update_at_usize use_specialized_update_at (sz 2) 0uy
+  in
+  let specialized_as_well:Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global =
+    Rust_primitives.Hax.Monomorphized_update_at.update_at_usize specialized_as_well (sz 2) 0uy
+  in
+  use_generic_update_at, use_specialized_update_at, specialized_as_well
+  <:
+  (t_MyArray & t_Slice u8 & Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global)
+
 [@@ FStar.Tactics.Typeclasses.tcinstance]
 let impl: Core.Ops.Index.t_Index t_MyArray usize =
   {

diff --git a/tests/Cargo.lock b/tests/Cargo.lock
diff --git a/tests/attributes/src/lib.rs b/tests/attributes/src/lib.rs
@@ -77,6 +77,23 @@ mod refined_indexes {
             &self[index]
         }
     }
+
+    #[hax::exclude]
+    impl std::ops::IndexMut<usize> for MyArray {
+        fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+            &mut self[index]
+        }
+    }
+
+    fn mutation_example(
+        use_generic_update_at: &mut MyArray,
+        use_specialized_update_at: &mut [u8],
+        specialized_as_well: &mut Vec<u8>,
+    ) {
+        use_generic_update_at[2] = 0;
+        use_specialized_update_at[2] = 0;
+        specialized_as_well[2] = 0;
+    }
 }
 mod newtype_pattern {
     use hax_lib_macros as hax;