From cc60caba1d2fc8c0594f7e09c5176022d08ba695 Mon Sep 17 00:00:00 2001
From: David Neto <dneto@google.com>
Date: Thu, 25 Aug 2016 16:42:36 -0400
Subject: [PATCH] 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    |  7 ++++---
 source/opt/function.cpp     | 14 ++++++++------
 source/opt/ir_loader.cpp    | 14 ++++++++++++++
 source/opt/ir_loader.h      |  5 +++--
 test/opt/test_ir_loader.cpp | 11 +++++++++++
 5 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/source/opt/basic_block.h b/source/opt/basic_block.h
index 52f9707..c8c5749 100644
--- a/source/opt/basic_block.h
+++ b/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);
diff --git a/source/opt/function.cpp b/source/opt/function.cpp
index 4327258..56f0ba1 100644
--- a/source/opt/function.cpp
+++ b/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
diff --git a/source/opt/ir_loader.cpp b/source/opt/ir_loader.cpp
index 47aa82d..b34e44f 100644
--- a/source/opt/ir_loader.cpp
+++ b/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_);
diff --git a/source/opt/ir_loader.h b/source/opt/ir_loader.h
index c4407b9..8f7dda1 100644
--- a/source/opt/ir_loader.h
+++ b/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:
diff --git a/test/opt/test_ir_loader.cpp b/test/opt/test_ir_loader.cpp
index 6066997..4447358 100644
--- a/test/opt/test_ir_loader.cpp
+++ b/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() {}
-- 
2.7.4