Properly handle JSWeakMap semantics.

Summary:
We were not properly handling the semantics of JSWeakMap.  The spec says
(http://www.ecma-international.org/ecma-262/#sec-weakmap-objects) that if an object used a key is only reachable from the the corresponding value, then the key should not be considered reachable, and the key/value pair should (eventually) be removed from the table.

This diff fixes that.  We modify marking to recognize JSWeakMaps.  When these are found reachable, we collect those in a vector, but do not mark through them.  When the marking transitive closure is otherwise complete, we consider the JSWeakMaps.  We find entries with reachable keys, and mark transitively from the corresponding values.  This is done carefully to ensure we reach a full transitive closure.  After this we clear the keys and values of entries with unreachable keys.  Normal weak ref processing will mark the table as having freeable slots.

A subsequent diffs may optimize this basic idea: if we consider a map multiple times, it is wasteful to scan all the entries.  So we should keep track of entries whose keys are not yet known to be unreachable, and consider only them.

Minor things:
* Modify HermesValue::getWeakSize to check if a JSWeakMap has freeable slots, and free those before getting the size, to get a consistent value.
* Factor out "CompleteMarkState::pushCell", now that there's a second use.  Also, assert that the cell pushed is valid, which found a problem in my debugging.

Reviewed By: dulinriley

Differential Revision: D18651375

fbshipit-source-id: 79c06875330147ecf15434bea2c92d3d3141ced9

CMakeLists.txt

@@ -542,6 +542,7 @@ set(HERMES_LIT_TEST_PARAMS
   serialize_enabled=${HERMESVM_SERIALIZE}
   profiler=${HERMES_PROFILER_MODE_IN_LIT_TEST}
   use_js_library_implementation=${HERMESVM_USE_JS_LIBRARY_IMPLEMENTATION}
+  gc=${HERMESVM_GCKIND}
   # TODO: Figure out how to tell if CMake is doing a UBSAN build.
   ubsan=OFF
   )

include/hermes/VM/CompleteMarkState.h

@@ -17,8 +17,13 @@
 namespace hermes {
 namespace vm {
 
+/// Forward declartions.
 struct FullMSCUpdateAcceptor;
 
+template <CellKind>
+class JSWeakMapImpl;
+using JSWeakMap = JSWeakMapImpl<CellKind::WeakMapKind>;
+
 /// Intermediate state from marking.
 struct CompleteMarkState {
   /// Forward declaration of the Acceptor used to mark fields of marked objects.
@@ -32,6 +37,10 @@ struct CompleteMarkState {
   /// traversal).
   void markTransitive(void *ptr);
 
+  /// The pointer \p cell is assumed to point to a reachable object.
+  /// Push it on the appropriate markStack.
+  void pushCell(GCCell *cell);
+
   /// Continually pops elements from the mark stack and scans their pointer
   /// fields.  If such a field points to an unmarked object, mark it and push it
   /// on the mark stack.  If such a push would overflow the mark stack, sets a
@@ -82,6 +91,9 @@ struct CompleteMarkState {
   /// Stores the current object whose fields are being scanned
   /// to be marked if needed.
   GCCell *currentParPointer = nullptr;
+
+  /// The WeakMap objects that have been discovered to be reachable.
+  std::vector<JSWeakMap *> reachableWeakMaps_;
 };
 
 /// Returns a heap acceptor for mark-sweep-compact pointer update.

include/hermes/VM/GenGCNC.h

@@ -666,6 +666,18 @@ class GenGC final : public GCBase {
   /// close the mark bits.
   void completeMarking();
 
+  /// In the first phase of marking, before this is called, we treat
+  /// JSWeakMaps specially: when we mark a reachable JSWeakMap, we do
+  /// not mark from it, but rather save a pointer to it in a vector.
+  /// Then we call this method, which finds the keys that are
+  /// reachable, and marks transitively from the corresponding value.
+  /// This is done carefully, to reach a correct global transitive
+  /// closure, in cases where keys are reachable only via values of
+  /// other keys.  When this marking is done, entries with unreachable
+  /// keys are cleared.  Normal WeakRef processing at the end of GC
+  /// will delete the cleared entries from the map.
+  void completeWeakMapMarking();
+
   /// Does any work necessary for GC stats at the end of collection.
   /// Returns the number of allocated objects before collection starts.
   /// (In optimized builds, does nothing, and returns zero.)

include/hermes/VM/JSWeakMapImpl.h

@@ -30,6 +30,11 @@ struct WeakRefKey {
   uint32_t hash;
 
   WeakRefKey(WeakRef<JSObject> ref, uint32_t hash) : ref(ref), hash(hash) {}
+
+  /// Returns the object reference of ref; returns null if ref is not valid.
+  /// Should only be called during GC; the \param gc argument is used only to
+  /// verify this.
+  JSObject *getObject(GC *gc) const;
 };
 
 /// Enable using WeakRef<JSObject> in DenseMap.
@@ -159,6 +164,48 @@ class JSWeakMapImplBase : public JSObject {
       PointerBase *base,
       JSWeakMapImplBase *self);
 
+  /// An iterator over the keys of the map.
+  struct KeyIterator {
+    DenseMapT::iterator mapIter;
+
+    KeyIterator &operator++(int /*dummy*/) {
+      mapIter++;
+      return *this;
+    }
+
+    WeakRefKey &operator*() {
+      return mapIter->first;
+    }
+
+    WeakRefKey *operator->() {
+      return &mapIter->first;
+    }
+
+    bool operator!=(const KeyIterator &other) {
+      return mapIter != other.mapIter;
+    }
+  };
+
+  // Return begin and end iterators for the keys of the map.
+  KeyIterator keys_begin();
+  KeyIterator keys_end();
+
+  /// Returns the value corresponding to the given \p key.  Returns
+  /// the empty HermesValue if \p key is not in the map.  May only be
+  /// called during GC.
+  /// \param gc Used to verify that the call is during GC, and provides
+  /// a PointerBase.
+  HermesValue getValueDirect(GC *gc, const WeakRefKey &key);
+
+  /// If the given \p key is in the map, clears the entry
+  /// corresponding to \p key -- clears the slot of the WeakRef in
+  /// key, and sets the value to the empty HermesValue.  May only be
+  /// called during GC.
+  /// \param gc Used to verify that the call is during GC, and provides
+  /// a PointerBase.
+  /// \return whether the key was in the map.
+  bool clearEntryDirect(GC *gc, const WeakRefKey &key);
+
  protected:
   static void _finalizeImpl(GCCell *cell, GC *gc) {
     auto *self = vmcast<JSWeakMapImplBase>(cell);

include/hermes/VM/WeakRef.h

@@ -61,7 +61,8 @@ class WeakRef : public WeakRefBase {
   /// This is an unsafe function since the referenced object may be freed any
   /// time that GC occurs.
   OptValue<typename Traits::value_type> unsafeGetOptional() const {
-    return isValid() ? Traits::decode(slot_->value()) : llvm::None;
+    return isValid() ? Traits::decode(slot_->value())
+                     : OptValue<typename Traits::value_type>(llvm::None);
   }
 
   /// Before calling this function the user must check whether weak reference is
@@ -84,6 +85,11 @@ class WeakRef : public WeakRefBase {
     assert(vmisa<T>(other.unsafeGetSlot()->value()) && "invalid WeakRef cast");
     return WeakRef<T>(other.unsafeGetSlot());
   }
+
+  /// Clear the slot to which the WeakRef refers.
+  void clear() {
+    unsafeGetSlot()->clearPointer();
+  }
 };
 
 } // namespace vm

lib/VM/JSWeakMapImpl.cpp

@@ -76,6 +76,29 @@ bool JSWeakMapImplBase::deleteValue(
   return true;
 }
 
+// Only during GC.
+bool JSWeakMapImplBase::clearEntryDirect(GC *gc, const WeakRefKey &key) {
+  assert(gc->inGC() && "Should only be used by the GC implementation.");
+  DenseMapT::iterator it = map_.find(key);
+  if (it == map_.end()) {
+    return false;
+  }
+  it->first.ref.clear();
+  valueStorage_.get(gc->getPointerBase())
+      ->at(it->second)
+      .setNonPtr(HermesValue::encodeEmptyValue());
+  return true;
+}
+
+HermesValue JSWeakMapImplBase::getValueDirect(GC *gc, const WeakRefKey &key) {
+  assert(gc->inGC() && "Should only be used by the GC implementation.");
+  DenseMapT::iterator it = map_.find(key);
+  if (it == map_.end()) {
+    return HermesValue::encodeEmptyValue();
+  }
+  return valueStorage_.get(gc->getPointerBase())->at(it->second);
+}
+
 /// \return true if the \p key exists in the map.
 bool JSWeakMapImplBase::hasValue(
     Handle<JSWeakMapImplBase> self,
@@ -111,6 +134,23 @@ uint32_t JSWeakMapImplBase::debugFreeSlotsAndGetSize(
   return self->map_.size();
 }
 
+JSWeakMapImplBase::KeyIterator JSWeakMapImplBase::keys_begin() {
+  return KeyIterator{map_.begin()};
+}
+
+JSWeakMapImplBase::KeyIterator JSWeakMapImplBase::keys_end() {
+  return KeyIterator{map_.end()};
+}
+
+JSObject *detail::WeakRefKey::getObject(GC *gc) const {
+  assert(gc->inGC() && "Should only be used by the GC implementation.");
+  if (auto val = ref.unsafeGetOptional()) {
+    return *val;
+  } else {
+    return nullptr;
+  }
+}
+
 void JSWeakMapImplBase::_markWeakImpl(GCCell *cell, WeakRefAcceptor &acceptor) {
   auto *self = reinterpret_cast<JSWeakMapImplBase *>(cell);
   for (auto it = self->map_.begin(); it != self->map_.end(); ++it) {
@@ -145,7 +185,7 @@ void JSWeakMapImplBase::deleteInternal(
     PointerBase *base,
     JSWeakMapImplBase::DenseMapT::iterator it) {
   assert(it != map_.end() && "Invalid iterator to deleteInternal");
-  valueStorage_.get(base)
+  valueStorage_.getNonNull(base)
       ->at(it->second)
       .setNonPtr(HermesValue::encodeNativeUInt32(freeListHead_));
   freeListHead_ = it->second;

lib/VM/gcs/CompleteMarkState.cpp

@@ -6,9 +6,11 @@
  */
 
 #include "hermes/VM/CompleteMarkState.h"
+#include "hermes/VM/Casting.h"
 #include "hermes/VM/CompleteMarkState-inline.h"
 #include "hermes/VM/GCBase-inline.h"
 #include "hermes/VM/GCBase.h"
+#include "hermes/VM/JSWeakMapImpl.h"
 
 namespace hermes {
 namespace vm {
@@ -48,16 +50,22 @@ void CompleteMarkState::markTransitive(void *ptr) {
   // wouldn't make it past the earlier check for a false value if the
   // indices were equal.
   if (ptr < reinterpret_cast<void *>(currentParPointer)) {
-    GCCell *cell = reinterpret_cast<GCCell *>(ptr);
-    // Push cell to the correct mark stack.
-    std::vector<GCCell *> *stack =
-        (cell->isVariableSize() ? &varSizeMarkStack_ : &markStack_);
-    if (stack->size() == kMarkStackLimit) {
-      markStackOverflow_ = true;
-    } else {
-      stack->push_back(cell);
-      numPtrsPushedByParent++;
-    }
+    pushCell(reinterpret_cast<GCCell *>(ptr));
+  }
+}
+
+void CompleteMarkState::pushCell(GCCell *cell) {
+  // Push cell to the correct mark stack.
+#ifdef HERMES_SLOW_DEBUG
+  assert(cell->isValid());
+#endif
+  std::vector<GCCell *> *stack =
+      (cell->isVariableSize() ? &varSizeMarkStack_ : &markStack_);
+  if (stack->size() == kMarkStackLimit) {
+    markStackOverflow_ = true;
+  } else {
+    stack->push_back(cell);
+    numPtrsPushedByParent++;
   }
 }
 
@@ -78,7 +86,15 @@ void CompleteMarkState::drainMarkStack(
     // skipped for fixed sized cells.
     numPtrsPushedByParent = 0;
 
-    GCBase::markCell(cell, gc, acceptor);
+    // Don't mark through JSWeakMaps.  (Presently, this check is very
+    // specific.  If there are ever other cell kinds that need special
+    // weak marking handling, we could put a boolean in the vtable or
+    // metadata table).
+    if (cell->getKind() == CellKind::WeakMapKind) {
+      reachableWeakMaps_.push_back(vmcast<JSWeakMap>(cell));
+    } else {
+      GCBase::markCell(cell, gc, acceptor);
+    }
 
     // All fields of a fixed-sized cell should be marked by this point, but var
     // sized GCCells may not. Pop if the last round of marking pushed nothing,

lib/VM/gcs/GenGCNC.cpp

@@ -25,6 +25,7 @@
 #include "hermes/VM/GCPointer-inline.h"
 #include "hermes/VM/HeapSnapshot.h"
 #include "hermes/VM/HermesValue-inline.h"
+#include "hermes/VM/JSWeakMapImpl.h"
 #include "hermes/VM/Serializer.h"
 #include "hermes/VM/SlotAcceptorDefault-inline.h"
 #include "hermes/VM/StringPrimitive.h"
@@ -630,9 +631,102 @@ void GenGC::completeMarking() {
       if (markState_.markStackOverflow_)
         break;
     }
+    // The marking loop above will have accumulated WeakMaps;
+    // find things reachable from values of reachable keys.
+    // (Note that this can also set markStackOverflow_).
+    completeWeakMapMarking();
   } while (markState_.markStackOverflow_);
 }
 
+/// For all reachable keys in \p weakMap, mark from the corresponding value
+/// using \p markState and the given \p acceptor, reaching a
+/// transitive closure (or setting \p markState.markStackOverflow_).
+/// Returns whether any previously-unmarked values were marked.
+static bool markFromReachableWeakMapKeys(
+    GC *gc,
+    JSWeakMap *weakMap,
+    CompleteMarkState *markState,
+    CompleteMarkState::FullMSCMarkTransitiveAcceptor &acceptor) {
+  bool newlyMarkedValue = false;
+  for (auto iter = weakMap->keys_begin(), end = weakMap->keys_end();
+       iter != end;
+       iter++) {
+    GCCell *cell = iter->getObject(gc);
+    if (!cell || !AlignedHeapSegment::getCellMarkBit(cell)) {
+      continue;
+    }
+    HermesValue val = weakMap->getValueDirect(gc, *iter);
+    if (val.isPointer()) {
+      GCCell *valCell = reinterpret_cast<GCCell *>(val.getPointer());
+      if (!AlignedHeapSegment::getCellMarkBit(valCell)) {
+        AlignedHeapSegment::setCellMarkBit(valCell);
+        markState->pushCell(valCell);
+        markState->drainMarkStack(gc, acceptor);
+        newlyMarkedValue = true;
+      }
+    }
+  }
+  return newlyMarkedValue;
+}
+
+/// For all non-null keys in \p weakMap that are unreachable, clear
+/// the key (clear the pointer in the WeakRefSlot) and value (set it
+/// to undefined).
+static void clearEntriesWithUnreachableKeys(GenGC *gc, JSWeakMap *weakMap) {
+  for (auto iter = weakMap->keys_begin(), end = weakMap->keys_end();
+       iter != end;
+       iter++) {
+    JSObject *keyObj = iter->getObject(gc);
+    if (keyObj && !AlignedHeapSegment::getCellMarkBit(keyObj)) {
+      weakMap->clearEntryDirect(gc, *iter);
+    }
+  }
+}
+
+void GenGC::completeWeakMapMarking() {
+  CompleteMarkState::FullMSCMarkTransitiveAcceptor acceptor(*this, &markState_);
+
+  // Set the currentParPointer to a maximal value, so all pointers scanned
+  // will be pushed on the mark stack.
+  markState_.currentParPointer =
+      reinterpret_cast<GCCell *>(static_cast<intptr_t>(-1));
+  // Must declare this outside the loop, but the initial value doesn't matter:
+  // we make it false at the start of each loop iteration.
+  bool newReachableValueFound = true;
+  do {
+    newReachableValueFound = false;
+    // Note that new reachable weak maps may be discovered during the loop, so
+    // markState_.reachableWeakMaps_.size() may increase during the loop.
+    for (unsigned i = 0; i < markState_.reachableWeakMaps_.size(); i++) {
+      JSWeakMap *weakMap = markState_.reachableWeakMaps_[i];
+      if (markFromReachableWeakMapKeys(this, weakMap, &markState_, acceptor)) {
+        newReachableValueFound = true;
+      }
+    }
+  } while (newReachableValueFound);
+  for (auto *weakMap : markState_.reachableWeakMaps_) {
+    clearEntriesWithUnreachableKeys(this, weakMap);
+    // The argument for why this works is delicate.  We need to call
+    // markCell, because it marks fields of the WeakMap that are not
+    // part of the table -- in particular, the pointer to the
+    // valueStorage_ array.  There would be a problem, however, if
+    // this markCell/drain pair could cause any keys of any WeakMaps
+    // to become newly reachable.  If so, we'd need to scan their
+    // corresponding values.  This cannot happen, however.  The
+    // marking loop above ensured that all values in the valueStorage_
+    // corresponding to reachable keys have been marked.  The call
+    // just above ensures that values corresponding to unreachable
+    // keys have been cleared.  So we'll mark *only* the valueStorage_
+    // array; the drainMarkStack call should not do any work.  Perhaps
+    // we should assert this.
+    GCBase::markCell(weakMap, this, acceptor);
+    markState_.drainMarkStack(this, acceptor);
+  }
+
+  markState_.currentParPointer = nullptr;
+  markState_.reachableWeakMaps_.clear();
+}
+
 void GenGC::finalizeUnreachableObjects() {
   youngGen_.finalizeUnreachableObjects();
   oldGen_.finalizeUnreachableObjects();