[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

@@ -59,6 +59,30 @@ enum BlockExit {
     },
 }
 
+/// What code generation did with a provided [`BlockExit`] value.
+///
+/// A function that accepts a [`BlockExit`] argument should return a value of
+/// this type, to indicate whether the code it generated ended up using the
+/// provided exit, or ignored it and did a non-local exit of some other kind
+/// (say, [`Break`] or [`Continue`]). Some callers must use this information to
+/// decide whether to generate the target block at all.
+///
+/// [`Break`]: Statement::Break
+/// [`Continue`]: Statement::Continue
+#[must_use]
+enum BlockExitDisposition {
+    /// The generated code used the provided `BlockExit` value. If it included a
+    /// block label, the caller should be sure to actually emit the block it
+    /// refers to.
+    Used,
+
+    /// The generated code did not use the provided `BlockExit` value. If it
+    /// included a block label, the caller should not bother to actually emit
+    /// the block it refers to, unless it knows the block is needed for
+    /// something else.
+    Discarded,
+}
+
 #[derive(Clone, Copy, Default)]
 struct LoopContext {
     continuing_id: Option<Word>,
@@ -2065,14 +2089,30 @@ 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`. This function's return value indicates whether it acted on
+    /// this parameter or not; see [`BlockExitDisposition`].
+    ///
+    /// 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.
+    ///
+    /// [`Break`]: Statement::Break
+    /// [`Continue`]: Statement::Continue
     fn write_block(
         &mut self,
         label_id: Word,
         naga_block: &crate::Block,
         exit: BlockExit,
         loop_context: LoopContext,
         debug_info: Option<&DebugInfoInner>,
-    ) -> Result<(), Error> {
+    ) -> Result<BlockExitDisposition, Error> {
         let mut block = Block::new(label_id);
         for (statement, span) in naga_block.span_iter() {
             if let (Some(debug_info), false) = (
@@ -2108,15 +2148,22 @@ 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,
                     )?;
 
-                    block = Block::new(merge_id);
+                    match merge_used {
+                        BlockExitDisposition::Used => {
+                            block = Block::new(merge_id);
+                        }
+                        BlockExitDisposition::Discarded => {
+                            return Ok(BlockExitDisposition::Discarded);
+                        }
+                    }
                 }
                 Statement::If {
                     condition,
@@ -2152,7 +2199,11 @@ impl<'w> BlockContext<'w> {
                     );
 
                     if let Some(block_id) = accept_id {
-                        self.write_block(
+                        // We can ignore the `BlockExitDisposition` returned here because,
+                        // even if `merge_id` is not actually reachable, it is always
+                        // referred to by the `OpSelectionMerge` instruction we emitted
+                        // earlier.
+                        let _ = self.write_block(
                             block_id,
                             accept,
                             BlockExit::Branch { target: merge_id },
@@ -2161,7 +2212,11 @@ impl<'w> BlockContext<'w> {
                         )?;
                     }
                     if let Some(block_id) = reject_id {
-                        self.write_block(
+                        // We can ignore the `BlockExitDisposition` returned here because,
+                        // even if `merge_id` is not actually reachable, it is always
+                        // referred to by the `OpSelectionMerge` instruction we emitted
+                        // earlier.
+                        let _ = self.write_block(
                             block_id,
                             reject,
                             BlockExit::Branch { target: merge_id },
@@ -2239,7 +2294,15 @@ impl<'w> BlockContext<'w> {
                         } else {
                             merge_id
                         };
-                        self.write_block(
+                        // We can ignore the `BlockExitDisposition` returned here because
+                        // `case_finish_id` is always referred to by either:
+                        //
+                        // - the `OpSwitch`, if it's the next case's label for a
+                        //   fall-through, or
+                        //
+                        // - the `OpSelectionMerge`, if it's the switch's overall merge
+                        //   block because there's no fall-through.
+                        let _ = self.write_block(
                             *label_id,
                             &case.body,
                             BlockExit::Branch {
@@ -2285,7 +2348,10 @@ impl<'w> BlockContext<'w> {
                     ));
                     self.function.consume(block, Instruction::branch(body_id));
 
-                    self.write_block(
+                    // We can ignore the `BlockExitDisposition` returned here because,
+                    // even if `continuing_id` is not actually reachable, it is always
+                    // referred to by the `OpLoopMerge` instruction we emitted earlier.
+                    let _ = self.write_block(
                         body_id,
                         body,
                         BlockExit::Branch {
@@ -2308,7 +2374,10 @@ impl<'w> BlockContext<'w> {
                         },
                     };
 
-                    self.write_block(
+                    // We can ignore the `BlockExitDisposition` returned here because,
+                    // even if `merge_id` is not actually reachable, it is always referred
+                    // to by the `OpLoopMerge` instruction we emitted earlier.
+                    let _ = self.write_block(
                         continuing_id,
                         continuing,
                         exit,
@@ -2324,14 +2393,14 @@ impl<'w> BlockContext<'w> {
                 Statement::Break => {
                     self.function
                         .consume(block, Instruction::branch(loop_context.break_id.unwrap()));
-                    return Ok(());
+                    return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Continue => {
                     self.function.consume(
                         block,
                         Instruction::branch(loop_context.continuing_id.unwrap()),
                     );
-                    return Ok(());
+                    return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Return { value: Some(value) } => {
                     let value_id = self.cached[value];
@@ -2350,15 +2419,15 @@ impl<'w> BlockContext<'w> {
                         None => Instruction::return_value(value_id),
                     };
                     self.function.consume(block, instruction);
-                    return Ok(());
+                    return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Return { value: None } => {
                     self.function.consume(block, Instruction::return_void());
-                    return Ok(());
+                    return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Kill => {
                     self.function.consume(block, Instruction::kill());
-                    return Ok(());
+                    return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Barrier(flags) => {
                     self.writer.write_barrier(flags, &mut block);
@@ -2724,15 +2793,17 @@ impl<'w> BlockContext<'w> {
         };
 
         self.function.consume(block, termination);
-        Ok(())
+        Ok(BlockExitDisposition::Used)
     }
 
     pub(super) fn write_function_body(
         &mut self,
         entry_id: Word,
         debug_info: Option<&DebugInfoInner>,
     ) -> Result<(), Error> {
-        self.write_block(
+        // We can ignore the `BlockExitDisposition` returned here because
+        // `BlockExit::Return` doesn't refer to a block.
+        let _ = self.write_block(
             entry_id,
             &self.ir_function.body,
             super::block::BlockExit::Return,

naga/tests/in/6220-break-from-loop.param.ron

@@ -0,0 +1,2 @@
+(
+)

naga/tests/in/6220-break-from-loop.wgsl

@@ -0,0 +1,43 @@
+// #6220: Don't generate unreachable SPIR-V blocks that branch into
+// structured control flow constructs.
+//
+// Suppose we have Naga code like this:
+//
+//     Block {
+//       ... prelude
+//       Block { ... nested }
+//       ... postlude
+//     }
+//
+// The SPIR-V back end used to always generate three separate SPIR-V
+// blocks for the sections labeled "prelude", "nested", and
+// "postlude", each block ending with a branch to the next, even if
+// they were empty.
+//
+// However, the function below generates code that includes the
+// following structure:
+//
+//     Loop {
+//       body: Block {
+//         ... prelude
+//         Block { Break }
+//         ... postlude
+//       }
+//       continuing: ...
+//     }
+//
+// In this case, even though the `Break` renders the "postlude"
+// unreachable, we used to generate a SPIR-V block for it anyway,
+// ending with a branch to the `Loop`'s "continuing" block. However,
+// SPIR-V's structured control flow rules forbid branches to a loop
+// construct's continue target from outside the loop, so the SPIR-V
+// module containing the unreachable block didn't pass validation.
+//
+// One might assume that unreachable blocks shouldn't affect
+// validation, but the spec doesn't clearly agree, and this doesn't
+// seem to be the way validation has been implemented.
+fn break_from_loop() {
+    for (var i = 0; i < 4; i += 1) {
+      break;
+    }
+}

naga/tests/out/spv/6220-break-from-loop.spvasm

@@ -0,0 +1,44 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 26
+OpCapability Shader
+OpCapability Linkage
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+%2 = OpTypeVoid
+%3 = OpTypeInt 32 1
+%6 = OpTypeFunction %2
+%7 = OpConstant  %3  0
+%8 = OpConstant  %3  4
+%9 = OpConstant  %3  1
+%11 = OpTypePointer Function %3
+%18 = OpTypeBool
+%5 = OpFunction  %2  None %6
+%4 = OpLabel
+%10 = OpVariable  %11  Function %7
+OpBranch %12
+%12 = OpLabel
+OpBranch %13
+%13 = OpLabel
+OpLoopMerge %14 %16 None
+OpBranch %15
+%15 = OpLabel
+%17 = OpLoad  %3  %10
+%19 = OpSLessThan  %18  %17 %8
+OpSelectionMerge %20 None
+OpBranchConditional %19 %20 %21
+%21 = OpLabel
+OpBranch %14
+%20 = OpLabel
+OpBranch %22
+%22 = OpLabel
+OpBranch %14
+%16 = OpLabel
+%24 = OpLoad  %3  %10
+%25 = OpIAdd  %3  %24 %9
+OpStore %10 %25
+OpBranch %13
+%14 = OpLabel
+OpReturn
+OpFunctionEnd

naga/tests/snapshots.rs

@@ -932,6 +932,7 @@ fn convert_wgsl() {
             "phony_assignment",
             Targets::SPIRV | Targets::METAL | Targets::GLSL | Targets::HLSL | Targets::WGSL,
         ),
+        ("6220-break-from-loop", Targets::SPIRV),
     ];
 
     for &(name, targets) in inputs.iter() {