Scala.js: Support js.async and js.await, including JSPI on Wasm.

This is forward port of the compiler changes in the two commits of
the Scala.js PR scala-js/scala-js#5130

---

We add support for a new pair of primitive methods, `js.async` and
`js.await`. They correspond to JavaScript `async` functions and
`await` expressions.

`js.await(p)` awaits a `Promise`, but it must be directly scoped
within a `js.async { ... }` block.

At the IR level, `js.await(p)` directly translates to a dedicated
IR node `JSAwait(arg)`. `js.async` blocks don't have a direct
representation. Instead, the IR models `async function`s and
`async =>`functions, as `Closure`s with an additionnal `async`
flag. This corresponds to the JavaScript model for `async/await`.

A `js.async { body }` block therefore corresponds to an
immediately-applied `async Closure`, which in JavaScript would be
written as `(async () => body)()`.

---

We then optionally allow orphan `js.await(p)` on WebAssembly.

With the JavaScript Promise Integration (JSPI), there can be as
many frames as we want between a `WebAssembly.promising` function
and the corresponding `WebAssembly.Suspending` calls. The former
is introduced by our `js.async` blocks, and the latter by calls to
`js.await`.

Normally, `js.await` must be directly enclosed within a `js.async`
block. This ensures that it can be compiled to a JavaScript `async`
function and `await` expression.

We introduce a sort of "language import" to allow "orphan awaits".
This way, we can decouple the `js.await` calls from their
`js.async` blocks. The generated IR will then only link when
targeting WebAssembly. Technically, `js.await` requires an
implicit `js.AwaitPermit`. The default one only allows `js.await`
directly inside `js.async`, and we can import a stronger one that
allows orphan awaits.

There must still be a `js.async` block *dynamically* enclosing any
`js.await` call (on the call stack), without intervening JS frame.
If that dynamic property is not satisfied, a
`WebAssembly.SuspendError` gets thrown. This last point is not yet
implemented in Node.js at the time of this commit; instead such a
situation traps. The corresponding test is therefore currently
ignored.

---

Since these features fundamentally depend on a) the target ES
version and b) the WebAssembly backend, we add more configurations
of Scala.js to the Scala 3 CI.

Author: sjrd

.github/workflows/ci.yaml

@@ -155,7 +155,10 @@ jobs:
 
       - name: Scala.js Test
         run: |
-          ./project/scripts/sbt ";sjsSandbox/run ;sjsSandbox/test ;sjsJUnitTests/test ;set sjsJUnitTests/scalaJSLinkerConfig ~= switchToESModules ;sjsJUnitTests/test ;sjsCompilerTests/test"
+          ./project/scripts/sbt ";sjsSandbox/run ;sjsSandbox/test ;sjsJUnitTests/test ;set sjsJUnitTests/scalaJSLinkerConfig ~= switchToESModules ;sjsJUnitTests/test"
+          ./project/scripts/sbt ";sjsJUnitTests/clean ;set sjsJUnitTests/scalaJSLinkerConfig ~= switchToLatestESVersion ;sjsJUnitTests/test"
+          ./project/scripts/sbt ";sjsJUnitTests/clean ;set sjsJUnitTests/scalaJSLinkerConfig ~= switchToLatestESVersion ;set Global/enableWebAssembly := true; sjsJUnitTests/test"
+          ./project/scripts/sbt ";sjsCompilerTests/test"
 
   test_windows_fast:
     runs-on: [self-hosted, Windows]

compiler/src/dotty/tools/backend/sjs/JSCodeGen.scala

@@ -77,6 +77,7 @@ class JSCodeGen()(using genCtx: Context) {
 
   val currentClassSym = new ScopedVar[Symbol]
   private val delambdafyTargetDefDefs = new ScopedVar[MutableSymbolMap[DefDef]]
+  private val methodsAllowingJSAwait = new ScopedVar[mutable.Set[Symbol]]
   private val currentMethodSym = new ScopedVar[Symbol]
   private val localNames = new ScopedVar[LocalNameGenerator]
   private val thisLocalVarName = new ScopedVar[Option[LocalName]]
@@ -91,6 +92,7 @@ class JSCodeGen()(using genCtx: Context) {
     withScopedVars(
         currentClassSym := null,
         delambdafyTargetDefDefs := null,
+        methodsAllowingJSAwait := null,
         currentMethodSym := null,
         localNames := null,
         thisLocalVarName := null,
@@ -267,6 +269,7 @@ class JSCodeGen()(using genCtx: Context) {
         withScopedVars(
             currentClassSym := sym,
             delambdafyTargetDefDefs := MutableSymbolMap(),
+            methodsAllowingJSAwait := mutable.Set.empty,
         ) {
           val tree = if (sym.isJSType) {
             if (!sym.is(Trait) && sym.isNonNativeJSClass)
@@ -2424,6 +2427,7 @@ class JSCodeGen()(using genCtx: Context) {
       withScopedVars(
           currentClassSym := sym,
           delambdafyTargetDefDefs := MutableSymbolMap(),
+          methodsAllowingJSAwait := mutable.Set.empty,
       ) {
         genNonNativeJSClass(typeDef)
       }
@@ -4073,6 +4077,54 @@ class JSCodeGen()(using genCtx: Context) {
         // js.import.meta
         js.JSImportMeta()
 
+      case JS_ASYNC =>
+        // js.async(arg)
+        assert(args.size == 1,
+            s"Expected exactly 1 argument for JS primitive $code but got " +
+            s"${args.size} at $pos")
+
+        def extractStatsAndClosure(arg: Tree): (List[Tree], Closure) = (arg: @unchecked) match
+          case arg: Closure =>
+            (Nil, arg)
+          case Block(outerStats, expr) =>
+            val (innerStats, closure) = extractStatsAndClosure(expr)
+            (outerStats ::: innerStats, closure)
+
+        val (stats, fun @ Closure(_, target, _)) = extractStatsAndClosure(args.head)
+        methodsAllowingJSAwait += target.symbol
+        val genStats = stats.map(genStat(_))
+        val asyncExpr = genClosure(fun) match {
+          case js.NewLambda(_, closure: js.Closure)
+              if closure.params.isEmpty && closure.resultType == jstpe.AnyType =>
+            val newFlags = closure.flags.withTyped(false).withAsync(true)
+            js.JSFunctionApply(closure.copy(flags = newFlags), Nil)
+          case other =>
+            throw FatalError(
+                s"Unexpected tree generated for the Function0 argument to js.async at ${tree.sourcePos}: $other")
+        }
+        js.Block(genStats, asyncExpr)
+
+      case JS_AWAIT =>
+        // js.await(arg)
+        val (arg, permitValue) = genArgs2
+        if (!methodsAllowingJSAwait.contains(currentMethodSym)) {
+          // This is an orphan await
+          if (!(args(1).tpe <:< jsdefn.WasmJSPI_allowOrphanJSAwaitModuleClassRef)) {
+            report.error(
+                "Illegal use of js.await().\n" +
+                "It can only be used inside a js.async {...} block, without any lambda,\n" +
+                "by-name argument or nested method in-between.\n" +
+                "If you compile for WebAssembly, you can allow arbitrary js.await()\n" +
+                "calls by adding the following import:\n" +
+                "import scala.scalajs.js.wasm.JSPI.allowOrphanJSAwait",
+                tree.sourcePos)
+          }
+        }
+        /* In theory we should evaluate `permit` after `arg` but before the `JSAwait`.
+         * It *should* always be side-effect-free, though, so we just discard it.
+         */
+        js.JSAwait(arg)
+
       case DYNAMIC_IMPORT =>
         // runtime.dynamicImport
         assert(args.size == 1,

compiler/src/dotty/tools/backend/sjs/JSDefinitions.scala

@@ -42,6 +42,10 @@ final class JSDefinitions()(using Context) {
     def JSPackage_undefined(using Context) = JSPackage_undefinedR.symbol
     @threadUnsafe lazy val JSPackage_dynamicImportR = ScalaJSJSPackageClass.requiredMethodRef("dynamicImport")
     def JSPackage_dynamicImport(using Context) = JSPackage_dynamicImportR.symbol
+    @threadUnsafe lazy val JSPackage_asyncR = ScalaJSJSPackageClass.requiredMethodRef("async")
+    def JSPackage_async(using Context) = JSPackage_asyncR.symbol
+    @threadUnsafe lazy val JSPackage_awaitR = ScalaJSJSPackageClass.requiredMethodRef("await")
+    def JSPackage_await(using Context) = JSPackage_awaitR.symbol
 
   @threadUnsafe lazy val JSNativeAnnotType: TypeRef = requiredClassRef("scala.scalajs.js.native")
   def JSNativeAnnot(using Context) = JSNativeAnnotType.symbol.asClass
@@ -210,6 +214,10 @@ final class JSDefinitions()(using Context) {
     @threadUnsafe lazy val Special_unwrapFromThrowableR = SpecialPackageClass.requiredMethodRef("unwrapFromThrowable")
     def Special_unwrapFromThrowable(using Context) = Special_unwrapFromThrowableR.symbol
 
+  @threadUnsafe lazy val WasmJSPIModuleRef = requiredModuleRef("scala.scalajs.js.wasm.JSPI")
+    @threadUnsafe lazy val WasmJSPI_allowOrphanJSAwaitModuleRef = requiredModuleRef("scala.scalajs.js.wasm.JSPI.allowOrphanJSAwait")
+    def WasmJSPI_allowOrphanJSAwaitModuleClassRef(using Context) = WasmJSPIModuleRef.select(WasmJSPI_allowOrphanJSAwaitModuleRef.symbol)
+
   @threadUnsafe lazy val WrappedArrayType: TypeRef = requiredClassRef("scala.scalajs.js.WrappedArray")
   def WrappedArrayClass(using Context) = WrappedArrayType.symbol.asClass
 

compiler/src/dotty/tools/backend/sjs/JSPrimitives.scala

@@ -30,7 +30,10 @@ object JSPrimitives {
   inline val JS_IMPORT = JS_NEW_TARGET + 1  // js.import.apply(specifier)
   inline val JS_IMPORT_META = JS_IMPORT + 1 // js.import.meta
 
-  inline val CONSTRUCTOROF = JS_IMPORT_META + 1                         // runtime.constructorOf(clazz)
+  inline val JS_ASYNC = JS_IMPORT_META + 1 // js.async
+  inline val JS_AWAIT = JS_ASYNC + 1       // js.await
+
+  inline val CONSTRUCTOROF = JS_AWAIT + 1                               // runtime.constructorOf(clazz)
   inline val CREATE_INNER_JS_CLASS = CONSTRUCTOROF + 1                  // runtime.createInnerJSClass
   inline val CREATE_LOCAL_JS_CLASS = CREATE_INNER_JS_CLASS + 1          // runtime.createLocalJSClass
   inline val WITH_CONTEXTUAL_JS_CLASS_VALUE = CREATE_LOCAL_JS_CLASS + 1 // runtime.withContextualJSClassValue
@@ -110,6 +113,8 @@ class JSPrimitives(ictx: Context) extends DottyPrimitives(ictx) {
 
     addPrimitive(jsdefn.JSPackage_typeOf, TYPEOF)
     addPrimitive(jsdefn.JSPackage_native, JS_NATIVE)
+    addPrimitive(jsdefn.JSPackage_async, JS_ASYNC)
+    addPrimitive(jsdefn.JSPackage_await, JS_AWAIT)
 
     addPrimitive(defn.BoxedUnit_UNIT, UNITVAL)
 

project/Build.scala

@@ -33,6 +33,8 @@ import com.typesafe.tools.mima.plugin.MimaPlugin.autoImport._
 import org.scalajs.sbtplugin.ScalaJSPlugin
 import org.scalajs.sbtplugin.ScalaJSPlugin.autoImport._
 
+import org.scalajs.linker.interface.ESVersion
+
 import sbtbuildinfo.BuildInfoPlugin
 import sbtbuildinfo.BuildInfoPlugin.autoImport._
 import sbttastymima.TastyMiMaPlugin
@@ -2772,7 +2774,9 @@ object Build {
           case FullOptStage => (Test / fullLinkJS / scalaJSLinkerConfig).value
         }
 
-        if (linkerConfig.moduleKind != ModuleKind.NoModule && !linkerConfig.closureCompiler)
+        val isWebAssembly = linkerConfig.experimentalUseWebAssembly
+
+        if (linkerConfig.moduleKind != ModuleKind.NoModule && !linkerConfig.closureCompiler && !isWebAssembly)
           Seq(baseDirectory.value / "test-require-multi-modules")
         else
           Nil
@@ -2800,6 +2804,8 @@ object Build {
 
         val moduleKind = linkerConfig.moduleKind
         val hasModules = moduleKind != ModuleKind.NoModule
+        val hasAsyncAwait = linkerConfig.esFeatures.esVersion >= ESVersion.ES2017
+        val isWebAssembly = linkerConfig.experimentalUseWebAssembly
 
         def conditionally(cond: Boolean, subdir: String): Seq[File] =
           if (!cond) Nil
@@ -2829,9 +2835,12 @@ object Build {
 
           ++ conditionally(!hasModules, "js/src/test/require-no-modules")
           ++ conditionally(hasModules, "js/src/test/require-modules")
-          ++ conditionally(hasModules && !linkerConfig.closureCompiler, "js/src/test/require-multi-modules")
+          ++ conditionally(hasModules && !linkerConfig.closureCompiler && !isWebAssembly, "js/src/test/require-multi-modules")
           ++ conditionally(moduleKind == ModuleKind.ESModule, "js/src/test/require-dynamic-import")
           ++ conditionally(moduleKind == ModuleKind.ESModule, "js/src/test/require-esmodule")
+
+          ++ conditionally(hasAsyncAwait, "js/src/test/require-async-await")
+          ++ conditionally(hasAsyncAwait && isWebAssembly, "js/src/test/require-orphan-await")
         )
       },
 

project/DottyJSPlugin.scala

@@ -6,20 +6,34 @@ import sbt.Keys.*
 import org.scalajs.sbtplugin.ScalaJSPlugin
 import org.scalajs.sbtplugin.ScalaJSPlugin.autoImport._
 
-import org.scalajs.linker.interface.StandardConfig
+import org.scalajs.linker.interface.{ESVersion, StandardConfig}
+
+import org.scalajs.jsenv.nodejs.NodeJSEnv
 
 object DottyJSPlugin extends AutoPlugin {
 
   object autoImport {
     val switchToESModules: StandardConfig => StandardConfig =
       config => config.withModuleKind(ModuleKind.ESModule)
+
+    val switchToLatestESVersion: StandardConfig => StandardConfig =
+      config => config.withESFeatures(_.withESVersion(ESVersion.ES2021))
+
+    val enableWebAssembly: SettingKey[Boolean] =
+      settingKey("enable all the configuration items required for WebAssembly")
   }
 
+  import autoImport._
+
   val writePackageJSON = taskKey[Unit](
       "Write package.json to configure module type for Node.js")
 
   override def requires: Plugins = ScalaJSPlugin
 
+  override def globalSettings: Seq[Setting[_]] = Def.settings(
+    enableWebAssembly := false,
+  )
+
   override def projectSettings: Seq[Setting[_]] = Def.settings(
 
     /* #11709 Remove the dependency on scala3-library that ScalaJSPlugin adds.
@@ -40,6 +54,31 @@ object DottyJSPlugin extends AutoPlugin {
     // Typecheck the Scala.js IR found on the classpath
     scalaJSLinkerConfig ~= (_.withCheckIR(true)),
 
+    // Maybe configure WebAssembly
+    scalaJSLinkerConfig := {
+      val prev = scalaJSLinkerConfig.value
+      if (enableWebAssembly.value) {
+        prev
+          .withModuleKind(ModuleKind.ESModule)
+          .withExperimentalUseWebAssembly(true)
+      } else {
+        prev
+      }
+    },
+    jsEnv := {
+      val baseConfig = NodeJSEnv.Config()
+      val config = if (enableWebAssembly.value) {
+        baseConfig.withArgs(List(
+          "--experimental-wasm-exnref",
+          "--experimental-wasm-imported-strings", // for JS string builtins
+          "--experimental-wasm-jspi", // for JSPI, used by async/await
+        ))
+      } else {
+        baseConfig
+      }
+      new NodeJSEnv(config)
+    },
+
     Compile / jsEnvInput := (Compile / jsEnvInput).dependsOn(writePackageJSON).value,
     Test / jsEnvInput := (Test / jsEnvInput).dependsOn(writePackageJSON).value,