Allow symbols as weak collection keys

js/src/builtin/WeakMapObject.cpp

@@ -8,6 +8,7 @@
 #include "jsapi.h"
 #include "jscntxt.h"
 
+#include "builtin/WeakRefObject.h"
 #include "vm/SelfHosting.h"
 
 #include "vm/Interpreter-inl.h"
@@ -21,23 +22,82 @@ IsWeakMap(HandleValue v)
     return v.isObject() && v.toObject().is<WeakMapObject>();
 }
 
+struct WeakMapObject::Data
+{
+    ObjectValueMap objectMap;
+    SymbolValueMap symbolMap;
+
+    Data(JSContext* cx, JSObject* owner)
+      : objectMap(cx, owner),
+        symbolMap(cx, owner)
+    {}
+
+    bool init() {
+        return objectMap.init() && symbolMap.init();
+    }
+};
+
+ObjectValueMap*
+WeakMapObject::getMap()
+{
+    Data* data = getData();
+    return data ? &data->objectMap : nullptr;
+}
+
+SymbolValueMap*
+WeakMapObject::getSymbolMap()
+{
+    Data* data = getData();
+    return data ? &data->symbolMap : nullptr;
+}
+
+static bool
+EnsureWeakMapData(JSContext* cx, Handle<WeakMapObject*> mapObj)
+{
+    if (mapObj->getData())
+        return true;
+
+    auto data = cx->make_unique<WeakMapObject::Data>(cx, mapObj.get());
+    if (!data)
+        return false;
+
+    if (!data->init()) {
+        JS_ReportOutOfMemory(cx);
+        return false;
+    }
+
+    mapObj->setPrivate(data.release());
+    return true;
+}
+
 MOZ_ALWAYS_INLINE bool
 WeakMap_has_impl(JSContext* cx, const CallArgs& args)
 {
     MOZ_ASSERT(IsWeakMap(args.thisv()));
 
-    if (!args.get(0).isObject()) {
+    if (!CanBeHeldWeakly(args.get(0))) {
         args.rval().setBoolean(false);
         return true;
     }
 
-    if (ObjectValueMap* map = args.thisv().toObject().as<WeakMapObject>().getMap()) {
+    WeakMapObject& weakMap = args.thisv().toObject().as<WeakMapObject>();
+    if (args.get(0).isObject()) {
+        if (ObjectValueMap* map = weakMap.getMap()) {
             JSObject* key = &args[0].toObject();
             if (map->has(key)) {
                 args.rval().setBoolean(true);
                 return true;
             }
         }
+    } else {
+        if (SymbolValueMap* map = weakMap.getSymbolMap()) {
+            JS::Symbol* key = args[0].toSymbol();
+            if (map->has(key)) {
+                args.rval().setBoolean(true);
+                return true;
+            }
+        }
+    }
 
     args.rval().setBoolean(false);
     return true;
@@ -55,18 +115,29 @@ WeakMap_get_impl(JSContext* cx, const CallArgs& args)
 {
     MOZ_ASSERT(IsWeakMap(args.thisv()));
 
-    if (!args.get(0).isObject()) {
+    if (!CanBeHeldWeakly(args.get(0))) {
         args.rval().setUndefined();
         return true;
     }
 
-    if (ObjectValueMap* map = args.thisv().toObject().as<WeakMapObject>().getMap()) {
+    WeakMapObject& weakMap = args.thisv().toObject().as<WeakMapObject>();
+    if (args.get(0).isObject()) {
+        if (ObjectValueMap* map = weakMap.getMap()) {
             JSObject* key = &args[0].toObject();
             if (ObjectValueMap::Ptr ptr = map->lookup(key)) {
                 args.rval().set(ptr->value());
                 return true;
             }
         }
+    } else {
+        if (SymbolValueMap* map = weakMap.getSymbolMap()) {
+            JS::Symbol* key = args[0].toSymbol();
+            if (SymbolValueMap::Ptr ptr = map->lookup(key)) {
+                args.rval().set(ptr->value());
+                return true;
+            }
+        }
+    }
 
     args.rval().setUndefined();
     return true;
@@ -84,12 +155,14 @@ WeakMap_delete_impl(JSContext* cx, const CallArgs& args)
 {
     MOZ_ASSERT(IsWeakMap(args.thisv()));
 
-    if (!args.get(0).isObject()) {
+    if (!CanBeHeldWeakly(args.get(0))) {
         args.rval().setBoolean(false);
         return true;
     }
 
-    if (ObjectValueMap* map = args.thisv().toObject().as<WeakMapObject>().getMap()) {
+    WeakMapObject& weakMap = args.thisv().toObject().as<WeakMapObject>();
+    if (args.get(0).isObject()) {
+        if (ObjectValueMap* map = weakMap.getMap()) {
             JSObject* key = &args[0].toObject();
             if (ObjectValueMap::Ptr ptr = map->lookup(key)) {
                 map->remove(ptr);
@@ -97,6 +170,16 @@ WeakMap_delete_impl(JSContext* cx, const CallArgs& args)
                 return true;
             }
         }
+    } else {
+        if (SymbolValueMap* map = weakMap.getSymbolMap()) {
+            JS::Symbol* key = args[0].toSymbol();
+            if (SymbolValueMap::Ptr ptr = map->lookup(key)) {
+                map->remove(ptr);
+                args.rval().setBoolean(true);
+                return true;
+            }
+        }
+    }
 
     args.rval().setBoolean(false);
     return true;
@@ -126,22 +209,13 @@ TryPreserveReflector(JSContext* cx, HandleObject obj)
     return true;
 }
 
-static MOZ_ALWAYS_INLINE bool
+static bool
-SetWeakMapEntryInternal(JSContext* cx, Handle<WeakMapObject*> mapObj,
+SetWeakMapObjectEntryInternal(JSContext* cx, Handle<WeakMapObject*> mapObj,
                               HandleObject key, HandleValue value)
 {
+    if (!EnsureWeakMapData(cx, mapObj))
+        return false;
     ObjectValueMap* map = mapObj->getMap();
-    if (!map) {
-        auto newMap = cx->make_unique<ObjectValueMap>(cx, mapObj.get());
-        if (!newMap)
-            return false;
-        if (!newMap->init()) {
-            JS_ReportOutOfMemory(cx);
-            return false;
-        }
-        map = newMap.release();
-        mapObj->setPrivate(map);
-    }
 
     // Preserve wrapped native keys to prevent wrapper optimization.
     if (!TryPreserveReflector(cx, key))
@@ -162,13 +236,28 @@ SetWeakMapEntryInternal(JSContext* cx, Handle<WeakMapObject*> mapObj,
     return true;
 }
 
-MOZ_ALWAYS_INLINE bool
+static bool
-WeakMap_set_impl(JSContext* cx, const CallArgs& args)
+SetWeakMapSymbolEntryInternal(JSContext* cx, Handle<WeakMapObject*> mapObj,
+                              Handle<JS::Symbol*> key, HandleValue value)
 {
-    MOZ_ASSERT(IsWeakMap(args.thisv()));
+    if (!EnsureWeakMapData(cx, mapObj))
+        return false;
+    SymbolValueMap* map = mapObj->getSymbolMap();
 
-    if (!args.get(0).isObject()) {
+    MOZ_ASSERT_IF(value.isObject(), value.toObject().compartment() == mapObj->compartment());
-        UniqueChars bytes = DecompileValueGenerator(cx, JSDVG_SEARCH_STACK, args.get(0), nullptr);
+    if (!map->put(key, value)) {
+        JS_ReportOutOfMemory(cx);
+        return false;
+    }
+    return true;
+}
+
+static MOZ_ALWAYS_INLINE bool
+SetWeakMapEntryInternal(JSContext* cx, Handle<WeakMapObject*> mapObj,
+                        HandleValue key, HandleValue value)
+{
+    if (!CanBeHeldWeakly(key)) {
+        UniqueChars bytes = DecompileValueGenerator(cx, JSDVG_SEARCH_STACK, key, nullptr);
         if (!bytes)
             return false;
         JS_ReportErrorNumberLatin1(cx, GetErrorMessage, nullptr, JSMSG_NOT_NONNULL_OBJECT,
@@ -176,11 +265,24 @@ WeakMap_set_impl(JSContext* cx, const CallArgs& args)
         return false;
     }
 
-    RootedObject key(cx, &args[0].toObject());
+    if (key.isObject()) {
+        RootedObject objectKey(cx, &key.toObject());
+        return SetWeakMapObjectEntryInternal(cx, mapObj, objectKey, value);
+    }
+
+    Rooted<JS::Symbol*> symbolKey(cx, key.toSymbol());
+    return SetWeakMapSymbolEntryInternal(cx, mapObj, symbolKey, value);
+}
+
+MOZ_ALWAYS_INLINE bool
+WeakMap_set_impl(JSContext* cx, const CallArgs& args)
+{
+    MOZ_ASSERT(IsWeakMap(args.thisv()));
+
     Rooted<JSObject*> thisObj(cx, &args.thisv().toObject());
     Rooted<WeakMapObject*> map(cx, &thisObj->as<WeakMapObject>());
 
-    if (!SetWeakMapEntryInternal(cx, map, key, args.get(1)))
+    if (!SetWeakMapEntryInternal(cx, map, args.get(0), args.get(1)))
         return false;
     args.rval().set(args.thisv());
     return true;
@@ -193,6 +295,13 @@ js::WeakMap_set(JSContext* cx, unsigned argc, Value* vp)
     return CallNonGenericMethod<IsWeakMap, WeakMap_set_impl>(cx, args);
 }
 
+bool
+js::SetWeakMapEntryValue(JSContext* cx, HandleObject mapObj, HandleValue key, HandleValue val)
+{
+    Rooted<WeakMapObject*> rootedMap(cx, &mapObj->as<WeakMapObject>());
+    return SetWeakMapEntryInternal(cx, rootedMap, key, val);
+}
+
 JS_FRIEND_API(bool)
 JS_NondeterministicGetWeakMapKeys(JSContext* cx, HandleObject objArg, MutableHandleObject ret)
 {
@@ -205,11 +314,11 @@ JS_NondeterministicGetWeakMapKeys(JSContext* cx, HandleObject objArg, MutableHan
     RootedObject arr(cx, NewDenseEmptyArray(cx));
     if (!arr)
         return false;
-    ObjectValueMap* map = obj->as<WeakMapObject>().getMap();
+    WeakMapObject::Data* data = obj->as<WeakMapObject>().getData();
-    if (map) {
+    if (data) {
         // Prevent GC from mutating the weakmap while iterating.
         AutoSuppressGC suppress(cx);
-        for (ObjectValueMap::Base::Range r = map->all(); !r.empty(); r.popFront()) {
+        for (ObjectValueMap::Base::Range r = data->objectMap.all(); !r.empty(); r.popFront()) {
             JS::ExposeObjectToActiveJS(r.front().key());
             RootedObject key(cx, r.front().key());
             if (!cx->compartment()->wrap(cx, &key))
@@ -217,6 +326,14 @@ JS_NondeterministicGetWeakMapKeys(JSContext* cx, HandleObject objArg, MutableHan
             if (!NewbornArrayPush(cx, arr, ObjectValue(*key)))
                 return false;
         }
+        for (SymbolValueMap::Base::Range r = data->symbolMap.all(); !r.empty(); r.popFront()) {
+            gc::ExposeGCThingToActiveJS(JS::GCCellPtr(r.front().key().get()));
+            RootedValue key(cx, SymbolValue(r.front().key()));
+            if (!cx->compartment()->wrap(cx, &key))
+                return false;
+            if (!NewbornArrayPush(cx, arr, key))
+                return false;
+        }
     }
     ret.set(arr);
     return true;
@@ -225,21 +342,23 @@ JS_NondeterministicGetWeakMapKeys(JSContext* cx, HandleObject objArg, MutableHan
 static void
 WeakMap_mark(JSTracer* trc, JSObject* obj)
 {
-    if (ObjectValueMap* map = obj->as<WeakMapObject>().getMap())
+    if (WeakMapObject::Data* data = obj->as<WeakMapObject>().getData()) {
-        map->trace(trc);
+        data->objectMap.trace(trc);
+        data->symbolMap.trace(trc);
+    }
 }
 
 static void
 WeakMap_finalize(FreeOp* fop, JSObject* obj)
 {
     MOZ_ASSERT(fop->maybeOffMainThread());
-    if (ObjectValueMap* map = obj->as<WeakMapObject>().getMap()) {
+    if (WeakMapObject::Data* data = obj->as<WeakMapObject>().getData()) {
 #ifdef DEBUG
-        map->~ObjectValueMap();
+        data->~Data();
-        memset(static_cast<void*>(map), 0xdc, sizeof(*map));
+        memset(static_cast<void*>(data), 0xdc, sizeof(*data));
-        fop->free_(map);
+        fop->free_(data);
 #else
-        fop->delete_(map);
+        fop->delete_(data);
 #endif
     }
 }
@@ -282,7 +401,8 @@ JS::SetWeakMapEntry(JSContext* cx, HandleObject mapObj, HandleObject key,
     CHECK_REQUEST(cx);
     assertSameCompartment(cx, key, val);
     Rooted<WeakMapObject*> rootedMap(cx, &mapObj->as<WeakMapObject>());
-    return SetWeakMapEntryInternal(cx, rootedMap, key, val);
+    RootedValue keyValue(cx, ObjectValue(*key));
+    return SetWeakMapEntryInternal(cx, rootedMap, keyValue, val);
 }
 
 static bool
@@ -386,4 +506,3 @@ js::InitBareWeakMapCtor(JSContext* cx, HandleObject obj)
 {
     return InitWeakMapClass(cx, obj, false);
 }
-

js/src/builtin/WeakMapObject.h

@@ -16,7 +16,11 @@ class WeakMapObject : public NativeObject
   public:
     static const Class class_;
 
-    ObjectValueMap* getMap() { return static_cast<ObjectValueMap*>(getPrivate()); }
+    struct Data;
+
+    Data* getData() { return static_cast<Data*>(getPrivate()); }
+    ObjectValueMap* getMap();
+    SymbolValueMap* getSymbolMap();
 };
 
 } // namespace js

js/src/builtin/WeakSet.js

@@ -32,7 +32,7 @@ function WeakSet_add(value) {
         ThrowTypeError(JSMSG_INCOMPATIBLE_PROTO, "WeakSet", "add", typeof S);
 
     // Step 5.
-    if (!IsObject(value))
+    if (!CanBeHeldWeakly(value))
         ThrowTypeError(JSMSG_NOT_NONNULL_OBJECT, DecompileArg(0, value));
 
     // Steps 7-8.
@@ -55,7 +55,7 @@ function WeakSet_delete(value) {
         ThrowTypeError(JSMSG_INCOMPATIBLE_PROTO, "WeakSet", "delete", typeof S);
 
     // Step 5.
-    if (!IsObject(value))
+    if (!CanBeHeldWeakly(value))
         return false;
 
     // Steps 7-8.
@@ -75,7 +75,7 @@ function WeakSet_has(value) {
         ThrowTypeError(JSMSG_INCOMPATIBLE_PROTO, "WeakSet", "has", typeof S);
 
     // Step 6.
-    if (!IsObject(value))
+    if (!CanBeHeldWeakly(value))
         return false;
 
     // Steps 7-8.

js/src/builtin/WeakSetObject.cpp

@@ -12,6 +12,7 @@
 #include "builtin/MapObject.h"
 #include "builtin/SelfHostingDefines.h"
 #include "builtin/WeakMapObject.h"
+#include "builtin/WeakRefObject.h"
 #include "vm/GlobalObject.h"
 #include "vm/SelfHosting.h"
 
@@ -108,7 +109,6 @@ WeakSetObject::construct(JSContext* cx, unsigned argc, Value* vp)
 
         if (optimized) {
             RootedValue keyVal(cx);
-            RootedObject keyObject(cx);
             RootedValue placeholder(cx, BooleanValue(true));
             RootedObject map(cx, &obj->getReservedSlot(WEAKSET_MAP_SLOT).toObject());
             RootedArrayObject array(cx, &iterable.toObject().as<ArrayObject>());
@@ -116,7 +116,7 @@ WeakSetObject::construct(JSContext* cx, unsigned argc, Value* vp)
                 keyVal.set(array->getDenseElement(index));
                 MOZ_ASSERT(!keyVal.isMagic(JS_ELEMENTS_HOLE));
 
-                if (keyVal.isPrimitive()) {
+                if (!CanBeHeldWeakly(keyVal)) {
                     UniqueChars bytes =
                         DecompileValueGenerator(cx, JSDVG_SEARCH_STACK, keyVal, nullptr);
                     if (!bytes)
@@ -126,8 +126,7 @@ WeakSetObject::construct(JSContext* cx, unsigned argc, Value* vp)
                     return false;
                 }
 
-                keyObject = &keyVal.toObject();
+                if (!SetWeakMapEntryValue(cx, map, keyVal, placeholder))
-                if (!SetWeakMapEntry(cx, map, keyObject, placeholder))
                     return false;
             }
         } else {

js/src/jsweakmap.h

@@ -16,11 +16,25 @@
 #include "gc/Marking.h"
 #include "gc/StoreBuffer.h"
 #include "js/HashTable.h"
+#include "vm/Symbol.h"
 
 namespace js {
 
 class WeakMapBase;
 
+template <>
+struct MovableCellHasher<JS::Symbol*>
+{
+    using Key = JS::Symbol*;
+    using Lookup = JS::Symbol*;
+
+    static bool hasHash(const Lookup& l) { return true; }
+    static bool ensureHash(const Lookup& l) { return true; }
+    static HashNumber hash(const Lookup& l) { return l->hash(); }
+    static bool match(const Key& k, const Lookup& l) { return k == l; }
+    static void rekey(Key& k, const Key& newKey) { k = newKey; }
+};
+
 // A subclass template of js::HashMap whose keys and values may be garbage-collected. When
 // a key is collected, the table entry disappears, dropping its reference to the value.
 //
@@ -296,9 +310,16 @@ class WeakMap : public HashMap<Key, Value, HashPolicy, RuntimeAllocPolicy>,
         return nullptr;
     }
 
+    JSObject* getDelegate(JS::Symbol* sym) const {
+        return nullptr;
+    }
+
   private:
     void exposeGCThingToActiveJS(const JS::Value& v) const { JS::ExposeValueToActiveJS(v); }
     void exposeGCThingToActiveJS(JSObject* obj) const { JS::ExposeObjectToActiveJS(obj); }
+    void exposeGCThingToActiveJS(JS::Symbol* sym) const {
+        gc::ExposeGCThingToActiveJS(JS::GCCellPtr(sym));
+    }
 
     bool keyNeedsMark(JSObject* key) const {
         JSObject* delegate = getDelegate(key);
@@ -313,6 +334,10 @@ class WeakMap : public HashMap<Key, Value, HashPolicy, RuntimeAllocPolicy>,
         return false;
     }
 
+    bool keyNeedsMark(JS::Symbol* sym) const {
+        return false;
+    }
+
     bool findZoneEdges() override {
         // This is overridden by ObjectValueMap.
         return true;
@@ -378,6 +403,9 @@ WeakMap_set(JSContext* cx, unsigned argc, Value* vp);
 extern bool
 WeakMap_delete(JSContext* cx, unsigned argc, Value* vp);
 
+extern bool
+SetWeakMapEntryValue(JSContext* cx, HandleObject mapObj, HandleValue key, HandleValue val);
+
 extern JSObject*
 InitWeakMapClass(JSContext* cx, HandleObject obj);
 
@@ -394,6 +422,16 @@ class ObjectValueMap : public WeakMap<HeapPtr<JSObject*>, HeapPtr<Value>,
     virtual bool findZoneEdges();
 };
 
+class SymbolValueMap : public WeakMap<HeapPtr<JS::Symbol*>, HeapPtr<Value>,
+                                      MovableCellHasher<HeapPtr<JS::Symbol*>>>
+{
+  public:
+    SymbolValueMap(JSContext* cx, JSObject* obj)
+      : WeakMap<HeapPtr<JS::Symbol*>, HeapPtr<Value>,
+                MovableCellHasher<HeapPtr<JS::Symbol*>>>(cx, obj)
+    {}
+};
+
 
 // Generic weak map for mapping objects to other objects.
 class ObjectWeakMap

js/src/tests/non262/WeakMap/symbol-keys.js

@@ -0,0 +1,36 @@
+var key = Symbol("weak");
+var map = new WeakMap();
+assertEq(map.has(key), false);
+assertEq(map.get(key), undefined);
+assertEq(map.set(key, 13), map);
+assertEq(map.has(key), true);
+assertEq(map.get(key), 13);
+assertEq(map.delete(key), true);
+assertEq(map.has(key), false);
+
+var constructedKey = Symbol("constructed");
+var constructed = new WeakMap([[constructedKey, 7]]);
+assertEq(constructed.get(constructedKey), 7);
+
+var registered = Symbol.for("registered");
+assertEq(map.has(registered), false);
+assertEq(map.get(registered), undefined);
+assertEq(map.delete(registered), false);
+assertThrowsInstanceOf(() => map.set(registered, 1), TypeError);
+assertThrowsInstanceOf(() => new WeakMap([[registered, 1]]), TypeError);
+
+var setKey = Symbol("set");
+var set = new WeakSet([setKey]);
+assertEq(set.has(setKey), true);
+assertEq(set.delete(setKey), true);
+assertEq(set.has(setKey), false);
+assertEq(set.add(setKey), set);
+assertEq(set.has(setKey), true);
+
+assertEq(set.has(registered), false);
+assertEq(set.delete(registered), false);
+assertThrowsInstanceOf(() => set.add(registered), TypeError);
+assertThrowsInstanceOf(() => new WeakSet([registered]), TypeError);
+
+if (typeof reportCompare === "function")
+  reportCompare(0, 0);

js/src/vm/SelfHosting.cpp

@@ -38,6 +38,7 @@
 #include "builtin/SelfHostingDefines.h"
 #include "builtin/Stream.h"
 #include "builtin/TypedObject.h"
+#include "builtin/WeakRefObject.h"
 #include "builtin/WeakSetObject.h"
 #include "gc/Marking.h"
 #include "gc/Policy.h"
@@ -105,6 +106,14 @@ intrinsic_IsObject(JSContext* cx, unsigned argc, Value* vp)
     return true;
 }
 
+static bool
+intrinsic_CanBeHeldWeakly(JSContext* cx, unsigned argc, Value* vp)
+{
+    CallArgs args = CallArgsFromVp(argc, vp);
+    args.rval().setBoolean(CanBeHeldWeakly(args.get(0)));
+    return true;
+}
+
 static bool
 intrinsic_IsArray(JSContext* cx, unsigned argc, Value* vp)
 {
@@ -2333,6 +2342,7 @@ static const JSFunctionSpec intrinsic_functions[] = {
     // Helper funtions after this point.
     JS_INLINABLE_FN("ToObject",      intrinsic_ToObject,                1,0, IntrinsicToObject),
     JS_INLINABLE_FN("IsObject",      intrinsic_IsObject,                1,0, IntrinsicIsObject),
+    JS_FN("CanBeHeldWeakly",         intrinsic_CanBeHeldWeakly,         1,0),
     JS_INLINABLE_FN("IsArray",       intrinsic_IsArray,                 1,0, ArrayIsArray),
     JS_INLINABLE_FN("IsWrappedArrayConstructor", intrinsic_IsWrappedArrayConstructor, 1,0,
                     IntrinsicIsWrappedArrayConstructor),