Bug 1845728 - Ellide keys call in Object.keys(..).length. r=iain

This optimization is made to handle cases where Object.keys(..) is called only
to compute the number of owned properties as a short-cut to iterating over the
properties of the object.

However, this short-cut, which is well deployed on the Web, does not make sense
given that Object.keys(..) iterates over the owned properties, to allocate an
array, to store the names of the properties. This optimization ellide the
allocation of the array when the content of the array is not needed.

Note, this optimization skip the case of Proxy where the Object.keys(..)
function might be effectful, or if the object might be mutated while
Object.keys(..) value is kept alive.

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

js/src/builtin/Object.cpp

@@ -1889,6 +1889,112 @@ static bool TryEnumerableOwnPropertiesNative(JSContext* cx, HandleObject obj,
   return true;
 }
 
+// Optimization dedicated for `Object.keys(..).length` JS pattern. This function
+// replicates TryEnumerableOwnPropertiesNative code, except that instead of
+// generating an array we only return the length of the array that would have
+// been generated.
+//
+// As opposed to TryEnumerableOwnPropertiesNative, this function only support
+// EnumerableOwnPropertiesKind::Keys variant.
+static bool CountEnumerableOwnPropertiesNative(JSContext* cx, HandleObject obj,
+                                               int32_t& rval, bool* optimized) {
+  *optimized = false;
+
+  // Use the fast path if |obj| has neither extra indexed properties nor a
+  // newEnumerate hook. String objects need to be special-cased, because
+  // they're only marked as indexed after their enumerate hook ran. And
+  // because their enumerate hook is slowish, it's more performant to
+  // exclude them directly instead of executing the hook first.
+  if (!obj->is<NativeObject>() || obj->as<NativeObject>().isIndexed() ||
+      obj->getClass()->getNewEnumerate() || obj->is<StringObject>()) {
+    return true;
+  }
+
+#ifdef ENABLE_RECORD_TUPLE
+  // Skip the optimized path in case of record and tuples.
+  if (obj->is<TupleObject>() || obj->is<RecordObject>()) {
+    return true;
+  }
+#endif
+
+  Handle<NativeObject*> nobj = obj.as<NativeObject>();
+
+  // Resolve lazy properties on |nobj|.
+  if (JSEnumerateOp enumerate = nobj->getClass()->getEnumerate()) {
+    if (!enumerate(cx, nobj)) {
+      return false;
+    }
+
+    // Ensure no extra indexed properties were added through enumerate().
+    if (nobj->isIndexed()) {
+      return true;
+    }
+  }
+
+  *optimized = true;
+
+  int32_t num_properties = 0;
+
+  // If possible, attempt to use the shape's iterator cache.
+  Rooted<PropertyIteratorObject*> piter(cx,
+                                        LookupInShapeIteratorCache(cx, nobj));
+  if (piter) {
+    NativeIterator* ni = piter->getNativeIterator();
+    MOZ_ASSERT(ni->isReusable());
+
+    // Guard against indexes.
+    if (!ni->mayHavePrototypeProperties()) {
+      rval = ni->numKeys();
+      return true;
+    }
+  }
+
+  for (uint32_t i = 0, len = nobj->getDenseInitializedLength(); i < len; i++) {
+    if (nobj->getDenseElement(i).isMagic(JS_ELEMENTS_HOLE)) {
+      continue;
+    }
+
+    num_properties += 1;
+  }
+
+  if (obj->is<TypedArrayObject>()) {
+    Handle<TypedArrayObject*> tobj = obj.as<TypedArrayObject>();
+    size_t len = tobj->length();
+
+    // Fail early if the typed array contains too many elements for a
+    // dense array, because we likely OOM anyway when trying to allocate
+    // more than 2GB for the properties vector. This also means we don't
+    // need to handle indices greater than MAX_INT32 in the loop below.
+    if (len > NativeObject::MAX_DENSE_ELEMENTS_COUNT) {
+      ReportOutOfMemory(cx);
+      return false;
+    }
+
+    MOZ_ASSERT(num_properties == 0, "typed arrays cannot have dense elements");
+    num_properties = len;
+  }
+
+  // All enumerable properties with string property keys are data
+  // properties. This allows us to collect the property values while
+  // iterating over the shape hierarchy without worrying over accessors
+  // modifying any state.
+
+  if (nobj->hasEnumerableProperty()) {
+    for (ShapePropertyIter<AllowGC::NoGC> iter(obj.as<NativeObject>()->shape());
+         !iter.done(); iter++) {
+      jsid id = iter->key();
+      if (!iter->enumerable() || id.isSymbol()) {
+        continue;
+      }
+      MOZ_ASSERT(!id.isInt(), "Unexpected indexed property");
+      num_properties += 1;
+    }
+  }
+
+  rval = num_properties;
+  return true;
+}
+
 // ES2018 draft rev c164be80f7ea91de5526b33d54e5c9321ed03d3f
 // 7.3.21 EnumerableOwnProperties ( O, kind )
 template <EnumerableOwnPropertiesKind kind>
@@ -2006,7 +2112,7 @@ static bool EnumerableOwnProperties(JSContext* cx, const JS::CallArgs& args) {
 
 // ES2018 draft rev c164be80f7ea91de5526b33d54e5c9321ed03d3f
 // 19.1.2.16 Object.keys ( O )
-static bool obj_keys(JSContext* cx, unsigned argc, Value* vp) {
+bool js::obj_keys(JSContext* cx, unsigned argc, Value* vp) {
   AutoJSMethodProfilerEntry pseudoFrame(cx, "Object", "keys");
   CallArgs args = CallArgsFromVp(argc, vp);
 
@@ -2032,6 +2138,26 @@ static bool obj_keys(JSContext* cx, unsigned argc, Value* vp) {
   return GetOwnPropertyKeys(cx, obj, JSITER_OWNONLY, args.rval());
 }
 
+bool js::obj_keys_length(JSContext* cx, HandleObject obj, int32_t& length) {
+  bool optimized;
+  if (!CountEnumerableOwnPropertiesNative(cx, obj, length, &optimized)) {
+    return false;
+  }
+  if (optimized) {
+    return true;
+  }
+
+  // Object.keys: Steps 2-3.
+  // (GetOwnPropertyKeys / CountOwnPropertyKeys)
+  RootedIdVector keys(cx);
+  if (!GetPropertyKeys(cx, obj, JSITER_OWNONLY, &keys)) {
+    return false;
+  }
+
+  length = keys.length();
+  return true;
+}
+
 // ES2018 draft rev c164be80f7ea91de5526b33d54e5c9321ed03d3f
 // 19.1.2.21 Object.values ( O )
 static bool obj_values(JSContext* cx, unsigned argc, Value* vp) {
@@ -2331,7 +2457,7 @@ static const JSFunctionSpec object_static_methods[] = {
                       "ObjectGetOwnPropertyDescriptor", 2, 0),
     JS_SELF_HOSTED_FN("getOwnPropertyDescriptors",
                       "ObjectGetOwnPropertyDescriptors", 1, 0),
-    JS_FN("keys", obj_keys, 1, 0),
+    JS_INLINABLE_FN("keys", obj_keys, 1, 0, ObjectKeys),
     JS_FN("values", obj_values, 1, 0),
     JS_FN("entries", obj_entries, 1, 0),
     JS_INLINABLE_FN("is", obj_is, 2, 0, ObjectIs),

js/src/builtin/Object.h

@@ -34,6 +34,13 @@ PlainObject* ObjectCreateWithTemplate(JSContext* cx,
 
 [[nodiscard]] bool obj_create(JSContext* cx, unsigned argc, JS::Value* vp);
 
+[[nodiscard]] bool obj_keys(JSContext* cx, unsigned argc, JS::Value* vp);
+
+// Similar to calling obj_keys followed by asking the length property, except
+// that we do not materialize the keys array.
+[[nodiscard]] bool obj_keys_length(JSContext* cx, HandleObject obj,
+                                   int32_t& length);
+
 [[nodiscard]] bool obj_is(JSContext* cx, unsigned argc, JS::Value* vp);
 
 [[nodiscard]] bool obj_toString(JSContext* cx, unsigned argc, JS::Value* vp);

js/src/jit-test/tests/ion/object-keys-00.js

@@ -0,0 +1,31 @@
+// This test case is used to test the optimized code path where the computation
+// of `Object.keys(...).length` no longer compute `Object.keys(...)` as an
+// intermediate result.
+//
+// This test verifies that the result remain consistent after the optimization.
+
+function cmp_keys_length(a, b) {
+    return Object.keys(a).length == Object.keys(b).length;
+}
+
+let objs = [
+    {x:0, a: 1, b: 2},
+    {x:1, b: 1, c: 2},
+    {x:2, c: 1, d: 2},
+    {x:3, a: 1, b: 2, c: 3},
+    {x:4, b: 1, c: 2, d: 3},
+    {x:5, a: 1, b: 2, c: 3, d: 4}
+];
+
+function max_of(o) {
+    return o?.d ?? o?.c ?? o?.b ?? 0;
+}
+
+with({}) {}
+for (let i = 0; i < 1000; i++) {
+    for (let o1 of objs) {
+        for (let o2 of objs) {
+            assertEq(cmp_keys_length(o1, o2), max_of(o1) == max_of(o2));
+        }
+    }
+}

js/src/jit-test/tests/ion/object-keys-01.js

@@ -0,0 +1,36 @@
+
+// This test case is used to test one common configuration seen in the wild and
+// that we expect to be successful at optimizing properly.
+
+// Similar functions are part of popular framework such as React and Angular.
+function shallowEqual(o1, o2) {
+    var k1 = Object.keys(o1);
+    var k2 = Object.keys(o2);
+    if (k1.length != k2.length) {
+        return false;
+    }
+    for (var k = 0; k < k1.length; k++) {
+        if (!Object.hasOwnProperty.call(o2, k1[k]) || !Object.is(o1[k1[k]], o2[k1[k]])) {
+            return false;
+        }
+    }
+    return true;
+}
+
+let objs = [
+    {x:0, a: 1, b: 2},
+    {x:1, b: 1, c: 2},
+    {x:2, c: 1, d: 2},
+    {x:3, a: 1, b: 2, c: 3},
+    {x:4, b: 1, c: 2, d: 3},
+    {x:5, a: 1, b: 2, c: 3, d: 4}
+];
+
+with({}) {}
+for (let i = 0; i < 1000; i++) {
+    for (let o1 of objs) {
+        for (let o2 of objs) {
+            assertEq(shallowEqual(o1, o2), Object.is(o1, o2));
+        }
+    }
+}

js/src/jit-test/tests/ion/object-keys-02.js

@@ -0,0 +1,54 @@
+
+// This test case check some code frequently used in the wild, with some object
+// (proxy) for which the optimization we have deployed do not work, as the
+// assumption the Object.keys(...) can be elided without having noticeable
+// side-effects.
+
+// Similar functions are part of popular framework such as React and Angular.
+function shallowEqual(o1, o2) {
+    var k1 = Object.keys(o1);
+    var k2 = Object.keys(o2);
+    if (k1.length != k2.length) {
+        return false;
+    }
+    for (var k = 0; k < k1.length; k++) {
+        if (!Object.hasOwnProperty.call(o2, k1[k]) || !Object.is(o1[k1[k]], o2[k1[k]])) {
+            return false;
+        }
+    }
+    return true;
+}
+
+let sideEffectCounter = 0;
+const payload = {x: 5, a: 1, b: 2, c: 3, d: 4};
+const handler = {
+    ownKeys(target) {
+        // side-effect that should not be removed.
+        sideEffectCounter++;
+        // answer returned.
+        return Reflect.ownKeys(target);
+    },
+};
+const proxy = new Proxy(payload, handler);
+
+let objs = [
+    {x:0, a: 1, b: 2},
+    {x:1, b: 1, c: 2},
+    {x:2, c: 1, d: 2},
+    {x:3, a: 1, b: 2, c: 3},
+    {x:4, b: 1, c: 2, d: 3},
+    proxy
+];
+
+with({}) {}
+let iterMax = 1000;
+for (let i = 0; i < iterMax; i++) {
+    for (let o1 of objs) {
+        for (let o2 of objs) {
+            assertEq(shallowEqual(o1, o2), Object.is(o1, o2));
+        }
+    }
+}
+
+let expectedSideEffects = 2 * objs.length * iterMax;
+assertEq(sideEffectCounter, expectedSideEffects);

js/src/jit-test/tests/ion/object-keys-03.js

@@ -0,0 +1,60 @@
+
+// This test case checks that some code can be optimized for non-proxy object,
+// and than we can correctly fallback if a proxy object ever flow into this
+// code.
+
+// Similar functions are part of popular framework such as React and Angular.
+function shallowEqual(o1, o2) {
+    var k1 = Object.keys(o1);
+    var k2 = Object.keys(o2);
+    if (k1.length != k2.length) {
+        return false;
+    }
+    for (var k = 0; k < k1.length; k++) {
+        if (!Object.hasOwnProperty.call(o2, k1[k]) || !Object.is(o1[k1[k]], o2[k1[k]])) {
+            return false;
+        }
+    }
+    return true;
+}
+
+let sideEffectCounter = 0;
+const payload = {x: 5, a: 1, b: 2, c: 3, d: 4};
+const handler = {
+    ownKeys(target) {
+        // side-effect that should not be removed.
+        sideEffectCounter++;
+        // answer returned.
+        return Reflect.ownKeys(target);
+    },
+};
+const proxy = new Proxy(payload, handler);
+
+let objs = [
+    {x:0, a: 1, b: 2},
+    {x:1, b: 1, c: 2},
+    {x:2, c: 1, d: 2},
+    {x:3, a: 1, b: 2, c: 3},
+    {x:4, b: 1, c: 2, d: 3},
+    {x:5, a: 1, b: 2, c: 3, d: 4}
+];
+
+// Ion compile shallowEqual ...
+with({}) {}
+let iterMax = 1000;
+for (let i = 0; i < iterMax; i++) {
+    for (let o1 of objs) {
+        for (let o2 of objs) {
+            assertEq(shallowEqual(o1, o2), Object.is(o1, o2));
+        }
+    }
+}
+
+assertEq(sideEffectCounter, 0);
+
+// ... before calling it with a proxy.
+// This should bailout with a guard failure.
+shallowEqual(objs[0], proxy);
+
+// Assert that the proxy's ownKeys function has been called.
+assertEq(sideEffectCounter, 1);