Bug 1796081 - Part 2: Only check the key color once when populating the ephemeron table r=sfink

Currently the table population is done separately from marking the entry which
means we check the key color twice, potentially getting different results if
another thread is marking at the same time. This could result in us not adding
ephemeron table entries when necessary.

The patch consolidates these two paths.

Depends on D159677

Differential Revision: https://phabricator.services.mozilla.com/D159678

js/src/gc/WeakMap-inl.h

@@ -107,8 +107,15 @@ WeakMap<K, V>::WeakMap(JS::Zone* zone, JSObject* memOf)
 // If the entry is live, ensure its key and value are marked. Also make sure the
 // key is at least as marked as min(map, delegate), so it cannot get discarded
 // and then recreated by rewrapping the delegate.
+//
+// Optionally adds edges to the ephemeron edges table for any keys (or
+// delegates) where future changes to their mark color would require marking the
+// value (or the key).
 template <class K, class V>
-bool WeakMap<K, V>::markEntry(GCMarker* marker, K& key, V& value) {
+bool WeakMap<K, V>::markEntry(GCMarker* marker, K& key, V& value,
+                              bool populateWeakKeysTable) {
+  MOZ_ASSERT(mapColor);
+
   bool marked = false;
   CellColor markColor = marker->markColor();
   CellColor keyColor = gc::detail::GetEffectiveColor(marker, key);
@@ -116,7 +123,6 @@ bool WeakMap<K, V>::markEntry(GCMarker* marker, K& key, V& value) {
 
   if (delegate) {
     CellColor delegateColor = gc::detail::GetEffectiveColor(marker, delegate);
-    MOZ_ASSERT(mapColor);
     // The key needs to stay alive while both the delegate and map are live.
     CellColor proxyPreserveColor = std::min(delegateColor, mapColor);
     if (keyColor < proxyPreserveColor) {
@@ -131,8 +137,8 @@ bool WeakMap<K, V>::markEntry(GCMarker* marker, K& key, V& value) {
     }
   }
 
+  gc::Cell* cellValue = gc::ToMarkable(value);
   if (keyColor) {
-    gc::Cell* cellValue = gc::ToMarkable(value);
     if (cellValue) {
       CellColor targetColor = std::min(mapColor, keyColor);
       CellColor valueColor = gc::detail::GetEffectiveColor(marker, cellValue);
@@ -147,6 +153,28 @@ bool WeakMap<K, V>::markEntry(GCMarker* marker, K& key, V& value) {
     }
   }
 
+  if (populateWeakKeysTable) {
+    // Note that delegateColor >= keyColor because marking a key marks its
+    // delegate, so we only need to check whether keyColor < mapColor to tell
+    // this.
+
+    if (keyColor < mapColor) {
+      MOZ_ASSERT(marker->weakMapAction() == JS::WeakMapTraceAction::Expand);
+      // The final color of the key is not yet known. Record this weakmap and
+      // the lookup key in the list of weak keys. If the key has a delegate,
+      // then the lookup key is the delegate (because marking the key will end
+      // up marking the delegate and thereby mark the entry.)
+      gc::TenuredCell* tenuredValue = nullptr;
+      if (cellValue && cellValue->isTenured()) {
+        tenuredValue = &cellValue->asTenured();
+      }
+
+      if (!this->addImplicitEdges(key, delegate, tenuredValue)) {
+        marker->abortLinearWeakMarking();
+      }
+    }
+  }
+
   return marked;
 }
 
@@ -245,45 +273,16 @@ bool WeakMap<K, V>::markEntries(GCMarker* marker) {
   MOZ_ASSERT(mapColor);
   bool markedAny = false;
 
+  // If we don't populate the weak keys table now then we do it when we enter
+  // weak marking mode.
+  bool populateWeakKeysTable =
+      marker->incrementalWeakMapMarkingEnabled || marker->isWeakMarking();
+
   for (Enum e(*this); !e.empty(); e.popFront()) {
-    if (markEntry(marker, e.front().mutableKey(), e.front().value())) {
+    if (markEntry(marker, e.front().mutableKey(), e.front().value(),
+                  populateWeakKeysTable)) {
       markedAny = true;
     }
-    if (!marker->incrementalWeakMapMarkingEnabled && !marker->isWeakMarking()) {
-      // Populate weak keys table when we enter weak marking mode.
-      continue;
-    }
-
-    // Adds edges to the ephemeron edges table for any keys (or delegates) where
-    // future changes to their mark color would require marking the value (or
-    // the key).
-    //
-    // Note that delegateColor >= keyColor because marking a key marks its
-    // delegate, so we only need to check whether keyColor < mapColor to tell
-    // this.
-
-    CellColor keyColor =
-        gc::detail::GetEffectiveColor(marker, e.front().key().get());
-
-    if (keyColor < mapColor) {
-      MOZ_ASSERT(marker->weakMapAction() == JS::WeakMapTraceAction::Expand);
-      // The final color of the key is not yet known. Record this weakmap and
-      // the lookup key in the list of weak keys. If the key has a delegate,
-      // then the lookup key is the delegate (because marking the key will end
-      // up marking the delegate and thereby mark the entry.)
-      gc::Cell* weakKey = e.front().key();
-      gc::Cell* value = gc::ToMarkable(e.front().value());
-      gc::Cell* delegate = gc::detail::GetDelegate(e.front().key());
-
-      gc::TenuredCell* tenuredValue = nullptr;
-      if (value && value->isTenured()) {
-        tenuredValue = &value->asTenured();
-      }
-
-      if (!addImplicitEdges(weakKey, delegate, tenuredValue)) {
-        marker->abortLinearWeakMarking();
-      }
-    }
   }
 
   return markedAny;

js/src/gc/WeakMap.h

@@ -277,7 +277,8 @@ class WeakMap
   }
 #endif
 
-  bool markEntry(GCMarker* marker, Key& key, Value& value);
+  bool markEntry(GCMarker* marker, Key& key, Value& value,
+                 bool populateWeakKeysTable);
 
   void trace(JSTracer* trc) override;