spirv-opt: Fix memory leaks on ID overflow (#6887)

Several passes clone a basic block or an instruction into a raw pointer
and only transfer ownership to the module later. When ID allocation
fails part way through, these passes bail out before ownership is
transferred, and the clone is leaked. The solution is to take ownership
immediately.

These leaks are found in the presubmits for #6881.
diff --git a/source/opt/fold_spec_constant_op_and_composite_pass.cpp b/source/opt/fold_spec_constant_op_and_composite_pass.cpp
index edcf551..1ae3b41 100644
--- a/source/opt/fold_spec_constant_op_and_composite_pass.cpp
+++ b/source/opt/fold_spec_constant_op_and_composite_pass.cpp
@@ -208,11 +208,14 @@
   }
 
   if (need_to_clone) {
-    new_const_inst = new_const_inst->Clone(context());
     uint32_t new_id = TakeNextId();
     if (new_id == 0) {
       return nullptr;
     }
+    new_const_inst = new_const_inst->Clone(context());
+    if (!new_const_inst) {
+      return nullptr;
+    }
     new_const_inst->SetResultId(new_id);
     new_const_inst->InsertAfter(insert_pos);
     get_def_use_mgr()->AnalyzeInstDefUse(new_const_inst);
diff --git a/source/opt/loop_unroller.cpp b/source/opt/loop_unroller.cpp
index 16a0ea9..759cce9 100644
--- a/source/opt/loop_unroller.cpp
+++ b/source/opt/loop_unroller.cpp
@@ -406,6 +406,9 @@
     return false;
   }
 
+  // The merge block created by |DuplicateLoop| is not owned by the function.
+  std::unique_ptr<BasicBlock> new_loop_merge_block(new_loop->GetMergeBlock());
+
   // Add the blocks to the function.
   AddBlocksToFunction(loop->GetMergeBlock());
   blocks_to_add_.clear();
@@ -428,8 +431,7 @@
 
   // Add the new merge block to the back of the list of blocks to be added. It
   // needs to be the last block added to maintain dominator order in the binary.
-  blocks_to_add_.push_back(
-      std::unique_ptr<BasicBlock>(new_loop->GetMergeBlock()));
+  blocks_to_add_.push_back(std::move(new_loop_merge_block));
 
   // Add the blocks to the function.
   AddBlocksToFunction(loop->GetMergeBlock());
@@ -651,16 +653,16 @@
 bool LoopUnrollerUtilsImpl::CopyBasicBlock(Loop* loop, const BasicBlock* itr,
                                            bool preserve_instructions) {
   // Clone the block exactly, including the IDs.
-  BasicBlock* basic_block = itr->Clone(context_);
+  std::unique_ptr<BasicBlock> basic_block(itr->Clone(context_));
   if (!basic_block) return false;
   basic_block->SetParent(itr->GetParent());
 
   // We do not want to duplicate DebugDeclare.
-  KillDebugDeclares(basic_block);
+  KillDebugDeclares(basic_block.get());
 
   // Assign each result a new unique ID and keep a mapping of the old ids to
   // the new ones.
-  if (!AssignNewResultIds(basic_block)) {
+  if (!AssignNewResultIds(basic_block.get())) {
     return false;
   }
 
@@ -673,12 +675,12 @@
       context_->UpdateDefUse(merge_inst);
     }
 
-    state_.new_continue_block = basic_block;
+    state_.new_continue_block = basic_block.get();
   }
 
   // If this is the header block we are copying.
   if (itr == loop->GetHeaderBlock()) {
-    state_.new_header_block = basic_block;
+    state_.new_header_block = basic_block.get();
 
     if (!preserve_instructions) {
       // Remove the loop merge instruction if it exists.
@@ -688,19 +690,19 @@
   }
 
   // If this is the latch block being copied, record it in the state.
-  if (itr == loop->GetLatchBlock()) state_.new_latch_block = basic_block;
+  if (itr == loop->GetLatchBlock()) state_.new_latch_block = basic_block.get();
 
   // If this is the condition block we are copying.
   if (itr == loop_condition_block_) {
-    state_.new_condition_block = basic_block;
+    state_.new_condition_block = basic_block.get();
   }
 
+  // Keep tracking the old block via a map.
+  state_.new_blocks[itr->id()] = basic_block.get();
+
   // Add this block to the list of blocks to add to the function at the end of
   // the unrolling process.
-  blocks_to_add_.push_back(std::unique_ptr<BasicBlock>(basic_block));
-
-  // Keep tracking the old block via a map.
-  state_.new_blocks[itr->id()] = basic_block;
+  blocks_to_add_.push_back(std::move(basic_block));
   return true;
 }
 
@@ -863,13 +865,14 @@
   }
 
   // Clone the merge block, give it a new id and record it in the state.
-  BasicBlock* new_merge = old_loop->GetMergeBlock()->Clone(context_);
+  std::unique_ptr<BasicBlock> new_merge(
+      old_loop->GetMergeBlock()->Clone(context_));
   if (!new_merge) return false;
   new_merge->SetParent(old_loop->GetMergeBlock()->GetParent());
-  if (!AssignNewResultIds(new_merge)) {
+  if (!AssignNewResultIds(new_merge.get())) {
     return false;
   }
-  state_.new_blocks[old_loop->GetMergeBlock()->id()] = new_merge;
+  state_.new_blocks[old_loop->GetMergeBlock()->id()] = new_merge.get();
 
   // Remap the operands of every instruction in the loop to point to the new
   // copies.
@@ -884,7 +887,7 @@
   new_loop->SetHeaderBlock(state_.new_header_block);
   new_loop->SetContinueBlock(state_.new_continue_block);
   new_loop->SetLatchBlock(state_.new_latch_block);
-  new_loop->SetMergeBlock(new_merge);
+  new_loop->SetMergeBlock(new_merge.release());
   return true;
 }
 
diff --git a/source/opt/loop_unswitch_pass.cpp b/source/opt/loop_unswitch_pass.cpp
index b49f243..fef3c77 100644
--- a/source/opt/loop_unswitch_pass.cpp
+++ b/source/opt/loop_unswitch_pass.cpp
@@ -184,15 +184,15 @@
       bool ok = true;
       if_merge_block->ForEachPhiInst(
           [loop_merge_block, &ok, &builder, this](Instruction* phi) -> bool {
-            Instruction* cloned = phi->Clone(context_);
+            std::unique_ptr<Instruction> cloned(phi->Clone(context_));
             uint32_t new_id = TakeNextId();
             if (new_id == 0) {
               ok = false;
               return false;
             }
             cloned->SetResultId(new_id);
-            builder.AddInstruction(std::unique_ptr<Instruction>(cloned));
-            phi->SetInOperand(0, {cloned->result_id()});
+            builder.AddInstruction(std::move(cloned));
+            phi->SetInOperand(0, {new_id});
             phi->SetInOperand(1, {loop_merge_block->id()});
             for (uint32_t j = phi->NumInOperands() - 1; j > 1; j--)
               phi->RemoveInOperand(j);
diff --git a/source/opt/loop_utils.cpp b/source/opt/loop_utils.cpp
index ced68de..dffb47f 100644
--- a/source/opt/loop_utils.cpp
+++ b/source/opt/loop_utils.cpp
@@ -598,7 +598,7 @@
   for (BasicBlock* old_bb : ordered_loop_blocks) {
     // For each basic block in the loop, we clone it and register the mapping
     // between old and new ids.
-    BasicBlock* new_bb = old_bb->Clone(context_);
+    std::unique_ptr<BasicBlock> new_bb(old_bb->Clone(context_));
     if (!new_bb) return nullptr;
     new_bb->SetParent(&function_);
     uint32_t new_label_id = context_->TakeNextId();
@@ -607,14 +607,13 @@
     }
     new_bb->GetLabelInst()->SetResultId(new_label_id);
     def_use_mgr->AnalyzeInstDef(new_bb->GetLabelInst());
-    context_->set_instr_block(new_bb->GetLabelInst(), new_bb);
-    cloning_result->cloned_bb_.emplace_back(new_bb);
+    context_->set_instr_block(new_bb->GetLabelInst(), new_bb.get());
 
-    cloning_result->old_to_new_bb_[old_bb->id()] = new_bb;
+    cloning_result->old_to_new_bb_[old_bb->id()] = new_bb.get();
     cloning_result->new_to_old_bb_[new_bb->id()] = old_bb;
     cloning_result->value_map_[old_bb->id()] = new_bb->id();
 
-    if (loop_->IsInsideLoop(old_bb)) new_loop->AddBasicBlock(new_bb);
+    if (loop_->IsInsideLoop(old_bb)) new_loop->AddBasicBlock(new_bb.get());
 
     for (auto new_inst = new_bb->begin(), old_inst = old_bb->begin();
          new_inst != new_bb->end(); ++new_inst, ++old_inst) {
@@ -632,6 +631,8 @@
         def_use_mgr->AnalyzeInstDef(&*new_inst);
       }
     }
+
+    cloning_result->cloned_bb_.emplace_back(std::move(new_bb));
   }
 
   // All instructions (including all labels) have been cloned,
diff --git a/source/opt/replace_desc_array_access_using_var_index.cpp b/source/opt/replace_desc_array_access_using_var_index.cpp
index 007eeff..a3d1a06 100644
--- a/source/opt/replace_desc_array_access_using_var_index.cpp
+++ b/source/opt/replace_desc_array_access_using_var_index.cpp
@@ -396,10 +396,10 @@
 BasicBlock* ReplaceDescArrayAccessUsingVarIndex::CreateDefaultBlock(
     bool null_const_for_phi_is_needed, std::vector<uint32_t>* phi_operands,
     uint32_t merge_block_id) const {
-  auto* default_block = CreateNewBlock();
+  std::unique_ptr<BasicBlock> default_block(CreateNewBlock());
   if (!default_block) return nullptr;
-  AddBranchToBlock(default_block, merge_block_id);
-  if (!null_const_for_phi_is_needed) return default_block;
+  AddBranchToBlock(default_block.get(), merge_block_id);
+  if (!null_const_for_phi_is_needed) return default_block.release();
 
   // Create null value for OpPhi
   Instruction* inst = context()->get_def_use_mgr()->GetDef((*phi_operands)[0]);
@@ -408,7 +408,7 @@
     return nullptr;
   }
   phi_operands->push_back(null_const_inst->result_id());
-  return default_block;
+  return default_block.release();
 }
 
 Instruction* ReplaceDescArrayAccessUsingVarIndex::GetConstNull(
diff --git a/source/opt/split_combined_image_sampler_pass.cpp b/source/opt/split_combined_image_sampler_pass.cpp
index 47296fa..0d2af3b 100644
--- a/source/opt/split_combined_image_sampler_pass.cpp
+++ b/source/opt/split_combined_image_sampler_pass.cpp
@@ -580,9 +580,10 @@
 
   // Rewite OpFunctionParameter in function definitions.
   for (Function& fn : *context()->module()) {
-    // Rewrite the function parameters and record their replacements.
+    // Rewrite the function parameters and record their replacements. A
+    // rewritten parameter is no longer owned by the function.
     struct Replacement {
-      Instruction* combined;
+      std::unique_ptr<Instruction> combined;
       Instruction* image;
       Instruction* sampler;
     };
@@ -601,7 +602,7 @@
           }
 
           // Replace this parameter with two new parameters.
-          auto* combined_inst = param.release();
+          std::unique_ptr<Instruction> combined_inst = std::move(param);
           auto* combined_type = def_use_mgr_->GetDef(combined_inst->type_id());
           auto [image_type, sampler_type] = SplitType(*combined_type);
           if (!image_type || !sampler_type) {
@@ -626,8 +627,8 @@
               context(), spv::Op::OpFunctionParameter,
               sampler_type->result_id(), sampler_param_id,
               Instruction::OperandList{});
-          replacements.push_back(
-              {combined_inst, image_param.get(), sampler_param.get()});
+          replacements.push_back({std::move(combined_inst), image_param.get(),
+                                  sampler_param.get()});
           appender = std::move(image_param);
           appender = std::move(sampler_param);
         };
@@ -641,7 +642,12 @@
       modified_ = true;
       def_use_mgr_->AnalyzeInstDefUse(r.image);
       def_use_mgr_->AnalyzeInstDefUse(r.sampler);
-      CHECK_STATUS(RemapUses(r.combined, r.image, r.sampler));
+      spv_result_t status = RemapUses(r.combined.get(), r.image, r.sampler);
+      if (status == SPV_SUCCESS) {
+        // |RemapUses| killed the combined instruction, which deleted it.
+        (void)r.combined.release();
+      }
+      CHECK_STATUS(status);
     }
   }
   return SPV_SUCCESS;