Add support for WGSL's atomicCompareExchangeWeak with the `__atomic…

…_compare_exchange_result` struct, and add SPIR-V codegen for it.

Partially addresses #2113, #1755.

src/back/spv/block.rs

@@ -1954,6 +1954,19 @@ impl<'w> BlockContext<'w> {
                 } => {
                     let id = self.gen_id();
                     let result_type_id = self.get_expression_type_id(&self.fun_info[result].ty);
+                    // TODO: look this up from the atomic expression's scalar type so that it works with i32 as well
+                    let scalar_u32 = self.get_type_id(LookupType::Local(LocalType::Value {
+                        vector_size: None,
+                        kind: crate::ScalarKind::Uint,
+                        width: 4,
+                        pointer_space: None,
+                    }));
+                    let bool_type_id = self.get_type_id(LookupType::Local(LocalType::Value {
+                        vector_size: None,
+                        kind: crate::ScalarKind::Bool,
+                        width: crate::BOOL_WIDTH,
+                        pointer_space: None,
+                    }));
 
                     self.cached[result] = id;
 
@@ -2079,8 +2092,31 @@ impl<'w> BlockContext<'w> {
                                 value_id,
                             )
                         }
-                        crate::AtomicFunction::Exchange { compare: Some(_) } => {
-                            return Err(Error::FeatureNotImplemented("atomic CompareExchange"));
+                        crate::AtomicFunction::Exchange { compare: Some(cmp) } => {
+                            let cas_result_id = self.gen_id();
+                            let equality_result_id = self.gen_id();
+                            block.body.push(Instruction::atomic_ternary(
+                                spirv::Op::AtomicCompareExchange,
+                                scalar_u32,
+                                cas_result_id,
+                                pointer_id,
+                                scope_constant_id,
+                                semantics_id,
+                                value_id,
+                                self.cached[cmp],
+                            ));
+                            block.body.push(Instruction::binary(
+                                spirv::Op::IEqual,
+                                bool_type_id,
+                                equality_result_id,
+                                cas_result_id,
+                                self.cached[cmp],
+                            ));
+                            Instruction::composite_construct(
+                                result_type_id,
+                                id,
+                                &[cas_result_id, equality_result_id],
+                            )
                         }
                     };
 

src/back/spv/instructions.rs

@@ -793,6 +793,28 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn atomic_ternary(
+        op: Op,
+        result_type_id: Word,
+        id: Word,
+        pointer: Word,
+        scope_id: Word,
+        semantics_id: Word,
+        value1: Word,
+        value2: Word,
+    ) -> Self {
+        let mut instruction = Self::new(op);
+        instruction.set_type(result_type_id);
+        instruction.set_result(id);
+        instruction.add_operand(pointer);
+        instruction.add_operand(scope_id);
+        instruction.add_operand(semantics_id);
+        instruction.add_operand(semantics_id);
+        instruction.add_operand(value1);
+        instruction.add_operand(value2);
+        instruction
+    }
+
     //
     // Bit Instructions
     //

src/front/wgsl/mod.rs

@@ -1628,7 +1628,7 @@ impl Parser {
             crate::TypeInner::Scalar { kind, width } => crate::Expression::AtomicResult {
                 kind,
                 width,
-                comparison: false,
+                comparison: None,
             },
             _ => return Err(Error::InvalidAtomicOperandType(value_span)),
         };
@@ -1857,10 +1857,50 @@ impl Parser {
 
                     let expression = match *ctx.resolve_type(value)? {
                         crate::TypeInner::Scalar { kind, width } => {
+                            let bool_ty = ctx.types.insert(
+                                crate::Type {
+                                    name: None,
+                                    inner: crate::TypeInner::Scalar {
+                                        kind: crate::ScalarKind::Bool,
+                                        width: crate::BOOL_WIDTH,
+                                    },
+                                },
+                                NagaSpan::UNDEFINED,
+                            );
+                            let scalar_ty = ctx.types.insert(
+                                crate::Type {
+                                    name: None,
+                                    inner: crate::TypeInner::Scalar { kind, width },
+                                },
+                                NagaSpan::UNDEFINED,
+                            );
+                            let struct_ty = ctx.types.insert(
+                                crate::Type {
+                                    name: Some("__atomic_compare_exchange_result".to_string()),
+                                    inner: crate::TypeInner::Struct {
+                                        members: vec![
+                                            crate::StructMember {
+                                                name: Some("old_value".to_string()),
+                                                ty: scalar_ty,
+                                                binding: None,
+                                                offset: 0,
+                                            },
+                                            crate::StructMember {
+                                                name: Some("exchanged".to_string()),
+                                                ty: bool_ty,
+                                                binding: None,
+                                                offset: 4,
+                                            },
+                                        ],
+                                        span: 8,
+                                    },
+                                },
+                                NagaSpan::UNDEFINED,
+                            );
                             crate::Expression::AtomicResult {
                                 kind,
                                 width,
-                                comparison: true,
+                                comparison: Some(struct_ty),
                             }
                         }
                         _ => return Err(Error::InvalidAtomicOperandType(value_span)),

src/lib.rs

@@ -1404,7 +1404,7 @@ pub enum Expression {
     AtomicResult {
         kind: ScalarKind,
         width: Bytes,
-        comparison: bool,
+        comparison: Option<Handle<Type>>,
     },
     /// Get the length of an array.
     /// The expression must resolve to a pointer to an array with a dynamic size.

src/proc/typifier.rs

@@ -649,12 +649,8 @@ impl<'a> ResolveContext<'a> {
                 width,
                 comparison,
             } => {
-                if comparison {
-                    TypeResolution::Value(Ti::Vector {
-                        size: crate::VectorSize::Bi,
-                        kind,
-                        width,
-                    })
+                if let Some(struct_ty) = comparison {
+                    TypeResolution::Handle(struct_ty)
                 } else {
                     TypeResolution::Value(Ti::Scalar { kind, width })
                 }

src/valid/function.rs

@@ -363,11 +363,17 @@ impl super::Validator {
                 .into_other());
         }
         match context.expressions[result] {
-            //TODO: support atomic result with comparison
+            //TODO: does the result of an atomicCompareExchange need additional validation, or does the existing validation for
+            // the struct type it returns suffice?
             crate::Expression::AtomicResult {
                 kind,
                 width,
-                comparison: false,
+                comparison: Some(_),
+            } if kind == ptr_kind && width == ptr_width => {}
+            crate::Expression::AtomicResult {
+                kind,
+                width,
+                comparison: None,
             } if kind == ptr_kind && width == ptr_width => {}
             _ => {
                 return Err(AtomicError::ResultTypeMismatch(result)