[spv-out] Track block termination statically.

Rather than giving `Block` an optional `terminator` field in addition to its
`body`, track block termination statically, with two types:

- `Block` is a block without a termination instruction. This is what most code
  generation functions operate on.

- `TerminatedBlock` is a block with a termination instruction. This is what
  `Function::blocks` holds.

The `Function::consume` method takes a `Block` by value, together with a
termination instruction, and turns it into a `TerminatedBlock`.

This lets us remove some unwraps and awkward conditions.

As part of this change, `Writer::write_block` no longer hits an `unimplemented!`
for Naga statements following `Break`, `Return`, and so on. Instead, it simply
doesn't emit code for them, which is a correct translation. If we want to forbid
these, we should handle that in validation instead.

src/back/spv/mod.rs

@@ -72,10 +72,26 @@ impl IdGenerator {
     }
 }
 
+/// A SPIR-V block to which we are still adding instructions.
+///
+/// A `Block` represents a SPIR-V block that does not yet have a termination
+/// instruction like `OpBranch` or `OpReturn`.
+///
+/// The `OpLabel` that starts the block is implicit. It will be emitted based on
+/// `label_id` when we write the block to a `LogicalLayout`.
+///
+/// To terminate a `Block`, pass the block and the termination instruction to
+/// `Function::consume`. This takes ownership of the `Block` and transforms it
+/// into a `TerminatedBlock`.
 struct Block {
     label_id: Word,
     body: Vec<Instruction>,
-    termination: Option<Instruction>,
+}
+
+/// A SPIR-V block that ends with a termination instruction.
+struct TerminatedBlock {
+    label_id: Word,
+    body: Vec<Instruction>,
 }
 
 impl Block {
@@ -83,7 +99,6 @@ impl Block {
         Block {
             label_id,
             body: Vec::new(),
-            termination: None,
         }
     }
 }
@@ -109,14 +124,17 @@ struct Function {
     signature: Option<Instruction>,
     parameters: Vec<Instruction>,
     variables: crate::FastHashMap<Handle<crate::LocalVariable>, LocalVariable>,
-    blocks: Vec<Block>,
+    blocks: Vec<TerminatedBlock>,
     entry_point_context: Option<EntryPointContext>,
 }
 
 impl Function {
     fn consume(&mut self, mut block: Block, termination: Instruction) {
-        block.termination = Some(termination);
-        self.blocks.push(block);
+        block.body.push(termination);
+        self.blocks.push(TerminatedBlock {
+            label_id: block.label_id,
+            body: block.body,
+        })
     }
 
     fn parameter_id(&self, index: u32) -> Word {

src/back/spv/writer.rs

@@ -47,7 +47,6 @@ impl Function {
             for instruction in block.body.iter() {
                 instruction.to_words(sink);
             }
-            block.termination.as_ref().unwrap().to_words(sink);
         }
     }
 }
@@ -2369,9 +2368,6 @@ impl Writer {
         let mut block = Block::new(label_id);
 
         for statement in statements {
-            if block.termination.is_some() {
-                unimplemented!("No statements are expected after block termination");
-            }
             match *statement {
                 crate::Statement::Emit(ref range) => {
                     for handle in range.clone() {
@@ -2576,11 +2572,15 @@ impl Writer {
                     block = Block::new(merge_id);
                 }
                 crate::Statement::Break => {
-                    block.termination = Some(Instruction::branch(loop_context.break_id.unwrap()));
+                    function.consume(block, Instruction::branch(loop_context.break_id.unwrap()));
+                    return Ok(());
                 }
                 crate::Statement::Continue => {
-                    block.termination =
-                        Some(Instruction::branch(loop_context.continuing_id.unwrap()));
+                    function.consume(
+                        block,
+                        Instruction::branch(loop_context.continuing_id.unwrap()),
+                    );
+                    return Ok(());
                 }
                 crate::Statement::Return { value: Some(value) } => {
                     let value_id = self.cached[value];
@@ -2598,13 +2598,16 @@ impl Writer {
                         }
                         None => Instruction::return_value(value_id),
                     };
-                    block.termination = Some(instruction);
+                    function.consume(block, instruction);
+                    return Ok(());
                 }
                 crate::Statement::Return { value: None } => {
-                    block.termination = Some(Instruction::return_void());
+                    function.consume(block, Instruction::return_void());
+                    return Ok(());
                 }
                 crate::Statement::Kill => {
-                    block.termination = Some(Instruction::kill());
+                    function.consume(block, Instruction::kill());
+                    return Ok(());
                 }
                 crate::Statement::Barrier(flags) => {
                     let memory_scope = if flags.contains(crate::Barrier::STORAGE) {
@@ -2734,24 +2737,22 @@ impl Writer {
             }
         }
 
-        if block.termination.is_none() {
-            block.termination = Some(match exit_id {
-                Some(id) => Instruction::branch(id),
-                // This can happen if the last branch had all the paths
-                // leading out of the graph (i.e. returning).
-                // Or it may be the end of the function.
-                None => match ir_function.result {
-                    Some(ref result) if function.entry_point_context.is_none() => {
-                        let type_id = self.get_type_id(LookupType::Handle(result.ty))?;
-                        let null_id = self.write_constant_null(type_id);
-                        Instruction::return_value(null_id)
-                    }
-                    _ => Instruction::return_void(),
-                },
-            });
-        }
+        let termination = match exit_id {
+            Some(id) => Instruction::branch(id),
+            // This can happen if the last branch had all the paths
+            // leading out of the graph (i.e. returning).
+            // Or it may be the end of the function.
+            None => match ir_function.result {
+                Some(ref result) if function.entry_point_context.is_none() => {
+                    let type_id = self.get_type_id(LookupType::Handle(result.ty))?;
+                    let null_id = self.write_constant_null(type_id);
+                    Instruction::return_value(null_id)
+                }
+                _ => Instruction::return_void(),
+            },
+        };
 
-        function.blocks.push(block);
+        function.consume(block, termination);
         Ok(())
     }
 

src/back/wgsl/writer.rs

@@ -528,7 +528,7 @@ impl<W: Write> Writer<W> {
                 if let Some(storage_class) = storage_class_str(class) {
                     write!(self.out, "<{}>", storage_class)?;
                 }
-            },
+            }
             _ => {
                 return Err(Error::Unimplemented(format!(
                     "write_value_type {:?}",