[identity-map] Remove double-lookups in IdentityMap

Remove the pattern of calling 'Find' followed by 'Set' for IdentityMap,
with a single 'FindOrInsert' that explicitly returns whether an existing
entry was found, or the entry was inserted. This replaces 'Get', which
would return either an initialised or uninitialised entry (and callers
would rely on default initialisation to check this).

Also replace 'Set' with 'Insert', which explicitly requires that the
element didn't exist before. This matches expectations where it was
used (where those weren't replaced wholesale with 'FindOrInsert'), and
makes the naming consistent with 'FindOrInsert'.

Change-Id: I8fb76f4ac14fb92b88474965aafb1ace5fb79145
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2443135
Auto-Submit: Leszek Swirski <leszeks@chromium.org>
Reviewed-by: Dominik Inführ <dinfuehr@chromium.org>
Reviewed-by: Maya Lekova <mslekova@chromium.org>
Reviewed-by: Ross McIlroy <rmcilroy@chromium.org>
Commit-Queue: Dominik Inführ <dinfuehr@chromium.org>
Commit-Queue: Maya Lekova <mslekova@chromium.org>
Cr-Commit-Position: refs/heads/master@{#70300}

src/builtins/constants-table-builder.cc

@@ -46,15 +46,12 @@ uint32_t BuiltinsConstantsTableBuilder::AddObject(Handle<Object> object) {
   DCHECK(!object->IsCode());
 #endif
 
-  uint32_t* maybe_key = map_.Find(object);
-  if (maybe_key == nullptr) {
+  auto find_result = map_.FindOrInsert(object);
+  if (!find_result.already_exists) {
     DCHECK(object->IsHeapObject());
-    uint32_t index = map_.size();
-    map_.Set(object, index);
-    return index;
-  } else {
-    return *maybe_key;
+    *find_result.entry = map_.size() - 1;
   }
+  return *find_result.entry;
 }
 
 namespace {
@@ -85,7 +82,7 @@ void BuiltinsConstantsTableBuilder::PatchSelfReference(
   uint32_t key;
   if (map_.Delete(self_reference, &key)) {
     DCHECK(code_object->IsCode());
-    map_.Set(code_object, key);
+    map_.Insert(code_object, key);
   }
 }
 
@@ -96,7 +93,7 @@ void BuiltinsConstantsTableBuilder::PatchBasicBlockCountersReference(
   uint32_t key;
   if (map_.Delete(ReadOnlyRoots(isolate_).basic_block_counters_marker(),
                   &key)) {
-    map_.Set(counters, key);
+    map_.Insert(counters, key);
   }
 }
 

src/compiler-dispatcher/compiler-dispatcher.cc

@@ -114,7 +114,7 @@ void CompilerDispatcher::RegisterSharedFunctionInfo(
   auto job_it = jobs_.find(job_id);
   DCHECK_NE(job_it, jobs_.end());
   Job* job = job_it->second.get();
-  shared_to_unoptimized_job_id_.Set(function_handle, job_id);
+  shared_to_unoptimized_job_id_.Insert(function_handle, job_id);
 
   {
     base::MutexGuard lock(&mutex_);

src/compiler/js-heap-broker.cc

@@ -2441,7 +2441,7 @@ void JSHeapBroker::CopyCanonicalHandlesForTesting(
   for (auto it = it_scope.begin(); it != it_scope.end(); ++it) {
     Address* entry = *it.entry();
     Object key = it.key();
-    canonical_handles_->Set(key, entry);
+    canonical_handles_->Insert(key, entry);
   }
 }
 

src/compiler/js-heap-broker.h

@@ -250,13 +250,13 @@ class V8_EXPORT_PRIVATE JSHeapBroker {
       }
 
       Object obj(address);
-      Address** entry = canonical_handles_->Get(obj);
-      if (*entry == nullptr) {
+      auto find_result = canonical_handles_->FindOrInsert(obj);
+      if (!find_result.already_exists) {
         // Allocate new PersistentHandle if one wasn't created before.
         DCHECK(local_heap_);
-        *entry = local_heap_->NewPersistentHandle(obj).location();
+        *find_result.entry = local_heap_->NewPersistentHandle(obj).location();
       }
-      return Handle<T>(*entry);
+      return Handle<T>(*find_result.entry);
     } else {
       return Handle<T>(object, isolate());
     }

src/handles/handles.cc

@@ -175,12 +175,12 @@ Address* CanonicalHandleScope::Lookup(Address object) {
       return isolate_->root_handle(root_index).location();
     }
   }
-  Address** entry = identity_map_->Get(Object(object));
-  if (*entry == nullptr) {
+  auto find_result = identity_map_->FindOrInsert(Object(object));
+  if (!find_result.already_exists) {
     // Allocate new handle location.
-    *entry = HandleScope::CreateHandle(isolate_, object);
+    *find_result.entry = HandleScope::CreateHandle(isolate_, object);
   }
-  return *entry;
+  return *find_result.entry;
 }
 
 std::unique_ptr<CanonicalHandlesMap>

src/objects/value-serializer.cc

@@ -384,7 +384,7 @@ void ValueSerializer::TransferArrayBuffer(uint32_t transfer_id,
                                           Handle<JSArrayBuffer> array_buffer) {
   DCHECK(!array_buffer_transfer_map_.Find(array_buffer));
   DCHECK(!array_buffer->is_shared());
-  array_buffer_transfer_map_.Set(array_buffer, transfer_id);
+  array_buffer_transfer_map_.Insert(array_buffer, transfer_id);
 }
 
 Maybe<bool> ValueSerializer::WriteObject(Handle<Object> object) {
@@ -500,16 +500,16 @@ void ValueSerializer::WriteString(Handle<String> string) {
 
 Maybe<bool> ValueSerializer::WriteJSReceiver(Handle<JSReceiver> receiver) {
   // If the object has already been serialized, just write its ID.
-  uint32_t* id_map_entry = id_map_.Get(receiver);
-  if (uint32_t id = *id_map_entry) {
+  auto find_result = id_map_.FindOrInsert(receiver);
+  if (find_result.already_exists) {
     WriteTag(SerializationTag::kObjectReference);
-    WriteVarint(id - 1);
+    WriteVarint(*find_result.entry - 1);
     return ThrowIfOutOfMemory();
   }
 
   // Otherwise, allocate an ID for it.
   uint32_t id = next_id_++;
-  *id_map_entry = id + 1;
+  *find_result.entry = id + 1;
 
   // Eliminate callable and exotic objects, which should not be serialized.
   InstanceType instance_type = receiver->map().instance_type();

src/snapshot/read-only-serializer.cc

@@ -61,7 +61,7 @@ void ReadOnlySerializer::SerializeObjectImpl(Handle<HeapObject> obj) {
     CHECK_NULL(serialized_objects_.Find(obj));
     // There's no "IdentitySet", so use an IdentityMap with a value that is
     // later ignored.
-    serialized_objects_.Set(obj, 0);
+    serialized_objects_.Insert(obj, 0);
   }
 #endif
 }

src/snapshot/references.h

@@ -121,7 +121,7 @@ class SerializerReferenceMap {
 
   void Add(HeapObject object, SerializerReference reference) {
     DCHECK_NULL(LookupReference(object));
-    map_.Set(object, reference);
+    map_.Insert(object, reference);
   }
 
   void AddBackingStore(void* backing_store, SerializerReference reference) {
@@ -132,7 +132,7 @@ class SerializerReferenceMap {
   SerializerReference AddAttachedReference(HeapObject object) {
     SerializerReference reference =
         SerializerReference::AttachedReference(attached_reference_index_++);
-    map_.Set(object, reference);
+    map_.Insert(object, reference);
     return reference;
   }
 

src/snapshot/serializer.cc

@@ -298,21 +298,16 @@ void Serializer::ResolvePendingForwardReference(int forward_reference_id) {
 void Serializer::RegisterObjectIsPending(Handle<HeapObject> obj) {
   if (*obj == ReadOnlyRoots(isolate()).not_mapped_symbol()) return;
 
-#ifdef DEBUG
-  bool already_exists = forward_refs_per_pending_object_.Find(obj) != nullptr;
-#endif
-
   // Add the given object to the pending objects -> forward refs map.
-  PendingObjectReferences* refs = forward_refs_per_pending_object_.Get(obj);
-
-  DCHECK_EQ(already_exists, *refs != nullptr);
-  USE(refs);
+  auto find_result = forward_refs_per_pending_object_.FindOrInsert(obj);
+  USE(find_result);
 
   // If the above emplace didn't actually add the object, then the object must
   // already have been registered pending by deferring. It might not be in the
   // deferred objects queue though, since it may be the very object we just
   // popped off that queue, so just check that it can be deferred.
-  DCHECK_IMPLIES(already_exists, CanBeDeferred(*obj));
+  DCHECK_IMPLIES(find_result.already_exists, *find_result.entry != nullptr);
+  DCHECK_IMPLIES(find_result.already_exists, CanBeDeferred(*obj));
 }
 
 void Serializer::ResolvePendingObject(Handle<HeapObject> obj) {

src/snapshot/serializer.h

@@ -139,14 +139,12 @@ class ObjectCacheIndexMap {
   // map and return false. In either case set |*index_out| to the index
   // associated with the map.
   bool LookupOrInsert(Handle<HeapObject> obj, int* index_out) {
-    int* maybe_index = map_.Find(obj);
-    if (maybe_index) {
-      *index_out = *maybe_index;
-      return true;
+    auto find_result = map_.FindOrInsert(obj);
+    if (!find_result.already_exists) {
+      *find_result.entry = next_index_++;
     }
-    *index_out = next_index_;
-    map_.Set(obj, next_index_++);
-    return false;
+    *index_out = *find_result.entry;
+    return find_result.already_exists;
   }
 
  private:

src/utils/identity-map.cc

@@ -61,19 +61,19 @@ int IdentityMapBase::ScanKeysFor(Address address) const {
   return -1;
 }
 
-int IdentityMapBase::InsertKey(Address address) {
+std::pair<int, bool> IdentityMapBase::InsertKey(Address address) {
   Address not_mapped = ReadOnlyRoots(heap_).not_mapped_symbol().ptr();
   while (true) {
     int start = Hash(address) & mask_;
     int limit = capacity_ / 2;
     // Search up to {limit} entries.
     for (int index = start; --limit > 0; index = (index + 1) & mask_) {
-      if (keys_[index] == address) return index;  // Found.
+      if (keys_[index] == address) return {index, true};  // Found.
       if (keys_[index] == not_mapped) {           // Free entry.
         size_++;
         DCHECK_LE(size_, capacity_);
         keys_[index] = address;
-        return index;
+        return {index, false};
       }
     }
     // Should only have to resize once, since we grow 4x.
@@ -132,16 +132,19 @@ int IdentityMapBase::Lookup(Address key) const {
   return index;
 }
 
-int IdentityMapBase::LookupOrInsert(Address key) {
+std::pair<int, bool> IdentityMapBase::LookupOrInsert(Address key) {
   // Perform an optimistic lookup.
   int index = ScanKeysFor(key);
+  bool already_exists;
   if (index < 0) {
     // Miss; rehash if there was a GC, then insert.
     if (gc_counter_ != heap_->gc_count()) Rehash();
-    index = InsertKey(key);
+    std::tie(index, already_exists) = InsertKey(key);
+  } else {
+    already_exists = true;
   }
   DCHECK_GE(index, 0);
-  return index;
+  return {index, already_exists};
 }
 
 int IdentityMapBase::Hash(Address address) const {
@@ -151,10 +154,35 @@ int IdentityMapBase::Hash(Address address) const {
 
 // Searches this map for the given key using the object's address
 // as the identity, returning:
-//    found => a pointer to the storage location for the value
-//    not found => a pointer to a new storage location for the value
-IdentityMapBase::RawEntry IdentityMapBase::GetEntry(Address key) {
+//    found => a pointer to the storage location for the value, true
+//    not found => a pointer to a new storage location for the value, false
+IdentityMapFindResult<uintptr_t> IdentityMapBase::FindOrInsertEntry(
+    Address key) {
   CHECK(!is_iterable());  // Don't allow insertion while iterable.
+  if (capacity_ == 0) {
+    return {InsertEntry(key), false};
+  }
+  auto lookup_result = LookupOrInsert(key);
+  return {&values_[lookup_result.first], lookup_result.second};
+}
+
+// Searches this map for the given key using the object's address
+// as the identity, returning:
+//    found => a pointer to the storage location for the value
+//    not found => {nullptr}
+IdentityMapBase::RawEntry IdentityMapBase::FindEntry(Address key) const {
+  // Don't allow find by key while iterable (might rehash).
+  CHECK(!is_iterable());
+  if (size_ == 0) return nullptr;
+  int index = Lookup(key);
+  return index >= 0 ? &values_[index] : nullptr;
+}
+
+// Inserts the given key using the object's address as the identity, returning
+// a pointer to the new storage location for the value.
+IdentityMapBase::RawEntry IdentityMapBase::InsertEntry(Address key) {
+  // Don't allow find by key while iterable (might rehash).
+  CHECK(!is_iterable());
   if (capacity_ == 0) {
     // Allocate the initial storage for keys and values.
     capacity_ = kInitialIdentityMapSize;
@@ -169,22 +197,16 @@ IdentityMapBase::RawEntry IdentityMapBase::GetEntry(Address key) {
 
     strong_roots_entry_ = heap_->RegisterStrongRoots(
         FullObjectSlot(keys_), FullObjectSlot(keys_ + capacity_));
+  } else {
+    // Rehash if there was a GC, then insert.
+    if (gc_counter_ != heap_->gc_count()) Rehash();
   }
-  int index = LookupOrInsert(key);
-  return &values_[index];
-}
 
-// Searches this map for the given key using the object's address
-// as the identity, returning:
-//    found => a pointer to the storage location for the value
-//    not found => {nullptr}
-IdentityMapBase::RawEntry IdentityMapBase::FindEntry(Address key) const {
-  // Don't allow find by key while iterable (might rehash).
-  CHECK(!is_iterable());
-  if (size_ == 0) return nullptr;
-  // Remove constness since lookup might have to rehash.
-  int index = Lookup(key);
-  return index >= 0 ? &values_[index] : nullptr;
+  int index;
+  bool already_exists;
+  std::tie(index, already_exists) = InsertKey(key);
+  DCHECK(!already_exists);
+  return &values_[index];
 }
 
 // Deletes the given key from the map using the object's address as the
@@ -254,7 +276,7 @@ void IdentityMapBase::Rehash() {
   }
   // Reinsert all the key/value pairs that were in the wrong place.
   for (auto pair : reinsert) {
-    int index = InsertKey(pair.first);
+    int index = InsertKey(pair.first).first;
     DCHECK_GE(index, 0);
     values_[index] = pair.second;
   }
@@ -281,7 +303,7 @@ void IdentityMapBase::Resize(int new_capacity) {
 
   for (int i = 0; i < old_capacity; i++) {
     if (old_keys[i] == not_mapped) continue;
-    int index = InsertKey(old_keys[i]);
+    int index = InsertKey(old_keys[i]).first;
     DCHECK_GE(index, 0);
     values_[index] = old_values[i];
   }

src/utils/identity-map.h

@@ -18,6 +18,12 @@ namespace internal {
 class Heap;
 class StrongRootsEntry;
 
+template <typename T>
+struct IdentityMapFindResult {
+  T* entry;
+  bool already_exists;
+};
+
 // Base class of identity maps contains shared code for all template
 // instantions.
 class V8_EXPORT_PRIVATE IdentityMapBase {
@@ -46,8 +52,9 @@ class V8_EXPORT_PRIVATE IdentityMapBase {
         is_iterable_(false) {}
   virtual ~IdentityMapBase();
 
-  RawEntry GetEntry(Address key);
+  IdentityMapFindResult<uintptr_t> FindOrInsertEntry(Address key);
   RawEntry FindEntry(Address key) const;
+  RawEntry InsertEntry(Address key);
   bool DeleteEntry(Address key, uintptr_t* deleted_value);
   void Clear();
 
@@ -65,9 +72,9 @@ class V8_EXPORT_PRIVATE IdentityMapBase {
  private:
   // Internal implementation should not be called directly by subclasses.
   int ScanKeysFor(Address address) const;
-  int InsertKey(Address address);
+  std::pair<int, bool> InsertKey(Address address);
   int Lookup(Address key) const;
-  int LookupOrInsert(Address key);
+  std::pair<int, bool> LookupOrInsert(Address key);
   bool DeleteIndex(int index, uintptr_t* deleted_value);
   void Rehash();
   void Resize(int new_capacity);
@@ -107,10 +114,15 @@ class IdentityMap : public IdentityMapBase {
 
   // Searches this map for the given key using the object's address
   // as the identity, returning:
-  //    found => a pointer to the storage location for the value
-  //    not found => a pointer to a new storage location for the value
-  V* Get(Handle<Object> key) { return Get(*key); }
-  V* Get(Object key) { return reinterpret_cast<V*>(GetEntry(key.ptr())); }
+  //    found => a pointer to the storage location for the value, true
+  //    not found => a pointer to a new storage location for the value, false
+  IdentityMapFindResult<V> FindOrInsert(Handle<Object> key) {
+    return FindOrInsert(*key);
+  }
+  IdentityMapFindResult<V> FindOrInsert(Object key) {
+    auto raw = FindOrInsertEntry(key.ptr());
+    return {reinterpret_cast<V*>(raw.entry), raw.already_exists};
+  }
 
   // Searches this map for the given key using the object's address
   // as the identity, returning:
@@ -121,10 +133,11 @@ class IdentityMap : public IdentityMapBase {
     return reinterpret_cast<V*>(FindEntry(key.ptr()));
   }
 
-  // Set the value for the given key.
-  void Set(Handle<Object> key, V v) { Set(*key, v); }
-  void Set(Object key, V v) {
-    *(reinterpret_cast<V*>(GetEntry(key.ptr()))) = v;
+  // Insert the value for the given key. The key must not have previously
+  // existed.
+  void Insert(Handle<Object> key, V v) { Insert(*key, v); }
+  void Insert(Object key, V v) {
+    *reinterpret_cast<V*>(InsertEntry(key.ptr())) = v;
   }
 
   bool Delete(Handle<Object> key, V* deleted_value) {