Bug 1436400 - Part 8: Add code path for bytecode encoding module scripts. r=nbp

The module scripts are scanned immediately before executing, and if the
module script is not yet executed, and it matches the requirement for encoding
(size, fetch count, etc), it's marked as encode,
and they're scanned again immediately after executing, and marked scripts are
queued for bytecode encoding.

This patch doesn't actually enable the encoding, given ShouldCacheBytecode
returns false because the caching is not enabled for the channel and
GetCacheTokenFetchCount fails.

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

dom/script/ModuleLoader.cpp

@@ -460,6 +460,12 @@ nsresult ModuleLoader::CompileOrFinishModuleScript(
       return NS_ERROR_FAILURE;
     }
 
+    if (ScriptLoader::ShouldCacheBytecode(aRequest)) {
+      if (!JS::StartIncrementalEncoding(aCx, std::move(stencil))) {
+        return NS_ERROR_FAILURE;
+      }
+    }
+
     return NS_OK;
   }
 
@@ -487,6 +493,12 @@ nsresult ModuleLoader::CompileOrFinishModuleScript(
     return NS_ERROR_FAILURE;
   }
 
+  if (ScriptLoader::ShouldCacheBytecode(aRequest)) {
+    if (!JS::StartIncrementalEncoding(aCx, std::move(stencil))) {
+      return NS_ERROR_FAILURE;
+    }
+  }
+
   return NS_OK;
 }
 

dom/script/ScriptLoader.cpp

@@ -20,6 +20,7 @@
 #include "js/loader/LoadedScript.h"
 #include "js/loader/ModuleLoadRequest.h"
 #include "js/MemoryFunctions.h"
+#include "js/Modules.h"
 #include "js/OffThreadScriptCompilation.h"
 #include "js/PropertyAndElement.h"  // JS_DefineProperty
 #include "js/Realm.h"
@@ -589,6 +590,7 @@ nsresult ScriptLoader::StartLoadInternal(ScriptLoadRequest* aRequest,
   nsCOMPtr<nsICacheInfoChannel> cic(do_QueryInterface(channel));
   if (cic && StaticPrefs::dom_script_loader_bytecode_cache_enabled() &&
       // Bug 1436400: no bytecode cache support for modules yet.
+      // TODO: Remove this also to enable bytecode encoding.
       !aRequest->IsModuleRequest()) {
     MOZ_ASSERT(!aRequest->GetLoadContext()->GetWebExtGlobal(),
                "Can not bytecode cache WebExt code");
@@ -2127,6 +2129,9 @@ nsresult ScriptLoader::CompileOrDecodeClassicScript(
 
 /* static */
 nsCString& ScriptLoader::BytecodeMimeTypeFor(ScriptLoadRequest* aRequest) {
+  if (aRequest->IsModuleRequest()) {
+    return nsContentUtils::JSModuleBytecodeMimeType();
+  }
   return nsContentUtils::JSScriptBytecodeMimeType();
 }
 
@@ -2148,6 +2153,7 @@ nsresult ScriptLoader::MaybePrepareForBytecodeEncodingAfterExecute(
     //       first encode.
     MOZ_ASSERT(aRequest->mBytecodeOffset == aRequest->mScriptBytecode.length());
     RegisterForBytecodeEncoding(aRequest);
+    MOZ_ASSERT(IsAlreadyHandledForBytecodeEncodingPreparation(aRequest));
 
     return aRv;
   }
@@ -2157,6 +2163,54 @@ nsresult ScriptLoader::MaybePrepareForBytecodeEncodingAfterExecute(
   TRACE_FOR_TEST_NONE(aRequest->GetLoadContext()->GetScriptElement(),
                       "scriptloader_no_encode");
   aRequest->mCacheInfo = nullptr;
+  MOZ_ASSERT(IsAlreadyHandledForBytecodeEncodingPreparation(aRequest));
+
+  return aRv;
+}
+
+bool ScriptLoader::IsAlreadyHandledForBytecodeEncodingPreparation(
+    ScriptLoadRequest* aRequest) {
+  return aRequest->isInList() || !aRequest->mCacheInfo;
+}
+
+void ScriptLoader::MaybePrepareModuleForBytecodeEncodingBeforeExecute(
+    JSContext* aCx, ModuleLoadRequest* aRequest) {
+  {
+    ModuleScript* moduleScript = aRequest->mModuleScript;
+    JS::Rooted<JSObject*> module(aCx, moduleScript->ModuleRecord());
+
+    if (JS::IsModuleEvaluated(module)) {
+      // This module is already evaluated, and all imported modules should also
+      // already be evaluated.
+      return;
+    }
+
+    if (aRequest->IsMarkedForBytecodeEncoding()) {
+      // This module is imported multiple times, and already marked.
+      return;
+    }
+
+    JS::Rooted<JSScript*> script(aCx, JS::GetModuleScript(module));
+    MaybePrepareForBytecodeEncodingBeforeExecute(aRequest, script);
+  }
+
+  for (ModuleLoadRequest* childRequest : aRequest->mImports) {
+    MaybePrepareModuleForBytecodeEncodingBeforeExecute(aCx, childRequest);
+  }
+}
+
+nsresult ScriptLoader::MaybePrepareModuleForBytecodeEncodingAfterExecute(
+    ModuleLoadRequest* aRequest, nsresult aRv) {
+  if (IsAlreadyHandledForBytecodeEncodingPreparation(aRequest)) {
+    // This module is imported multiple times and already handled.
+    return aRv;
+  }
+
+  aRv = MaybePrepareForBytecodeEncodingAfterExecute(aRequest, aRv);
+
+  for (ModuleLoadRequest* childRequest : aRequest->mImports) {
+    aRv = MaybePrepareModuleForBytecodeEncodingAfterExecute(childRequest, aRv);
+  }
 
   return aRv;
 }
@@ -2338,7 +2392,6 @@ void ScriptLoader::EncodeBytecode() {
   RefPtr<ScriptLoadRequest> request;
   while (!mBytecodeEncodingQueue.isEmpty()) {
     request = mBytecodeEncodingQueue.StealFirst();
-    MOZ_ASSERT(!request->IsModuleRequest());
     MOZ_ASSERT(!request->GetLoadContext()->GetWebExtGlobal(),
                "Not handling global above");
     EncodeRequestBytecode(aes.cx(), request);
@@ -2441,7 +2494,6 @@ void ScriptLoader::GiveUpBytecodeEncoding() {
     LOG(("ScriptLoadRequest (%p): Cannot serialize bytecode", request.get()));
     TRACE_FOR_TEST_NONE(request->GetLoadContext()->GetScriptElement(),
                         "scriptloader_bytecode_failed");
-    MOZ_ASSERT(!request->IsModuleRequest());
     MOZ_ASSERT(!request->GetLoadContext()->GetWebExtGlobal());
 
     if (aes.isSome()) {

dom/script/ScriptLoader.h

@@ -573,6 +573,17 @@ class ScriptLoader final : public JS::loader::ScriptLoaderInterface {
   nsresult MaybePrepareForBytecodeEncodingAfterExecute(
       ScriptLoadRequest* aRequest, nsresult aRv);
 
+  // Returns true if MaybePrepareForBytecodeEncodingAfterExecute is called
+  // for given script load request.
+  bool IsAlreadyHandledForBytecodeEncodingPreparation(
+      ScriptLoadRequest* aRequest);
+
+  void MaybePrepareModuleForBytecodeEncodingBeforeExecute(
+      JSContext* aCx, ModuleLoadRequest* aRequest) override;
+
+  nsresult MaybePrepareModuleForBytecodeEncodingAfterExecute(
+      ModuleLoadRequest* aRequest, nsresult aRv) override;
+
   // Implements https://html.spec.whatwg.org/#run-a-classic-script
   nsresult EvaluateScript(nsIGlobalObject* aGlobalObject,
                           ScriptLoadRequest* aRequest);
@@ -590,7 +601,7 @@ class ScriptLoader final : public JS::loader::ScriptLoaderInterface {
    * no more script have to be processed.  If all conditions are met, queue an
    * event to encode all the bytecode and save them on the cache.
    */
-  void MaybeTriggerBytecodeEncoding();
+  void MaybeTriggerBytecodeEncoding() override;
 
   /**
    * Iterate over all script load request and save the bytecode of executed

js/loader/ModuleLoaderBase.cpp

@@ -777,6 +777,8 @@ nsresult ModuleLoaderBase::EvaluateModule(nsIGlobalObject* aGlobalObject,
 
   JS::Rooted<JS::Value> rval(cx);
 
+  mLoader->MaybePrepareModuleForBytecodeEncodingBeforeExecute(cx, request);
+
   rv = nsJSUtils::ModuleEvaluate(cx, module, &rval);
 
   if (NS_SUCCEEDED(rv)) {
@@ -814,11 +816,10 @@ nsresult ModuleLoaderBase::EvaluateModule(nsIGlobalObject* aGlobalObject,
     }
   }
 
-  if (aRequest->HasLoadContext()) {
-    TRACE_FOR_TEST_NONE(aRequest->GetLoadContext()->GetScriptElement(),
-                        "scriptloader_no_encode");
-  }
-  aRequest->mCacheInfo = nullptr;
+  rv = mLoader->MaybePrepareModuleForBytecodeEncodingAfterExecute(request, rv);
+
+  mLoader->MaybeTriggerBytecodeEncoding();
+
   return rv;
 }
 

js/loader/ModuleLoaderBase.h

@@ -70,6 +70,14 @@ class ScriptLoaderInterface : public nsISupports {
       JSContext* cx, ScriptLoadRequest* aRequest,
       JS::Handle<JSObject*> aScopeChain, JS::CompileOptions* aOptions,
       JS::MutableHandle<JSScript*> aIntroductionScript) = 0;
+
+  virtual void MaybePrepareModuleForBytecodeEncodingBeforeExecute(
+      JSContext* aCx, ModuleLoadRequest* aRequest) = 0;
+
+  virtual nsresult MaybePrepareModuleForBytecodeEncodingAfterExecute(
+      ModuleLoadRequest* aRequest, nsresult aRv) = 0;
+
+  virtual void MaybeTriggerBytecodeEncoding() = 0;
 };
 
 class ModuleLoaderBase : public nsISupports {