Remove depth counter in the continue context in favor of pushing/popping

from the stack.

naga/src/back/continue_forward.rs

@@ -59,32 +59,40 @@
 //! [`SwitchCase`]: crate::SwitchCase
 
 use crate::proc::Namer;
+use std::rc::Rc;
 
 /// 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).
+    /// * When entering a nested switch, add a [`Switch`] state to the stack with a new variable
+    ///   name.
+    /// * When entering an inner loop, add a [`Loop`] state to the stack.
+    /// * When exiting the loop, pop from the stack.
     ///
     /// [`Switch`]: Nesting::Switch
-    Loop { depth: u32 },
+    /// [`Switch`]: Nesting::Loop
+    Loop,
     /// 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).
+    /// * When entering a nested loop, add a [`Loop`] state to the stack.
+    /// * When entering a nested switch, add a [`Switch`] state to the stack with a clone of the
+    ///   variable name.
+    /// * When encountering a continue statement, record that the switch contains a continue.
+    /// * When exiting the switch, pop from the stack. If there is a [`Switch`] on top of the new
+    ///   stack, propagate if a continue was encountered.
+    ///
+    /// [`Switch`]: Nesting::Switch
+    /// [`Switch`]: Nesting::Loop
     Switch {
-        depth: u32,
-        variable: String,
+        variable: Rc<String>,
         // Used to skip generating checks of the continue variable if no continue statements are
         // found nested in this switch.
         //
@@ -98,7 +106,7 @@ pub(crate) enum ExitControlFlow {
     None,
     /// Emit `if (continue_variable) { continue; }`
     Continue {
-        variable: String,
+        variable: Rc<String>,
     },
     /// Emit `if (continue_variable) { break; }`
     ///
@@ -107,7 +115,7 @@ pub(crate) enum ExitControlFlow {
     /// 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: String,
+        variable: Rc<String>,
     },
 }
 
@@ -130,27 +138,17 @@ impl ContinueCtx {
 
     /// 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,
-        }
+        self.stack.push(Nesting::Loop);
     }
 
     /// Updates internal state to record exiting a loop.
     pub fn exit_loop(&mut self) {
-        match self.stack.last_mut() {
+        match self.stack.pop() {
             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 { .. }) => {
+            Some(Nesting::Loop) => {}
+            Some(Nesting::Switch { .. }) => {
                 log::error!("Unexpected switch state when exiting loop");
             }
         }
@@ -161,22 +159,24 @@ impl ContinueCtx {
     ///
     /// `variable` is a unique variable name for storing if a continue is encountered in this
     /// switch.
-    pub fn enter_switch(&mut self, namer: &mut Namer) -> Option<String> {
-        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.
+    pub fn enter_switch(&mut self, namer: &mut Namer) -> Option<Rc<String>> {
+        match self.stack.last() {
+            // If stack is empty we are not in loop. So we don't need a variable for forwarding
+            // continue statements when writing a switch.
             None => None,
-            Some(&mut Nesting::Loop { .. }) => {
-                let variable = namer.call("should_continue");
+            Some(&Nesting::Loop { .. }) => {
+                let variable = Rc::new(namer.call("should_continue"));
                 self.stack.push(Nesting::Switch {
-                    depth: 1,
-                    variable: variable.clone(),
+                    variable: Rc::clone(&variable),
                     continue_encountered: false,
                 });
                 Some(variable)
             }
-            Some(&mut Nesting::Switch { ref mut depth, .. }) => {
-                *depth += 1;
+            Some(&Nesting::Switch { ref variable, .. }) => {
+                self.stack.push(Nesting::Switch {
+                    variable: Rc::clone(variable),
+                    continue_encountered: false,
+                });
                 // We already have a variable we can use.
                 None
             }
@@ -186,34 +186,33 @@ impl ContinueCtx {
     /// 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() {
+        match self.stack.pop() {
             None => ExitControlFlow::None,
-            Some(&mut Nesting::Loop { .. }) => {
+            Some(Nesting::Loop { .. }) => {
                 log::error!("Unexpected loop state when exiting switch");
                 ExitControlFlow::None
             }
-            Some(&mut Nesting::Switch {
-                ref mut depth,
-                ref variable,
-                continue_encountered,
+            Some(Nesting::Switch {
+                variable,
+                continue_encountered: inner_continue,
             }) => {
-                *depth -= 1;
-                if *depth == 0 {
-                    if let Some(Nesting::Switch { variable, .. }) = self.stack.pop() {
-                        if continue_encountered {
-                            ExitControlFlow::Continue { variable }
-                        } else {
-                            ExitControlFlow::None
-                        }
-                    } else {
-                        unreachable!()
-                    }
-                } else if continue_encountered {
-                    // Switch was nested in another switch that we need to break out of first
-                    // before we can continue.
-                    ExitControlFlow::Break {
-                        variable: variable.clone(),
-                    }
+                // Check if this is nested in another switch.
+                let control_flow = if let Some(&mut Nesting::Switch {
+                    ref mut continue_encountered,
+                    ..
+                }) = self.stack.last_mut()
+                {
+                    // Propagate up that there is a continue statement present.
+                    *continue_encountered |= inner_continue;
+                    ExitControlFlow::Break { variable }
+                } else {
+                    ExitControlFlow::Continue { variable }
+                };
+
+                // We only need to emit checks and control flow if a continue statement is actualy
+                // present,
+                if inner_continue {
+                    control_flow
                 } else {
                     ExitControlFlow::None
                 }
@@ -231,11 +230,10 @@ impl ContinueCtx {
         if let Some(&mut Nesting::Switch {
             ref variable,
             ref mut continue_encountered,
-            ..
         }) = self.stack.last_mut()
         {
             *continue_encountered = true;
-            Some(&*variable)
+            Some(variable)
         } else {
             None
         }

naga/tests/out/hlsl/control-flow.hlsl

@@ -149,9 +149,6 @@ void loop_switch_omit_continue_variable_checks(int x_2, int y_1, int z_1, int w)
                                 break;
                             }
                         }
-                        if (should_continue_6) {
-                            break;
-                        }
                         break;
                     }
                 }