Repository: gfx-rs/wgpu (rust)

Issue:
Unable to break out of the loop in SPIR-V
<!-- Thank you for filing this! Please read the [debugging tips](https://github.com/gfx-rs/wgpu/wiki/Debugging-wgpu-Applications).
That may let you investigate on your own, or provide additional information that helps us to assist.-->

**Description**
Getting invalid SPIR-V out of Naga:
>error: line 36: Block '23[%23]' branches to the loop continue target '16[%16]', but is not contained in the associated loop construct '13[%13]'
  OpBranch %16


**Repro steps**
Program:
```wgsl
fn main() {
    for (var i = 0; i < 4; i += 1) {
      break;
    }
}
```

**Expected vs observed behavior**
Expected to produce valid SPIR-V

Pull request description:
[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 #6220.


Patch:
diff --git a/naga/src/back/spv/block.rs b/naga/src/back/spv/block.rs
index 97db1cb54..c6adba85b 100644
--- a/naga/src/back/spv/block.rs
+++ b/naga/src/back/spv/block.rs
@@ -4,8 +4,7 @@ Implementations for `BlockContext` methods.
 
 use super::{
     helpers, index::BoundsCheckResult, make_local, selection::Selection, Block, BlockContext,
-    Dimension, Error, Instruction, LocalType, LookupType, LoopContext, ResultMember, Writer,
-    WriterFlags,
+    Dimension, Error, Instruction, LocalType, LookupType, ResultMember, Writer, WriterFlags,
 };
 use crate::{arena::Handle, proc::TypeResolution, Statement};
 use spirv::Word;
@@ -39,7 +38,7 @@ enum ExpressionPointer {
 }
 
 /// The termination statement to be added to the end of the block
-pub enum BlockExit {
+enum BlockExit {
     /// Generates an OpReturn (void return)
     Return,
     /// Generates an OpBranch to the specified block
@@ -60,6 +59,36 @@ pub 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>,
+    break_id: Option<Word>,
+}
+
 #[derive(Debug)]
 pub(crate) struct DebugInfoInner<'a> {
     pub source_code: &'a str,
@@ -2060,14 +2089,30 @@ impl<'w> BlockContext<'w> {
         }
     }
 
-    pub(super) fn write_block(
+    /// 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) = (
@@ -2103,7 +2148,7 @@ 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 },
@@ -2111,7 +2156,14 @@ impl<'w> BlockContext<'w> {
                         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,
@@ -2147,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 },
@@ -2156,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 },
@@ -2234,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 {
@@ -2280,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 {
@@ -2303,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,
@@ -2319,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];
@@ -2345,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);
@@ -2719,6 +2793,24 @@ impl<'w> BlockContext<'w> {
         };
 
         self.function.consume(block, termination);
+        Ok(BlockExitDisposition::Used)
+    }
+
+    pub(super) fn write_function_body(
+        &mut self,
+        entry_id: Word,
+        debug_info: Option<&DebugInfoInner>,
+    ) -> Result<(), Error> {
+        // 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,
+            LoopContext::default(),
+            debug_info,
+        )?;
+
         Ok(())
     }
 }
diff --git a/naga/src/back/spv/mod.rs b/naga/src/back/spv/mod.rs
index 764d06c90..e6397017c 100644
--- a/naga/src/back/spv/mod.rs
+++ b/naga/src/back/spv/mod.rs
@@ -665,12 +665,6 @@ impl BlockContext<'_> {
     }
 }
 
-#[derive(Clone, Copy, Default)]
-struct LoopContext {
-    continuing_id: Option<Word>,
-    break_id: Option<Word>,
-}
-
 pub struct Writer {
     physical_layout: PhysicalLayout,
     logical_layout: LogicalLayout,
diff --git a/naga/src/back/spv/writer.rs b/naga/src/back/spv/writer.rs
index 7f41caf29..14f1fc002 100644
--- a/naga/src/back/spv/writer.rs
+++ b/naga/src/back/spv/writer.rs
@@ -3,7 +3,7 @@ use super::{
     helpers::{contains_builtin, global_needs_wrapper, map_storage_class},
     make_local, Block, BlockContext, CachedConstant, CachedExpressions, DebugInfo,
     EntryPointContext, Error, Function, FunctionArgument, GlobalVariable, IdGenerator, Instruction,
-    LocalType, LocalVariable, LogicalLayout, LookupFunctionType, LookupType, LoopContext, Options,
+    LocalType, LocalVariable, LogicalLayout, LookupFunctionType, LookupType, Options,
     PhysicalLayout, PipelineOptions, ResultMember, Writer, WriterFlags, BITS_PER_BYTE,
 };
 use crate::{
@@ -756,13 +756,7 @@ impl Writer {
             next_id
         };
 
-        context.write_block(
-            main_id,
-            &ir_function.body,
-            super::block::BlockExit::Return,
-            LoopContext::default(),
-            debug_info.as_ref(),
-        )?;
+        context.write_function_body(main_id, debug_info.as_ref())?;
 
         // Consume the `BlockContext`, ending its borrows and letting the
         // `Writer` steal back its cached expression table and temp_list.


Test patch:
diff --git a/naga/tests/in/6220-break-from-loop.param.ron b/naga/tests/in/6220-break-from-loop.param.ron
new file mode 100644
index 000000000..72873dd66
--- /dev/null
+++ b/naga/tests/in/6220-break-from-loop.param.ron
@@ -0,0 +1,2 @@
+(
+)
diff --git a/naga/tests/in/6220-break-from-loop.wgsl b/naga/tests/in/6220-break-from-loop.wgsl
new file mode 100644
index 000000000..424886a75
--- /dev/null
+++ b/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;
+    }
+}
diff --git a/naga/tests/out/spv/6220-break-from-loop.spvasm b/naga/tests/out/spv/6220-break-from-loop.spvasm
new file mode 100644
index 000000000..9dabbc207
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/naga/tests/snapshots.rs b/naga/tests/snapshots.rs
index 596e4cea1..adf67f833 100644
--- a/naga/tests/snapshots.rs
+++ b/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() {