Fix incorrect atomic bounds check on metal back-end (gfx-rs#2099)

* Fix incorrect atomic bounds check on metal back-end

Generalize put_atomic_fetch to handle `exchange` as well, rather than special-cased code which didn't do the bounds check (the check handling as fixed in gfx-rs#1703 but only for the fetch cases, exchange was skipped).

Fixes #1848

* Add tests for atomic exchange

src/back/msl/writer.rs

@@ -1187,6 +1187,17 @@ impl<W: Write> Writer<W> {
         key: &str,
         value: Handle<crate::Expression>,
         context: &ExpressionContext,
+    ) -> BackendResult {
+        self.put_atomic_operation(pointer, "fetch_", key, value, context)
+    }
+
+    fn put_atomic_operation(
+        &mut self,
+        pointer: Handle<crate::Expression>,
+        key1: &str,
+        key2: &str,
+        value: Handle<crate::Expression>,
+        context: &ExpressionContext,
     ) -> BackendResult {
         // If the pointer we're passing to the atomic operation needs to be conditional
         // for `ReadZeroSkipWrite`, the condition needs to *surround* the atomic op, and
@@ -1202,8 +1213,8 @@ impl<W: Write> Writer<W> {
 
         write!(
             self.out,
-            "{}::atomic_fetch_{}_explicit({}",
-            NAMESPACE, key, ATOMIC_REFERENCE
+            "{}::atomic_{}{}_explicit({}",
+            NAMESPACE, key1, key2, ATOMIC_REFERENCE
         )?;
         self.put_access_chain(pointer, policy, context)?;
         write!(self.out, ", ")?;
@@ -2725,15 +2736,13 @@ impl<W: Write> Writer<W> {
                             self.put_atomic_fetch(pointer, "max", value, &context.expression)?;
                         }
                         crate::AtomicFunction::Exchange { compare: None } => {
-                            write!(
-                                self.out,
-                                "{}::atomic_exchange_explicit({}",
-                                NAMESPACE, ATOMIC_REFERENCE,
+                            self.put_atomic_operation(
+                                pointer,
+                                "exchange",
+                                "",
+                                value,
+                                &context.expression,
                             )?;
-                            self.put_expression(pointer, &context.expression, true)?;
-                            write!(self.out, ", ")?;
-                            self.put_expression(value, &context.expression, true)?;
-                            write!(self.out, ", {}::memory_order_relaxed)", NAMESPACE)?;
                         }
                         crate::AtomicFunction::Exchange { .. } => {
                             return Err(Error::FeatureNotImplemented(

tests/in/bounds-check-zero-atomic.wgsl

@@ -23,3 +23,16 @@ fn fetch_add_atomic_static_sized_array(i: i32) -> u32 {
 fn fetch_add_atomic_dynamic_sized_array(i: i32) -> u32 {
    return atomicAdd(&globals.c[i], 1u);
 }
+
+fn exchange_atomic() -> u32 {
+   return atomicExchange(&globals.a, 1u);
+}
+
+fn exchange_atomic_static_sized_array(i: i32) -> u32 {
+   return atomicExchange(&globals.b[i], 1u);
+}
+
+fn exchange_atomic_dynamic_sized_array(i: i32) -> u32 {
+   return atomicExchange(&globals.c[i], 1u);
+}
+

tests/out/msl/bounds-check-zero-atomic.msl

@@ -49,3 +49,29 @@ uint fetch_add_atomic_dynamic_sized_array(
     uint _e5 = uint(i_1) < 1 + (_buffer_sizes.size0 - 44 - 4) / 4 ? metal::atomic_fetch_add_explicit(&globals.c[i_1], 1u, metal::memory_order_relaxed) : DefaultConstructible();
     return _e5;
 }
+
+uint exchange_atomic(
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    uint _e3 = metal::atomic_exchange_explicit(&globals.a, 1u, metal::memory_order_relaxed);
+    return _e3;
+}
+
+uint exchange_atomic_static_sized_array(
+    int i_2,
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    uint _e5 = uint(i_2) < 10 ? metal::atomic_exchange_explicit(&globals.b.inner[i_2], 1u, metal::memory_order_relaxed) : DefaultConstructible();
+    return _e5;
+}
+
+uint exchange_atomic_dynamic_sized_array(
+    int i_3,
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    uint _e5 = uint(i_3) < 1 + (_buffer_sizes.size0 - 44 - 4) / 4 ? metal::atomic_exchange_explicit(&globals.c[i_3], 1u, metal::memory_order_relaxed) : DefaultConstructible();
+    return _e5;
+}