naga/back/wgsl/
writer.rs

1use alloc::{
2    format,
3    string::{String, ToString},
4    vec,
5    vec::Vec,
6};
7use core::fmt::Write;
8
9use super::Error;
10use super::ToWgslIfImplemented as _;
11use crate::{back::wgsl::polyfill::InversePolyfill, common::wgsl::TypeContext};
12use crate::{
13    back::{self, Baked},
14    common::{
15        self,
16        wgsl::{address_space_str, ToWgsl, TryToWgsl},
17    },
18    proc::{self, NameKey},
19    valid, Handle, Module, ShaderStage, TypeInner,
20};
21
22/// Shorthand result used internally by the backend
23type BackendResult = Result<(), Error>;
24
25/// WGSL [attribute](https://gpuweb.github.io/gpuweb/wgsl/#attributes)
26enum Attribute {
27    Binding(u32),
28    BuiltIn(crate::BuiltIn),
29    Group(u32),
30    Invariant,
31    Interpolate(Option<crate::Interpolation>, Option<crate::Sampling>),
32    Location(u32),
33    BlendSrc(u32),
34    Stage(ShaderStage),
35    WorkGroupSize([u32; 3]),
36}
37
38/// The WGSL form that `write_expr_with_indirection` should use to render a Naga
39/// expression.
40///
41/// Sometimes a Naga `Expression` alone doesn't provide enough information to
42/// choose the right rendering for it in WGSL. For example, one natural WGSL
43/// rendering of a Naga `LocalVariable(x)` expression might be `&x`, since
44/// `LocalVariable` produces a pointer to the local variable's storage. But when
45/// rendering a `Store` statement, the `pointer` operand must be the left hand
46/// side of a WGSL assignment, so the proper rendering is `x`.
47///
48/// The caller of `write_expr_with_indirection` must provide an `Expected` value
49/// to indicate how ambiguous expressions should be rendered.
50#[derive(Clone, Copy, Debug)]
51enum Indirection {
52    /// Render pointer-construction expressions as WGSL `ptr`-typed expressions.
53    ///
54    /// This is the right choice for most cases. Whenever a Naga pointer
55    /// expression is not the `pointer` operand of a `Load` or `Store`, it
56    /// must be a WGSL pointer expression.
57    Ordinary,
58
59    /// Render pointer-construction expressions as WGSL reference-typed
60    /// expressions.
61    ///
62    /// For example, this is the right choice for the `pointer` operand when
63    /// rendering a `Store` statement as a WGSL assignment.
64    Reference,
65}
66
67bitflags::bitflags! {
68    #[cfg_attr(feature = "serialize", derive(serde::Serialize))]
69    #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
70    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
71    pub struct WriterFlags: u32 {
72        /// Always annotate the type information instead of inferring.
73        const EXPLICIT_TYPES = 0x1;
74    }
75}
76
77pub struct Writer<W> {
78    out: W,
79    flags: WriterFlags,
80    names: crate::FastHashMap<NameKey, String>,
81    namer: proc::Namer,
82    named_expressions: crate::NamedExpressions,
83    required_polyfills: crate::FastIndexSet<InversePolyfill>,
84}
85
86impl<W: Write> Writer<W> {
87    pub fn new(out: W, flags: WriterFlags) -> Self {
88        Writer {
89            out,
90            flags,
91            names: crate::FastHashMap::default(),
92            namer: proc::Namer::default(),
93            named_expressions: crate::NamedExpressions::default(),
94            required_polyfills: crate::FastIndexSet::default(),
95        }
96    }
97
98    fn reset(&mut self, module: &Module) {
99        self.names.clear();
100        self.namer.reset(
101            module,
102            &crate::keywords::wgsl::RESERVED_SET,
103            // an identifier must not start with two underscore
104            proc::CaseInsensitiveKeywordSet::empty(),
105            &["__", "_naga"],
106            &mut self.names,
107        );
108        self.named_expressions.clear();
109        self.required_polyfills.clear();
110    }
111
112    /// Determine if `ty` is the Naga IR presentation of a WGSL builtin type.
113    ///
114    /// Return true if `ty` refers to the Naga IR form of a WGSL builtin type
115    /// like `__atomic_compare_exchange_result`.
116    ///
117    /// Even though the module may use the type, the WGSL backend should avoid
118    /// emitting a definition for it, since it is [predeclared] in WGSL.
119    ///
120    /// This also covers types like [`NagaExternalTextureParams`], which other
121    /// backends use to lower WGSL constructs like external textures to their
122    /// implementations. WGSL can express these directly, so the types need not
123    /// be emitted.
124    ///
125    /// [predeclared]: https://www.w3.org/TR/WGSL/#predeclared
126    /// [`NagaExternalTextureParams`]: crate::ir::SpecialTypes::external_texture_params
127    fn is_builtin_wgsl_struct(&self, module: &Module, ty: Handle<crate::Type>) -> bool {
128        module
129            .special_types
130            .predeclared_types
131            .values()
132            .any(|t| *t == ty)
133            || Some(ty) == module.special_types.external_texture_params
134            || Some(ty) == module.special_types.external_texture_transfer_function
135    }
136
137    pub fn write(&mut self, module: &Module, info: &valid::ModuleInfo) -> BackendResult {
138        if !module.overrides.is_empty() {
139            return Err(Error::Unimplemented(
140                "Pipeline constants are not yet supported for this back-end".to_string(),
141            ));
142        }
143
144        self.reset(module);
145
146        // Write all `enable` declarations
147        self.write_enable_declarations(module)?;
148
149        // Write all structs
150        for (handle, ty) in module.types.iter() {
151            if let TypeInner::Struct { ref members, .. } = ty.inner {
152                {
153                    if !self.is_builtin_wgsl_struct(module, handle) {
154                        self.write_struct(module, handle, members)?;
155                        writeln!(self.out)?;
156                    }
157                }
158            }
159        }
160
161        // Write all named constants
162        let mut constants = module
163            .constants
164            .iter()
165            .filter(|&(_, c)| c.name.is_some())
166            .peekable();
167        while let Some((handle, _)) = constants.next() {
168            self.write_global_constant(module, handle)?;
169            // Add extra newline for readability on last iteration
170            if constants.peek().is_none() {
171                writeln!(self.out)?;
172            }
173        }
174
175        // Write all globals
176        for (ty, global) in module.global_variables.iter() {
177            self.write_global(module, global, ty)?;
178        }
179
180        if !module.global_variables.is_empty() {
181            // Add extra newline for readability
182            writeln!(self.out)?;
183        }
184
185        // Write all regular functions
186        for (handle, function) in module.functions.iter() {
187            let fun_info = &info[handle];
188
189            let func_ctx = back::FunctionCtx {
190                ty: back::FunctionType::Function(handle),
191                info: fun_info,
192                expressions: &function.expressions,
193                named_expressions: &function.named_expressions,
194            };
195
196            // Write the function
197            self.write_function(module, function, &func_ctx)?;
198
199            writeln!(self.out)?;
200        }
201
202        // Write all entry points
203        for (index, ep) in module.entry_points.iter().enumerate() {
204            let attributes = match ep.stage {
205                ShaderStage::Vertex | ShaderStage::Fragment => vec![Attribute::Stage(ep.stage)],
206                ShaderStage::Compute => vec![
207                    Attribute::Stage(ShaderStage::Compute),
208                    Attribute::WorkGroupSize(ep.workgroup_size),
209                ],
210                ShaderStage::Task | ShaderStage::Mesh => unreachable!(),
211            };
212
213            self.write_attributes(&attributes)?;
214            // Add a newline after attribute
215            writeln!(self.out)?;
216
217            let func_ctx = back::FunctionCtx {
218                ty: back::FunctionType::EntryPoint(index as u16),
219                info: info.get_entry_point(index),
220                expressions: &ep.function.expressions,
221                named_expressions: &ep.function.named_expressions,
222            };
223            self.write_function(module, &ep.function, &func_ctx)?;
224
225            if index < module.entry_points.len() - 1 {
226                writeln!(self.out)?;
227            }
228        }
229
230        // Write any polyfills that were required.
231        for polyfill in &self.required_polyfills {
232            writeln!(self.out)?;
233            write!(self.out, "{}", polyfill.source)?;
234            writeln!(self.out)?;
235        }
236
237        Ok(())
238    }
239
240    /// Helper method which writes all the `enable` declarations
241    /// needed for a module.
242    fn write_enable_declarations(&mut self, module: &Module) -> BackendResult {
243        let mut needs_f16 = false;
244        let mut needs_dual_source_blending = false;
245        let mut needs_clip_distances = false;
246
247        // Determine which `enable` declarations are needed
248        for (_, ty) in module.types.iter() {
249            match ty.inner {
250                TypeInner::Scalar(scalar)
251                | TypeInner::Vector { scalar, .. }
252                | TypeInner::Matrix { scalar, .. } => {
253                    needs_f16 |= scalar == crate::Scalar::F16;
254                }
255                TypeInner::Struct { ref members, .. } => {
256                    for binding in members.iter().filter_map(|m| m.binding.as_ref()) {
257                        match *binding {
258                            crate::Binding::Location {
259                                blend_src: Some(_), ..
260                            } => {
261                                needs_dual_source_blending = true;
262                            }
263                            crate::Binding::BuiltIn(crate::BuiltIn::ClipDistance) => {
264                                needs_clip_distances = true;
265                            }
266                            _ => {}
267                        }
268                    }
269                }
270                _ => {}
271            }
272        }
273
274        // Write required declarations
275        let mut any_written = false;
276        if needs_f16 {
277            writeln!(self.out, "enable f16;")?;
278            any_written = true;
279        }
280        if needs_dual_source_blending {
281            writeln!(self.out, "enable dual_source_blending;")?;
282            any_written = true;
283        }
284        if needs_clip_distances {
285            writeln!(self.out, "enable clip_distances;")?;
286            any_written = true;
287        }
288        if any_written {
289            // Empty line for readability
290            writeln!(self.out)?;
291        }
292
293        Ok(())
294    }
295
296    /// Helper method used to write
297    /// [functions](https://gpuweb.github.io/gpuweb/wgsl/#functions)
298    ///
299    /// # Notes
300    /// Ends in a newline
301    fn write_function(
302        &mut self,
303        module: &Module,
304        func: &crate::Function,
305        func_ctx: &back::FunctionCtx<'_>,
306    ) -> BackendResult {
307        let func_name = match func_ctx.ty {
308            back::FunctionType::EntryPoint(index) => &self.names[&NameKey::EntryPoint(index)],
309            back::FunctionType::Function(handle) => &self.names[&NameKey::Function(handle)],
310        };
311
312        // Write function name
313        write!(self.out, "fn {func_name}(")?;
314
315        // Write function arguments
316        for (index, arg) in func.arguments.iter().enumerate() {
317            // Write argument attribute if a binding is present
318            if let Some(ref binding) = arg.binding {
319                self.write_attributes(&map_binding_to_attribute(binding))?;
320            }
321            // Write argument name
322            let argument_name = &self.names[&func_ctx.argument_key(index as u32)];
323
324            write!(self.out, "{argument_name}: ")?;
325            // Write argument type
326            self.write_type(module, arg.ty)?;
327            if index < func.arguments.len() - 1 {
328                // Add a separator between args
329                write!(self.out, ", ")?;
330            }
331        }
332
333        write!(self.out, ")")?;
334
335        // Write function return type
336        if let Some(ref result) = func.result {
337            write!(self.out, " -> ")?;
338            if let Some(ref binding) = result.binding {
339                self.write_attributes(&map_binding_to_attribute(binding))?;
340            }
341            self.write_type(module, result.ty)?;
342        }
343
344        write!(self.out, " {{")?;
345        writeln!(self.out)?;
346
347        // Write function local variables
348        for (handle, local) in func.local_variables.iter() {
349            // Write indentation (only for readability)
350            write!(self.out, "{}", back::INDENT)?;
351
352            // Write the local name
353            // The leading space is important
354            write!(self.out, "var {}: ", self.names[&func_ctx.name_key(handle)])?;
355
356            // Write the local type
357            self.write_type(module, local.ty)?;
358
359            // Write the local initializer if needed
360            if let Some(init) = local.init {
361                // Put the equal signal only if there's a initializer
362                // The leading and trailing spaces aren't needed but help with readability
363                write!(self.out, " = ")?;
364
365                // Write the constant
366                // `write_constant` adds no trailing or leading space/newline
367                self.write_expr(module, init, func_ctx)?;
368            }
369
370            // Finish the local with `;` and add a newline (only for readability)
371            writeln!(self.out, ";")?
372        }
373
374        if !func.local_variables.is_empty() {
375            writeln!(self.out)?;
376        }
377
378        // Write the function body (statement list)
379        for sta in func.body.iter() {
380            // The indentation should always be 1 when writing the function body
381            self.write_stmt(module, sta, func_ctx, back::Level(1))?;
382        }
383
384        writeln!(self.out, "}}")?;
385
386        self.named_expressions.clear();
387
388        Ok(())
389    }
390
391    /// Helper method to write a attribute
392    fn write_attributes(&mut self, attributes: &[Attribute]) -> BackendResult {
393        for attribute in attributes {
394            match *attribute {
395                Attribute::Location(id) => write!(self.out, "@location({id}) ")?,
396                Attribute::BlendSrc(blend_src) => write!(self.out, "@blend_src({blend_src}) ")?,
397                Attribute::BuiltIn(builtin_attrib) => {
398                    let builtin = builtin_attrib.to_wgsl_if_implemented()?;
399                    write!(self.out, "@builtin({builtin}) ")?;
400                }
401                Attribute::Stage(shader_stage) => {
402                    let stage_str = match shader_stage {
403                        ShaderStage::Vertex => "vertex",
404                        ShaderStage::Fragment => "fragment",
405                        ShaderStage::Compute => "compute",
406                        ShaderStage::Task | ShaderStage::Mesh => unreachable!(),
407                    };
408                    write!(self.out, "@{stage_str} ")?;
409                }
410                Attribute::WorkGroupSize(size) => {
411                    write!(
412                        self.out,
413                        "@workgroup_size({}, {}, {}) ",
414                        size[0], size[1], size[2]
415                    )?;
416                }
417                Attribute::Binding(id) => write!(self.out, "@binding({id}) ")?,
418                Attribute::Group(id) => write!(self.out, "@group({id}) ")?,
419                Attribute::Invariant => write!(self.out, "@invariant ")?,
420                Attribute::Interpolate(interpolation, sampling) => {
421                    if sampling.is_some() && sampling != Some(crate::Sampling::Center) {
422                        let interpolation = interpolation
423                            .unwrap_or(crate::Interpolation::Perspective)
424                            .to_wgsl();
425                        let sampling = sampling.unwrap_or(crate::Sampling::Center).to_wgsl();
426                        write!(self.out, "@interpolate({interpolation}, {sampling}) ")?;
427                    } else if interpolation.is_some()
428                        && interpolation != Some(crate::Interpolation::Perspective)
429                    {
430                        let interpolation = interpolation
431                            .unwrap_or(crate::Interpolation::Perspective)
432                            .to_wgsl();
433                        write!(self.out, "@interpolate({interpolation}) ")?;
434                    }
435                }
436            };
437        }
438        Ok(())
439    }
440
441    /// Helper method used to write structs
442    /// Write the full declaration of a struct type.
443    ///
444    /// Write out a definition of the struct type referred to by
445    /// `handle` in `module`. The output will be an instance of the
446    /// `struct_decl` production in the WGSL grammar.
447    ///
448    /// Use `members` as the list of `handle`'s members. (This
449    /// function is usually called after matching a `TypeInner`, so
450    /// the callers already have the members at hand.)
451    fn write_struct(
452        &mut self,
453        module: &Module,
454        handle: Handle<crate::Type>,
455        members: &[crate::StructMember],
456    ) -> BackendResult {
457        write!(self.out, "struct {}", self.names[&NameKey::Type(handle)])?;
458        write!(self.out, " {{")?;
459        writeln!(self.out)?;
460        for (index, member) in members.iter().enumerate() {
461            // The indentation is only for readability
462            write!(self.out, "{}", back::INDENT)?;
463            if let Some(ref binding) = member.binding {
464                self.write_attributes(&map_binding_to_attribute(binding))?;
465            }
466            // Write struct member name and type
467            let member_name = &self.names[&NameKey::StructMember(handle, index as u32)];
468            write!(self.out, "{member_name}: ")?;
469            self.write_type(module, member.ty)?;
470            write!(self.out, ",")?;
471            writeln!(self.out)?;
472        }
473
474        writeln!(self.out, "}}")?;
475
476        Ok(())
477    }
478
479    fn write_type(&mut self, module: &Module, ty: Handle<crate::Type>) -> BackendResult {
480        // This actually can't be factored out into a nice constructor method,
481        // because the borrow checker needs to be able to see that the borrows
482        // of `self.names` and `self.out` are disjoint.
483        let type_context = WriterTypeContext {
484            module,
485            names: &self.names,
486        };
487        type_context.write_type(ty, &mut self.out)?;
488
489        Ok(())
490    }
491
492    fn write_type_resolution(
493        &mut self,
494        module: &Module,
495        resolution: &proc::TypeResolution,
496    ) -> BackendResult {
497        // This actually can't be factored out into a nice constructor method,
498        // because the borrow checker needs to be able to see that the borrows
499        // of `self.names` and `self.out` are disjoint.
500        let type_context = WriterTypeContext {
501            module,
502            names: &self.names,
503        };
504        type_context.write_type_resolution(resolution, &mut self.out)?;
505
506        Ok(())
507    }
508
509    /// Helper method used to write statements
510    ///
511    /// # Notes
512    /// Always adds a newline
513    fn write_stmt(
514        &mut self,
515        module: &Module,
516        stmt: &crate::Statement,
517        func_ctx: &back::FunctionCtx<'_>,
518        level: back::Level,
519    ) -> BackendResult {
520        use crate::{Expression, Statement};
521
522        match *stmt {
523            Statement::Emit(ref range) => {
524                for handle in range.clone() {
525                    let info = &func_ctx.info[handle];
526                    let expr_name = if let Some(name) = func_ctx.named_expressions.get(&handle) {
527                        // Front end provides names for all variables at the start of writing.
528                        // But we write them to step by step. We need to recache them
529                        // Otherwise, we could accidentally write variable name instead of full expression.
530                        // Also, we use sanitized names! It defense backend from generating variable with name from reserved keywords.
531                        Some(self.namer.call(name))
532                    } else {
533                        let expr = &func_ctx.expressions[handle];
534                        let min_ref_count = expr.bake_ref_count();
535                        // Forcefully creating baking expressions in some cases to help with readability
536                        let required_baking_expr = match *expr {
537                            Expression::ImageLoad { .. }
538                            | Expression::ImageQuery { .. }
539                            | Expression::ImageSample { .. } => true,
540                            _ => false,
541                        };
542                        if min_ref_count <= info.ref_count || required_baking_expr {
543                            Some(Baked(handle).to_string())
544                        } else {
545                            None
546                        }
547                    };
548
549                    if let Some(name) = expr_name {
550                        write!(self.out, "{level}")?;
551                        self.start_named_expr(module, handle, func_ctx, &name)?;
552                        self.write_expr(module, handle, func_ctx)?;
553                        self.named_expressions.insert(handle, name);
554                        writeln!(self.out, ";")?;
555                    }
556                }
557            }
558            // TODO: copy-paste from glsl-out
559            Statement::If {
560                condition,
561                ref accept,
562                ref reject,
563            } => {
564                write!(self.out, "{level}")?;
565                write!(self.out, "if ")?;
566                self.write_expr(module, condition, func_ctx)?;
567                writeln!(self.out, " {{")?;
568
569                let l2 = level.next();
570                for sta in accept {
571                    // Increase indentation to help with readability
572                    self.write_stmt(module, sta, func_ctx, l2)?;
573                }
574
575                // If there are no statements in the reject block we skip writing it
576                // This is only for readability
577                if !reject.is_empty() {
578                    writeln!(self.out, "{level}}} else {{")?;
579
580                    for sta in reject {
581                        // Increase indentation to help with readability
582                        self.write_stmt(module, sta, func_ctx, l2)?;
583                    }
584                }
585
586                writeln!(self.out, "{level}}}")?
587            }
588            Statement::Return { value } => {
589                write!(self.out, "{level}")?;
590                write!(self.out, "return")?;
591                if let Some(return_value) = value {
592                    // The leading space is important
593                    write!(self.out, " ")?;
594                    self.write_expr(module, return_value, func_ctx)?;
595                }
596                writeln!(self.out, ";")?;
597            }
598            // TODO: copy-paste from glsl-out
599            Statement::Kill => {
600                write!(self.out, "{level}")?;
601                writeln!(self.out, "discard;")?
602            }
603            Statement::Store { pointer, value } => {
604                write!(self.out, "{level}")?;
605
606                let is_atomic_pointer = func_ctx
607                    .resolve_type(pointer, &module.types)
608                    .is_atomic_pointer(&module.types);
609
610                if is_atomic_pointer {
611                    write!(self.out, "atomicStore(")?;
612                    self.write_expr(module, pointer, func_ctx)?;
613                    write!(self.out, ", ")?;
614                    self.write_expr(module, value, func_ctx)?;
615                    write!(self.out, ")")?;
616                } else {
617                    self.write_expr_with_indirection(
618                        module,
619                        pointer,
620                        func_ctx,
621                        Indirection::Reference,
622                    )?;
623                    write!(self.out, " = ")?;
624                    self.write_expr(module, value, func_ctx)?;
625                }
626                writeln!(self.out, ";")?
627            }
628            Statement::Call {
629                function,
630                ref arguments,
631                result,
632            } => {
633                write!(self.out, "{level}")?;
634                if let Some(expr) = result {
635                    let name = Baked(expr).to_string();
636                    self.start_named_expr(module, expr, func_ctx, &name)?;
637                    self.named_expressions.insert(expr, name);
638                }
639                let func_name = &self.names[&NameKey::Function(function)];
640                write!(self.out, "{func_name}(")?;
641                for (index, &argument) in arguments.iter().enumerate() {
642                    if index != 0 {
643                        write!(self.out, ", ")?;
644                    }
645                    self.write_expr(module, argument, func_ctx)?;
646                }
647                writeln!(self.out, ");")?
648            }
649            Statement::Atomic {
650                pointer,
651                ref fun,
652                value,
653                result,
654            } => {
655                write!(self.out, "{level}")?;
656                if let Some(result) = result {
657                    let res_name = Baked(result).to_string();
658                    self.start_named_expr(module, result, func_ctx, &res_name)?;
659                    self.named_expressions.insert(result, res_name);
660                }
661
662                let fun_str = fun.to_wgsl();
663                write!(self.out, "atomic{fun_str}(")?;
664                self.write_expr(module, pointer, func_ctx)?;
665                if let crate::AtomicFunction::Exchange { compare: Some(cmp) } = *fun {
666                    write!(self.out, ", ")?;
667                    self.write_expr(module, cmp, func_ctx)?;
668                }
669                write!(self.out, ", ")?;
670                self.write_expr(module, value, func_ctx)?;
671                writeln!(self.out, ");")?
672            }
673            Statement::ImageAtomic {
674                image,
675                coordinate,
676                array_index,
677                ref fun,
678                value,
679            } => {
680                write!(self.out, "{level}")?;
681                let fun_str = fun.to_wgsl();
682                write!(self.out, "textureAtomic{fun_str}(")?;
683                self.write_expr(module, image, func_ctx)?;
684                write!(self.out, ", ")?;
685                self.write_expr(module, coordinate, func_ctx)?;
686                if let Some(array_index_expr) = array_index {
687                    write!(self.out, ", ")?;
688                    self.write_expr(module, array_index_expr, func_ctx)?;
689                }
690                write!(self.out, ", ")?;
691                self.write_expr(module, value, func_ctx)?;
692                writeln!(self.out, ");")?;
693            }
694            Statement::WorkGroupUniformLoad { pointer, result } => {
695                write!(self.out, "{level}")?;
696                // TODO: Obey named expressions here.
697                let res_name = Baked(result).to_string();
698                self.start_named_expr(module, result, func_ctx, &res_name)?;
699                self.named_expressions.insert(result, res_name);
700                write!(self.out, "workgroupUniformLoad(")?;
701                self.write_expr(module, pointer, func_ctx)?;
702                writeln!(self.out, ");")?;
703            }
704            Statement::ImageStore {
705                image,
706                coordinate,
707                array_index,
708                value,
709            } => {
710                write!(self.out, "{level}")?;
711                write!(self.out, "textureStore(")?;
712                self.write_expr(module, image, func_ctx)?;
713                write!(self.out, ", ")?;
714                self.write_expr(module, coordinate, func_ctx)?;
715                if let Some(array_index_expr) = array_index {
716                    write!(self.out, ", ")?;
717                    self.write_expr(module, array_index_expr, func_ctx)?;
718                }
719                write!(self.out, ", ")?;
720                self.write_expr(module, value, func_ctx)?;
721                writeln!(self.out, ");")?;
722            }
723            // TODO: copy-paste from glsl-out
724            Statement::Block(ref block) => {
725                write!(self.out, "{level}")?;
726                writeln!(self.out, "{{")?;
727                for sta in block.iter() {
728                    // Increase the indentation to help with readability
729                    self.write_stmt(module, sta, func_ctx, level.next())?
730                }
731                writeln!(self.out, "{level}}}")?
732            }
733            Statement::Switch {
734                selector,
735                ref cases,
736            } => {
737                // Start the switch
738                write!(self.out, "{level}")?;
739                write!(self.out, "switch ")?;
740                self.write_expr(module, selector, func_ctx)?;
741                writeln!(self.out, " {{")?;
742
743                let l2 = level.next();
744                let mut new_case = true;
745                for case in cases {
746                    if case.fall_through && !case.body.is_empty() {
747                        // TODO: we could do the same workaround as we did for the HLSL backend
748                        return Err(Error::Unimplemented(
749                            "fall-through switch case block".into(),
750                        ));
751                    }
752
753                    match case.value {
754                        crate::SwitchValue::I32(value) => {
755                            if new_case {
756                                write!(self.out, "{l2}case ")?;
757                            }
758                            write!(self.out, "{value}")?;
759                        }
760                        crate::SwitchValue::U32(value) => {
761                            if new_case {
762                                write!(self.out, "{l2}case ")?;
763                            }
764                            write!(self.out, "{value}u")?;
765                        }
766                        crate::SwitchValue::Default => {
767                            if new_case {
768                                if case.fall_through {
769                                    write!(self.out, "{l2}case ")?;
770                                } else {
771                                    write!(self.out, "{l2}")?;
772                                }
773                            }
774                            write!(self.out, "default")?;
775                        }
776                    }
777
778                    new_case = !case.fall_through;
779
780                    if case.fall_through {
781                        write!(self.out, ", ")?;
782                    } else {
783                        writeln!(self.out, ": {{")?;
784                    }
785
786                    for sta in case.body.iter() {
787                        self.write_stmt(module, sta, func_ctx, l2.next())?;
788                    }
789
790                    if !case.fall_through {
791                        writeln!(self.out, "{l2}}}")?;
792                    }
793                }
794
795                writeln!(self.out, "{level}}}")?
796            }
797            Statement::Loop {
798                ref body,
799                ref continuing,
800                break_if,
801            } => {
802                write!(self.out, "{level}")?;
803                writeln!(self.out, "loop {{")?;
804
805                let l2 = level.next();
806                for sta in body.iter() {
807                    self.write_stmt(module, sta, func_ctx, l2)?;
808                }
809
810                // The continuing is optional so we don't need to write it if
811                // it is empty, but the `break if` counts as a continuing statement
812                // so even if `continuing` is empty we must generate it if a
813                // `break if` exists
814                if !continuing.is_empty() || break_if.is_some() {
815                    writeln!(self.out, "{l2}continuing {{")?;
816                    for sta in continuing.iter() {
817                        self.write_stmt(module, sta, func_ctx, l2.next())?;
818                    }
819
820                    // The `break if` is always the last
821                    // statement of the `continuing` block
822                    if let Some(condition) = break_if {
823                        // The trailing space is important
824                        write!(self.out, "{}break if ", l2.next())?;
825                        self.write_expr(module, condition, func_ctx)?;
826                        // Close the `break if` statement
827                        writeln!(self.out, ";")?;
828                    }
829
830                    writeln!(self.out, "{l2}}}")?;
831                }
832
833                writeln!(self.out, "{level}}}")?
834            }
835            Statement::Break => {
836                writeln!(self.out, "{level}break;")?;
837            }
838            Statement::Continue => {
839                writeln!(self.out, "{level}continue;")?;
840            }
841            Statement::ControlBarrier(barrier) | Statement::MemoryBarrier(barrier) => {
842                if barrier.contains(crate::Barrier::STORAGE) {
843                    writeln!(self.out, "{level}storageBarrier();")?;
844                }
845
846                if barrier.contains(crate::Barrier::WORK_GROUP) {
847                    writeln!(self.out, "{level}workgroupBarrier();")?;
848                }
849
850                if barrier.contains(crate::Barrier::SUB_GROUP) {
851                    writeln!(self.out, "{level}subgroupBarrier();")?;
852                }
853
854                if barrier.contains(crate::Barrier::TEXTURE) {
855                    writeln!(self.out, "{level}textureBarrier();")?;
856                }
857            }
858            Statement::RayQuery { .. } => unreachable!(),
859            Statement::SubgroupBallot { result, predicate } => {
860                write!(self.out, "{level}")?;
861                let res_name = Baked(result).to_string();
862                self.start_named_expr(module, result, func_ctx, &res_name)?;
863                self.named_expressions.insert(result, res_name);
864
865                write!(self.out, "subgroupBallot(")?;
866                if let Some(predicate) = predicate {
867                    self.write_expr(module, predicate, func_ctx)?;
868                }
869                writeln!(self.out, ");")?;
870            }
871            Statement::SubgroupCollectiveOperation {
872                op,
873                collective_op,
874                argument,
875                result,
876            } => {
877                write!(self.out, "{level}")?;
878                let res_name = Baked(result).to_string();
879                self.start_named_expr(module, result, func_ctx, &res_name)?;
880                self.named_expressions.insert(result, res_name);
881
882                match (collective_op, op) {
883                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::All) => {
884                        write!(self.out, "subgroupAll(")?
885                    }
886                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Any) => {
887                        write!(self.out, "subgroupAny(")?
888                    }
889                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Add) => {
890                        write!(self.out, "subgroupAdd(")?
891                    }
892                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Mul) => {
893                        write!(self.out, "subgroupMul(")?
894                    }
895                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Max) => {
896                        write!(self.out, "subgroupMax(")?
897                    }
898                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Min) => {
899                        write!(self.out, "subgroupMin(")?
900                    }
901                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::And) => {
902                        write!(self.out, "subgroupAnd(")?
903                    }
904                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Or) => {
905                        write!(self.out, "subgroupOr(")?
906                    }
907                    (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Xor) => {
908                        write!(self.out, "subgroupXor(")?
909                    }
910                    (crate::CollectiveOperation::ExclusiveScan, crate::SubgroupOperation::Add) => {
911                        write!(self.out, "subgroupExclusiveAdd(")?
912                    }
913                    (crate::CollectiveOperation::ExclusiveScan, crate::SubgroupOperation::Mul) => {
914                        write!(self.out, "subgroupExclusiveMul(")?
915                    }
916                    (crate::CollectiveOperation::InclusiveScan, crate::SubgroupOperation::Add) => {
917                        write!(self.out, "subgroupInclusiveAdd(")?
918                    }
919                    (crate::CollectiveOperation::InclusiveScan, crate::SubgroupOperation::Mul) => {
920                        write!(self.out, "subgroupInclusiveMul(")?
921                    }
922                    _ => unimplemented!(),
923                }
924                self.write_expr(module, argument, func_ctx)?;
925                writeln!(self.out, ");")?;
926            }
927            Statement::SubgroupGather {
928                mode,
929                argument,
930                result,
931            } => {
932                write!(self.out, "{level}")?;
933                let res_name = Baked(result).to_string();
934                self.start_named_expr(module, result, func_ctx, &res_name)?;
935                self.named_expressions.insert(result, res_name);
936
937                match mode {
938                    crate::GatherMode::BroadcastFirst => {
939                        write!(self.out, "subgroupBroadcastFirst(")?;
940                    }
941                    crate::GatherMode::Broadcast(_) => {
942                        write!(self.out, "subgroupBroadcast(")?;
943                    }
944                    crate::GatherMode::Shuffle(_) => {
945                        write!(self.out, "subgroupShuffle(")?;
946                    }
947                    crate::GatherMode::ShuffleDown(_) => {
948                        write!(self.out, "subgroupShuffleDown(")?;
949                    }
950                    crate::GatherMode::ShuffleUp(_) => {
951                        write!(self.out, "subgroupShuffleUp(")?;
952                    }
953                    crate::GatherMode::ShuffleXor(_) => {
954                        write!(self.out, "subgroupShuffleXor(")?;
955                    }
956                    crate::GatherMode::QuadBroadcast(_) => {
957                        write!(self.out, "quadBroadcast(")?;
958                    }
959                    crate::GatherMode::QuadSwap(direction) => match direction {
960                        crate::Direction::X => {
961                            write!(self.out, "quadSwapX(")?;
962                        }
963                        crate::Direction::Y => {
964                            write!(self.out, "quadSwapY(")?;
965                        }
966                        crate::Direction::Diagonal => {
967                            write!(self.out, "quadSwapDiagonal(")?;
968                        }
969                    },
970                }
971                self.write_expr(module, argument, func_ctx)?;
972                match mode {
973                    crate::GatherMode::BroadcastFirst => {}
974                    crate::GatherMode::Broadcast(index)
975                    | crate::GatherMode::Shuffle(index)
976                    | crate::GatherMode::ShuffleDown(index)
977                    | crate::GatherMode::ShuffleUp(index)
978                    | crate::GatherMode::ShuffleXor(index)
979                    | crate::GatherMode::QuadBroadcast(index) => {
980                        write!(self.out, ", ")?;
981                        self.write_expr(module, index, func_ctx)?;
982                    }
983                    crate::GatherMode::QuadSwap(_) => {}
984                }
985                writeln!(self.out, ");")?;
986            }
987        }
988
989        Ok(())
990    }
991
992    /// Return the sort of indirection that `expr`'s plain form evaluates to.
993    ///
994    /// An expression's 'plain form' is the most general rendition of that
995    /// expression into WGSL, lacking `&` or `*` operators:
996    ///
997    /// - The plain form of `LocalVariable(x)` is simply `x`, which is a reference
998    ///   to the local variable's storage.
999    ///
1000    /// - The plain form of `GlobalVariable(g)` is simply `g`, which is usually a
1001    ///   reference to the global variable's storage. However, globals in the
1002    ///   `Handle` address space are immutable, and `GlobalVariable` expressions for
1003    ///   those produce the value directly, not a pointer to it. Such
1004    ///   `GlobalVariable` expressions are `Ordinary`.
1005    ///
1006    /// - `Access` and `AccessIndex` are `Reference` when their `base` operand is a
1007    ///   pointer. If they are applied directly to a composite value, they are
1008    ///   `Ordinary`.
1009    ///
1010    /// Note that `FunctionArgument` expressions are never `Reference`, even when
1011    /// the argument's type is `Pointer`. `FunctionArgument` always evaluates to the
1012    /// argument's value directly, so any pointer it produces is merely the value
1013    /// passed by the caller.
1014    fn plain_form_indirection(
1015        &self,
1016        expr: Handle<crate::Expression>,
1017        module: &Module,
1018        func_ctx: &back::FunctionCtx<'_>,
1019    ) -> Indirection {
1020        use crate::Expression as Ex;
1021
1022        // Named expressions are `let` expressions, which apply the Load Rule,
1023        // so if their type is a Naga pointer, then that must be a WGSL pointer
1024        // as well.
1025        if self.named_expressions.contains_key(&expr) {
1026            return Indirection::Ordinary;
1027        }
1028
1029        match func_ctx.expressions[expr] {
1030            Ex::LocalVariable(_) => Indirection::Reference,
1031            Ex::GlobalVariable(handle) => {
1032                let global = &module.global_variables[handle];
1033                match global.space {
1034                    crate::AddressSpace::Handle => Indirection::Ordinary,
1035                    _ => Indirection::Reference,
1036                }
1037            }
1038            Ex::Access { base, .. } | Ex::AccessIndex { base, .. } => {
1039                let base_ty = func_ctx.resolve_type(base, &module.types);
1040                match *base_ty {
1041                    TypeInner::Pointer { .. } | TypeInner::ValuePointer { .. } => {
1042                        Indirection::Reference
1043                    }
1044                    _ => Indirection::Ordinary,
1045                }
1046            }
1047            _ => Indirection::Ordinary,
1048        }
1049    }
1050
1051    fn start_named_expr(
1052        &mut self,
1053        module: &Module,
1054        handle: Handle<crate::Expression>,
1055        func_ctx: &back::FunctionCtx,
1056        name: &str,
1057    ) -> BackendResult {
1058        // Write variable name
1059        write!(self.out, "let {name}")?;
1060        if self.flags.contains(WriterFlags::EXPLICIT_TYPES) {
1061            write!(self.out, ": ")?;
1062            // Write variable type
1063            self.write_type_resolution(module, &func_ctx.info[handle].ty)?;
1064        }
1065
1066        write!(self.out, " = ")?;
1067        Ok(())
1068    }
1069
1070    /// Write the ordinary WGSL form of `expr`.
1071    ///
1072    /// See `write_expr_with_indirection` for details.
1073    fn write_expr(
1074        &mut self,
1075        module: &Module,
1076        expr: Handle<crate::Expression>,
1077        func_ctx: &back::FunctionCtx<'_>,
1078    ) -> BackendResult {
1079        self.write_expr_with_indirection(module, expr, func_ctx, Indirection::Ordinary)
1080    }
1081
1082    /// Write `expr` as a WGSL expression with the requested indirection.
1083    ///
1084    /// In terms of the WGSL grammar, the resulting expression is a
1085    /// `singular_expression`. It may be parenthesized. This makes it suitable
1086    /// for use as the operand of a unary or binary operator without worrying
1087    /// about precedence.
1088    ///
1089    /// This does not produce newlines or indentation.
1090    ///
1091    /// The `requested` argument indicates (roughly) whether Naga
1092    /// `Pointer`-valued expressions represent WGSL references or pointers. See
1093    /// `Indirection` for details.
1094    fn write_expr_with_indirection(
1095        &mut self,
1096        module: &Module,
1097        expr: Handle<crate::Expression>,
1098        func_ctx: &back::FunctionCtx<'_>,
1099        requested: Indirection,
1100    ) -> BackendResult {
1101        // If the plain form of the expression is not what we need, emit the
1102        // operator necessary to correct that.
1103        let plain = self.plain_form_indirection(expr, module, func_ctx);
1104        match (requested, plain) {
1105            (Indirection::Ordinary, Indirection::Reference) => {
1106                write!(self.out, "(&")?;
1107                self.write_expr_plain_form(module, expr, func_ctx, plain)?;
1108                write!(self.out, ")")?;
1109            }
1110            (Indirection::Reference, Indirection::Ordinary) => {
1111                write!(self.out, "(*")?;
1112                self.write_expr_plain_form(module, expr, func_ctx, plain)?;
1113                write!(self.out, ")")?;
1114            }
1115            (_, _) => self.write_expr_plain_form(module, expr, func_ctx, plain)?,
1116        }
1117
1118        Ok(())
1119    }
1120
1121    fn write_const_expression(
1122        &mut self,
1123        module: &Module,
1124        expr: Handle<crate::Expression>,
1125        arena: &crate::Arena<crate::Expression>,
1126    ) -> BackendResult {
1127        self.write_possibly_const_expression(module, expr, arena, |writer, expr| {
1128            writer.write_const_expression(module, expr, arena)
1129        })
1130    }
1131
1132    fn write_possibly_const_expression<E>(
1133        &mut self,
1134        module: &Module,
1135        expr: Handle<crate::Expression>,
1136        expressions: &crate::Arena<crate::Expression>,
1137        write_expression: E,
1138    ) -> BackendResult
1139    where
1140        E: Fn(&mut Self, Handle<crate::Expression>) -> BackendResult,
1141    {
1142        use crate::Expression;
1143
1144        match expressions[expr] {
1145            Expression::Literal(literal) => match literal {
1146                crate::Literal::F16(value) => write!(self.out, "{value}h")?,
1147                crate::Literal::F32(value) => write!(self.out, "{value}f")?,
1148                crate::Literal::U32(value) => write!(self.out, "{value}u")?,
1149                crate::Literal::I32(value) => {
1150                    // `-2147483648i` is not valid WGSL. The most negative `i32`
1151                    // value can only be expressed in WGSL using AbstractInt and
1152                    // a unary negation operator.
1153                    if value == i32::MIN {
1154                        write!(self.out, "i32({value})")?;
1155                    } else {
1156                        write!(self.out, "{value}i")?;
1157                    }
1158                }
1159                crate::Literal::Bool(value) => write!(self.out, "{value}")?,
1160                crate::Literal::F64(value) => write!(self.out, "{value:?}lf")?,
1161                crate::Literal::I64(value) => {
1162                    // `-9223372036854775808li` is not valid WGSL. Nor can we simply use the
1163                    // AbstractInt trick above, as AbstractInt also cannot represent
1164                    // `9223372036854775808`. Instead construct the second most negative
1165                    // AbstractInt, subtract one from it, then cast to i64.
1166                    if value == i64::MIN {
1167                        write!(self.out, "i64({} - 1)", value + 1)?;
1168                    } else {
1169                        write!(self.out, "{value}li")?;
1170                    }
1171                }
1172                crate::Literal::U64(value) => write!(self.out, "{value:?}lu")?,
1173                crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => {
1174                    return Err(Error::Custom(
1175                        "Abstract types should not appear in IR presented to backends".into(),
1176                    ));
1177                }
1178            },
1179            Expression::Constant(handle) => {
1180                let constant = &module.constants[handle];
1181                if constant.name.is_some() {
1182                    write!(self.out, "{}", self.names[&NameKey::Constant(handle)])?;
1183                } else {
1184                    self.write_const_expression(module, constant.init, &module.global_expressions)?;
1185                }
1186            }
1187            Expression::ZeroValue(ty) => {
1188                self.write_type(module, ty)?;
1189                write!(self.out, "()")?;
1190            }
1191            Expression::Compose { ty, ref components } => {
1192                self.write_type(module, ty)?;
1193                write!(self.out, "(")?;
1194                for (index, component) in components.iter().enumerate() {
1195                    if index != 0 {
1196                        write!(self.out, ", ")?;
1197                    }
1198                    write_expression(self, *component)?;
1199                }
1200                write!(self.out, ")")?
1201            }
1202            Expression::Splat { size, value } => {
1203                let size = common::vector_size_str(size);
1204                write!(self.out, "vec{size}(")?;
1205                write_expression(self, value)?;
1206                write!(self.out, ")")?;
1207            }
1208            _ => unreachable!(),
1209        }
1210
1211        Ok(())
1212    }
1213
1214    /// Write the 'plain form' of `expr`.
1215    ///
1216    /// An expression's 'plain form' is the most general rendition of that
1217    /// expression into WGSL, lacking `&` or `*` operators. The plain forms of
1218    /// `LocalVariable(x)` and `GlobalVariable(g)` are simply `x` and `g`. Such
1219    /// Naga expressions represent both WGSL pointers and references; it's the
1220    /// caller's responsibility to distinguish those cases appropriately.
1221    fn write_expr_plain_form(
1222        &mut self,
1223        module: &Module,
1224        expr: Handle<crate::Expression>,
1225        func_ctx: &back::FunctionCtx<'_>,
1226        indirection: Indirection,
1227    ) -> BackendResult {
1228        use crate::Expression;
1229
1230        if let Some(name) = self.named_expressions.get(&expr) {
1231            write!(self.out, "{name}")?;
1232            return Ok(());
1233        }
1234
1235        let expression = &func_ctx.expressions[expr];
1236
1237        // Write the plain WGSL form of a Naga expression.
1238        //
1239        // The plain form of `LocalVariable` and `GlobalVariable` expressions is
1240        // simply the variable name; `*` and `&` operators are never emitted.
1241        //
1242        // The plain form of `Access` and `AccessIndex` expressions are WGSL
1243        // `postfix_expression` forms for member/component access and
1244        // subscripting.
1245        match *expression {
1246            Expression::Literal(_)
1247            | Expression::Constant(_)
1248            | Expression::ZeroValue(_)
1249            | Expression::Compose { .. }
1250            | Expression::Splat { .. } => {
1251                self.write_possibly_const_expression(
1252                    module,
1253                    expr,
1254                    func_ctx.expressions,
1255                    |writer, expr| writer.write_expr(module, expr, func_ctx),
1256                )?;
1257            }
1258            Expression::Override(_) => unreachable!(),
1259            Expression::FunctionArgument(pos) => {
1260                let name_key = func_ctx.argument_key(pos);
1261                let name = &self.names[&name_key];
1262                write!(self.out, "{name}")?;
1263            }
1264            Expression::Binary { op, left, right } => {
1265                write!(self.out, "(")?;
1266                self.write_expr(module, left, func_ctx)?;
1267                write!(self.out, " {} ", back::binary_operation_str(op))?;
1268                self.write_expr(module, right, func_ctx)?;
1269                write!(self.out, ")")?;
1270            }
1271            Expression::Access { base, index } => {
1272                self.write_expr_with_indirection(module, base, func_ctx, indirection)?;
1273                write!(self.out, "[")?;
1274                self.write_expr(module, index, func_ctx)?;
1275                write!(self.out, "]")?
1276            }
1277            Expression::AccessIndex { base, index } => {
1278                let base_ty_res = &func_ctx.info[base].ty;
1279                let mut resolved = base_ty_res.inner_with(&module.types);
1280
1281                self.write_expr_with_indirection(module, base, func_ctx, indirection)?;
1282
1283                let base_ty_handle = match *resolved {
1284                    TypeInner::Pointer { base, space: _ } => {
1285                        resolved = &module.types[base].inner;
1286                        Some(base)
1287                    }
1288                    _ => base_ty_res.handle(),
1289                };
1290
1291                match *resolved {
1292                    TypeInner::Vector { .. } => {
1293                        // Write vector access as a swizzle
1294                        write!(self.out, ".{}", back::COMPONENTS[index as usize])?
1295                    }
1296                    TypeInner::Matrix { .. }
1297                    | TypeInner::Array { .. }
1298                    | TypeInner::BindingArray { .. }
1299                    | TypeInner::ValuePointer { .. } => write!(self.out, "[{index}]")?,
1300                    TypeInner::Struct { .. } => {
1301                        // This will never panic in case the type is a `Struct`, this is not true
1302                        // for other types so we can only check while inside this match arm
1303                        let ty = base_ty_handle.unwrap();
1304
1305                        write!(
1306                            self.out,
1307                            ".{}",
1308                            &self.names[&NameKey::StructMember(ty, index)]
1309                        )?
1310                    }
1311                    ref other => return Err(Error::Custom(format!("Cannot index {other:?}"))),
1312                }
1313            }
1314            Expression::ImageSample {
1315                image,
1316                sampler,
1317                gather: None,
1318                coordinate,
1319                array_index,
1320                offset,
1321                level,
1322                depth_ref,
1323                clamp_to_edge,
1324            } => {
1325                use crate::SampleLevel as Sl;
1326
1327                let suffix_cmp = match depth_ref {
1328                    Some(_) => "Compare",
1329                    None => "",
1330                };
1331                let suffix_level = match level {
1332                    Sl::Auto => "",
1333                    Sl::Zero if clamp_to_edge => "BaseClampToEdge",
1334                    Sl::Zero | Sl::Exact(_) => "Level",
1335                    Sl::Bias(_) => "Bias",
1336                    Sl::Gradient { .. } => "Grad",
1337                };
1338
1339                write!(self.out, "textureSample{suffix_cmp}{suffix_level}(")?;
1340                self.write_expr(module, image, func_ctx)?;
1341                write!(self.out, ", ")?;
1342                self.write_expr(module, sampler, func_ctx)?;
1343                write!(self.out, ", ")?;
1344                self.write_expr(module, coordinate, func_ctx)?;
1345
1346                if let Some(array_index) = array_index {
1347                    write!(self.out, ", ")?;
1348                    self.write_expr(module, array_index, func_ctx)?;
1349                }
1350
1351                if let Some(depth_ref) = depth_ref {
1352                    write!(self.out, ", ")?;
1353                    self.write_expr(module, depth_ref, func_ctx)?;
1354                }
1355
1356                match level {
1357                    Sl::Auto => {}
1358                    Sl::Zero => {
1359                        // Level 0 is implied for depth comparison and BaseClampToEdge
1360                        if depth_ref.is_none() && !clamp_to_edge {
1361                            write!(self.out, ", 0.0")?;
1362                        }
1363                    }
1364                    Sl::Exact(expr) => {
1365                        write!(self.out, ", ")?;
1366                        self.write_expr(module, expr, func_ctx)?;
1367                    }
1368                    Sl::Bias(expr) => {
1369                        write!(self.out, ", ")?;
1370                        self.write_expr(module, expr, func_ctx)?;
1371                    }
1372                    Sl::Gradient { x, y } => {
1373                        write!(self.out, ", ")?;
1374                        self.write_expr(module, x, func_ctx)?;
1375                        write!(self.out, ", ")?;
1376                        self.write_expr(module, y, func_ctx)?;
1377                    }
1378                }
1379
1380                if let Some(offset) = offset {
1381                    write!(self.out, ", ")?;
1382                    self.write_const_expression(module, offset, func_ctx.expressions)?;
1383                }
1384
1385                write!(self.out, ")")?;
1386            }
1387
1388            Expression::ImageSample {
1389                image,
1390                sampler,
1391                gather: Some(component),
1392                coordinate,
1393                array_index,
1394                offset,
1395                level: _,
1396                depth_ref,
1397                clamp_to_edge: _,
1398            } => {
1399                let suffix_cmp = match depth_ref {
1400                    Some(_) => "Compare",
1401                    None => "",
1402                };
1403
1404                write!(self.out, "textureGather{suffix_cmp}(")?;
1405                match *func_ctx.resolve_type(image, &module.types) {
1406                    TypeInner::Image {
1407                        class: crate::ImageClass::Depth { multi: _ },
1408                        ..
1409                    } => {}
1410                    _ => {
1411                        write!(self.out, "{}, ", component as u8)?;
1412                    }
1413                }
1414                self.write_expr(module, image, func_ctx)?;
1415                write!(self.out, ", ")?;
1416                self.write_expr(module, sampler, func_ctx)?;
1417                write!(self.out, ", ")?;
1418                self.write_expr(module, coordinate, func_ctx)?;
1419
1420                if let Some(array_index) = array_index {
1421                    write!(self.out, ", ")?;
1422                    self.write_expr(module, array_index, func_ctx)?;
1423                }
1424
1425                if let Some(depth_ref) = depth_ref {
1426                    write!(self.out, ", ")?;
1427                    self.write_expr(module, depth_ref, func_ctx)?;
1428                }
1429
1430                if let Some(offset) = offset {
1431                    write!(self.out, ", ")?;
1432                    self.write_const_expression(module, offset, func_ctx.expressions)?;
1433                }
1434
1435                write!(self.out, ")")?;
1436            }
1437            Expression::ImageQuery { image, query } => {
1438                use crate::ImageQuery as Iq;
1439
1440                let texture_function = match query {
1441                    Iq::Size { .. } => "textureDimensions",
1442                    Iq::NumLevels => "textureNumLevels",
1443                    Iq::NumLayers => "textureNumLayers",
1444                    Iq::NumSamples => "textureNumSamples",
1445                };
1446
1447                write!(self.out, "{texture_function}(")?;
1448                self.write_expr(module, image, func_ctx)?;
1449                if let Iq::Size { level: Some(level) } = query {
1450                    write!(self.out, ", ")?;
1451                    self.write_expr(module, level, func_ctx)?;
1452                };
1453                write!(self.out, ")")?;
1454            }
1455
1456            Expression::ImageLoad {
1457                image,
1458                coordinate,
1459                array_index,
1460                sample,
1461                level,
1462            } => {
1463                write!(self.out, "textureLoad(")?;
1464                self.write_expr(module, image, func_ctx)?;
1465                write!(self.out, ", ")?;
1466                self.write_expr(module, coordinate, func_ctx)?;
1467                if let Some(array_index) = array_index {
1468                    write!(self.out, ", ")?;
1469                    self.write_expr(module, array_index, func_ctx)?;
1470                }
1471                if let Some(index) = sample.or(level) {
1472                    write!(self.out, ", ")?;
1473                    self.write_expr(module, index, func_ctx)?;
1474                }
1475                write!(self.out, ")")?;
1476            }
1477            Expression::GlobalVariable(handle) => {
1478                let name = &self.names[&NameKey::GlobalVariable(handle)];
1479                write!(self.out, "{name}")?;
1480            }
1481
1482            Expression::As {
1483                expr,
1484                kind,
1485                convert,
1486            } => {
1487                let inner = func_ctx.resolve_type(expr, &module.types);
1488                match *inner {
1489                    TypeInner::Matrix {
1490                        columns,
1491                        rows,
1492                        scalar,
1493                    } => {
1494                        let scalar = crate::Scalar {
1495                            kind,
1496                            width: convert.unwrap_or(scalar.width),
1497                        };
1498                        let scalar_kind_str = scalar.to_wgsl_if_implemented()?;
1499                        write!(
1500                            self.out,
1501                            "mat{}x{}<{}>",
1502                            common::vector_size_str(columns),
1503                            common::vector_size_str(rows),
1504                            scalar_kind_str
1505                        )?;
1506                    }
1507                    TypeInner::Vector {
1508                        size,
1509                        scalar: crate::Scalar { width, .. },
1510                    } => {
1511                        let scalar = crate::Scalar {
1512                            kind,
1513                            width: convert.unwrap_or(width),
1514                        };
1515                        let vector_size_str = common::vector_size_str(size);
1516                        let scalar_kind_str = scalar.to_wgsl_if_implemented()?;
1517                        if convert.is_some() {
1518                            write!(self.out, "vec{vector_size_str}<{scalar_kind_str}>")?;
1519                        } else {
1520                            write!(self.out, "bitcast<vec{vector_size_str}<{scalar_kind_str}>>")?;
1521                        }
1522                    }
1523                    TypeInner::Scalar(crate::Scalar { width, .. }) => {
1524                        let scalar = crate::Scalar {
1525                            kind,
1526                            width: convert.unwrap_or(width),
1527                        };
1528                        let scalar_kind_str = scalar.to_wgsl_if_implemented()?;
1529                        if convert.is_some() {
1530                            write!(self.out, "{scalar_kind_str}")?
1531                        } else {
1532                            write!(self.out, "bitcast<{scalar_kind_str}>")?
1533                        }
1534                    }
1535                    _ => {
1536                        return Err(Error::Unimplemented(format!(
1537                            "write_expr expression::as {inner:?}"
1538                        )));
1539                    }
1540                };
1541                write!(self.out, "(")?;
1542                self.write_expr(module, expr, func_ctx)?;
1543                write!(self.out, ")")?;
1544            }
1545            Expression::Load { pointer } => {
1546                let is_atomic_pointer = func_ctx
1547                    .resolve_type(pointer, &module.types)
1548                    .is_atomic_pointer(&module.types);
1549
1550                if is_atomic_pointer {
1551                    write!(self.out, "atomicLoad(")?;
1552                    self.write_expr(module, pointer, func_ctx)?;
1553                    write!(self.out, ")")?;
1554                } else {
1555                    self.write_expr_with_indirection(
1556                        module,
1557                        pointer,
1558                        func_ctx,
1559                        Indirection::Reference,
1560                    )?;
1561                }
1562            }
1563            Expression::LocalVariable(handle) => {
1564                write!(self.out, "{}", self.names[&func_ctx.name_key(handle)])?
1565            }
1566            Expression::ArrayLength(expr) => {
1567                write!(self.out, "arrayLength(")?;
1568                self.write_expr(module, expr, func_ctx)?;
1569                write!(self.out, ")")?;
1570            }
1571
1572            Expression::Math {
1573                fun,
1574                arg,
1575                arg1,
1576                arg2,
1577                arg3,
1578            } => {
1579                use crate::MathFunction as Mf;
1580
1581                enum Function {
1582                    Regular(&'static str),
1583                    InversePolyfill(InversePolyfill),
1584                }
1585
1586                let function = match fun.try_to_wgsl() {
1587                    Some(name) => Function::Regular(name),
1588                    None => match fun {
1589                        Mf::Inverse => {
1590                            let ty = func_ctx.resolve_type(arg, &module.types);
1591                            let Some(overload) = InversePolyfill::find_overload(ty) else {
1592                                return Err(Error::unsupported("math function", fun));
1593                            };
1594
1595                            Function::InversePolyfill(overload)
1596                        }
1597                        _ => return Err(Error::unsupported("math function", fun)),
1598                    },
1599                };
1600
1601                match function {
1602                    Function::Regular(fun_name) => {
1603                        write!(self.out, "{fun_name}(")?;
1604                        self.write_expr(module, arg, func_ctx)?;
1605                        for arg in IntoIterator::into_iter([arg1, arg2, arg3]).flatten() {
1606                            write!(self.out, ", ")?;
1607                            self.write_expr(module, arg, func_ctx)?;
1608                        }
1609                        write!(self.out, ")")?
1610                    }
1611                    Function::InversePolyfill(inverse) => {
1612                        write!(self.out, "{}(", inverse.fun_name)?;
1613                        self.write_expr(module, arg, func_ctx)?;
1614                        write!(self.out, ")")?;
1615                        self.required_polyfills.insert(inverse);
1616                    }
1617                }
1618            }
1619
1620            Expression::Swizzle {
1621                size,
1622                vector,
1623                pattern,
1624            } => {
1625                self.write_expr(module, vector, func_ctx)?;
1626                write!(self.out, ".")?;
1627                for &sc in pattern[..size as usize].iter() {
1628                    self.out.write_char(back::COMPONENTS[sc as usize])?;
1629                }
1630            }
1631            Expression::Unary { op, expr } => {
1632                let unary = match op {
1633                    crate::UnaryOperator::Negate => "-",
1634                    crate::UnaryOperator::LogicalNot => "!",
1635                    crate::UnaryOperator::BitwiseNot => "~",
1636                };
1637
1638                write!(self.out, "{unary}(")?;
1639                self.write_expr(module, expr, func_ctx)?;
1640
1641                write!(self.out, ")")?
1642            }
1643
1644            Expression::Select {
1645                condition,
1646                accept,
1647                reject,
1648            } => {
1649                write!(self.out, "select(")?;
1650                self.write_expr(module, reject, func_ctx)?;
1651                write!(self.out, ", ")?;
1652                self.write_expr(module, accept, func_ctx)?;
1653                write!(self.out, ", ")?;
1654                self.write_expr(module, condition, func_ctx)?;
1655                write!(self.out, ")")?
1656            }
1657            Expression::Derivative { axis, ctrl, expr } => {
1658                use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl};
1659                let op = match (axis, ctrl) {
1660                    (Axis::X, Ctrl::Coarse) => "dpdxCoarse",
1661                    (Axis::X, Ctrl::Fine) => "dpdxFine",
1662                    (Axis::X, Ctrl::None) => "dpdx",
1663                    (Axis::Y, Ctrl::Coarse) => "dpdyCoarse",
1664                    (Axis::Y, Ctrl::Fine) => "dpdyFine",
1665                    (Axis::Y, Ctrl::None) => "dpdy",
1666                    (Axis::Width, Ctrl::Coarse) => "fwidthCoarse",
1667                    (Axis::Width, Ctrl::Fine) => "fwidthFine",
1668                    (Axis::Width, Ctrl::None) => "fwidth",
1669                };
1670                write!(self.out, "{op}(")?;
1671                self.write_expr(module, expr, func_ctx)?;
1672                write!(self.out, ")")?
1673            }
1674            Expression::Relational { fun, argument } => {
1675                use crate::RelationalFunction as Rf;
1676
1677                let fun_name = match fun {
1678                    Rf::All => "all",
1679                    Rf::Any => "any",
1680                    _ => return Err(Error::UnsupportedRelationalFunction(fun)),
1681                };
1682                write!(self.out, "{fun_name}(")?;
1683
1684                self.write_expr(module, argument, func_ctx)?;
1685
1686                write!(self.out, ")")?
1687            }
1688            // Not supported yet
1689            Expression::RayQueryGetIntersection { .. }
1690            | Expression::RayQueryVertexPositions { .. } => unreachable!(),
1691            // Nothing to do here, since call expression already cached
1692            Expression::CallResult(_)
1693            | Expression::AtomicResult { .. }
1694            | Expression::RayQueryProceedResult
1695            | Expression::SubgroupBallotResult
1696            | Expression::SubgroupOperationResult { .. }
1697            | Expression::WorkGroupUniformLoadResult { .. } => {}
1698        }
1699
1700        Ok(())
1701    }
1702
1703    /// Helper method used to write global variables
1704    /// # Notes
1705    /// Always adds a newline
1706    fn write_global(
1707        &mut self,
1708        module: &Module,
1709        global: &crate::GlobalVariable,
1710        handle: Handle<crate::GlobalVariable>,
1711    ) -> BackendResult {
1712        // Write group and binding attributes if present
1713        if let Some(ref binding) = global.binding {
1714            self.write_attributes(&[
1715                Attribute::Group(binding.group),
1716                Attribute::Binding(binding.binding),
1717            ])?;
1718            writeln!(self.out)?;
1719        }
1720
1721        // First write global name and address space if supported
1722        write!(self.out, "var")?;
1723        let (address, maybe_access) = address_space_str(global.space);
1724        if let Some(space) = address {
1725            write!(self.out, "<{space}")?;
1726            if let Some(access) = maybe_access {
1727                write!(self.out, ", {access}")?;
1728            }
1729            write!(self.out, ">")?;
1730        }
1731        write!(
1732            self.out,
1733            " {}: ",
1734            &self.names[&NameKey::GlobalVariable(handle)]
1735        )?;
1736
1737        // Write global type
1738        self.write_type(module, global.ty)?;
1739
1740        // Write initializer
1741        if let Some(init) = global.init {
1742            write!(self.out, " = ")?;
1743            self.write_const_expression(module, init, &module.global_expressions)?;
1744        }
1745
1746        // End with semicolon
1747        writeln!(self.out, ";")?;
1748
1749        Ok(())
1750    }
1751
1752    /// Helper method used to write global constants
1753    ///
1754    /// # Notes
1755    /// Ends in a newline
1756    fn write_global_constant(
1757        &mut self,
1758        module: &Module,
1759        handle: Handle<crate::Constant>,
1760    ) -> BackendResult {
1761        let name = &self.names[&NameKey::Constant(handle)];
1762        // First write only constant name
1763        write!(self.out, "const {name}: ")?;
1764        self.write_type(module, module.constants[handle].ty)?;
1765        write!(self.out, " = ")?;
1766        let init = module.constants[handle].init;
1767        self.write_const_expression(module, init, &module.global_expressions)?;
1768        writeln!(self.out, ";")?;
1769
1770        Ok(())
1771    }
1772
1773    // See https://github.com/rust-lang/rust-clippy/issues/4979.
1774    #[allow(clippy::missing_const_for_fn)]
1775    pub fn finish(self) -> W {
1776        self.out
1777    }
1778}
1779
1780struct WriterTypeContext<'m> {
1781    module: &'m Module,
1782    names: &'m crate::FastHashMap<NameKey, String>,
1783}
1784
1785impl TypeContext for WriterTypeContext<'_> {
1786    fn lookup_type(&self, handle: Handle<crate::Type>) -> &crate::Type {
1787        &self.module.types[handle]
1788    }
1789
1790    fn type_name(&self, handle: Handle<crate::Type>) -> &str {
1791        self.names[&NameKey::Type(handle)].as_str()
1792    }
1793
1794    fn write_unnamed_struct<W: Write>(&self, _: &TypeInner, _: &mut W) -> core::fmt::Result {
1795        unreachable!("the WGSL back end should always provide type handles");
1796    }
1797
1798    fn write_override<W: Write>(&self, _: Handle<crate::Override>, _: &mut W) -> core::fmt::Result {
1799        unreachable!("overrides should be validated out");
1800    }
1801
1802    fn write_non_wgsl_inner<W: Write>(&self, _: &TypeInner, _: &mut W) -> core::fmt::Result {
1803        unreachable!("backends should only be passed validated modules");
1804    }
1805
1806    fn write_non_wgsl_scalar<W: Write>(&self, _: crate::Scalar, _: &mut W) -> core::fmt::Result {
1807        unreachable!("backends should only be passed validated modules");
1808    }
1809}
1810
1811fn map_binding_to_attribute(binding: &crate::Binding) -> Vec<Attribute> {
1812    match *binding {
1813        crate::Binding::BuiltIn(built_in) => {
1814            if let crate::BuiltIn::Position { invariant: true } = built_in {
1815                vec![Attribute::BuiltIn(built_in), Attribute::Invariant]
1816            } else {
1817                vec![Attribute::BuiltIn(built_in)]
1818            }
1819        }
1820        crate::Binding::Location {
1821            location,
1822            interpolation,
1823            sampling,
1824            blend_src: None,
1825        } => vec![
1826            Attribute::Location(location),
1827            Attribute::Interpolate(interpolation, sampling),
1828        ],
1829        crate::Binding::Location {
1830            location,
1831            interpolation,
1832            sampling,
1833            blend_src: Some(blend_src),
1834        } => vec![
1835            Attribute::Location(location),
1836            Attribute::BlendSrc(blend_src),
1837            Attribute::Interpolate(interpolation, sampling),
1838        ],
1839    }
1840}