Bug 798678 - Refactor wrapper preservation for weakmaps (r=mccr8)

js/src/jit-test/tests/basic/bug673468.js

@@ -0,0 +1,8 @@
+var g = newGlobal();
+var k = g.eval('var u = new Object(); u');
+var m = new WeakMap();
+m.set(k, {});
+k = null;
+gc();
+k = g.eval('u');
+assertEq(m.has(k), true);

js/src/jsapi.h

@@ -1883,6 +1883,9 @@ typedef void
 typedef JSBool
 (* JSEqualityOp)(JSContext *cx, JSHandleObject obj, JSHandleValue v, JSBool *bp);
 
+typedef JSRawObject
+(* JSWeakmapKeyDelegateOp)(JSRawObject obj);
+
 /*
  * Typedef for native functions called by the JS VM.
  *

js/src/jsclass.h

@@ -255,9 +255,22 @@ struct ClassExtension
      * WeakMaps use this to override the wrapper disposal optimization.
      */
     bool                isWrappedNative;
+
+    /*
+     * If an object is used as a key in a weakmap, it may be desirable for the
+     * garbage collector to keep that object around longer than it otherwise
+     * would. A common case is when the key is a wrapper around an object in
+     * another compartment, and we want to avoid collecting the wrapper (and
+     * removing the weakmap entry) as long as the wrapped object is alive. In
+     * that case, the wrapped object is returned by the wrapper's
+     * weakmapKeyDelegateOp hook. As long as the wrapper is used as a weakmap
+     * key, it will not be collected (and remain in the weakmap) until the
+     * wrapped object is collected.
+     */
+    JSWeakmapKeyDelegateOp weakmapKeyDelegateOp;
 };
 
-#define JS_NULL_CLASS_EXT   {NULL,NULL,NULL,NULL,NULL,false}
+#define JS_NULL_CLASS_EXT   {NULL,NULL,NULL,NULL,NULL,false,NULL}
 
 struct ObjectOps
 {

js/src/jsproxy.cpp

@@ -353,6 +353,12 @@ BaseProxyHandler::finalize(JSFreeOp *fop, JSObject *proxy)
 {
 }
 
+JSObject *
+BaseProxyHandler::weakmapKeyDelegate(JSObject *proxy)
+{
+    return NULL;
+}
+
 bool
 BaseProxyHandler::getPrototypeOf(JSContext *cx, JSObject *proxy, JSObject **proto)
 {
@@ -549,6 +555,12 @@ IndirectProxyHandler::iteratorNext(JSContext *cx, JSObject *proxy, Value *vp)
     return true;
 }
 
+JSObject *
+IndirectProxyHandler::weakmapKeyDelegate(JSObject *proxy)
+{
+    return UnwrapObject(proxy);
+}
+
 DirectProxyHandler::DirectProxyHandler(void *family)
   : IndirectProxyHandler(family)
 {
@@ -2857,6 +2869,13 @@ proxy_TraceFunction(JSTracer *trc, RawObject obj)
     proxy_TraceObject(trc, obj);
 }
 
+static JSObject *
+proxy_WeakmapKeyDelegate(RawObject obj)
+{
+    JS_ASSERT(obj->isProxy());
+    return GetProxyHandler(obj)->weakmapKeyDelegate(obj);
+}
+
 static JSBool
 proxy_Convert(JSContext *cx, HandleObject proxy, JSType hint, MutableHandleValue vp)
 {
@@ -2868,8 +2887,7 @@ static void
 proxy_Finalize(FreeOp *fop, RawObject obj)
 {
     JS_ASSERT(obj->isProxy());
-    if (!obj->getSlot(JSSLOT_PROXY_HANDLER).isUndefined())
-        GetProxyHandler(obj)->finalize(fop, obj);
+    GetProxyHandler(obj)->finalize(fop, obj);
 }
 
 static JSBool
@@ -2889,6 +2907,17 @@ proxy_TypeOf(JSContext *cx, HandleObject proxy)
     return Proxy::typeOf(cx, proxy);
 }
 
+#define PROXY_CLASS_EXT                             \
+    {                                               \
+        NULL,                /* equality */         \
+        NULL,                /* outerObject */      \
+        NULL,                /* innerObject */      \
+        NULL,                /* iteratorObject */   \
+        NULL,                /* unused */           \
+        false,               /* isWrappedNative */  \
+        proxy_WeakmapKeyDelegate                    \
+    }
+
 JS_FRIEND_DATA(Class) js::ObjectProxyClass = {
     "Proxy",
     Class::NON_NATIVE | JSCLASS_IMPLEMENTS_BARRIERS | JSCLASS_HAS_RESERVED_SLOTS(4),
@@ -2905,7 +2934,7 @@ JS_FRIEND_DATA(Class) js::ObjectProxyClass = {
     proxy_HasInstance,       /* hasInstance */
     NULL,                    /* construct   */
     proxy_TraceObject,       /* trace       */
-    JS_NULL_CLASS_EXT,
+    PROXY_CLASS_EXT,
     {
         proxy_LookupGeneric,
         proxy_LookupProperty,
@@ -2961,7 +2990,10 @@ JS_FRIEND_DATA(Class) js::OuterWindowProxyClass = {
         NULL,                /* equality    */
         NULL,                /* outerObject */
         proxy_innerObject,
-        NULL                 /* unused */
+        NULL,                /* iteratorObject */
+        NULL,                /* unused */
+        false,               /* isWrappedNative */
+        proxy_WeakmapKeyDelegate
     },
     {
         proxy_LookupGeneric,
@@ -3031,7 +3063,7 @@ JS_FRIEND_DATA(Class) js::FunctionProxyClass = {
     FunctionClass.hasInstance,
     proxy_Construct,
     proxy_TraceFunction,     /* trace       */
-    JS_NULL_CLASS_EXT,
+    PROXY_CLASS_EXT,
     {
         proxy_LookupGeneric,
         proxy_LookupProperty,

js/src/jsproxy.h

@@ -124,6 +124,9 @@ class JS_FRIEND_API(BaseProxyHandler) {
     virtual bool getElementIfPresent(JSContext *cx, JSObject *obj, JSObject *receiver,
                                      uint32_t index, Value *vp, bool *present);
     virtual bool getPrototypeOf(JSContext *cx, JSObject *proxy, JSObject **protop);
+
+    /* See comment for weakmapKeyDelegateOp in jsclass.h. */
+    virtual JSObject *weakmapKeyDelegate(JSObject *proxy);
 };
 
 /*
@@ -176,6 +179,7 @@ class JS_PUBLIC_API(IndirectProxyHandler) : public BaseProxyHandler {
                               Value *vp) MOZ_OVERRIDE;
     virtual bool iteratorNext(JSContext *cx, JSObject *proxy,
                               Value *vp) MOZ_OVERRIDE;
+    virtual JSObject *weakmapKeyDelegate(JSObject *proxy);
 };
 
 /*

js/src/jsweakmap.cpp

@@ -108,14 +108,7 @@ GetKeyArg(JSContext *cx, CallArgs &args)
         JS_ReportErrorNumber(cx, js_GetErrorMessage, NULL, JSMSG_NOT_NONNULL_OBJECT);
         return NULL;
     }
-    JSObject &key = vp->toObject();
-
-    // If the key is from another compartment, and we store the wrapper as the key
-    // the wrapper might be GC-ed since it is not strong referenced (Bug 673468).
-    // To avoid this we always use the unwrapped object as the key instead of its
-    // security wrapper. This also means that if the keys are ever exposed they must
-    // be re-wrapped (see: JS_NondeterministicGetWeakMapKeys).
-    return JS_UnwrapObject(&key);
+    return &vp->toObject();
 }
 
 JS_ALWAYS_INLINE bool
@@ -251,7 +244,7 @@ WeakMap_set_impl(JSContext *cx, CallArgs args)
 
     // Preserve wrapped native keys to prevent wrapper optimization.
     if (key->getClass()->ext.isWrappedNative) {
-        MOZ_ASSERT(cx->runtime->preserveWrapperCallback, "wrapped native weak map key needs preserveWrapperCallback");
+        JS_ASSERT(cx->runtime->preserveWrapperCallback);
         if (!cx->runtime->preserveWrapperCallback(cx, key)) {
             JS_ReportErrorNumber(cx, js_GetErrorMessage, NULL, JSMSG_BAD_WEAKMAP_KEY);
             return false;
@@ -288,12 +281,7 @@ JS_NondeterministicGetWeakMapKeys(JSContext *cx, JSObject *obj, JSObject **ret)
     ObjectValueMap *map = GetObjectMap(obj);
     if (map) {
         for (ObjectValueMap::Base::Range r = map->all(); !r.empty(); r.popFront()) {
-            RootedObject key(cx, r.front().key);
-            // Re-wrapping the key (see comment of GetKeyArg)
-            if (!JS_WrapObject(cx, key.address()))
-                return false;
-
-            if (!js_NewbornArrayPush(cx, arr, ObjectValue(*key)))
+            if (!js_NewbornArrayPush(cx, arr, ObjectValue(*r.front().key)))
                 return false;
         }
     }

js/src/jsweakmap.h

@@ -27,42 +27,10 @@ namespace js {
 //     is collected. If an entry is not collected, it remains in the WeakMap and it has a
 //     strong reference to the value.
 //
-// You must call this table's 'mark' method when the object of which it is a part is
+// You must call this table's 'trace' method when the object of which it is a part is
 // reached by the garbage collection tracer. Once a table is known to be live, the
 // implementation takes care of the iterative marking needed for weak tables and removing
 // table entries when collection is complete.
-//
-// You may provide your own MarkPolicy class to specify how keys and values are marked; a
-// policy template provides default definitions for some common key/value type
-// combinations.
-//
-// Details:
-//
-// The interface is as for a js::HashMap, with the following additions:
-//
-// - You must call the WeakMap's 'trace' member function when you discover that the map is
-//   part of a live object. (You'll typically call this from the containing type's 'trace'
-//   function.)
-//
-// - There is no AllocPolicy parameter; these are used with our garbage collector, so
-//   RuntimeAllocPolicy is hard-wired.
-//
-// - Optional fourth and fifth parameters are the MarkPolicies for the key and value type.
-//   A MarkPolicy has the constructor:
-//
-//     MarkPolicy(JSTracer *)
-//
-//   and the following member functions:
-//
-//     bool isMarked(const Type &x)
-//        Return true if x has been marked as live by the garbage collector.
-//
-//     bool mark(Type &x)
-//        Return false if x is already marked. Otherwise, mark x and return true.
-//
-//   If omitted, the MarkPolicy parameter defaults to js::DefaultMarkPolicy<Type>,
-//   a policy template with the obvious definitions for some typical
-//   SpiderMonkey type combinations.
 
 // The value for the next pointer for maps not in the map list.
 static WeakMapBase * const WeakMapNotInList = reinterpret_cast<WeakMapBase *>(1);
@@ -170,6 +138,23 @@ class WeakMap : public HashMap<Key, Value, HashPolicy, RuntimeAllocPolicy>, publ
             markValue(trc, &r.front().value);
     }
 
+    bool keyNeedsMark(JSObject *key) {
+        if (JSWeakmapKeyDelegateOp op = key->getClass()->ext.weakmapKeyDelegateOp) {
+            JSObject *delegate = op(key);
+            /*
+             * Check if the delegate is marked with any color to properly handle
+             * gray marking when the key's delegate is black and the map is
+             * gray.
+             */
+            return delegate && gc::IsObjectMarked(&delegate);
+        }
+        return false;
+    }
+
+    bool keyNeedsMark(gc::Cell *cell) {
+        return false;
+    }
+
     bool markIteratively(JSTracer *trc) {
         bool markedAny = false;
         for (Enum e(*this); !e.empty(); e.popFront()) {
@@ -180,6 +165,12 @@ class WeakMap : public HashMap<Key, Value, HashPolicy, RuntimeAllocPolicy>, publ
                     markedAny = true;
                 if (prior != e.front().key)
                     e.rekeyFront(e.front().key);
+            } else if (keyNeedsMark(e.front().key)) {
+                gc::Mark(trc, const_cast<Key *>(&e.front().key), "proxy-preserved WeakMap key");
+                if (prior != e.front().key)
+                    e.rekeyFront(e.front().key);
+                gc::Mark(trc, &e.front().value, "WeakMap entry");
+                markedAny = true;
             }
         }
         return markedAny;

js/xpconnect/tests/chrome/Makefile.in

@@ -62,6 +62,7 @@ MOCHITEST_CHROME_FILES = \
 		test_precisegc.xul \
 		test_sandboxImport.xul \
 		test_weakmaps.xul \
+		test_weakmap_keys_preserved.xul \
 		test_weakref.xul \
 		test_wrappers.xul \
 		$(NULL)

js/xpconnect/tests/chrome/test_weakmap_keys_preserved.xul

@@ -0,0 +1,37 @@
+<?xml version="1.0"?>
+<?xml-stylesheet type="text/css" href="chrome://global/skin"?>
+<?xml-stylesheet type="text/css" href="chrome://mochikit/content/tests/SimpleTest/test.css"?>
+<!--
+https://bugzilla.mozilla.org/show_bug.cgi?id=673468
+-->
+<window title="Mozilla Bug "
+        xmlns="http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul">
+  <script type="application/javascript" src="chrome://mochikit/content/tests/SimpleTest/SimpleTest.js"/>
+
+  <!-- test results are displayed in the html:body -->
+  <body xmlns="http://www.w3.org/1999/xhtml">
+  <a href="https://bugzilla.mozilla.org/show_bug.cgi?id="
+     target="_blank">Mozilla Bug 673468</a>
+  </body>
+
+  <!-- test code goes here -->
+  <script type="application/javascript">
+  <![CDATA[
+  /** Test for Bug 673468 **/
+
+  let Cc = Components.classes;
+  let Cu = Components.utils;
+  let Ci = Components.interfaces;
+
+  let system = Cc["@mozilla.org/systemprincipal;1"].createInstance();
+  let sandbox = Cu.Sandbox(system);
+  let map = sandbox.WeakMap();
+  let obj = {};
+  map.set(obj, {});
+
+  Cu.forceGC();
+
+  ok(map.has(obj), "Weakmap still contains our wrapper!");
+  ]]>
+  </script>
+</window>