Check for use of InstIds from the wrong SemIR::File

toolchain/base/BUILD

@@ -66,6 +66,7 @@ cc_library(
     hdrs = ["fixed_size_value_store.h"],
     deps = [
         ":mem_usage",
+        ":value_store",
         ":value_store_types",
         "//common:check",
         "@llvm-project//llvm:Support",

toolchain/base/fixed_size_value_store.h

@@ -13,6 +13,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
 #include "toolchain/base/mem_usage.h"
+#include "toolchain/base/value_store.h"
 #include "toolchain/base/value_store_types.h"
 
 namespace Carbon {
@@ -47,9 +48,11 @@ class FixedSizeValueStore {
   }
 
   // Makes a ValueStore of the specified size, initialized to a default.
-  static auto MakeWithExplicitSize(size_t size, ConstRefType default_value)
+  static auto MakeWithExplicitSize(CheckIRIdTag tag, size_t size,
+                                   ConstRefType default_value)
       -> FixedSizeValueStore {
     FixedSizeValueStore store;
+    store.tag_ = tag;
     store.values_.resize(size, default_value);
     return store;
   }
@@ -73,7 +76,8 @@ class FixedSizeValueStore {
   template <typename ValueStoreT>
     requires std::same_as<IdT, typename ValueStoreT::IdType>
   explicit FixedSizeValueStore(const ValueStoreT& size_source,
-                               ConstRefType default_value) {
+                               ConstRefType default_value)
+      : tag_(GetCheckIRIdTag(size_source)) {
     values_.resize(size_source.size(), default_value);
   }
 
@@ -86,7 +90,8 @@ class FixedSizeValueStore {
       llvm::function_ref<
           auto(IdT, typename ValueStoreT::ConstRefType)->ValueType>
           factory_fn)
-      : values_(llvm::map_range(source.enumerate(), factory_fn)) {}
+      : values_(llvm::map_range(source.enumerate(), factory_fn)),
+        tag_(GetCheckIRIdTag(source)) {}
 
   // Move-only.
   FixedSizeValueStore(FixedSizeValueStore&&) noexcept = default;
@@ -95,20 +100,23 @@ class FixedSizeValueStore {
 
   // Sets the value for an ID.
   auto Set(IdT id, ValueType value) -> void {
-    CARBON_DCHECK(id.index >= 0, "{0}", id);
-    values_[id.index] = value;
+    auto index = tag_.Remove(id.index);
+    CARBON_DCHECK(index >= 0, "{0}", id);
+    values_[index] = value;
   }
 
   // Returns a mutable value for an ID.
   auto Get(IdT id) -> RefType {
-    CARBON_DCHECK(id.index >= 0, "{0}", id);
-    return values_[id.index];
+    auto index = tag_.Remove(id.index);
+    CARBON_DCHECK(index >= 0, "{0}", id);
+    return values_[index];
   }
 
   // Returns the value for an ID.
   auto Get(IdT id) const -> ConstRefType {
-    CARBON_DCHECK(id.index >= 0, "{0}", id);
-    return values_[id.index];
+    auto index = tag_.Remove(id.index);
+    CARBON_DCHECK(index >= 0, "{0}", id);
+    return values_[index];
   }
 
   // Collects memory usage of the values.
@@ -135,6 +143,8 @@ class FixedSizeValueStore {
 
   // Storage for the `ValueT` objects, indexed by the id.
   llvm::SmallVector<ValueT, 0> values_;
+
+  CheckIRIdTag tag_;
 };
 
 }  // namespace Carbon

toolchain/base/value_store.h

@@ -32,6 +32,87 @@ class ValueStoreNotPrintable {};
 
 }  // namespace Internal
 
+struct CheckIRIdTag {
+  CheckIRIdTag()
+      : check_id_tag_(0),
+        initial_reserved_ids_(std::numeric_limits<int32_t>::max()) {}
+  explicit CheckIRIdTag(int32_t check_id_index, int32_t initial_reserved_ids)
+      : initial_reserved_ids_(initial_reserved_ids) {
+    // Shift down by 1 to get out of the high bit to avoid using any negative
+    // ids, since they have special uses.
+    // Shift down by another 1 to free up the second highest bit for a marker to
+    // indicate whether the index is tagged (& needs to be untagged) or not.
+    // Add one to the index so it's not zero-based, to make it a bit less likely
+    // this doesn't collide with anything else (though with the
+    // second-highest-bit-tagging this might not be needed).
+    check_id_tag_ = llvm::reverseBits((((check_id_index + 1) << 1) | 1) << 1);
+  }
+  auto GetCheckIRId() const -> int32_t {
+    return (llvm::reverseBits(check_id_tag_) >> 2) - 1;
+  }
+  auto Apply(int32_t index) const -> int32_t {
+    if (index < initial_reserved_ids_) {
+      return index;
+    }
+    // assert that check_id_tag_ doesn't have the second highest bit set
+    return index ^ check_id_tag_;
+  }
+  static auto DecomposeWithBestEffort(int32_t tagged_index)
+      -> std::pair<int32_t, int32_t> {
+    if (tagged_index < 0) {
+      return {-1, tagged_index};
+    }
+    if ((llvm::reverseBits(2) & tagged_index) == 0) {
+      return {-1, tagged_index};
+    }
+    int length = 0;
+    int location = 0;
+    for (int i = 0; i != 32; ++i) {
+      int current_run = 0;
+      int location_of_current_run = i;
+      while (i != 32 && (tagged_index & (1 << i)) == 0) {
+        ++current_run;
+        ++i;
+      }
+      if (current_run != 0) {
+        --i;
+      }
+      if (current_run > length) {
+        length = current_run;
+        location = location_of_current_run;
+      }
+    }
+    if (length < 8) {
+      return {-1, tagged_index};
+    }
+    auto index_mask = llvm::maskTrailingOnes<uint32_t>(location);
+    auto check_ir_id = (llvm::reverseBits(tagged_index & ~index_mask) >> 2) - 1;
+    auto index = tagged_index & index_mask;
+    return {check_ir_id, index};
+  }
+  auto Remove(int32_t tagged_index) const -> int32_t {
+    if ((llvm::reverseBits(2) & tagged_index) == 0) {
+      CARBON_DCHECK(tagged_index < initial_reserved_ids_,
+                    "This untagged index is outside the initial reserved ids "
+                    "and should have been tagged.");
+      return tagged_index;
+    }
+    auto index = tagged_index ^ check_id_tag_;
+    CARBON_DCHECK(index >= initial_reserved_ids_,
+                  "When removing tagging bits, found an index that "
+                  "shouldn't've been tagged in the first place.");
+    return index;
+  }
+  int32_t check_id_tag_;
+  int32_t initial_reserved_ids_;
+};
+
+template <typename ValueStoreT>
+auto GetCheckIRIdTag(const ValueStoreT& value_store) {
+  (void)value_store;
+  return CheckIRIdTag();
+}
+
 // A simple wrapper for accumulating values, providing IDs to later retrieve the
 // value. This does not do deduplication.
 template <typename IdT, typename ValueT>
@@ -74,6 +155,7 @@ class ValueStore
   };
 
   ValueStore() = default;
+  explicit ValueStore(CheckIRIdTag tag) : tag_(tag) {}
 
   // Stores the value and returns an ID to reference it.
   auto Add(ValueType value) -> IdType {
@@ -82,8 +164,8 @@ class ValueStore
     // tracking down issues easier.
     CARBON_DCHECK(size_ < std::numeric_limits<int32_t>::max(), "Id overflow");
 
-    IdType id(size_);
-    auto [chunk_index, pos] = IdToChunkIndices(id);
+    IdType id(tag_.Apply(size_));
+    auto [chunk_index, pos] = IdToChunkIndices(size_);
     ++size_;
 
     CARBON_DCHECK(static_cast<size_t>(chunk_index) <= chunks_.size(),
@@ -99,17 +181,15 @@ class ValueStore
 
   // Returns a mutable value for an ID.
   auto Get(IdType id) -> RefType {
-    CARBON_DCHECK(id.index >= 0, "{0}", id);
-    CARBON_DCHECK(id.index < size_, "{0}", id);
-    auto [chunk_index, pos] = IdToChunkIndices(id);
+    auto index = GetRawIndex(id);
+    auto [chunk_index, pos] = IdToChunkIndices(index);
     return chunks_[chunk_index].Get(pos);
   }
 
   // Returns the value for an ID.
   auto Get(IdType id) const -> ConstRefType {
-    CARBON_DCHECK(id.index >= 0, "{0}", id);
-    CARBON_DCHECK(id.index < size_, "{0}", id);
-    auto [chunk_index, pos] = IdToChunkIndices(id);
+    auto index = GetRawIndex(id);
+    auto [chunk_index, pos] = IdToChunkIndices(index);
     return chunks_[chunk_index].Get(pos);
   }
 
@@ -118,7 +198,7 @@ class ValueStore
     if (size <= size_) {
       return;
     }
-    auto [final_chunk_index, _] = IdToChunkIndices(IdType(size - 1));
+    auto [final_chunk_index, _] = IdToChunkIndices(size - 1);
     chunks_.resize(final_chunk_index + 1);
   }
 
@@ -128,10 +208,10 @@ class ValueStore
       return;
     }
 
-    auto [begin_chunk_index, begin_pos] = IdToChunkIndices(IdType(size_));
+    auto [begin_chunk_index, begin_pos] = IdToChunkIndices(size_);
     // Use an inclusive range so that if `size` would be the next chunk, we
     // don't try doing something with it.
-    auto [end_chunk_index, end_pos] = IdToChunkIndices(IdType(size - 1));
+    auto [end_chunk_index, end_pos] = IdToChunkIndices(size - 1);
     chunks_.resize(end_chunk_index + 1);
 
     // If the begin and end chunks are the same, we only fill from begin to end.
@@ -192,13 +272,34 @@ class ValueStore
     // `mapped_iterator` incorrectly infers the pointer type for `PointerProxy`.
     // NOLINTNEXTLINE(readability-const-return-type)
     auto index_to_id = [&](int32_t i) -> const std::pair<IdType, ConstRefType> {
-      return std::pair<IdType, ConstRefType>(IdType(i), Get(IdType(i)));
+      IdType id(tag_.Apply(i));
+      return std::pair<IdType, ConstRefType>(id, Get(id));
     };
     // Because indices into `ValueStore` are all sequential values from 0, we
     // can use llvm::seq to walk all indices in the store.
     return llvm::map_range(llvm::seq(size_), index_to_id);
   }
 
+  auto GetCheckIRIdTag() const -> CheckIRIdTag { return tag_; }
+  auto GetRawIndex(IdT id) const -> int32_t {
+    auto index = tag_.Remove(id.index);
+    CARBON_DCHECK(index >= 0, "{0}", index);
+    // Attempt to decompose id.index to include extra detail in the check here
+#ifndef NDEBUG
+    if (index >= size_) {
+      auto [check_ir_id, decomposed_index] =
+          CheckIRIdTag::DecomposeWithBestEffort(id.index);
+      CARBON_DCHECK(
+          index < size_,
+          "Untagged index was outside of container range. Possibly tagged "
+          "index {0}. Best-effort decomposition: CheckIRId: {1}, index: {2}. "
+          "The CheckIRIdTag for this container is: {3}",
+          id.index, check_ir_id, decomposed_index, tag_.GetCheckIRId());
+    }
+#endif
+    return index;
+  }
+
  private:
   // A chunk of `ValueType`s which has a fixed capacity, but variable size.
   // Tracks the size internally for verifying bounds.
@@ -314,7 +415,7 @@ class ValueStore
 
   // Converts an id into an index into the set of chunks, and an offset into
   // that specific chunk. Looks for index overflow in non-optimized builds.
-  static auto IdToChunkIndices(IdType id) -> std::pair<int32_t, int32_t> {
+  static auto IdToChunkIndices(int32_t index) -> std::pair<int32_t, int32_t> {
     constexpr auto LowBits = Chunk::IndexBits();
 
     // Verify there are no unused bits when indexing up to the `Capacity`. This
@@ -328,16 +429,18 @@ class ValueStore
     static_assert(LowBits < 30);
 
     // The index of the chunk is the high bits.
-    auto chunk = id.index >> LowBits;
+    auto chunk = index >> LowBits;
     // The index into the chunk is the low bits.
-    auto pos = id.index & ((1 << LowBits) - 1);
+    auto pos = index & ((1 << LowBits) - 1);
     return {chunk, pos};
   }
 
   // Number of elements added to the store. The number should never exceed what
   // fits in an `int32_t`, which is checked in non-optimized builds in Add().
   int32_t size_ = 0;
 
+  CheckIRIdTag tag_;
+
   // Storage for the `ValueType` objects, indexed by the id. We use a vector of
   // chunks of `ValueType` instead of just a vector of `ValueType` so that
   // addresses of `ValueType` objects are stable. This allows the rest of the

toolchain/check/check_unit.cpp

@@ -195,7 +195,7 @@ auto CheckUnit::CollectTransitiveImports(SemIR::InstId import_decl_id,
   // exported to this IR.
   auto ir_to_result_index =
       FixedSizeValueStore<SemIR::CheckIRId, int>::MakeWithExplicitSize(
-          unit_and_imports_->unit->total_ir_count, -1);
+          CheckIRIdTag(), unit_and_imports_->unit->total_ir_count, -1);
 
   // First add direct imports. This means that if an entity is imported both
   // directly and indirectly, the import path will reflect the direct import.

toolchain/check/context.cpp

@@ -32,9 +32,9 @@ Context::Context(DiagnosticEmitterBase* emitter,
       scope_stack_(sem_ir_),
       deferred_definition_worklist_(vlog_stream),
       vtable_stack_("vtable_stack_", *sem_ir, vlog_stream),
-      check_ir_map_(
-          FixedSizeValueStore<SemIR::CheckIRId, SemIR::ImportIRId>::
-              MakeWithExplicitSize(total_ir_count_, SemIR::ImportIRId::None)),
+      check_ir_map_(FixedSizeValueStore<SemIR::CheckIRId, SemIR::ImportIRId>::
+                        MakeWithExplicitSize(CheckIRIdTag(), total_ir_count_,
+                                             SemIR::ImportIRId::None)),
       global_init_(this),
       region_stack_([this](SemIR::LocId loc_id, std::string label) {
         TODO(loc_id, label);

toolchain/check/import_ref.cpp

@@ -36,8 +36,9 @@ namespace Carbon::Check {
 // Adds the ImportIR, excluding the update to the check_ir_map.
 static auto InternalAddImportIR(Context& context, SemIR::ImportIR import_ir)
     -> SemIR::ImportIRId {
-  context.import_ir_constant_values().push_back(
-      SemIR::ConstantValueStore(SemIR::ConstantId::None));
+  context.import_ir_constant_values().push_back(SemIR::ConstantValueStore(
+      SemIR::ConstantId::None,
+      import_ir.sem_ir ? &import_ir.sem_ir->insts() : nullptr));
   return context.import_irs().Add(import_ir);
 }