spirv-opt: Handle id overflow in MergeReturnPass (#6340)

This CL adds error handling to the MergeReturnPass to gracefully handle
cases where the pass runs out of IDs.

The following functions were modified to return a boolean indicating
success or failure:
- AddNewPhiNodes
- AddReturnFlag
- AddReturnValue
- BranchToBlock
- CreatePhiNodesForInst
- ProcessStructuredBlock
- RecordReturned
- UpdatePhiNodes

The callers of these functions were updated to check the return value
and propagate the failure. This prevents the pass from crashing when it
runs out of IDs.
diff --git a/source/opt/merge_return_pass.cpp b/source/opt/merge_return_pass.cpp
index 427b4ec..47b0763 100644
--- a/source/opt/merge_return_pass.cpp
+++ b/source/opt/merge_return_pass.cpp
@@ -134,7 +134,9 @@
       state_.pop_back();
     }
 
-    ProcessStructuredBlock(block);
+    if (!ProcessStructuredBlock(block)) {
+      return false;
+    }
 
     // Generate state for next block if warranted
     GenerateState(block);
@@ -169,7 +171,9 @@
   // We have not kept the dominator tree up-to-date.
   // Invalidate it at this point to make sure it will be rebuilt.
   context()->RemoveDominatorAnalysis(function);
-  AddNewPhiNodes();
+  if (!AddNewPhiNodes()) {
+    return false;
+  }
   return true;
 }
 
@@ -196,7 +200,9 @@
 }
 
 bool MergeReturnPass::CreateReturn(BasicBlock* block) {
-  AddReturnValue();
+  if (!AddReturnValue()) {
+    return false;
+  }
 
   if (return_value_) {
     // Load and return the final return value
@@ -229,12 +235,18 @@
   return true;
 }
 
-void MergeReturnPass::ProcessStructuredBlock(BasicBlock* block) {
+bool MergeReturnPass::ProcessStructuredBlock(BasicBlock* block) {
+  if (block->tail() == block->end()) {
+    return true;
+  }
+
   spv::Op tail_opcode = block->tail()->opcode();
   if (tail_opcode == spv::Op::OpReturn ||
       tail_opcode == spv::Op::OpReturnValue) {
     if (!return_flag_) {
-      AddReturnFlag();
+      if (!AddReturnFlag()) {
+        return false;
+      }
     }
   }
 
@@ -243,15 +255,20 @@
       tail_opcode == spv::Op::OpUnreachable) {
     assert(CurrentState().InBreakable() &&
            "Should be in the placeholder construct.");
-    BranchToBlock(block, CurrentState().BreakMergeId());
+    if (!BranchToBlock(block, CurrentState().BreakMergeId())) {
+      return false;
+    }
     return_blocks_.insert(block->id());
   }
+  return true;
 }
 
-void MergeReturnPass::BranchToBlock(BasicBlock* block, uint32_t target) {
+bool MergeReturnPass::BranchToBlock(BasicBlock* block, uint32_t target) {
   if (block->tail()->opcode() == spv::Op::OpReturn ||
       block->tail()->opcode() == spv::Op::OpReturnValue) {
-    RecordReturned(block);
+    if (!RecordReturned(block)) {
+      return false;
+    }
     RecordReturnValue(block);
   }
 
@@ -259,7 +276,9 @@
   if (target_block->GetLoopMergeInst()) {
     cfg()->SplitLoopHeader(target_block);
   }
-  UpdatePhiNodes(block, target_block);
+  if (!UpdatePhiNodes(block, target_block)) {
+    return false;
+  }
 
   Instruction* return_inst = block->terminator();
   return_inst->SetOpcode(spv::Op::OpBranch);
@@ -267,19 +286,26 @@
   context()->get_def_use_mgr()->AnalyzeInstDefUse(return_inst);
   new_edges_[target_block].insert(block->id());
   cfg()->AddEdge(block->id(), target);
+  return true;
 }
 
-void MergeReturnPass::UpdatePhiNodes(BasicBlock* new_source,
+bool MergeReturnPass::UpdatePhiNodes(BasicBlock* new_source,
                                      BasicBlock* target) {
-  target->ForEachPhiInst([this, new_source](Instruction* inst) {
+  bool succeeded = true;
+  target->ForEachPhiInst([this, new_source, &succeeded](Instruction* inst) {
     uint32_t undefId = Type2Undef(inst->type_id());
+    if (undefId == 0) {
+      succeeded = false;
+      return;
+    }
     inst->AddOperand({SPV_OPERAND_TYPE_ID, {undefId}});
     inst->AddOperand({SPV_OPERAND_TYPE_ID, {new_source->id()}});
     context()->UpdateDefUse(inst);
   });
+  return succeeded;
 }
 
-void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
+bool MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
                                             Instruction& inst) {
   DominatorAnalysis* dom_tree =
       context()->GetDominatorAnalysis(merge_block->GetParent());
@@ -313,7 +339,7 @@
         });
 
     if (users_to_update.empty()) {
-      return;
+      return true;
     }
 
     // There is at least one values that needs to be replaced.
@@ -357,6 +383,9 @@
     if (regenerateInstruction) {
       std::unique_ptr<Instruction> regen_inst(inst.Clone(context()));
       uint32_t new_id = TakeNextId();
+      if (new_id == 0) {
+        return false;
+      }
       regen_inst->SetResultId(new_id);
       Instruction* insert_pos = &*merge_block->begin();
       while (insert_pos->opcode() == spv::Op::OpPhi) {
@@ -366,19 +395,31 @@
       get_def_use_mgr()->AnalyzeInstDefUse(new_phi);
       context()->set_instr_block(new_phi, merge_block);
 
-      new_phi->ForEachInId([dom_tree, merge_block, this](uint32_t* use_id) {
+      bool succeeded = true;
+      new_phi->ForEachInId([dom_tree, merge_block, this,
+                            &succeeded](uint32_t* use_id) {
+        if (!succeeded) {
+          return;
+        }
         Instruction* use = get_def_use_mgr()->GetDef(*use_id);
         BasicBlock* use_bb = context()->get_instr_block(use);
         if (use_bb != nullptr && !dom_tree->Dominates(use_bb, merge_block)) {
-          CreatePhiNodesForInst(merge_block, *use);
+          if (!CreatePhiNodesForInst(merge_block, *use)) {
+            succeeded = false;
+          }
         }
       });
+      if (!succeeded) {
+        return false;
+      }
     } else {
       InstructionBuilder builder(
           context(), &*merge_block->begin(),
           IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
-      // TODO(1841): Handle id overflow.
       new_phi = builder.AddPhi(inst.type_id(), phi_operands);
+      if (new_phi == nullptr) {
+        return false;
+      }
     }
     uint32_t result_of_phi = new_phi->result_id();
 
@@ -392,6 +433,7 @@
       context()->AnalyzeUses(user);
     }
   }
+  return true;
 }
 
 bool MergeReturnPass::PredicateBlocks(
@@ -484,6 +526,9 @@
   cfg()->RemoveSuccessorEdges(block);
 
   auto old_body_id = TakeNextId();
+  if (old_body_id == 0) {
+    return false;
+  }
   BasicBlock* old_body = block->SplitBasicBlock(context(), old_body_id, iter);
   predicated->insert(old_body);
 
@@ -520,9 +565,11 @@
   analysis::Bool bool_type;
   uint32_t bool_id = context()->get_type_mgr()->GetId(&bool_type);
   assert(bool_id != 0);
-  // TODO(1841): Handle id overflow.
-  uint32_t load_id =
-      builder.AddLoad(bool_id, return_flag_->result_id())->result_id();
+  Instruction* load_inst = builder.AddLoad(bool_id, return_flag_->result_id());
+  if (load_inst == nullptr) {
+    return false;
+  }
+  uint32_t load_id = load_inst->result_id();
 
   // 2. Branch to |merge_block| (true) or |old_body| (false)
   builder.AddConditionalBranch(load_id, merge_block->id(), old_body->id(),
@@ -535,7 +582,9 @@
   }
 
   // 3. Update OpPhi instructions in |merge_block|.
-  UpdatePhiNodes(block, merge_block);
+  if (!UpdatePhiNodes(block, merge_block)) {
+    return false;
+  }
 
   // 4. Update the CFG.  We do this after updating the OpPhi instructions
   // because |UpdatePhiNodes| assumes the edge from |block| has not been added
@@ -548,10 +597,10 @@
   return true;
 }
 
-void MergeReturnPass::RecordReturned(BasicBlock* block) {
+bool MergeReturnPass::RecordReturned(BasicBlock* block) {
   if (block->tail()->opcode() != spv::Op::OpReturn &&
       block->tail()->opcode() != spv::Op::OpReturnValue)
-    return;
+    return true;
 
   assert(return_flag_ && "Did not generate the return flag variable.");
 
@@ -564,6 +613,9 @@
     const analysis::Constant* true_const =
         const_mgr->GetConstant(bool_type, {true});
     constant_true_ = const_mgr->GetDefiningInstruction(true_const);
+    if (!constant_true_) {
+      return false;
+    }
     context()->UpdateDefUse(constant_true_);
   }
 
@@ -577,6 +629,7 @@
       &*block->tail().InsertBefore(std::move(return_store));
   context()->set_instr_block(store_inst, block);
   context()->AnalyzeDefUse(store_inst);
+  return true;
 }
 
 void MergeReturnPass::RecordReturnValue(BasicBlock* block) {
@@ -600,18 +653,21 @@
   context()->AnalyzeDefUse(store_inst);
 }
 
-void MergeReturnPass::AddReturnValue() {
-  if (return_value_) return;
+bool MergeReturnPass::AddReturnValue() {
+  if (return_value_) return true;
 
   uint32_t return_type_id = function_->type_id();
   if (get_def_use_mgr()->GetDef(return_type_id)->opcode() ==
       spv::Op::OpTypeVoid)
-    return;
+    return true;
 
   uint32_t return_ptr_type = context()->get_type_mgr()->FindPointerToType(
       return_type_id, spv::StorageClass::Function);
 
   uint32_t var_id = TakeNextId();
+  if (var_id == 0) {
+    return false;
+  }
   std::unique_ptr<Instruction> returnValue(
       new Instruction(context(), spv::Op::OpVariable, return_ptr_type, var_id,
                       std::initializer_list<Operand>{
@@ -627,27 +683,44 @@
 
   context()->get_decoration_mgr()->CloneDecorations(
       function_->result_id(), var_id, {spv::Decoration::RelaxedPrecision});
+  return true;
 }
 
-void MergeReturnPass::AddReturnFlag() {
-  if (return_flag_) return;
+bool MergeReturnPass::AddReturnFlag() {
+  if (return_flag_) return true;
 
   analysis::TypeManager* type_mgr = context()->get_type_mgr();
   analysis::ConstantManager* const_mgr = context()->get_constant_mgr();
 
   analysis::Bool temp;
   uint32_t bool_id = type_mgr->GetTypeInstruction(&temp);
+  if (bool_id == 0) {
+    return false;
+  }
   analysis::Bool* bool_type = type_mgr->GetType(bool_id)->AsBool();
 
   const analysis::Constant* false_const =
       const_mgr->GetConstant(bool_type, {false});
-  uint32_t const_false_id =
-      const_mgr->GetDefiningInstruction(false_const)->result_id();
+  Instruction* false_inst = const_mgr->GetDefiningInstruction(false_const);
+  if (false_inst == nullptr) {
+    return false;
+  }
+  uint32_t const_false_id = false_inst->result_id();
 
   uint32_t bool_ptr_id =
       type_mgr->FindPointerToType(bool_id, spv::StorageClass::Function);
 
+  if (bool_ptr_id == 0) {
+    return false;
+    ;
+  }
+
   uint32_t var_id = TakeNextId();
+
+  if (var_id == 0) {
+    return false;
+  }
+
   std::unique_ptr<Instruction> returnFlag(new Instruction(
       context(), spv::Op::OpVariable, bool_ptr_id, var_id,
       std::initializer_list<Operand>{{SPV_OPERAND_TYPE_STORAGE_CLASS,
@@ -661,6 +734,7 @@
   return_flag_ = &*entry_block->begin();
   context()->AnalyzeDefUse(return_flag_);
   context()->set_instr_block(return_flag_, entry_block);
+  return true;
 }
 
 std::vector<BasicBlock*> MergeReturnPass::CollectReturnBlocks(
@@ -739,16 +813,19 @@
   return true;
 }
 
-void MergeReturnPass::AddNewPhiNodes() {
+bool MergeReturnPass::AddNewPhiNodes() {
   std::list<BasicBlock*> order;
   cfg()->ComputeStructuredOrder(function_, &*function_->begin(), &order);
 
   for (BasicBlock* bb : order) {
-    AddNewPhiNodes(bb);
+    if (!AddNewPhiNodes(bb)) {
+      return false;
+    }
   }
+  return true;
 }
 
-void MergeReturnPass::AddNewPhiNodes(BasicBlock* bb) {
+bool MergeReturnPass::AddNewPhiNodes(BasicBlock* bb) {
   // New phi nodes are needed for any id whose definition used to dominate |bb|,
   // but no longer dominates |bb|.  These are found by walking the dominator
   // tree starting at the original immediate dominator of |bb| and ending at its
@@ -766,16 +843,19 @@
 
   BasicBlock* dominator = dom_tree->ImmediateDominator(bb);
   if (dominator == nullptr) {
-    return;
+    return true;
   }
 
   BasicBlock* current_bb = context()->get_instr_block(original_dominator_[bb]);
   while (current_bb != nullptr && current_bb != dominator) {
     for (Instruction& inst : *current_bb) {
-      CreatePhiNodesForInst(bb, inst);
+      if (!CreatePhiNodesForInst(bb, inst)) {
+        return false;
+      }
     }
     current_bb = dom_tree->ImmediateDominator(current_bb);
   }
+  return true;
 }
 
 void MergeReturnPass::RecordImmediateDominators(Function* function) {
@@ -859,8 +939,12 @@
     ++split_pos;
   }
 
+  uint32_t new_block_id = TakeNextId();
+  if (new_block_id == 0) {
+    return false;
+  }
   BasicBlock* old_block =
-      start_block->SplitBasicBlock(context(), TakeNextId(), split_pos);
+      start_block->SplitBasicBlock(context(), new_block_id, split_pos);
 
   // Find DebugFunctionDefinition inst in the old block, and if we can find it,
   // move it to the entry block. Since DebugFunctionDefinition is not necessary
diff --git a/source/opt/merge_return_pass.h b/source/opt/merge_return_pass.h
index d39c1d9..d83ffc3 100644
--- a/source/opt/merge_return_pass.h
+++ b/source/opt/merge_return_pass.h
@@ -173,21 +173,22 @@
   //
   // Note this will break the semantics.  To fix this, PredicateBlock will have
   // to be called on the merge block the branch targets.
-  void ProcessStructuredBlock(BasicBlock* block);
+  bool ProcessStructuredBlock(BasicBlock* block);
 
   // Creates a variable used to store whether or not the control flow has
   // traversed a block that used to have a return.  A pointer to the instruction
-  // declaring the variable is stored in |return_flag_|.
-  void AddReturnFlag();
+  // declaring the variable is stored in |return_flag_|. Returns true if it
+  // succeeds.
+  bool AddReturnFlag();
 
   // Creates the variable used to store the return value when passing through
-  // a block that use to contain an OpReturnValue.
-  void AddReturnValue();
+  // a block that use to contain an OpReturnValue. Returns true if it succeeds.
+  bool AddReturnValue();
 
-  // Adds a store that stores true to |return_flag_| immediately before the
-  // terminator of |block|. It is assumed that |AddReturnFlag| has already been
-  // called.
-  void RecordReturned(BasicBlock* block);
+  // Records that |block| used to be a return.  This is done by adding an
+  // instruction to store true to the |return_flag_|. Returns true if it
+  // succeeds.
+  bool RecordReturned(BasicBlock* block);
 
   // Adds an instruction that stores the value being returned in the
   // OpReturnValue in |block|.  The value is stored to |return_value_|, and the
@@ -198,10 +199,10 @@
   // have already been called to create the variable to store to.
   void RecordReturnValue(BasicBlock* block);
 
-  // Adds an unconditional branch in |block| that branches to |target|.  It also
-  // adds stores to |return_flag_| and |return_value_| as needed.
-  // |AddReturnFlag| and |AddReturnValue| must have already been called.
-  void BranchToBlock(BasicBlock* block, uint32_t target);
+  // Replaces the terminator of |block| with a branch to |target|.  If the
+  // terminator was a return, it will first call RecordReturned and
+  // RecordReturnValue. Returns true if it succeeds.
+  bool BranchToBlock(BasicBlock* block, uint32_t target);
 
   // For every basic block that is reachable from |return_block|, extra code is
   // added to jump around any code that should not be executed because the
@@ -239,32 +240,28 @@
   // return block at the end of the pass.
   bool CreateReturnBlock();
 
-  // Creates a Phi node in |merge_block| for the result of |inst|.
-  // Any uses of the result of |inst| that are no longer
-  // dominated by |inst|, are replaced with the result of the new |OpPhi|
-  // instruction.
-  void CreatePhiNodesForInst(BasicBlock* merge_block, Instruction& inst);
+  // For each use of |inst| that is no longer dominated by |inst|, a phi node
+  // is created in |merge_block|.  The original use is replaced by the result
+  // of the phi node. Returns true if it succeeds.
+  bool CreatePhiNodesForInst(BasicBlock* merge_block, Instruction& inst);
 
-  // Add new phi nodes for any id that no longer dominate all of it uses.  A phi
-  // node is added to a block |bb| for an id if the id is defined between the
-  // original immediate dominator of |bb| and its new immediate dominator.  It
-  // is assumed that at this point there are no unreachable blocks in the
-  // control flow graph.
-  void AddNewPhiNodes();
+  // Adds new phi nodes as needed to the function.  This is necessary because
+  // adding the predication code can change the dominator tree.  Returns false
+  // if there is a failure.
+  bool AddNewPhiNodes();
 
-  // Creates any new phi nodes that are needed in |bb|.  |AddNewPhiNodes| must
-  // have already been called on the original dominators of |bb|.
-  void AddNewPhiNodes(BasicBlock* bb);
+  // Adds new phi nodes to |bb| as needed.  This is necessary because adding
+  // the predication code can change the dominator tree.  Returns false if
+  // there is a failure.
+  bool AddNewPhiNodes(BasicBlock* bb);
 
   // Records the terminator of immediate dominator for every basic block in
   // |function|.
   void RecordImmediateDominators(Function* function);
 
-  // Modifies existing OpPhi instruction in |target| block to account for the
-  // new edge from |new_source|.  The value for that edge will be an Undef.
-  //
-  // The CFG must not include the edge from |new_source| to |target| yet.
-  void UpdatePhiNodes(BasicBlock* new_source, BasicBlock* target);
+  // For each OpPhi instruction in |target|, this function adds an operand for
+  // |new_source|.  The value will be OpUndef. Returns true if it succeeds.
+  bool UpdatePhiNodes(BasicBlock* new_source, BasicBlock* target);
 
   StructuredControlState& CurrentState() { return state_.back(); }
 
@@ -334,4 +331,4 @@
 }  // namespace opt
 }  // namespace spvtools
 
-#endif  // SOURCE_OPT_MERGE_RETURN_PASS_H_
\ No newline at end of file
+#endif  // SOURCE_OPT_MERGE_RETURN_PASS_H_