From ff595c7e2cf3476cb36bc1a0dd5f829186f6ef2e Mon Sep 17 00:00:00 2001 From: Jim Blandy Date: Tue, 4 Mar 2025 14:40:48 -0800 Subject: [PATCH] [naga wgsl-out] New `TypeContext`, for writing Naga types as WGSL. In `naga::back::wgsl`, introduce a new trait, `TypeContext`, which provides all the context necessary to generate WGSL code for a type. `Writer::write_type` and `write_type_inner` become default methods on `TypeContext`. `Writer` replaces them with functions that simply create a `TypeContext` from the `Writer` and forward the call. For simplicity's sake, rather than trying to return errors for Naga IR that we can't generate WGSL for, just panic. Validation should have rejected any such values, and it is not practical to expect Naga backends when given invalid modules: backends need to be free to assume that handles are valid, that arenas are free of cycles, and so on. --- naga/src/back/wgsl/writer.rs | 190 ++++++++++++++++++++++++----------- 1 file changed, 129 insertions(+), 61 deletions(-) diff --git a/naga/src/back/wgsl/writer.rs b/naga/src/back/wgsl/writer.rs index a255e1a336..df0b38f961 100644 --- a/naga/src/back/wgsl/writer.rs +++ b/naga/src/back/wgsl/writer.rs @@ -431,18 +431,43 @@ impl Writer { } } -impl Writer { +/// A context for printing Naga IR types as WGSL. +/// +/// This trait's default methods [`write_type`] and +/// [`write_type_inner`] do the work of formatting types as WGSL. +/// Implementors must provide the remaining methods, to customize +/// behavior for the context at hand. +/// +/// For example, the WGSL backend would provide an implementation of +/// [`type_name`] that handles hygienic renaming, whereas the WGSL +/// front end would simply show the name that was given in the source. +/// +/// [`write_type`]: TypeContext::write_type +/// [`write_type_inner`]: TypeContext::write_type_inner +/// [`type_name`]: TypeContext::type_name +trait TypeContext { + /// Return the [`Type`] referred to by `handle`. + /// + /// [`Type`]: crate::Type + fn lookup_type(&self, handle: Handle) -> &crate::Type; + + /// Return the name to be used for the type referred to by + /// `handle`. + fn type_name(&self, handle: Handle) -> &str; + + /// Write the WGSL form of `override` to `out`. + fn write_override(&self, r#override: Handle, out: &mut W) + -> core::fmt::Result; + /// Write the type `ty` as it would appear in a value's declaration. /// /// Write the type referred to by `ty` in `module` as it would appear in /// a `var`, `let`, etc. declaration, or in a function's argument list. - fn write_type(&mut self, module: &Module, ty: Handle) -> BackendResult { - let inner = &module.types[ty].inner; - match *inner { - TypeInner::Struct { .. } => { - write!(self.out, "{}", self.names[&NameKey::Type(ty)])?; - } - ref other => self.write_type_inner(module, other)?, + fn write_type(&self, handle: Handle, out: &mut W) -> core::fmt::Result { + let ty = self.lookup_type(handle); + match ty.inner { + TypeInner::Struct { .. } => out.write_str(self.type_name(handle))?, + ref other => self.write_type_inner(other, out)?, } Ok(()) @@ -458,19 +483,28 @@ impl Writer { /// [`TypeInner`]. /// /// [`Struct`]: TypeInner::Struct - fn write_type_inner(&mut self, module: &Module, inner: &TypeInner) -> BackendResult { + fn write_type_inner(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result { + fn unwrap_to_wgsl(value: T) -> &'static str { + value.try_to_wgsl().unwrap_or_else(|| { + unreachable!( + "validation should have forbidden {}: {value:?}", + T::DESCRIPTION + ); + }) + } + match *inner { TypeInner::Vector { size, scalar } => write!( - self.out, + out, "vec{}<{}>", common::vector_size_str(size), - scalar.to_wgsl_if_implemented()?, + unwrap_to_wgsl(scalar), )?, TypeInner::Sampler { comparison: false } => { - write!(self.out, "sampler")?; + write!(out, "sampler")?; } TypeInner::Sampler { comparison: true } => { - write!(self.out, "sampler_comparison")?; + write!(out, "sampler_comparison")?; } TypeInner::Image { dim, @@ -486,7 +520,7 @@ impl Writer { Ic::Sampled { kind, multi } => ( "", if multi { "multisampled_" } else { "" }, - crate::Scalar { kind, width: 4 }.to_wgsl_if_implemented()?, + unwrap_to_wgsl(crate::Scalar { kind, width: 4 }), "", ), Ic::Depth { multi } => { @@ -510,19 +544,19 @@ impl Writer { ), }; write!( - self.out, + out, "texture_{class_str}{multisampled_str}{dim_str}{arrayed_str}" )?; if !format_str.is_empty() { - write!(self.out, "<{format_str}{storage_str}>")?; + write!(out, "<{format_str}{storage_str}>")?; } } TypeInner::Scalar(scalar) => { - write!(self.out, "{}", scalar.to_wgsl_if_implemented()?)?; + write!(out, "{}", unwrap_to_wgsl(scalar))?; } TypeInner::Atomic(scalar) => { - write!(self.out, "atomic<{}>", scalar.to_wgsl_if_implemented()?)?; + write!(out, "atomic<{}>", unwrap_to_wgsl(scalar))?; } TypeInner::Array { base, @@ -532,37 +566,37 @@ impl Writer { // More info https://gpuweb.github.io/gpuweb/wgsl/#array-types // array -- Constant array // array -- Dynamic array - write!(self.out, "array<")?; + write!(out, "array<")?; match size { crate::ArraySize::Constant(len) => { - self.write_type(module, base)?; - write!(self.out, ", {len}")?; + self.write_type(base, out)?; + write!(out, ", {len}")?; } - crate::ArraySize::Pending(_) => { - unreachable!(); + crate::ArraySize::Pending(r#override) => { + self.write_override(r#override, out)?; } crate::ArraySize::Dynamic => { - self.write_type(module, base)?; + self.write_type(base, out)?; } } - write!(self.out, ">")?; + write!(out, ">")?; } TypeInner::BindingArray { base, size } => { // More info https://github.com/gpuweb/gpuweb/issues/2105 - write!(self.out, "binding_array<")?; + write!(out, "binding_array<")?; match size { crate::ArraySize::Constant(len) => { - self.write_type(module, base)?; - write!(self.out, ", {len}")?; + self.write_type(base, out)?; + write!(out, ", {len}")?; } - crate::ArraySize::Pending(_) => { - unreachable!(); + crate::ArraySize::Pending(r#override) => { + self.write_override(r#override, out)?; } crate::ArraySize::Dynamic => { - self.write_type(module, base)?; + self.write_type(base, out)?; } } - write!(self.out, ">")?; + write!(out, ">")?; } TypeInner::Matrix { columns, @@ -570,11 +604,11 @@ impl Writer { scalar, } => { write!( - self.out, + out, "mat{}x{}<{}>", common::vector_size_str(columns), common::vector_size_str(rows), - scalar.to_wgsl_if_implemented()? + unwrap_to_wgsl(scalar) )?; } TypeInner::Pointer { base, space } => { @@ -583,14 +617,14 @@ impl Writer { // Naga IR never produces pointers to handles, so it doesn't matter much // how we write such a type. Just write it as the base type alone. if let Some(space) = address { - write!(self.out, "ptr<{space}, ")?; + write!(out, "ptr<{space}, ")?; } - self.write_type(module, base)?; + self.write_type(base, out)?; if address.is_some() { if let Some(access) = maybe_access { - write!(self.out, ", {access}")?; + write!(out, ", {access}")?; } - write!(self.out, ">")?; + write!(out, ">")?; } } TypeInner::ValuePointer { @@ -600,20 +634,13 @@ impl Writer { } => { let (address, maybe_access) = address_space_str(space); if let Some(space) = address { - write!( - self.out, - "ptr<{}, {}", - space, - scalar.to_wgsl_if_implemented()? - )?; + write!(out, "ptr<{}, {}", space, unwrap_to_wgsl(scalar))?; if let Some(access) = maybe_access { - write!(self.out, ", {access}")?; + write!(out, ", {access}")?; } - write!(self.out, ">")?; + write!(out, ">")?; } else { - return Err(Error::Unimplemented(format!( - "ValuePointer to AddressSpace::Handle {inner:?}" - ))); + unreachable!("ValuePointer to AddressSpace::Handle {inner:?}"); } } TypeInner::ValuePointer { @@ -624,37 +651,78 @@ impl Writer { let (address, maybe_access) = address_space_str(space); if let Some(space) = address { write!( - self.out, + out, "ptr<{}, vec{}<{}>", space, common::vector_size_str(size), - scalar.to_wgsl_if_implemented()? + unwrap_to_wgsl(scalar) )?; if let Some(access) = maybe_access { - write!(self.out, ", {access}")?; + write!(out, ", {access}")?; } - write!(self.out, ">")?; + write!(out, ">")?; } else { - return Err(Error::Unimplemented(format!( - "ValuePointer to AddressSpace::Handle {inner:?}" - ))); + unreachable!("ValuePointer to AddressSpace::Handle {inner:?}"); } - write!(self.out, ">")?; + write!(out, ">")?; } TypeInner::AccelerationStructure { vertex_return } => { let caps = if vertex_return { "" } else { "" }; - write!(self.out, "acceleration_structure{}", caps)? - } - _ => { - return Err(Error::Unimplemented(format!("write_type_inner {inner:?}"))); + write!(out, "acceleration_structure{}", caps)? } + _ => unreachable!("invalid TypeInner"), } Ok(()) } } +struct WriterTypeContext<'m> { + module: &'m Module, + names: &'m crate::FastHashMap, +} + +impl TypeContext for WriterTypeContext<'_> { + fn lookup_type(&self, handle: Handle) -> &crate::Type { + &self.module.types[handle] + } + + fn type_name(&self, handle: Handle) -> &str { + self.names[&NameKey::Type(handle)].as_str() + } + + fn write_override(&self, _: Handle, _: &mut W) -> core::fmt::Result { + unreachable!("overrides should be validated out"); + } +} + impl Writer { + fn write_type(&mut self, module: &Module, ty: Handle) -> BackendResult { + // This actually can't be factored out into a nice constructor method, + // because the borrow checker needs to be able to see that the borrows + // of `self.names` and `self.out` are disjoint. + let type_context = WriterTypeContext { + module, + names: &self.names, + }; + type_context.write_type(ty, &mut self.out)?; + + Ok(()) + } + + fn write_type_inner(&mut self, module: &Module, inner: &TypeInner) -> BackendResult { + // This actually can't be factored out into a nice constructor method, + // because the borrow checker needs to be able to see that the borrows + // of `self.names` and `self.out` are disjoint. + let type_context = WriterTypeContext { + module, + names: &self.names, + }; + type_context.write_type_inner(inner, &mut self.out)?; + + Ok(()) + } + /// Helper method used to write statements /// /// # Notes