[naga spv-out] Don't emit unreachable blocks that jump into loops.

When generating SPIR-V, avoid generating unreachable blocks following
statements like `break`, `return`, and so on that cause non-local
exits. These unreachable blocks can cause SPIR-V validation to fail.

Fixes gfx-rs#6220.

naga/src/back/spv/block.rs

@@ -2034,14 +2034,33 @@ impl<'w> BlockContext<'w> {
         }
     }
 
+    /// Generate one or more SPIR-V blocks for `naga_block`.
+    ///
+    /// Use `label_id` as the label for the SPIR-V entry point block.
+    ///
+    /// If control reaches the end of the SPIR-V block, terminate it according
+    /// to `exit`.
+    ///
+    /// If the block contains [`Break`] or [`Continue`] statements,
+    /// `loop_context` supplies the labels of the SPIR-V blocks to jump to. If
+    /// either of these labels are `None`, then it should have been a Naga
+    /// validation error for the corresponding statement to occur in this
+    /// context.
+    ///
+    /// Return `true` if control can reach the end of the block. This means that
+    /// `exit` was used. Otherwise, return `false`, indicating that `exit` was
+    /// unused.
+    ///
+    /// [`Break`]: Statement::Break
+    /// [`Continue`]: Statement::Continue
     pub(super) fn write_block(
         &mut self,
         label_id: Word,
         naga_block: &crate::Block,
         exit: BlockExit,
         loop_context: LoopContext,
         debug_info: Option<&DebugInfoInner>,
-    ) -> Result<(), Error> {
+    ) -> Result<bool, Error> {
         let mut block = Block::new(label_id);
         for (statement, span) in naga_block.span_iter() {
             if let (Some(debug_info), false) = (
@@ -2077,14 +2096,60 @@ impl<'w> BlockContext<'w> {
                     self.function.consume(block, Instruction::branch(scope_id));
 
                     let merge_id = self.gen_id();
-                    self.write_block(
+                    let merge_used = self.write_block(
                         scope_id,
                         block_statements,
                         BlockExit::Branch { target: merge_id },
                         loop_context,
                         debug_info,
                     )?;
 
+                    if !merge_used {
+                        // Avoid generating unreachable SPIR-V blocks, because they might
+                        // try to jump back into loops in violation of SPIR-V's structured
+                        // control flow rules.
+                        //
+                        // For example, the WGSL front end generates code like this:
+                        //
+                        //     Loop {
+                        //       body: Block {
+                        //         If i < 4 {
+                        //           accept: Block { }
+                        //           reject: Block { Break }
+                        //         }
+                        //         Block { Break }
+                        //       }
+                        //       continuing: Block {
+                        //         i += 1;
+                        //       }
+                        //     }
+                        //
+                        // Suppose this call to `write_block` is handling the `Loop` body.
+                        // Thus, our `exit` argument is a `BlockExit::Branch`, jumping to
+                        // the `Loop`'s continuing block.
+                        //
+                        // When this call reaches the `Block { Break }`, the call to
+                        // `write_block` above will handle it. However, that recursive
+                        // call will return without ever using the `BlockExit` we pass to
+                        // it, since the block has a `Break`.
+                        //
+                        // If we ignore `merge_used` and naively start a new block, this
+                        // new block will be unreachable. However, when we finish the
+                        // current Naga `Block`, we will terminate this new unreachable
+                        // block according to our `exit` value, and try to branch to the
+                        // loop's continuing block.
+                        //
+                        // The result is that we will introduce an unreachable SPIR-V
+                        // block that branches into the loop construct.
+                        //
+                        // It seems to me that if a block is unreachable, its behavior
+                        // shouldn't affect the validity of the SPIR-V module. The newer
+                        // versions of the SPIR-V specification talk about "structural
+                        // reachability" that I would have expected to address this. But
+                        // the SPIR-V validator still complains.
+                        return Ok(false);
+                    }
+
                     block = Block::new(merge_id);
                 }
                 Statement::If {
@@ -2293,14 +2358,14 @@ impl<'w> BlockContext<'w> {
                 Statement::Break => {
                     self.function
                         .consume(block, Instruction::branch(loop_context.break_id.unwrap()));
-                    return Ok(());
+                    return Ok(false);
                 }
                 Statement::Continue => {
                     self.function.consume(
                         block,
                         Instruction::branch(loop_context.continuing_id.unwrap()),
                     );
-                    return Ok(());
+                    return Ok(false);
                 }
                 Statement::Return { value: Some(value) } => {
                     let value_id = self.cached[value];
@@ -2319,15 +2384,15 @@ impl<'w> BlockContext<'w> {
                         None => Instruction::return_value(value_id),
                     };
                     self.function.consume(block, instruction);
-                    return Ok(());
+                    return Ok(false);
                 }
                 Statement::Return { value: None } => {
                     self.function.consume(block, Instruction::return_void());
-                    return Ok(());
+                    return Ok(false);
                 }
                 Statement::Kill => {
                     self.function.consume(block, Instruction::kill());
-                    return Ok(());
+                    return Ok(false);
                 }
                 Statement::Barrier(flags) => {
                     self.writer.write_barrier(flags, &mut block);
@@ -2693,6 +2758,6 @@ impl<'w> BlockContext<'w> {
         };
 
         self.function.consume(block, termination);
-        Ok(())
+        Ok(true)
     }
 }