From 8ec2e751b4b75d6479fa8de955353c74be784ec0 Mon Sep 17 00:00:00 2001
From: Jim Blandy <jimb@red-bean.com>
Date: Tue, 4 Mar 2025 17:56:01 -0800
Subject: [PATCH] [naga] Implement `TryToWgsl` for `Scalar`, and use as
 appropriate.

Replace `naga::back::wgsl::writer::scalar_kind_str` with a `TryToWgsl`
implementation for `Scalar` in `common::wgsl`. Use this where needed
in the WGSL backend.
---
 naga/src/back/wgsl/writer.rs | 62 +++++++++---------------------------
 naga/src/common/wgsl.rs      | 19 +++++++++++
 2 files changed, 34 insertions(+), 47 deletions(-)

diff --git a/naga/src/back/wgsl/writer.rs b/naga/src/back/wgsl/writer.rs
index c56fcb5572..e2d99be13f 100644
--- a/naga/src/back/wgsl/writer.rs
+++ b/naga/src/back/wgsl/writer.rs
@@ -423,7 +423,7 @@ impl<W: Write> Writer<W> {
                 self.out,
                 "vec{}<{}>",
                 common::vector_size_str(size),
-                scalar_kind_str(scalar),
+                scalar.to_wgsl_if_implemented()?,
             )?,
             TypeInner::Sampler { comparison: false } => {
                 write!(self.out, "sampler")?;
@@ -445,7 +445,7 @@ impl<W: Write> Writer<W> {
                     Ic::Sampled { kind, multi } => (
                         "",
                         if multi { "multisampled_" } else { "" },
-                        scalar_kind_str(crate::Scalar { kind, width: 4 }),
+                        crate::Scalar { kind, width: 4 }.to_wgsl_if_implemented()?,
                         "",
                     ),
                     Ic::Depth { multi } => {
@@ -478,10 +478,10 @@ impl<W: Write> Writer<W> {
                 }
             }
             TypeInner::Scalar(scalar) => {
-                write!(self.out, "{}", scalar_kind_str(scalar))?;
+                write!(self.out, "{}", scalar.to_wgsl_if_implemented()?)?;
             }
             TypeInner::Atomic(scalar) => {
-                write!(self.out, "atomic<{}>", scalar_kind_str(scalar))?;
+                write!(self.out, "atomic<{}>", scalar.to_wgsl_if_implemented()?)?;
             }
             TypeInner::Array {
                 base,
@@ -533,7 +533,7 @@ impl<W: Write> Writer<W> {
                     "mat{}x{}<{}>",
                     common::vector_size_str(columns),
                     common::vector_size_str(rows),
-                    scalar_kind_str(scalar)
+                    scalar.to_wgsl_if_implemented()?
                 )?;
             }
             TypeInner::Pointer { base, space } => {
@@ -559,7 +559,12 @@ impl<W: Write> Writer<W> {
             } => {
                 let (address, maybe_access) = address_space_str(space);
                 if let Some(space) = address {
-                    write!(self.out, "ptr<{}, {}", space, scalar_kind_str(scalar))?;
+                    write!(
+                        self.out,
+                        "ptr<{}, {}",
+                        space,
+                        scalar.to_wgsl_if_implemented()?
+                    )?;
                     if let Some(access) = maybe_access {
                         write!(self.out, ", {access}")?;
                     }
@@ -582,7 +587,7 @@ impl<W: Write> Writer<W> {
                         "ptr<{}, vec{}<{}>",
                         space,
                         common::vector_size_str(size),
-                        scalar_kind_str(scalar)
+                        scalar.to_wgsl_if_implemented()?
                     )?;
                     if let Some(access) = maybe_access {
                         write!(self.out, ", {access}")?;
@@ -1590,7 +1595,7 @@ impl<W: Write> Writer<W> {
                             kind,
                             width: convert.unwrap_or(scalar.width),
                         };
-                        let scalar_kind_str = scalar_kind_str(scalar);
+                        let scalar_kind_str = scalar.to_wgsl_if_implemented()?;
                         write!(
                             self.out,
                             "mat{}x{}<{}>",
@@ -1608,7 +1613,7 @@ impl<W: Write> Writer<W> {
                             width: convert.unwrap_or(width),
                         };
                         let vector_size_str = common::vector_size_str(size);
-                        let scalar_kind_str = scalar_kind_str(scalar);
+                        let scalar_kind_str = scalar.to_wgsl_if_implemented()?;
                         if convert.is_some() {
                             write!(self.out, "vec{vector_size_str}<{scalar_kind_str}>")?;
                         } else {
@@ -1620,7 +1625,7 @@ impl<W: Write> Writer<W> {
                             kind,
                             width: convert.unwrap_or(width),
                         };
-                        let scalar_kind_str = scalar_kind_str(scalar);
+                        let scalar_kind_str = scalar.to_wgsl_if_implemented()?;
                         if convert.is_some() {
                             write!(self.out, "{scalar_kind_str}")?
                         } else {
@@ -1883,43 +1888,6 @@ const fn image_dimension_str(dim: crate::ImageDimension) -> &'static str {
     }
 }
 
-const fn scalar_kind_str(scalar: crate::Scalar) -> &'static str {
-    use crate::Scalar;
-    use crate::ScalarKind as Sk;
-
-    match scalar {
-        Scalar {
-            kind: Sk::Float,
-            width: 8,
-        } => "f64",
-        Scalar {
-            kind: Sk::Float,
-            width: 4,
-        } => "f32",
-        Scalar {
-            kind: Sk::Sint,
-            width: 4,
-        } => "i32",
-        Scalar {
-            kind: Sk::Uint,
-            width: 4,
-        } => "u32",
-        Scalar {
-            kind: Sk::Sint,
-            width: 8,
-        } => "i64",
-        Scalar {
-            kind: Sk::Uint,
-            width: 8,
-        } => "u64",
-        Scalar {
-            kind: Sk::Bool,
-            width: 1,
-        } => "bool",
-        _ => unreachable!(),
-    }
-}
-
 const fn address_space_str(
     space: crate::AddressSpace,
 ) -> (Option<&'static str>, Option<&'static str>) {
diff --git a/naga/src/common/wgsl.rs b/naga/src/common/wgsl.rs
index a0de0465d3..8b98e8d0c6 100644
--- a/naga/src/common/wgsl.rs
+++ b/naga/src/common/wgsl.rs
@@ -305,3 +305,22 @@ impl ToWgsl for crate::StorageFormat {
         }
     }
 }
+
+impl TryToWgsl for crate::Scalar {
+    const DESCRIPTION: &'static str = "scalar type";
+
+    fn try_to_wgsl(self) -> Option<&'static str> {
+        use crate::Scalar;
+
+        Some(match self {
+            Scalar::F64 => "f64",
+            Scalar::F32 => "f32",
+            Scalar::I32 => "i32",
+            Scalar::U32 => "u32",
+            Scalar::I64 => "i64",
+            Scalar::U64 => "u64",
+            Scalar::BOOL => "bool",
+            _ => return None,
+        })
+    }
+}