Workaround shader compiler bugs with degenerate switch statements

CHANGELOG.md

@@ -176,6 +176,11 @@ By @atlv24 in [#5383](https://github.com/gfx-rs/wgpu/pull/5383)
 
 - In spv-out don't decorate a `BindingArray`'s type with `Block` if the type is a struct with a runtime array by @Vecvec in [#5776](https://github.com/gfx-rs/wgpu/pull/5776)
 
+#### Naga
+
+- Work around shader consumers that have bugs handling `switch` statements with a single body for all cases. These are now written as `do {} while(false);` loops in hlsl-out and glsl-out. By @Imberflur in [#5654](https://github.com/gfx-rs/wgpu/pull/5654)
+- In hlsl-out, defer `continue` statements in switches by setting a flag and breaking from the switch. This allows such constructs to work with FXC which does not support `continue` within a switch. By @Imberflur in [#5654](https://github.com/gfx-rs/wgpu/pull/5654)
+
 ## v0.20.0 (2024-04-28)
 
 ### Major Changes

naga/src/back/continue_forward.rs

@@ -0,0 +1,207 @@
+//! Workarounds for platform bugs in switches and loops.
+//!
+//! In these docs, we use CamelCase links for Naga IR concepts, and ordinary
+//! `code` formatting for HLSL or GLSL concepts.
+//!
+//! ## Avoiding `continue` within `switch`
+//!
+//! As described in <https://github.com/gfx-rs/wgpu/issues/4485>, the FXC HLSL
+//! compiler doesn't allow `continue` statements within `switch` statements, but
+//! Naga IR does. We work around this by introducing synthetic boolean local
+//! variables and branches.
+//!
+//! Specifically:
+//!
+//! - We generate code for [`Continue`] statements that sets a `bool` local to
+//!   `true` and does a `break`.
+//!
+//! - When generating code for a [`Switch`] statement, we conservatively assume
+//!   it might contain such a [`Continue`] statement, so:
+//!
+//!   - If it's the outermost such [`Switch`] within the targeted [`Loop`], we
+//!     declare the `bool` local, initialized to `false`, and immediately after
+//!     the `switch`, we check the local and do a `continue` if it's set.
+//!
+//!   - If the [`Switch`] is nested within other [`Switch`]es, then after the
+//!     generated `switch`, we check the local (which we know was declared
+//!     before the surrounding `switch`) and do a `break` if it's set.
+//!
+//! So while we "weaken" the [`Continue`] statement into a `break` statement, we
+//! also place checks immediately at the locations to which those `break`
+//! statements will jump, until we can be sure we've reached the intended target
+//! of the original [`Continue`].
+//!
+//! In the case of nested [`Loop`] and [`Switch`] statements, there may be
+//! multiple `bool` locals in scope, but there's no problem knowing which one to
+//! operate on. At any point, there is at most one [`Loop`] statement that could
+//! be targeted by a [`Continue`] statement, so the correct `bool` local to set
+//! and test is always the one introduced for the innermost enclosing [`Loop`]'s
+//! outermost [`Switch`].
+//!
+//! # Avoiding single body `switch` statements
+//!
+//! As described in <https://github.com/gfx-rs/wgpu/issues/4514>, some language
+//! front ends miscompile `switch` statements where all cases branch to the same
+//! body. Our HLSL and GLSL backends write [`Switch`] statements with a single
+//! [`SwitchCase`] as `do {} while(false);` loops.
+//!
+//! However, this rewriting introduces a new loop that could "capture"
+//! `continue` statements in its body. To avoid doing so, we apply the
+//! [`Continue`]-to-`break` transformation described above.
+//!
+//! [`Continue`]: crate::Statement::Continue
+//! [`Loop`]: crate::Statement::Loop
+//! [`Switch`]: crate::Statement::Switch
+//! [`SwitchCase`]: crate::SwitchCase
+
+/// A summary of the code surrounding a statement.
+enum Nesting {
+    /// Currently nested in at least one loop.
+    ///
+    /// `continue` should apply to the current loop.
+    ///
+    /// * When entering a nested switch, add a [`Switch`] state to the stack.
+    /// * When entering an inner loop, increment the depth.
+    /// * When exiting the loop, decrement the depth (and pop if it reaches 0).
+    ///
+    /// [`Switch`]: Nesting::Switch
+    Loop { depth: u32 },
+    /// Currently nested in at least one switch that needs to forward continues.
+    ///
+    /// This includes switches transformed into `do {} while(false)` loops, but doesn't need to
+    /// include regular switches in backends that can support `continue` within switches.
+    ///
+    /// `continue` should be forwarded to surrounding loop.
+    ///
+    /// * When entering a nested loop, add a `Loop` state to the stack.
+    /// * When entering an inner switch, increment the depth.
+    /// * When exiting the switch, decrement the depth (and pop if it reaches 0).
+    Switch { depth: u32, variable_id: u32 },
+}
+
+pub(crate) enum ExitControlFlow {
+    None,
+    /// Emit `if (continue_variable) { continue; }`
+    Continue {
+        variable_id: u32,
+    },
+    /// Emit `if (continue_variable) { break; }`
+    ///
+    /// Used when nesting switches.
+    ///
+    /// Outer switch will be exited by the break, and then its associated check will check this
+    /// same variable and see that it is set.
+    Break {
+        variable_id: u32,
+    },
+}
+
+/// Utility for tracking nesting of loops and switches to orchestrate forwarding of continue
+/// statements inside of a switch to the enclosing loop.
+///
+/// See [module docs](self) for why we need this.
+#[derive(Default)]
+pub(crate) struct ContinueCtx {
+    stack: Vec<Nesting>,
+    next_id: u32,
+}
+
+impl ContinueCtx {
+    /// Resets internal state completely including the ID generation used for unique variable
+    /// names.
+    ///
+    /// Use this to reuse memory between writing sessions.
+    pub fn clear(&mut self) {
+        self.next_id = 0;
+        self.stack.clear();
+    }
+
+    /// Updates internal state to record entering a loop.
+    pub fn enter_loop(&mut self) {
+        match self.stack.last_mut() {
+            None | Some(&mut Nesting::Switch { .. }) => {
+                self.stack.push(Nesting::Loop { depth: 1 });
+            }
+            Some(&mut Nesting::Loop { ref mut depth }) => *depth += 1,
+        }
+    }
+
+    /// Updates internal state to record exiting a loop.
+    pub fn exit_loop(&mut self) {
+        match self.stack.last_mut() {
+            None => {
+                log::error!("Unexpected empty stack when exiting loop");
+            }
+            Some(&mut Nesting::Loop { ref mut depth }) => {
+                *depth -= 1;
+                if *depth == 0 {
+                    self.stack.pop();
+                }
+            }
+            Some(&mut Nesting::Switch { .. }) => {
+                log::error!("Unexpected switch state when exiting loop");
+            }
+        }
+    }
+
+    /// Updates internal state and returns `Some(variable_id)` if nested in a loop and a new
+    /// variable needs to be declared for forwarding continues.
+    ///
+    /// `variable_id` can be used to derive a unique variable name.
+    pub fn enter_switch(&mut self) -> Option<u32> {
+        match self.stack.last_mut() {
+            // If stack is empty we are not in loop. So we need a variable for forwarding continue
+            // statements when writing a switch.
+            None => None,
+            Some(&mut Nesting::Loop { .. }) => {
+                let variable_id = self.next_id;
+                // Always increment to avoid conflicting variable names from adjacent switches.
+                self.next_id += 1;
+                self.stack.push(Nesting::Switch {
+                    depth: 1,
+                    variable_id,
+                });
+                Some(variable_id)
+            }
+            Some(&mut Nesting::Switch { ref mut depth, .. }) => {
+                *depth += 1;
+                // We already have a variable we can use.
+                None
+            }
+        }
+    }
+
+    /// Updates internal state and returns whether this switch needs to be followed by a statement
+    /// to forward continues.
+    pub fn exit_switch(&mut self) -> ExitControlFlow {
+        match self.stack.last_mut() {
+            None => ExitControlFlow::None,
+            Some(&mut Nesting::Loop { .. }) => {
+                log::error!("Unexpected loop state when exiting switch");
+                ExitControlFlow::None
+            }
+            Some(&mut Nesting::Switch {
+                ref mut depth,
+                variable_id,
+            }) => {
+                *depth -= 1;
+                if *depth == 0 {
+                    self.stack.pop();
+                    ExitControlFlow::Continue { variable_id }
+                } else {
+                    ExitControlFlow::Break { variable_id }
+                }
+            }
+        }
+    }
+
+    /// Checks if a continue statement can be emitted directly (i.e. not in a switch) or if it
+    /// needs to be forwarded via setting a variable to `true` and breaking out of the switch.
+    pub fn needs_forwarding(&self) -> Option<u32> {
+        if let Some(&Nesting::Switch { variable_id, .. }) = self.stack.last() {
+            Some(variable_id)
+        } else {
+            None
+        }
+    }
+}