naga/common/wgsl/
types.rs

1//! Code for formatting Naga IR types as WGSL source code.
2
3use super::{address_space_str, ToWgsl, TryToWgsl};
4use crate::common;
5use crate::proc::TypeResolution;
6use crate::{Handle, Scalar, TypeInner};
7
8use alloc::string::String;
9use core::fmt::Write;
10
11/// A context for printing Naga IR types as WGSL.
12///
13/// This trait's default methods [`write_type`] and
14/// [`write_type_inner`] do the work of formatting types as WGSL.
15/// Implementors must provide the remaining methods, to customize
16/// behavior for the context at hand.
17///
18/// For example, the WGSL backend would provide an implementation of
19/// [`type_name`] that handles hygienic renaming, whereas the WGSL
20/// front end would simply show the name that was given in the source.
21///
22/// [`write_type`]: TypeContext::write_type
23/// [`write_type_inner`]: TypeContext::write_type_inner
24/// [`type_name`]: TypeContext::type_name
25pub trait TypeContext {
26    /// Return the [`Type`] referred to by `handle`.
27    ///
28    /// [`Type`]: crate::Type
29    fn lookup_type(&self, handle: Handle<crate::Type>) -> &crate::Type;
30
31    /// Return the name to be used for the type referred to by
32    /// `handle`.
33    fn type_name(&self, handle: Handle<crate::Type>) -> &str;
34
35    /// Write the WGSL form of `override` to `out`.
36    fn write_override<W: Write>(
37        &self,
38        r#override: Handle<crate::Override>,
39        out: &mut W,
40    ) -> core::fmt::Result;
41
42    /// Write a [`TypeInner::Struct`] for which we are unable to find a name.
43    ///
44    /// The names of struct types are only available if we have `Handle<Type>`,
45    /// not from [`TypeInner`]. For logging and debugging, it's fine to just
46    /// write something helpful to the developer, but for generating WGSL,
47    /// this should be unreachable.
48    fn write_unnamed_struct<W: Write>(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result;
49
50    /// Write a [`TypeInner`] that has no representation as WGSL source,
51    /// even including Naga extensions.
52    ///
53    /// A backend might implement this with a call to the [`unreachable!`]
54    /// macro, since backends are allowed to assume that the module has passed
55    /// validation.
56    ///
57    /// The default implementation is appropriate for generating type names to
58    /// appear in error messages. It punts to `TypeInner`'s [`core::fmt::Debug`]
59    /// implementation, since it's probably best to show the user something they
60    /// can act on.
61    fn write_non_wgsl_inner<W: Write>(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result {
62        write!(out, "{{non-WGSL Naga type {inner:?}}}")
63    }
64
65    /// Write a [`Scalar`] that has no representation as WGSL source,
66    /// even including Naga extensions.
67    ///
68    /// A backend might implement this with a call to the [`unreachable!`]
69    /// macro, since backends are allowed to assume that the module has passed
70    /// validation.
71    ///
72    /// The default implementation is appropriate for generating type names to
73    /// appear in error messages. It punts to `Scalar`'s [`core::fmt::Debug`]
74    /// implementation, since it's probably best to show the user something they
75    /// can act on.
76    fn write_non_wgsl_scalar<W: Write>(&self, scalar: Scalar, out: &mut W) -> core::fmt::Result {
77        match scalar.kind {
78            crate::ScalarKind::Sint
79            | crate::ScalarKind::Uint
80            | crate::ScalarKind::Float
81            | crate::ScalarKind::Bool => write!(out, "{{non-WGSL Naga scalar {scalar:?}}}"),
82
83            // The abstract types are kind of an odd quasi-WGSL category:
84            // they are definitely part of the spec, but they are not expressible
85            // in WGSL itself. So we want to call them out by name in error messages,
86            // but the WGSL backend should never generate these.
87            crate::ScalarKind::AbstractInt => out.write_str("{AbstractInt}"),
88            crate::ScalarKind::AbstractFloat => out.write_str("{AbstractFloat}"),
89        }
90    }
91
92    /// Write the type `ty` as it would appear in a value's declaration.
93    ///
94    /// Write the type referred to by `ty` in `module` as it would appear in
95    /// a `var`, `let`, etc. declaration, or in a function's argument list.
96    fn write_type<W: Write>(&self, handle: Handle<crate::Type>, out: &mut W) -> core::fmt::Result {
97        let ty = self.lookup_type(handle);
98        match ty.inner {
99            TypeInner::Struct { .. } => out.write_str(self.type_name(handle))?,
100            ref other => self.write_type_inner(other, out)?,
101        }
102
103        Ok(())
104    }
105
106    /// Write the [`TypeInner`] `inner` as it would appear in a value's declaration.
107    ///
108    /// Write `inner` as it would appear in a `var`, `let`, etc.
109    /// declaration, or in a function's argument list.
110    ///
111    /// Note that this cannot handle writing [`Struct`] types: those
112    /// must be referred to by name, but the name isn't available in
113    /// [`TypeInner`].
114    ///
115    /// [`Struct`]: TypeInner::Struct
116    fn write_type_inner<W: Write>(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result {
117        match try_write_type_inner(self, inner, out) {
118            Ok(()) => Ok(()),
119            Err(WriteTypeError::Format(err)) => Err(err),
120            Err(WriteTypeError::NonWgsl) => self.write_non_wgsl_inner(inner, out),
121        }
122    }
123
124    /// Write the [`Scalar`] `scalar` as a WGSL type.
125    fn write_scalar<W: Write>(&self, scalar: Scalar, out: &mut W) -> core::fmt::Result {
126        match scalar.try_to_wgsl() {
127            Some(string) => out.write_str(string),
128            None => self.write_non_wgsl_scalar(scalar, out),
129        }
130    }
131
132    /// Write the [`TypeResolution`] `resolution` as a WGSL type.
133    fn write_type_resolution<W: Write>(
134        &self,
135        resolution: &TypeResolution,
136        out: &mut W,
137    ) -> core::fmt::Result {
138        match *resolution {
139            TypeResolution::Handle(handle) => self.write_type(handle, out),
140            TypeResolution::Value(ref inner) => self.write_type_inner(inner, out),
141        }
142    }
143
144    fn write_type_conclusion<W: Write>(
145        &self,
146        conclusion: &crate::proc::Conclusion,
147        out: &mut W,
148    ) -> core::fmt::Result {
149        use crate::proc::Conclusion as Co;
150
151        match *conclusion {
152            Co::Value(ref inner) => self.write_type_inner(inner, out),
153            Co::Predeclared(ref predeclared) => out.write_str(&predeclared.struct_name()),
154        }
155    }
156
157    fn write_type_rule<W: Write>(
158        &self,
159        name: &str,
160        rule: &crate::proc::Rule,
161        out: &mut W,
162    ) -> core::fmt::Result {
163        write!(out, "fn {name}(")?;
164        for (i, arg) in rule.arguments.iter().enumerate() {
165            if i > 0 {
166                out.write_str(", ")?;
167            }
168            self.write_type_resolution(arg, out)?
169        }
170        out.write_str(") -> ")?;
171        self.write_type_conclusion(&rule.conclusion, out)?;
172        Ok(())
173    }
174
175    fn type_to_string(&self, handle: Handle<crate::Type>) -> String {
176        let mut buf = String::new();
177        self.write_type(handle, &mut buf).unwrap();
178        buf
179    }
180
181    fn type_resolution_to_string(&self, resolution: &TypeResolution) -> String {
182        let mut buf = String::new();
183        self.write_type_resolution(resolution, &mut buf).unwrap();
184        buf
185    }
186
187    fn type_rule_to_string(&self, name: &str, rule: &crate::proc::Rule) -> String {
188        let mut buf = String::new();
189        self.write_type_rule(name, rule, &mut buf).unwrap();
190        buf
191    }
192}
193
194fn try_write_type_inner<C, W>(ctx: &C, inner: &TypeInner, out: &mut W) -> Result<(), WriteTypeError>
195where
196    C: TypeContext + ?Sized,
197    W: Write,
198{
199    match *inner {
200        TypeInner::Vector { size, scalar } => {
201            write!(out, "vec{}<", common::vector_size_str(size))?;
202            ctx.write_scalar(scalar, out)?;
203            out.write_str(">")?;
204        }
205        TypeInner::Sampler { comparison: false } => {
206            write!(out, "sampler")?;
207        }
208        TypeInner::Sampler { comparison: true } => {
209            write!(out, "sampler_comparison")?;
210        }
211        TypeInner::Image {
212            dim,
213            arrayed,
214            class,
215        } => {
216            // More about texture types: https://gpuweb.github.io/gpuweb/wgsl/#sampled-texture-type
217            use crate::ImageClass as Ic;
218
219            let dim_str = dim.to_wgsl();
220            let arrayed_str = if arrayed { "_array" } else { "" };
221            match class {
222                Ic::Sampled { kind, multi } => {
223                    let multisampled_str = if multi { "multisampled_" } else { "" };
224                    write!(out, "texture_{multisampled_str}{dim_str}{arrayed_str}<")?;
225                    ctx.write_scalar(Scalar { kind, width: 4 }, out)?;
226                    out.write_str(">")?;
227                }
228                Ic::Depth { multi } => {
229                    let multisampled_str = if multi { "multisampled_" } else { "" };
230                    write!(
231                        out,
232                        "texture_depth_{multisampled_str}{dim_str}{arrayed_str}"
233                    )?;
234                }
235                Ic::Storage { format, access } => {
236                    let format_str = format.to_wgsl();
237                    let access_str = if access.contains(crate::StorageAccess::ATOMIC) {
238                        ",atomic"
239                    } else if access
240                        .contains(crate::StorageAccess::LOAD | crate::StorageAccess::STORE)
241                    {
242                        ",read_write"
243                    } else if access.contains(crate::StorageAccess::LOAD) {
244                        ",read"
245                    } else {
246                        ",write"
247                    };
248                    write!(
249                        out,
250                        "texture_storage_{dim_str}{arrayed_str}<{format_str}{access_str}>"
251                    )?;
252                }
253                Ic::External => {
254                    write!(out, "texture_external")?;
255                }
256            }
257        }
258        TypeInner::Scalar(scalar) => {
259            ctx.write_scalar(scalar, out)?;
260        }
261        TypeInner::Atomic(scalar) => {
262            out.write_str("atomic<")?;
263            ctx.write_scalar(scalar, out)?;
264            out.write_str(">")?;
265        }
266        TypeInner::Array {
267            base,
268            size,
269            stride: _,
270        } => {
271            // More info https://gpuweb.github.io/gpuweb/wgsl/#array-types
272            // array<A, 3> -- Constant array
273            // array<A> -- Dynamic array
274            write!(out, "array<")?;
275            match size {
276                crate::ArraySize::Constant(len) => {
277                    ctx.write_type(base, out)?;
278                    write!(out, ", {len}")?;
279                }
280                crate::ArraySize::Pending(r#override) => {
281                    ctx.write_override(r#override, out)?;
282                }
283                crate::ArraySize::Dynamic => {
284                    ctx.write_type(base, out)?;
285                }
286            }
287            write!(out, ">")?;
288        }
289        TypeInner::BindingArray { base, size } => {
290            // More info https://github.com/gpuweb/gpuweb/issues/2105
291            write!(out, "binding_array<")?;
292            match size {
293                crate::ArraySize::Constant(len) => {
294                    ctx.write_type(base, out)?;
295                    write!(out, ", {len}")?;
296                }
297                crate::ArraySize::Pending(r#override) => {
298                    ctx.write_override(r#override, out)?;
299                }
300                crate::ArraySize::Dynamic => {
301                    ctx.write_type(base, out)?;
302                }
303            }
304            write!(out, ">")?;
305        }
306        TypeInner::Matrix {
307            columns,
308            rows,
309            scalar,
310        } => {
311            write!(
312                out,
313                "mat{}x{}<",
314                common::vector_size_str(columns),
315                common::vector_size_str(rows),
316            )?;
317            ctx.write_scalar(scalar, out)?;
318            out.write_str(">")?;
319        }
320        TypeInner::Pointer { base, space } => {
321            let (address, maybe_access) = address_space_str(space);
322            // Everything but `AddressSpace::Handle` gives us a `address` name, but
323            // Naga IR never produces pointers to handles, so it doesn't matter much
324            // how we write such a type. Just write it as the base type alone.
325            if let Some(space) = address {
326                write!(out, "ptr<{space}, ")?;
327            }
328            ctx.write_type(base, out)?;
329            if address.is_some() {
330                if let Some(access) = maybe_access {
331                    write!(out, ", {access}")?;
332                }
333                write!(out, ">")?;
334            }
335        }
336        TypeInner::ValuePointer {
337            size: None,
338            scalar,
339            space,
340        } => {
341            let (address, maybe_access) = address_space_str(space);
342            if let Some(space) = address {
343                write!(out, "ptr<{space}, ")?;
344                ctx.write_scalar(scalar, out)?;
345                if let Some(access) = maybe_access {
346                    write!(out, ", {access}")?;
347                }
348                write!(out, ">")?;
349            } else {
350                return Err(WriteTypeError::NonWgsl);
351            }
352        }
353        TypeInner::ValuePointer {
354            size: Some(size),
355            scalar,
356            space,
357        } => {
358            let (address, maybe_access) = address_space_str(space);
359            if let Some(space) = address {
360                write!(out, "ptr<{}, vec{}<", space, common::vector_size_str(size),)?;
361                ctx.write_scalar(scalar, out)?;
362                out.write_str(">")?;
363                if let Some(access) = maybe_access {
364                    write!(out, ", {access}")?;
365                }
366                write!(out, ">")?;
367            } else {
368                return Err(WriteTypeError::NonWgsl);
369            }
370            write!(out, ">")?;
371        }
372        TypeInner::AccelerationStructure { vertex_return } => {
373            let caps = if vertex_return { "<vertex_return>" } else { "" };
374            write!(out, "acceleration_structure{caps}")?
375        }
376        TypeInner::Struct { .. } => {
377            ctx.write_unnamed_struct(inner, out)?;
378        }
379        TypeInner::RayQuery { vertex_return } => {
380            let caps = if vertex_return { "<vertex_return>" } else { "" };
381            write!(out, "ray_query{caps}")?
382        }
383    }
384
385    Ok(())
386}
387
388/// Error type returned by `try_write_type_inner`.
389///
390/// This type is private to the module.
391enum WriteTypeError {
392    Format(core::fmt::Error),
393    NonWgsl,
394}
395
396impl From<core::fmt::Error> for WriteTypeError {
397    fn from(err: core::fmt::Error) -> Self {
398        Self::Format(err)
399    }
400}
401
402/// Format types as WGSL based on a [`GlobalCtx`].
403///
404/// This is probably good enough for diagnostic output, but it has some
405/// limitations:
406///
407/// - It does not apply [`Namer`] renamings, to avoid collisions.
408///
409/// - It generates invalid WGSL for anonymous struct types.
410///
411/// - It doesn't write the lengths of override-expression-sized arrays
412///   correctly, unless the expression is just the override identifier.
413///
414/// [`GlobalCtx`]: crate::proc::GlobalCtx
415/// [`Namer`]: crate::proc::Namer
416impl TypeContext for crate::proc::GlobalCtx<'_> {
417    fn lookup_type(&self, handle: Handle<crate::Type>) -> &crate::Type {
418        &self.types[handle]
419    }
420
421    fn type_name(&self, handle: Handle<crate::Type>) -> &str {
422        self.types[handle]
423            .name
424            .as_deref()
425            .unwrap_or("{anonymous type}")
426    }
427
428    fn write_unnamed_struct<W: Write>(&self, _: &TypeInner, out: &mut W) -> core::fmt::Result {
429        write!(out, "{{unnamed struct}}")
430    }
431
432    fn write_override<W: Write>(
433        &self,
434        handle: Handle<crate::Override>,
435        out: &mut W,
436    ) -> core::fmt::Result {
437        match self.overrides[handle].name {
438            Some(ref name) => out.write_str(name),
439            None => write!(out, "{{anonymous override {handle:?}}}"),
440        }
441    }
442}
443
444/// Format types as WGSL based on a `UniqueArena<Type>`.
445///
446/// This is probably only good enough for logging:
447///
448/// - It does not apply any kind of [`Namer`] renamings.
449///
450/// - It generates invalid WGSL for anonymous struct types.
451///
452/// - It doesn't write override-sized arrays properly.
453///
454/// [`Namer`]: crate::proc::Namer
455impl TypeContext for crate::UniqueArena<crate::Type> {
456    fn lookup_type(&self, handle: Handle<crate::Type>) -> &crate::Type {
457        &self[handle]
458    }
459
460    fn type_name(&self, handle: Handle<crate::Type>) -> &str {
461        self[handle].name.as_deref().unwrap_or("{anonymous type}")
462    }
463
464    fn write_unnamed_struct<W: Write>(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result {
465        write!(out, "{{unnamed struct {inner:?}}}")
466    }
467
468    fn write_override<W: Write>(
469        &self,
470        handle: Handle<crate::Override>,
471        out: &mut W,
472    ) -> core::fmt::Result {
473        write!(out, "{{override {handle:?}}}")
474    }
475}