Bug 1317481 - Optimize away Generator/Promise handling for await in the topmost JS frame with already resolved/rejected Promise. r=anba,smaug

dom/workers/RuntimeService.cpp

@@ -1139,6 +1139,7 @@ public:
       microTaskQueue = &GetDebuggerMicroTaskQueue();
     }
 
+    JS::JobQueueMayNotBeEmpty(cx);
     microTaskQueue->push(runnable.forget());
   }
 

dom/worklet/WorkletThread.cpp

@@ -171,7 +171,6 @@ public:
   {
     RefPtr<MicroTaskRunnable> runnable(aRunnable);
 
-#ifdef DEBUG
     MOZ_ASSERT(!NS_IsMainThread());
     MOZ_ASSERT(runnable);
 
@@ -181,10 +180,12 @@ public:
     JSContext* cx = workletThread->GetJSContext();
     MOZ_ASSERT(cx);
 
+#ifdef DEBUG
     JS::Rooted<JSObject*> global(cx, JS::CurrentGlobalOrNull(cx));
     MOZ_ASSERT(global);
 #endif
 
+    JS::JobQueueMayNotBeEmpty(cx);
     GetMicroTaskQueue().push(runnable.forget());
   }
 

js/public/Value.h

@@ -222,6 +222,9 @@ enum JSWhyMagic
     /** standard constructors are not created for off-thread parsing. */
     JS_OFF_THREAD_CONSTRUCTOR,
 
+    /** used in jit::TrySkipAwait */
+    JS_CANNOT_SKIP_AWAIT,
+
     /** for local use */
     JS_GENERIC_MAGIC,
 

js/src/builtin/Promise.cpp

@@ -21,6 +21,8 @@
 #include "vm/Debugger.h"
 #include "vm/Iteration.h"
 #include "vm/JSContext.h"
+#include "vm/JSObject.h"
+#include "vm/SelfHosting.h"
 
 #include "vm/Compartment-inl.h"
 #include "vm/Debugger-inl.h"
@@ -4390,6 +4392,97 @@ js::PromiseLookup::isDefaultInstance(JSContext* cx, PromiseObject* promise,
     return hasDefaultProtoAndNoShadowedProperties(cx, promise);
 }
 
+// We can skip `await` with an already resolved value only if the current frame
+// is the topmost JS frame and the current job is the last job in the job queue.
+// This guarantees that any new job enqueued in the current turn will be
+// executed immediately after the current job.
+//
+// Currently we only support skipping jobs when the async function is resumed
+// at least once.
+static MOZ_MUST_USE bool
+IsTopMostAsyncFunctionCall(JSContext* cx)
+{
+    FrameIter iter(cx);
+
+    // The current frame should be the async function.
+    if (iter.done())
+        return false;
+    if (!iter.calleeTemplate())
+        return false;
+    MOZ_ASSERT(iter.calleeTemplate()->isAsync());
+
+    ++iter;
+
+    // The parent frame should be the `next` function of the generator that is
+    // internally called in AsyncFunctionResume.
+    if (iter.done())
+        return false;
+    if (!iter.calleeTemplate())
+        return false;
+
+    if (!IsSelfHostedFunctionWithName(iter.calleeTemplate(), cx->names().GeneratorNext))
+        return false;
+
+    ++iter;
+
+    // There should be no more frames.
+    if (iter.done())
+        return true;
+
+    return false;
+}
+
+MOZ_MUST_USE bool
+js::TrySkipAwait(JSContext* cx, HandleValue val, bool* canSkip, MutableHandleValue resolved)
+{
+    if (!cx->canSkipEnqueuingJobs) {
+        *canSkip = false;
+        return true;
+    }
+
+    if (!IsTopMostAsyncFunctionCall(cx)) {
+        *canSkip = false;
+        return true;
+    }
+
+    // Primitive values cannot be 'thenables', so we can trivially skip the
+    // await operation.
+    if (!val.isObject()) {
+        resolved.set(val);
+        *canSkip = true;
+        return true;
+    }
+
+    JSObject* obj = &val.toObject();
+    if (!obj->is<PromiseObject>()) {
+        *canSkip = false;
+        return true;
+    }
+
+    PromiseObject* promise = &obj->as<PromiseObject>();
+
+    if (promise->state() == JS::PromiseState::Pending) {
+        *canSkip = false;
+        return true;
+    }
+
+    PromiseLookup& promiseLookup = cx->realm()->promiseLookup;
+    if (!promiseLookup.isDefaultInstance(cx, promise)) {
+        *canSkip = false;
+        return true;
+    }
+
+    if (promise->state() == JS::PromiseState::Rejected) {
+        // We don't optimize rejected Promises for now.
+        *canSkip = false;
+        return true;
+    }
+
+    resolved.set(promise->value());
+    *canSkip = true;
+    return true;
+}
+
 OffThreadPromiseTask::OffThreadPromiseTask(JSContext* cx, Handle<PromiseObject*> promise)
   : runtime_(cx->runtime()),
     promise_(cx, promise),

js/src/builtin/Promise.h

@@ -159,6 +159,12 @@ AsyncFunctionThrown(JSContext* cx, Handle<PromiseObject*> resultPromise);
 MOZ_MUST_USE bool
 AsyncFunctionAwait(JSContext* cx, Handle<PromiseObject*> resultPromise, HandleValue value);
 
+// If the await operation can be skipped and the resolution value for `val` can
+// be acquired, stored the resolved value to `resolved` and `true` to
+// `*canSkip`.  Otherwise, stores `false` to `*canSkip`.
+MOZ_MUST_USE bool
+TrySkipAwait(JSContext* cx, HandleValue val, bool* canSkip, MutableHandleValue resolved);
+
 class AsyncGeneratorObject;
 
 MOZ_MUST_USE bool

js/src/frontend/BytecodeEmitter.cpp

@@ -6068,10 +6068,26 @@ BytecodeEmitter::emitAwaitInInnermostScope(ParseNode* pn)
 bool
 BytecodeEmitter::emitAwaitInScope(EmitterScope& currentScope)
 {
+    if (!emit1(JSOP_TRYSKIPAWAIT))              // VALUE_OR_RESOLVED CANSKIP
+        return false;
+
+    if (!emit1(JSOP_NOT))                       // VALUE_OR_RESOLVED !CANSKIP
+        return false;
+
+    InternalIfEmitter ifCanSkip(this);
+    if (!ifCanSkip.emitThen())                  // VALUE_OR_RESOLVED
+        return false;
+
     if (!emitGetDotGeneratorInScope(currentScope))
+        return false;                           // VALUE GENERATOR
+    if (!emitYieldOp(JSOP_AWAIT))               // RESOLVED
         return false;
-    if (!emitYieldOp(JSOP_AWAIT))
+
+    if (!ifCanSkip.emitEnd())
         return false;
+
+    MOZ_ASSERT(ifCanSkip.popped() == 0);
+
     return true;
 }
 

js/src/jit/BaselineCompiler.cpp

@@ -1032,7 +1032,6 @@ BaselineCompiler::emitBody()
             // Run-once opcode during self-hosting initialization.
           case JSOP_UNUSED126:
           case JSOP_UNUSED206:
-          case JSOP_UNUSED223:
           case JSOP_LIMIT:
             // === !! WARNING WARNING WARNING !! ===
             // Do you really want to sacrifice performance by not implementing
@@ -4122,6 +4121,38 @@ BaselineCompiler::emit_JSOP_TOASYNCITER()
     return true;
 }
 
+typedef bool (*TrySkipAwaitFn)(JSContext*, HandleValue, MutableHandleValue);
+static const VMFunction TrySkipAwaitInfo = FunctionInfo<TrySkipAwaitFn>(jit::TrySkipAwait, "TrySkipAwait");
+
+bool
+BaselineCompiler::emit_JSOP_TRYSKIPAWAIT()
+{
+    frame.syncStack(0);
+    masm.loadValue(frame.addressOfStackValue(frame.peek(-1)), R0);
+
+    prepareVMCall();
+    pushArg(R0);
+
+    if (!callVM(TrySkipAwaitInfo))
+        return false;
+
+    Label cannotSkip, done;
+    masm.branchTestMagicValue(Assembler::Equal, R0, JS_CANNOT_SKIP_AWAIT, &cannotSkip);
+    masm.moveValue(BooleanValue(true), R1);
+    masm.jump(&done);
+
+    masm.bind(&cannotSkip);
+    masm.loadValue(frame.addressOfStackValue(frame.peek(-1)), R0);
+    masm.moveValue(BooleanValue(false), R1);
+
+    masm.bind(&done);
+
+    frame.pop();
+    frame.push(R0);
+    frame.push(R1);
+    return true;
+}
+
 typedef bool (*ThrowObjectCoercibleFn)(JSContext*, HandleValue);
 static const VMFunction ThrowObjectCoercibleInfo =
     FunctionInfo<ThrowObjectCoercibleFn>(ThrowObjectCoercible, "ThrowObjectCoercible");

js/src/jit/BaselineCompiler.h

@@ -217,6 +217,7 @@ namespace jit {
     _(JSOP_INITIALYIELD)       \
     _(JSOP_YIELD)              \
     _(JSOP_AWAIT)              \
+    _(JSOP_TRYSKIPAWAIT)       \
     _(JSOP_DEBUGAFTERYIELD)    \
     _(JSOP_FINALYIELDRVAL)     \
     _(JSOP_RESUME)             \

js/src/jit/IonBuilder.cpp

@@ -2453,6 +2453,7 @@ IonBuilder::inspectOpcode(JSOp op)
       case JSOP_RESUME:
       case JSOP_DEBUGAFTERYIELD:
       case JSOP_AWAIT:
+      case JSOP_TRYSKIPAWAIT:
       case JSOP_GENERATOR:
 
       // Misc
@@ -2474,7 +2475,6 @@ IonBuilder::inspectOpcode(JSOp op)
 
       case JSOP_UNUSED126:
       case JSOP_UNUSED206:
-      case JSOP_UNUSED223:
       case JSOP_LIMIT:
         break;
     }

js/src/jit/VMFunctions.cpp

@@ -6,6 +6,7 @@
 
 #include "jit/VMFunctions.h"
 
+#include "builtin/Promise.h"
 #include "builtin/TypedObject.h"
 #include "frontend/BytecodeCompiler.h"
 #include "jit/arm/Simulator-arm.h"
@@ -1951,5 +1952,18 @@ typedef bool (*DoConcatStringObjectFn)(JSContext*, HandleValue, HandleValue,
 const VMFunction DoConcatStringObjectInfo =
     FunctionInfo<DoConcatStringObjectFn>(DoConcatStringObject, "DoConcatStringObject", TailCall, PopValues(2));
 
+MOZ_MUST_USE bool
+TrySkipAwait(JSContext* cx, HandleValue val, MutableHandleValue resolved)
+{
+    bool canSkip;
+    if (!TrySkipAwait(cx, val, &canSkip, resolved))
+        return false;
+
+    if (!canSkip)
+        resolved.setMagic(JS_CANNOT_SKIP_AWAIT);
+
+    return true;
+}
+
 } // namespace jit
 } // namespace js

js/src/jit/VMFunctions.h

@@ -957,6 +957,13 @@ bool
 DoConcatStringObject(JSContext* cx, HandleValue lhs, HandleValue rhs,
                      MutableHandleValue res);
 
+// Wrapper for js::TrySkipAwait.
+// If the await operation can be skipped and the resolution value for `val` can
+// be acquired, stored the resolved value to `resolved`.  Otherwise, stores
+// the JS_CANNOT_SKIP_AWAIT magic value to `resolved`.
+MOZ_MUST_USE bool
+TrySkipAwait(JSContext* cx, HandleValue val, MutableHandleValue resolved);
+
 // This is the tailcall version of DoConcatStringObject
 extern const VMFunction DoConcatStringObjectInfo;
 

js/src/jsapi.cpp

@@ -4995,6 +4995,18 @@ JS::SetPromiseRejectionTrackerCallback(JSContext* cx, JSPromiseRejectionTrackerC
     cx->promiseRejectionTrackerCallbackData = data;
 }
 
+extern JS_PUBLIC_API(void)
+JS::JobQueueIsEmpty(JSContext* cx)
+{
+    cx->canSkipEnqueuingJobs = true;
+}
+
+extern JS_PUBLIC_API(void)
+JS::JobQueueMayNotBeEmpty(JSContext* cx)
+{
+    cx->canSkipEnqueuingJobs = false;
+}
+
 JS_PUBLIC_API(JSObject*)
 JS::NewPromiseObject(JSContext* cx, HandleObject executor, HandleObject proto /* = nullptr */)
 {

js/src/jsapi.h

@@ -4060,6 +4060,34 @@ extern JS_PUBLIC_API(void)
 SetPromiseRejectionTrackerCallback(JSContext* cx, JSPromiseRejectionTrackerCallback callback,
                                    void* data = nullptr);
 
+/**
+ * Inform the runtime that the job queue is empty and the embedding is going to
+ * execute its last promise job. The runtime may now choose to skip creating
+ * promise jobs for asynchronous execution and instead continue execution
+ * synchronously. More specifically, this optimization is used to skip the
+ * standard job queuing behavior for `await` operations in async functions.
+ *
+ * This function may be called before executing the last job in the job queue.
+ * When it was called, JobQueueMayNotBeEmpty must be called in order to restore
+ * the default job queuing behavior before the embedding enqueues its next job
+ * into the job queue.
+ */
+extern JS_PUBLIC_API(void)
+JobQueueIsEmpty(JSContext* cx);
+
+/**
+ * Inform the runtime that job queue is no longer empty. The runtime can now no
+ * longer skip creating promise jobs for asynchronous execution, because
+ * pending jobs in the job queue must be executed first to preserve the FIFO
+ * (first in - first out) property of the queue. This effectively undoes
+ * JobQueueIsEmpty and re-enables the standard job queuing behavior.
+ *
+ * This function must be called whenever enqueuing a job to the job queue when
+ * JobQueueIsEmpty was called previously.
+ */
+extern JS_PUBLIC_API(void)
+JobQueueMayNotBeEmpty(JSContext* cx);
+
 /**
  * Returns a new instance of the Promise builtin class in the current
  * compartment, with the right slot layout.

js/src/vm/Interpreter.cpp

@@ -24,6 +24,7 @@
 #include "builtin/Array.h"
 #include "builtin/Eval.h"
 #include "builtin/ModuleObject.h"
+#include "builtin/Promise.h"
 #include "builtin/String.h"
 #include "jit/AtomicOperations.h"
 #include "jit/BaselineJIT.h"
@@ -2147,7 +2148,6 @@ CASE(JSOP_NOP_DESTRUCTURING)
 CASE(JSOP_TRY_DESTRUCTURING_ITERCLOSE)
 CASE(JSOP_UNUSED126)
 CASE(JSOP_UNUSED206)
-CASE(JSOP_UNUSED223)
 CASE(JSOP_CONDSWITCH)
 {
     MOZ_ASSERT(CodeSpec[*REGS.pc].length == 1);
@@ -3828,6 +3828,24 @@ CASE(JSOP_TOASYNCITER)
 }
 END_CASE(JSOP_TOASYNCITER)
 
+CASE(JSOP_TRYSKIPAWAIT)
+{
+    ReservedRooted<Value> val(&rootValue0, REGS.sp[-1]);
+    ReservedRooted<Value> resolved(&rootValue1);
+    bool canSkip;
+
+    if (!TrySkipAwait(cx, val, &canSkip, &resolved))
+        goto error;
+
+    if (canSkip) {
+        REGS.sp[-1] = resolved;
+        PUSH_BOOLEAN(true);
+    } else {
+        PUSH_BOOLEAN(false);
+    }
+}
+END_CASE(JSOP_TRYSKIPAWAIT)
+
 CASE(JSOP_SETFUNNAME)
 {
     MOZ_ASSERT(REGS.stackDepth() >= 2);

js/src/vm/JSContext.cpp

@@ -1045,6 +1045,7 @@ InternalEnqueuePromiseJobCallback(JSContext* cx, JS::HandleObject job,
                                   JS::HandleObject incumbentGlobal, void* data)
 {
     MOZ_ASSERT(job);
+    JS::JobQueueMayNotBeEmpty(cx);
     if (!cx->jobQueue->append(job)) {
         ReportOutOfMemory(cx);
         return false;
@@ -1098,6 +1099,7 @@ JS_FRIEND_API(bool)
 js::EnqueueJob(JSContext* cx, JS::HandleObject job)
 {
     MOZ_ASSERT(cx->jobQueue);
+    JS::JobQueueMayNotBeEmpty(cx);
     if (!cx->jobQueue->append(job)) {
         ReportOutOfMemory(cx);
         return false;
@@ -1153,6 +1155,12 @@ js::RunJobs(JSContext* cx)
                 continue;
 
             cx->jobQueue->get()[i] = nullptr;
+
+            // If the next job is the last job in the job queue, allow
+            // skipping the standard job queuing behavior.
+            if (i == cx->jobQueue->length() - 1)
+                JS::JobQueueIsEmpty(cx);
+
             AutoRealm ar(cx, job);
             {
                 if (!JS::Call(cx, UndefinedHandleValue, job, args, &rval)) {
@@ -1296,6 +1304,7 @@ JSContext::JSContext(JSRuntime* runtime, const JS::ContextOptions& options)
     jobQueue(nullptr),
     drainingJobQueue(false),
     stopDrainingJobQueue(false),
+    canSkipEnqueuingJobs(false),
     promiseRejectionTrackerCallback(nullptr),
     promiseRejectionTrackerCallbackData(nullptr)
 {

js/src/vm/JSContext.h

@@ -919,6 +919,7 @@ struct JSContext : public JS::RootingContext,
     js::ThreadData<JS::PersistentRooted<js::JobQueue>*> jobQueue;
     js::ThreadData<bool> drainingJobQueue;
     js::ThreadData<bool> stopDrainingJobQueue;
+    js::ThreadData<bool> canSkipEnqueuingJobs;
 
     js::ThreadData<JSPromiseRejectionTrackerCallback> promiseRejectionTrackerCallback;
     js::ThreadData<void*> promiseRejectionTrackerCallbackData;

js/src/vm/Opcodes.h

@@ -2268,7 +2268,18 @@
      *   Stack: val => val
      */ \
     macro(JSOP_ITERNEXT,      222, "iternext",   NULL,  1,  1,  1,  JOF_BYTE) \
-    macro(JSOP_UNUSED223,     223, "unused223",  NULL,  1,  0,  0,  JOF_BYTE) \
+    /*
+     * Pops the top of stack value as 'value', checks if the await for 'value'
+     * can be skipped.
+     * If the await operation can be skipped and the resolution value for
+     * 'value' can be acquired, pushes the resolution value and 'true' onto the
+     * stack.  Otherwise, pushes 'value' and 'false' on the stack.
+     *   Category: Statements
+     *   Type: Function
+     *   Operands:
+     *   Stack: value => value_or_resolved, canskip
+     */ \
+    macro(JSOP_TRYSKIPAWAIT,  223,"tryskipawait",  NULL,  1,  1,  2,  JOF_BYTE) \
     \
     /*
      * Creates rest parameter array for current function call, and pushes it

xpcom/base/CycleCollectedJSContext.cpp

@@ -24,6 +24,7 @@
 #include "mozilla/dom/PromiseBinding.h"
 #include "mozilla/dom/PromiseDebugging.h"
 #include "mozilla/dom/ScriptSettings.h"
+#include "jsapi.h"
 #include "js/Debug.h"
 #include "js/GCAPI.h"
 #include "js/Utility.h"
@@ -487,6 +488,7 @@ CycleCollectedJSContext::DispatchToMicroTask(
   MOZ_ASSERT(NS_IsMainThread());
   MOZ_ASSERT(runnable);
 
+  JS::JobQueueMayNotBeEmpty(Context());
   mPendingMicroTaskRunnables.push(runnable.forget());
 }
 
@@ -557,8 +559,13 @@ CycleCollectedJSContext::PerformMicroTaskCheckPoint(bool aForce)
       // Otherwise, mPendingMicroTaskRunnables will be replaced later with
       // all suppressed tasks in mDebuggerMicroTaskQueue unexpectedly.
       MOZ_ASSERT(NS_IsMainThread());
+      JS::JobQueueMayNotBeEmpty(Context());
       suppressed.push(runnable);
     } else {
+      if (mPendingMicroTaskRunnables.empty() &&
+          mDebuggerMicroTaskQueue.empty() && suppressed.empty()) {
+        JS::JobQueueIsEmpty(Context());
+      }
       didProcess = true;
       runnable->Run(aso);
     }
@@ -597,6 +604,10 @@ CycleCollectedJSContext::PerformDebuggerMicroTaskCheckpoint()
 
     // This function can re-enter, so we remove the element before calling.
     microtaskQueue->pop();
+
+    if (mPendingMicroTaskRunnables.empty() && mDebuggerMicroTaskQueue.empty()) {
+      JS::JobQueueIsEmpty(Context());
+    }
     runnable->Run(aso);
   }