IrLoader gracefully handles incomplete blocks and functions

This lets us write smaller test cases with the IrLoader, avoiding
boilerplate for function begin/end, and basic block begin/end.

Also ForEachInst is more forgiving of cases where a basic block
doesn't have a label, and when a function doesn't have a defining
or end instruction.

source/opt/basic_block.h

@@ -87,15 +87,16 @@ inline void BasicBlock::AddInstruction(std::unique_ptr<Instruction> i) {
 
 inline void BasicBlock::ForEachInst(const std::function<void(Instruction*)>& f,
                                     bool run_on_debug_line_insts) {
-  label_->ForEachInst(f, run_on_debug_line_insts);
+  if (label_) label_->ForEachInst(f, run_on_debug_line_insts);
   for (auto& inst : insts_) inst->ForEachInst(f, run_on_debug_line_insts);
 }
 
 inline void BasicBlock::ForEachInst(
     const std::function<void(const Instruction*)>& f,
     bool run_on_debug_line_insts) const {
-  static_cast<const Instruction*>(label_.get())
-      ->ForEachInst(f, run_on_debug_line_insts);
+  if (label_)
+    static_cast<const Instruction*>(label_.get())
+        ->ForEachInst(f, run_on_debug_line_insts);
   for (const auto& inst : insts_)
     static_cast<const Instruction*>(inst.get())
         ->ForEachInst(f, run_on_debug_line_insts);

source/opt/function.cpp

@@ -31,16 +31,17 @@ namespace ir {
 
 void Function::ForEachInst(const std::function<void(Instruction*)>& f,
                            bool run_on_debug_line_insts) {
-  def_inst_->ForEachInst(f, run_on_debug_line_insts);
+  if (def_inst_) def_inst_->ForEachInst(f, run_on_debug_line_insts);
   for (auto& param : params_) param->ForEachInst(f, run_on_debug_line_insts);
   for (auto& bb : blocks_) bb->ForEachInst(f, run_on_debug_line_insts);
-  end_inst_->ForEachInst(f, run_on_debug_line_insts);
+  if (end_inst_) end_inst_->ForEachInst(f, run_on_debug_line_insts);
 }
 
 void Function::ForEachInst(const std::function<void(const Instruction*)>& f,
                            bool run_on_debug_line_insts) const {
-  static_cast<const Instruction*>(def_inst_.get())
-      ->ForEachInst(f, run_on_debug_line_insts);
+  if (def_inst_)
+    static_cast<const Instruction*>(def_inst_.get())
+        ->ForEachInst(f, run_on_debug_line_insts);
 
   for (const auto& param : params_)
     static_cast<const Instruction*>(param.get())
@@ -50,8 +51,9 @@ void Function::ForEachInst(const std::function<void(const Instruction*)>& f,
     static_cast<const BasicBlock*>(bb.get())->ForEachInst(
         f, run_on_debug_line_insts);
 
-  static_cast<const Instruction*>(end_inst_.get())
-      ->ForEachInst(f, run_on_debug_line_insts);
+  if (end_inst_)
+    static_cast<const Instruction*>(end_inst_.get())
+        ->ForEachInst(f, run_on_debug_line_insts);
 }
 
 }  // namespace ir

source/opt/ir_loader.cpp

@@ -106,6 +106,20 @@ void IrLoader::AddInstruction(const spv_parsed_instruction_t* inst) {
 // Resolves internal references among the module, functions, basic blocks, etc.
 // This function should be called after adding all instructions.
 void IrLoader::EndModule() {
+  if (block_ && function_) {
+    // We're in the middle of a basic block, but the terminator is missing.
+    // Register the block anyway.  This lets us write tests with less
+    // boilerplate.
+    function_->AddBasicBlock(std::move(block_));
+    block_ = nullptr;
+  }
+  if (function_) {
+    // We're in the middle of a function, but the OpFunctionEnd is missing.
+    // Register the function anyway.  This lets us write tests with less
+    // boilerplate.
+    module_->AddFunction(std::move(function_));
+    function_ = nullptr;
+  }
   for (auto& function : *module_) {
     for (auto& bb : function) bb.SetParent(&function);
     function.SetParent(module_);

source/opt/ir_loader.h

@@ -60,8 +60,9 @@ class IrLoader {
   // returning.
   void AddInstruction(const spv_parsed_instruction_t* inst);
   // Finalizes the module construction. This must be called after the module
-  // header has been set and all instructions have been added.
-  // Resolves internal bookkeeping.
+  // header has been set and all instructions have been added.  This is
+  // forgiving in the case of a missing terminator instruction on a basic block,
+  // or a missing OpFunctionEnd.  Resolves internal bookkeeping.
   void EndModule();
 
  private:

test/opt/test_ir_loader.cpp

@@ -93,6 +93,17 @@ TEST(IrBuilder, RoundTrip) {
   // clang-format on
 }
 
+TEST(IrBuilder, RoundTripIncompleteBasicBlock) {
+  DoRoundTripCheck(
+      "%2 = OpFunction %1 None %3\n"
+      "%4 = OpLabel\n"
+      "OpNop\n");
+}
+
+TEST(IrBuilder, RoundTripIncompleteFunction) {
+  DoRoundTripCheck("%2 = OpFunction %1 None %3\n");
+}
+
 TEST(IrBuilder, KeepLineDebugInfo) {
   // #version 310 es
   // void main() {}