Add tests for WeakIdentifierMap and improve/simplify it

I added the tests to understand WeakIdentifieryMap behavior, e.g.
whether the map entry will be automatically removed when the object is
garbage collected, etc. The tests look also helpful for others.

Also simplify WeakIdentifierMap by removing the non-garbage-collecting
version which didn't look working. I tried that in test and the compiler
reported error requiring the type to be garbage collected.

Change-Id: Ic3f8d99d19d526d1a7a8e3940cf2e117323b9c20
Reviewed-on: https://chromium-review.googlesource.com/c/1435536
Commit-Queue: Xianzhu Wang <wangxianzhu@chromium.org>
Reviewed-by: Andrey Kosyakov <caseq@chromium.org>
Cr-Commit-Position: refs/heads/master@{#626145}

third_party/blink/renderer/core/BUILD.gn

@@ -1829,6 +1829,7 @@ jumbo_source_set("unit_tests") {
     "dom/tree_scope_adopter_test.cc",
     "dom/tree_scope_test.cc",
     "dom/user_gesture_indicator_test.cc",
+    "dom/weak_identifier_map_test.cc",
     "dom/whitespace_attacher_test.cc",
     "editing/caret_display_item_client_test.cc",
     "editing/finder/text_finder_test.cc",

third_party/blink/renderer/core/dom/weak_identifier_map.h

@@ -5,7 +5,9 @@
 #ifndef THIRD_PARTY_BLINK_RENDERER_CORE_DOM_WEAK_IDENTIFIER_MAP_H_
 #define THIRD_PARTY_BLINK_RENDERER_CORE_DOM_WEAK_IDENTIFIER_MAP_H_
 
+#include <limits>
 #include "third_party/blink/renderer/platform/heap/handle.h"
+#include "third_party/blink/renderer/platform/heap/persistent.h"
 #include "third_party/blink/renderer/platform/wtf/allocator.h"
 #include "third_party/blink/renderer/platform/wtf/hash_map.h"
 
@@ -14,53 +16,28 @@ namespace blink {
 // TODO(sof): WeakIdentifierMap<> belongs (out) in wtf/, but
 // cannot until GarbageCollected<> can be used from WTF.
 
-template <typename T, typename IdentifierType, bool isGarbageCollected>
-class WeakIdentifierMapBase {
-  USING_FAST_MALLOC(WeakIdentifierMapBase);
-
- protected:
-  using ObjectToIdentifier = HashMap<T*, IdentifierType>;
-  using IdentifierToObject = HashMap<IdentifierType, T*>;
-
-  ObjectToIdentifier object_to_identifier_;
-  IdentifierToObject identifier_to_object_;
-};
-
-template <typename T, typename IdentifierType>
-class WeakIdentifierMapBase<T, IdentifierType, true>
-    : public GarbageCollected<WeakIdentifierMapBase<T, IdentifierType, true>> {
+template <typename T, typename IdentifierType = int>
+class WeakIdentifierMap final
+    : public GarbageCollected<WeakIdentifierMap<T, IdentifierType>> {
  public:
+  WeakIdentifierMap() = default;
+
   void Trace(Visitor* visitor) {
     visitor->Trace(object_to_identifier_);
     visitor->Trace(identifier_to_object_);
   }
 
- protected:
-  using ObjectToIdentifier = HeapHashMap<WeakMember<T>, IdentifierType>;
-  using IdentifierToObject = HeapHashMap<IdentifierType, WeakMember<T>>;
-
-  ObjectToIdentifier object_to_identifier_;
-  IdentifierToObject identifier_to_object_;
-};
-
-template <typename T, typename IdentifierType = int>
-class WeakIdentifierMap final
-    : public WeakIdentifierMapBase<T,
-                                   IdentifierType,
-                                   IsGarbageCollectedType<T>::value> {
- public:
   static IdentifierType Identifier(T* object) {
     IdentifierType result = Instance().object_to_identifier_.at(object);
 
     if (WTF::IsHashTraitsEmptyValue<HashTraits<IdentifierType>>(result)) {
-      result = Next();
-      Instance().Put(object, result);
+      do {
+        result = Next();
+      } while (!LIKELY(Instance().Put(object, result)));
     }
     return result;
   }
 
-  WeakIdentifierMap() = default;
-
   static T* Lookup(IdentifierType identifier) {
     return Instance().identifier_to_object_.at(identifier);
   }
@@ -69,25 +46,49 @@ class WeakIdentifierMap final
     Instance().ObjectDestroyed(object);
   }
 
+  static void SetLastIdForTesting(IdentifierType i) { LastIdRef() = i; }
+
+  static size_t GetSizeForTesting() {
+    return Instance().object_to_identifier_.size();
+  }
+
  private:
   static WeakIdentifierMap<T, IdentifierType>& Instance();
 
   static IdentifierType Next() {
+    // On overflow, skip negative values for signed IdentifierType, and 0 which
+    // is not a valid key in HashMap by default.
+    if (UNLIKELY(LastIdRef() == std::numeric_limits<IdentifierType>::max()))
+      LastIdRef() = 0;
+    return ++LastIdRef();
+  }
+
+  static IdentifierType& LastIdRef() {
     static IdentifierType last_id = 0;
-    return ++last_id;
+    return last_id;
   }
 
-  void Put(T* object, IdentifierType identifier) {
-    DCHECK(object && !this->object_to_identifier_.Contains(object));
-    this->object_to_identifier_.Set(object, identifier);
-    this->identifier_to_object_.Set(identifier, object);
+  bool Put(T* object, IdentifierType identifier) {
+    if (!LIKELY(identifier_to_object_.insert(identifier, object).is_new_entry))
+      return false;
+    DCHECK(object && !object_to_identifier_.Contains(object));
+    object_to_identifier_.Set(object, identifier);
+    DCHECK_EQ(object_to_identifier_.size(), identifier_to_object_.size());
+    return true;
   }
 
   void ObjectDestroyed(T* object) {
-    IdentifierType identifier = this->object_to_identifier_.Take(object);
+    IdentifierType identifier = object_to_identifier_.Take(object);
     if (!WTF::IsHashTraitsEmptyValue<HashTraits<IdentifierType>>(identifier))
-      this->identifier_to_object_.erase(identifier);
+      identifier_to_object_.erase(identifier);
+    DCHECK_EQ(object_to_identifier_.size(), identifier_to_object_.size());
   }
+
+  using ObjectToIdentifier = HeapHashMap<WeakMember<T>, IdentifierType>;
+  using IdentifierToObject = HeapHashMap<IdentifierType, WeakMember<T>>;
+
+  ObjectToIdentifier object_to_identifier_;
+  IdentifierToObject identifier_to_object_;
 };
 
 #define DECLARE_WEAK_IDENTIFIER_MAP(T, ...)        \

third_party/blink/renderer/core/dom/weak_identifier_map_test.cc

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "third_party/blink/renderer/core/dom/weak_identifier_map.h"
+
+#include "testing/gtest/include/gtest/gtest.h"
+#include "third_party/blink/renderer/platform/heap/garbage_collected.h"
+
+namespace blink {
+
+class WeakIdentifierMapTest : public ::testing::Test {
+ public:
+  class TestClass final : public GarbageCollectedFinalized<TestClass> {
+   public:
+    virtual void Trace(Visitor*) {}
+  };
+
+  using TestMap = WeakIdentifierMap<TestClass>;
+
+  void SetUp() override;
+  void TearDown() override;
+
+  void CollectGarbage() {
+    ThreadState::Current()->CollectGarbage(
+        BlinkGC::kNoHeapPointersOnStack, BlinkGC::kAtomicMarking,
+        BlinkGC::kEagerSweeping, BlinkGC::GCReason::kForcedGC);
+  }
+};
+
+DECLARE_WEAK_IDENTIFIER_MAP(WeakIdentifierMapTest::TestClass);
+DEFINE_WEAK_IDENTIFIER_MAP(WeakIdentifierMapTest::TestClass);
+
+void WeakIdentifierMapTest::SetUp() {
+  EXPECT_EQ(0u, TestMap::GetSizeForTesting());
+}
+
+void WeakIdentifierMapTest::TearDown() {
+  CollectGarbage();
+  EXPECT_EQ(0u, TestMap::GetSizeForTesting());
+}
+
+TEST_F(WeakIdentifierMapTest, Basic) {
+  auto* a = MakeGarbageCollected<TestClass>();
+  auto* b = MakeGarbageCollected<TestClass>();
+
+  auto id_a = TestMap::Identifier(a);
+  EXPECT_NE(0, id_a);
+  EXPECT_EQ(id_a, TestMap::Identifier(a));
+  EXPECT_EQ(a, TestMap::Lookup(id_a));
+
+  auto id_b = TestMap::Identifier(b);
+  EXPECT_NE(0, id_b);
+  EXPECT_NE(id_a, id_b);
+  EXPECT_EQ(id_b, TestMap::Identifier(b));
+  EXPECT_EQ(b, TestMap::Lookup(id_b));
+
+  EXPECT_EQ(id_a, TestMap::Identifier(a));
+  EXPECT_EQ(a, TestMap::Lookup(id_a));
+
+  EXPECT_EQ(2u, TestMap::GetSizeForTesting());
+}
+
+TEST_F(WeakIdentifierMapTest, NotifyObjectDestroyed) {
+  auto* a = MakeGarbageCollected<TestClass>();
+  auto id_a = TestMap::Identifier(a);
+  TestMap::NotifyObjectDestroyed(a);
+  EXPECT_EQ(nullptr, TestMap::Lookup(id_a));
+
+  // Simulate that an object is newly allocated at the same address.
+  EXPECT_NE(id_a, TestMap::Identifier(a));
+}
+
+TEST_F(WeakIdentifierMapTest, GarbageCollected) {
+  auto* a = MakeGarbageCollected<TestClass>();
+  auto id_a = TestMap::Identifier(a);
+
+  a = nullptr;
+  CollectGarbage();
+  EXPECT_EQ(nullptr, TestMap::Lookup(id_a));
+}
+
+TEST_F(WeakIdentifierMapTest, UnusedID) {
+  auto* a = MakeGarbageCollected<TestClass>();
+  auto id_a = TestMap::Identifier(a);
+  EXPECT_EQ(nullptr, TestMap::Lookup(id_a + 1));
+}
+
+TEST_F(WeakIdentifierMapTest, Overflow) {
+  TestMap::SetLastIdForTesting(0);
+  auto* a = MakeGarbageCollected<TestClass>();
+  EXPECT_EQ(1, TestMap::Identifier(a));
+  EXPECT_EQ(a, TestMap::Lookup(1));
+
+  TestMap::SetLastIdForTesting(INT_MAX - 1);
+  auto* b = MakeGarbageCollected<TestClass>();
+  EXPECT_EQ(INT_MAX, TestMap::Identifier(b));
+  EXPECT_EQ(b, TestMap::Lookup(INT_MAX));
+
+  auto* c = MakeGarbageCollected<TestClass>();
+  EXPECT_EQ(2, TestMap::Identifier(c));
+  EXPECT_EQ(c, TestMap::Lookup(2));
+
+  DCHECK_EQ(3u, TestMap::GetSizeForTesting());
+}
+
+}  // namespace blink