Address comments

Author: RiverDave

naga/src/back/msl/keywords.rs

@@ -1,5 +1,8 @@
+use crate::common;
 use crate::proc::KeywordSet;
 use crate::racy_lock::RacyLock;
+use alloc::vec::Vec;
+use alloc::{boxed::Box, format};
 
 // MSLS - Metal Shading Language Specification:
 // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
@@ -346,7 +349,7 @@ pub const RESERVED: &[&str] = &[
     super::writer::MODF_FUNCTION,
     super::writer::ABS_FUNCTION,
     super::writer::DIV_FUNCTION,
-    super::writer::DOT_FUNCTION,
+    // DOT_FUNCTION_PREFIX variants are added dynamically below
     super::writer::MOD_FUNCTION,
     super::writer::NEG_FUNCTION,
     super::writer::F2I32_FUNCTION,
@@ -360,8 +363,65 @@ pub const RESERVED: &[&str] = &[
     super::writer::EXTERNAL_TEXTURE_WRAPPER_STRUCT,
 ];
 
+const CONCRETE_INTEGER_SCALARS: [crate::Scalar; 4] = [
+    crate::Scalar::I32,
+    crate::Scalar::U32,
+    crate::Scalar::I64,
+    crate::Scalar::U64,
+];
+
+const CONCRETE_VECTOR_SIZES: [crate::VectorSize; 3] = [
+    crate::VectorSize::Bi,
+    crate::VectorSize::Tri,
+    crate::VectorSize::Quad,
+];
+
 /// The above set of reserved keywords, turned into a cached HashSet. This saves
 /// significant time during [`Namer::reset`](crate::proc::Namer::reset).
 ///
 /// See <https://github.com/gfx-rs/wgpu/pull/7338> for benchmarks.
-pub static RESERVED_SET: RacyLock<KeywordSet> = RacyLock::new(|| KeywordSet::from_iter(RESERVED));
+pub static RESERVED_SET: RacyLock<KeywordSet> = RacyLock::new(|| {
+    let mut set = KeywordSet::from_iter(RESERVED);
+    // Add all concrete integer dot product function variants.
+    // These are generated to match the names produced by
+    // `Writer::get_dot_wrapper_function_helper_name`, ensuring they stay in sync.
+    let mut dot_function_names: Vec<&'static str> = Vec::new();
+    for scalar in CONCRETE_INTEGER_SCALARS {
+        // Map scalar to MSL type name (matching the logic in Scalar::to_msl_name)
+        let type_name = match scalar {
+            crate::Scalar {
+                kind: crate::ScalarKind::Sint,
+                width: 4,
+            } => "int",
+            crate::Scalar {
+                kind: crate::ScalarKind::Uint,
+                width: 4,
+            } => "uint",
+            crate::Scalar {
+                kind: crate::ScalarKind::Sint,
+                width: 8,
+            } => "long",
+            crate::Scalar {
+                kind: crate::ScalarKind::Uint,
+                width: 8,
+            } => "ulong",
+            _ => continue, // Skip non-integer or unsupported types
+        };
+
+        for size in CONCRETE_VECTOR_SIZES {
+            let size_suffix = common::vector_size_str(size);
+            let fun_name = format!(
+                "{}_{}{}",
+                super::writer::DOT_FUNCTION_PREFIX,
+                type_name,
+                size_suffix
+            );
+            // Convert to &'static str by leaking the String
+            // (In theory) This is safe because these are generated once and cached
+            let leaked = Box::leak(fun_name.into_boxed_str());
+            dot_function_names.push(leaked);
+        }
+    }
+    set.extend(dot_function_names.into_iter());
+    set
+});

naga/src/back/msl/writer.rs

@@ -55,7 +55,7 @@ pub(crate) const MODF_FUNCTION: &str = "naga_modf";
 pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
 pub(crate) const ABS_FUNCTION: &str = "naga_abs";
 pub(crate) const DIV_FUNCTION: &str = "naga_div";
-pub(crate) const DOT_FUNCTION: &str = "naga_dot";
+pub(crate) const DOT_FUNCTION_PREFIX: &str = "naga_dot";
 pub(crate) const MOD_FUNCTION: &str = "naga_mod";
 pub(crate) const NEG_FUNCTION: &str = "naga_neg";
 pub(crate) const F2I32_FUNCTION: &str = "naga_f2i32";
@@ -2338,12 +2338,14 @@ impl<W: Write> Writer<W> {
                                 },
                             ..
                         } => "dot",
-                        crate::TypeInner::Vector { size, scalar }
-                            if matches!(
-                                scalar.kind,
-                                crate::ScalarKind::Sint | crate::ScalarKind::Uint
-                            ) =>
-                        {
+                        crate::TypeInner::Vector {
+                            size,
+                            scalar:
+                                scalar @ crate::Scalar {
+                                    kind: crate::ScalarKind::Sint | crate::ScalarKind::Uint,
+                                    ..
+                                },
+                        } => {
                             // Integer vector dot: call our mangled helper `dot_{type}{N}(a, b)`.
                             let fun_name = self.get_dot_wrapper_function_helper_name(scalar, size);
                             write!(self.out, "{fun_name}(")?;
@@ -3368,6 +3370,15 @@ impl<W: Write> Writer<W> {
             } = *expr
             {
                 match fun {
+                    // WGSL's `dot` function works on any `vecN` type, but Metal's only
+                    // works on floating-point vectors, so we emit inline code for
+                    // integer vector `dot` calls. But that code uses each argument `N`
+                    // times, once for each component (see `put_dot_product`), so to
+                    // avoid duplicated evaluation, we must bake integer operands.
+                    // This applies both when using the polyfill (because of the duplicate
+                    // evaluation issue) and when we don't use the polyfill (because we
+                    // need them to be emitted before casting to packed chars -- see the
+                    // comment at the call to `put_casting_to_packed_chars`).
                     crate::MathFunction::Dot4U8Packed | crate::MathFunction::Dot4I8Packed => {
                         self.need_bake_expressions.insert(arg);
                         self.need_bake_expressions.insert(arg1.unwrap());
@@ -5785,15 +5796,15 @@ template <typename A>
     }
 
     /// Build the mangled helper name for integer vector dot products.
-    /// Result format: `{DOT_FUNCTION}_{type}{N}` (e.g., `naga_dot_int3`).
+    /// Result format: `{DOT_FUNCTION_PREFIX}_{type}{N}` (e.g., `naga_dot_int3`).
     fn get_dot_wrapper_function_helper_name(
         &self,
         scalar: crate::Scalar,
         size: crate::VectorSize,
     ) -> String {
         let type_name = scalar.to_msl_name();
         let size_suffix = common::vector_size_str(size);
-        format!("{DOT_FUNCTION}_{type_name}{size_suffix}")
+        format!("{DOT_FUNCTION_PREFIX}_{type_name}{size_suffix}")
     }
 
     #[allow(clippy::too_many_arguments)]