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.

Author: s-perron

source/opt/merge_return_pass.cpp

@@ -134,7 +134,9 @@ bool MergeReturnPass::ProcessStructured(
       state_.pop_back();
     }
 
-    ProcessStructuredBlock(block);
+    if (!ProcessStructuredBlock(block)) {
+      return false;
+    }
 
     // Generate state for next block if warranted
     GenerateState(block);
@@ -169,7 +171,9 @@ bool MergeReturnPass::ProcessStructured(
   // 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::CreateReturnBlock() {
 }
 
 bool MergeReturnPass::CreateReturn(BasicBlock* block) {
-  AddReturnValue();
+  if (!AddReturnValue()) {
+    return false;
+  }
 
   if (return_value_) {
     // Load and return the final return value
@@ -229,12 +235,18 @@ bool MergeReturnPass::CreateReturn(BasicBlock* block) {
   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,43 +255,57 @@ void MergeReturnPass::ProcessStructuredBlock(BasicBlock* block) {
       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);
   }
 
   BasicBlock* target_block = context()->get_instr_block(target);
   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);
   return_inst->ReplaceOperands({{SPV_OPERAND_TYPE_ID, {target}}});
   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 @@ void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
         });
 
     if (users_to_update.empty()) {
-      return;
+      return true;
     }
 
     // There is at least one values that needs to be replaced.
@@ -357,6 +383,9 @@ void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
     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 @@ void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
       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 @@ void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
       context()->AnalyzeUses(user);
     }
   }
+  return true;
 }
 
 bool MergeReturnPass::PredicateBlocks(
@@ -484,6 +526,9 @@ bool MergeReturnPass::BreakFromConstruct(
   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 @@ bool MergeReturnPass::BreakFromConstruct(
   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 @@ bool MergeReturnPass::BreakFromConstruct(
   }
 
   // 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 @@ bool MergeReturnPass::BreakFromConstruct(
   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 @@ void MergeReturnPass::RecordReturned(BasicBlock* block) {
     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 @@ void MergeReturnPass::RecordReturned(BasicBlock* block) {
       &*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 @@ void MergeReturnPass::RecordReturnValue(BasicBlock* block) {
   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 @@ void MergeReturnPass::AddReturnValue() {
 
   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 @@ void MergeReturnPass::AddReturnFlag() {
   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 @@ bool MergeReturnPass::MergeReturnBlocks(
   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 @@ void MergeReturnPass::AddNewPhiNodes(BasicBlock* bb) {
 
   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 @@ bool MergeReturnPass::CreateSingleCaseSwitch(BasicBlock* merge_target) {
     ++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