spirv-opt: Fail gracefully on ID overflow during spec constant folding

The instruction folder returns nullptr when it fails to create a new
instruction due to ID overflow. This is ambiguous because it also
returns nullptr when it simply cannot fold an instruction.

This change introduces an `id_overflow_` flag in `IRContext` that is
set by `TakeNextId()` when it runs out of IDs.

The `FoldSpecConstantOpAndCompositePass` is updated to check this flag
after attempting to fold a spec constant. If the flag is set, the pass
returns `Status::Failure`, preventing further processing and providing a
clear error.

Author: s-perron

source/opt/fold_spec_constant_op_and_composite_pass.cpp

@@ -96,17 +96,24 @@ Pass::Status FoldSpecConstantOpAndCompositePass::Process() {
       // Constants will be added to id_to_const_val_ and const_val_to_id_ so
       // that we can use the new Normal Constants when folding following Spec
       // Constants.
-      case spv::Op::OpSpecConstantOp:
-        modified |= ProcessOpSpecConstantOp(&inst_iter);
+      case spv::Op::OpSpecConstantOp: {
+        const auto status = ProcessOpSpecConstantOp(&inst_iter);
+        if (status == Status::Failure) {
+          return Status::Failure;
+        }
+        if (status == Status::SuccessWithChange) {
+          modified = true;
+        }
         break;
+      }
       default:
         break;
     }
   }
   return modified ? Status::SuccessWithChange : Status::SuccessWithoutChange;
 }
 
-bool FoldSpecConstantOpAndCompositePass::ProcessOpSpecConstantOp(
+Pass::Status FoldSpecConstantOpAndCompositePass::ProcessOpSpecConstantOp(
     Module::inst_iterator* pos) {
   Instruction* inst = &**pos;
   Instruction* folded_inst = nullptr;
@@ -116,18 +123,25 @@ bool FoldSpecConstantOpAndCompositePass::ProcessOpSpecConstantOp(
          "SPV_OPERAND_TYPE_SPEC_CONSTANT_OP_NUMBER type");
 
   folded_inst = FoldWithInstructionFolder(pos);
+  if (context()->id_overflow()) {
+    return Status::Failure;
+  }
+
   if (!folded_inst) {
     folded_inst = DoComponentWiseOperation(pos);
+    if (context()->id_overflow()) {
+      return Status::Failure;
+    }
   }
-  if (!folded_inst) return false;
+  if (!folded_inst) return Status::SuccessWithoutChange;
 
   // Replace the original constant with the new folded constant, kill the
   // original constant.
   uint32_t new_id = folded_inst->result_id();
   uint32_t old_id = inst->result_id();
   context()->ReplaceAllUsesWith(old_id, new_id);
   context()->KillDef(old_id);
-  return true;
+  return Status::SuccessWithChange;
 }
 
 Instruction* FoldSpecConstantOpAndCompositePass::FoldWithInstructionFolder(
@@ -186,7 +200,11 @@ Instruction* FoldSpecConstantOpAndCompositePass::FoldWithInstructionFolder(
 
   if (need_to_clone) {
     new_const_inst = new_const_inst->Clone(context());
-    new_const_inst->SetResultId(TakeNextId());
+    uint32_t new_id = TakeNextId();
+    if (new_id == 0) {
+      return nullptr;
+    }
+    new_const_inst->SetResultId(new_id);
     new_const_inst->InsertAfter(insert_pos);
     get_def_use_mgr()->AnalyzeInstDefUse(new_const_inst);
   }

source/opt/fold_spec_constant_op_and_composite_pass.h

@@ -19,6 +19,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "source/diagnostic.h"
 #include "source/opt/constants.h"
 #include "source/opt/def_use_manager.h"
 #include "source/opt/ir_context.h"
@@ -45,14 +46,14 @@ class FoldSpecConstantOpAndCompositePass : public Pass {
  private:
   // Processes the OpSpecConstantOp instruction pointed by the given
   // instruction iterator, folds it to normal constants if possible. Returns
-  // true if the spec constant is folded to normal constants. New instructions
-  // will be inserted before the OpSpecConstantOp instruction pointed by the
-  // instruction iterator. The instruction iterator, which is passed by
-  // pointer, will still point to the original OpSpecConstantOp instruction. If
-  // folding is done successfully, the original OpSpecConstantOp instruction
-  // will be changed to Nop and new folded instruction will be inserted before
-  // it.
-  bool ProcessOpSpecConstantOp(Module::inst_iterator* pos);
+  // kSuccess if the spec constant is folded to normal constants. New
+  // instructions will be inserted before the OpSpecConstantOp instruction
+  // pointed by the instruction iterator. The instruction iterator, which is
+  // passed by pointer, will still point to the original OpSpecConstantOp
+  // instruction. If folding is done successfully, the original OpSpecConstantOp
+  // instruction will be changed to Nop and new folded instruction will be
+  // inserted before it. Returns kFail if an id overflow occurs.
+  Status ProcessOpSpecConstantOp(Module::inst_iterator* pos);
 
   // Returns the result of folding the OpSpecConstantOp instruction
   // |inst_iter_ptr| using the instruction folder.
@@ -62,7 +63,24 @@ class FoldSpecConstantOpAndCompositePass : public Pass {
   // pointed by the given instruction iterator to a normal constant defining
   // instruction. Returns the pointer to the new constant defining instruction
   // if succeeded, otherwise return nullptr.
+  // instruction if succeeded, otherwise return nullptr.
   Instruction* DoComponentWiseOperation(Module::inst_iterator* inst_iter_ptr);
+
+  // Returns the next available id, or 0 if the id overflows.
+  uint32_t TakeNextId() {
+    uint32_t next_id = context()->TakeNextId();
+    if (next_id == 0) {
+      Fail() << "ID overflow. Try running compact-ids.";
+    }
+    return next_id;
+  }
+
+  // Records failure for the current module and returns a stream for printing
+  // diagnostics.
+  spvtools::DiagnosticStream Fail() {
+    return spvtools::DiagnosticStream({}, context()->consumer(), "",
+                                      SPV_ERROR_INTERNAL);
+  }
 };
 
 }  // namespace opt

source/opt/instruction.cpp

@@ -770,7 +770,7 @@ bool Instruction::IsFoldableByFoldScalar() const {
   // Even if the type of the instruction is foldable, its operands may not be
   // foldable (e.g., comparisons of 64bit types).  Check that all operand types
   // are foldable before accepting the instruction.
-  return WhileEachInOperand([&folder, this](const uint32_t* op_id) {
+  return WhileEachInId([&folder, this](const uint32_t* op_id) {
     Instruction* def_inst = context()->get_def_use_mgr()->GetDef(*op_id);
     Instruction* def_inst_type =
         context()->get_def_use_mgr()->GetDef(def_inst->type_id());
@@ -792,7 +792,7 @@ bool Instruction::IsFoldableByFoldVector() const {
   // Even if the type of the instruction is foldable, its operands may not be
   // foldable (e.g., comparisons of 64bit types).  Check that all operand types
   // are foldable before accepting the instruction.
-  return WhileEachInOperand([&folder, this](const uint32_t* op_id) {
+  return WhileEachInId([&folder, this](const uint32_t* op_id) {
     Instruction* def_inst = context()->get_def_use_mgr()->GetDef(*op_id);
     Instruction* def_inst_type =
         context()->get_def_use_mgr()->GetDef(def_inst->type_id());

source/opt/ir_context.h

@@ -109,7 +109,8 @@ class IRContext {
         id_to_name_(nullptr),
         max_id_bound_(kDefaultMaxIdBound),
         preserve_bindings_(false),
-        preserve_spec_constants_(false) {
+        preserve_spec_constants_(false),
+        id_overflow_(false) {
     SetContextMessageConsumer(syntax_context_, consumer_);
     module_->SetContext(this);
   }
@@ -127,7 +128,8 @@ class IRContext {
         id_to_name_(nullptr),
         max_id_bound_(kDefaultMaxIdBound),
         preserve_bindings_(false),
-        preserve_spec_constants_(false) {
+        preserve_spec_constants_(false),
+        id_overflow_(false) {
     SetContextMessageConsumer(syntax_context_, consumer_);
     module_->SetContext(this);
     InitializeCombinators();
@@ -563,6 +565,7 @@ class IRContext {
   inline uint32_t TakeNextId() {
     uint32_t next_id = module()->TakeNextIdBound();
     if (next_id == 0) {
+      id_overflow_ = true;
       if (consumer()) {
         std::string message = "ID overflow. Try running compact-ids.";
         consumer()(SPV_MSG_ERROR, "", {0, 0, 0}, message.c_str());
@@ -583,6 +586,13 @@ class IRContext {
     return next_id;
   }
 
+  // Returns true if an ID overflow has occurred since the last time the flag
+  // was cleared.
+  bool id_overflow() const { return id_overflow_; }
+
+  // Clears the ID overflow flag.
+  void clear_id_overflow() { id_overflow_ = false; }
+
   FeatureManager* get_feature_mgr() {
     if (!feature_mgr_.get()) {
       AnalyzeFeatures();
@@ -930,6 +940,9 @@ class IRContext {
   // Whether all specialization constants within |module_|
   // should be preserved.
   bool preserve_spec_constants_;
+
+  // Set to true if TakeNextId() fails.
+  bool id_overflow_;
 };
 
 inline IRContext::Analysis operator|(IRContext::Analysis lhs,