WIP

Author: dwblaikie

toolchain/base/BUILD

@@ -65,6 +65,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

@@ -10,6 +10,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 {
@@ -44,9 +45,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;
   }
@@ -57,7 +60,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);
   }
 
@@ -70,7 +74,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;
@@ -79,20 +84,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.
@@ -119,6 +127,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,42 @@ 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.
+    check_id_tag_ = llvm::reverseBits((((check_id_index + 1) << 1) | 1) << 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_;
+  }
+  auto Remove(int32_t tagged_index) const -> int32_t {
+    if ((llvm::reverseBits(2) & tagged_index) == 0) {
+      CARBON_CHECK(tagged_index < initial_reserved_ids_);
+      return tagged_index;
+    }
+    auto index = tagged_index ^ check_id_tag_;
+    CARBON_CHECK(index >= initial_reserved_ids_);
+    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 +110,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 +119,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 +136,19 @@ 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 = tag_.Remove(id.index);
+    CARBON_DCHECK(index >= 0, "{0}", index);
+    CARBON_DCHECK(index < size_, "{0}", index);
+    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 = tag_.Remove(id.index);
+    CARBON_DCHECK(index >= 0, "{0}", index);
+    CARBON_DCHECK(index < size_, "{0}", index);
+    auto [chunk_index, pos] = IdToChunkIndices(index);
     return chunks_[chunk_index].Get(pos);
   }
 
@@ -118,7 +157,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 +167,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,7 +231,8 @@ 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.
@@ -314,7 +354,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 +368,20 @@ 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;
 
+ public:
+  CheckIRIdTag tag_;
+
+ private:
   // 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(
-          total_ir_count_, -1);
+          CheckIRIdTag(), 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

@@ -31,9 +31,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

@@ -35,8 +35,11 @@ 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 ? CheckIRIdTag(import_ir.sem_ir->check_ir_id().index,
+                                      SemIR::SingletonInstKinds.size() + 1)
+                       : CheckIRIdTag(0, std::numeric_limits<int32_t>::max())));
   return context.import_irs().Add(import_ir);
 }
 
@@ -642,6 +645,7 @@ static auto AddImportRef(ImportContext& context, SemIR::InstId inst_id,
 static auto AddLoadedImportRef(ImportContext& context, SemIR::TypeId type_id,
                                SemIR::InstId inst_id,
                                SemIR::ConstantId const_id) -> SemIR::InstId {
+  context.import_insts().Get(inst_id);
   return AddLoadedImportRef(
       context.local_context(), type_id,
       SemIR::ImportIRInst(context.import_ir_id(), inst_id), const_id);
@@ -2088,7 +2092,8 @@ static auto TryResolveTypedInst(ImportRefResolver& resolver,
           resolver,
           GetSingletonType(resolver.local_context(),
                            SemIR::SpecificFunctionType::TypeInstId),
-          import_vtable_entry_inst_id, local_attached_constant_id);
+          /* bad: local_vtable_entry_inst_id */ import_vtable_entry_inst_id,
+          local_attached_constant_id);
     }
   }