Add edges to WeakMap values (facebook#1250)

Summary:
Pull Request resolved: facebook#1250

In new WeakMap implementation, we removed the value storage for
values, instead, they are stored in a native set. So the edges to
them aren't  added automatically. This diff adds them back in
`_snapshotAddEdgesImpl()`.

Reviewed By: neildhar

Differential Revision: D52740140

fbshipit-source-id: e5ca972ef3383f6bd8d230c62b2f7c72bb646536

include/hermes/VM/WeakRef.h

@@ -130,9 +130,11 @@ class WeakMapEntryRef {
       : slot_(
             runtime.getHeap().allocWeakMapEntrySlot(key, value, ownerMapPtr)) {}
 
-  /// \return a pointer to the key object.
-  GCCell *getKeyNoBarrierUnsafe(PointerBase &base) const {
-    return slot_->key.getNoBarrierUnsafe(base);
+  /// \return a pointer to the key object with read barrier. This should not be
+  /// called when the weak ref key object is collected.
+  GCCell *getKeyNonNull(PointerBase &base, GC &gc) const {
+    assert(isKeyValid() && "tried to access collected weak ref object");
+    return slot_->key.getNonNull(base, gc);
   }
 
   /// \return The mapped value by the the key object.

lib/VM/JSWeakMapImpl.cpp

@@ -145,7 +145,7 @@ void JSWeakMapImplBase::_snapshotAddEdgesImpl(
         HeapSnapshot::EdgeType::Internal, "map", self->getMapID(gc));
   }
 
-  // Add edges to objects pointed by WeakRef keys.
+  // Add edges to objects pointed by WeakRef keys and its mapped values.
   uint32_t edge_index = 0;
   for (const auto &key : self->set_) {
     // Skip if the ref is not valid.
@@ -156,7 +156,24 @@ void JSWeakMapImplBase::_snapshotAddEdgesImpl(
     snap.addNamedEdge(
         HeapSnapshot::EdgeType::Weak,
         indexName,
-        gc.getObjectID(key.ref.getKeyNoBarrierUnsafe(gc.getPointerBase())));
+        gc.getObjectID(key.ref.getKeyNonNull(gc.getPointerBase(), gc)));
+
+    auto mappedValue = key.ref.getMappedValue(gc);
+    // TODO(T175014649): nodes for numbers may not exist since they are not seen
+    // by PrimitiveNodeAcceptor. We can't simply add them in
+    // _snapshotAddNodesImpl() either, because PrimitiveNodeAcceptor may see the
+    // same number in a heap object and the assertion of not writing duplicate
+    // node will fail.
+    if (mappedValue.isNumber()) {
+      continue;
+    }
+    if (auto id = gc.getSnapshotID(mappedValue)) {
+      snap.addNamedEdge(
+          HeapSnapshot::EdgeType::Internal,
+          // Add a suffix to distinguish key and value.
+          indexName + "[value]",
+          id.getValue());
+    }
   }
 }
 

unittests/VMRuntime/HeapSnapshotTest.cpp

@@ -882,8 +882,8 @@ TEST_F(HeapSnapshotRuntimeTest, WeakMapTest) {
           "JSWeakMap",
           mapID,
           map->getAllocatedSize(),
-          firstNamed + 2));
-  EXPECT_EQ(nodesAndEdges.second.size(), firstNamed + 2);
+          firstNamed + 3));
+  EXPECT_EQ(nodesAndEdges.second.size(), firstNamed + 3);
 
   // Test the native edge.
   const auto nativeMapID = map->getMapID(runtime.getHeap());