[dart2wasm] Initial async/await support based on Promise integration

This implementation has a number of limitations:

- Returning a Future from an async function is not supported.
- Every async call switches to a new stack. A later optimization will
  stay on the same stack if the call is directly awaited.
- Exceptions that cross async boundaries do not receive full stack
  traces including the async suspension points. The stack trace of the
  final exception just reflects the stack trace for the outermost
  async call.

Change-Id: I2a5d5e30d6e955999ba55842c8b2ca3427cbf954
Cq-Include-Trybots: luci.dart.try:dart2wasm-linux-x64-d8-try
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/241010
Reviewed-by: William Hesse <whesse@google.com>
Reviewed-by: Joshua Litt <joshualitt@google.com>
diff --git a/pkg/dart2wasm/bin/run_wasm.js b/pkg/dart2wasm/bin/run_wasm.js
index 2d4e8f6..3e40f5a 100644
--- a/pkg/dart2wasm/bin/run_wasm.js
+++ b/pkg/dart2wasm/bin/run_wasm.js
@@ -6,7 +6,7 @@
 //
 // Run as follows:
 //
-// $> d8 --experimental-wasm-gc --wasm-gc-js-interop run_wasm.js -- <dart_module>.wasm [<ffi_module>.wasm]
+// $> d8 --experimental-wasm-gc --wasm-gc-js-interop --experimental-wasm-stack-switching --experimental-wasm-type-reflection run_wasm.js -- <dart_module>.wasm [<ffi_module>.wasm]
 //
 // If an FFI module is specified, it will be instantiated first, and its
 // exports will be supplied as imports to the Dart module under the 'ffi'
@@ -73,7 +73,7 @@
     return new factoryFunction();
 }
 
-// Imports for printing and event loop
+// Imports
 var dart2wasm = {
     printToConsole: function(string) {
         console.log(stringFromDartString(string))
@@ -83,6 +83,25 @@
             dartInstance.exports.$call0(closure);
         }, milliseconds);
     },
+    futurePromise: WebAssembly.suspendOnReturnedPromise(
+        new WebAssembly.Function(
+            {parameters: ['externref'], results: ['externref']},
+            function(future) {
+                return new Promise(function (resolve, reject) {
+                    dartInstance.exports.$awaitCallback(future, resolve);
+                });
+            })),
+    callResolve: function(resolve, result) {
+        // This trampoline is needed because [resolve] is a JS function that
+        // can't be called directly from Wasm.
+        resolve(result);
+    },
+    callAsyncBridge: function(args, completer) {
+        // This trampoline is needed because [asyncBridge] is a function wrapped
+        // by `returnPromiseOnSuspend`, and the stack-switching functionality of
+        // that wrapper is implemented as part of the export adapter.
+        asyncBridge(args, completer);
+    },
     getCurrentStackTrace: function() {
         // [Error] should be supported in most browsers.
         // A possible future optimization we could do is to just save the
@@ -331,10 +350,7 @@
     return new WebAssembly.Instance(module, imports);
 }
 
-// Import from the global scope.
-var importObject = (typeof window !== 'undefined')
-    ? window
-    : Realm.global(Realm.current());
+var importObject = globalThis;
 
 // Is an FFI module specified?
 if (arguments.length > 1) {
@@ -347,5 +363,10 @@
 // Instantiate the Dart module, importing from the global scope.
 var dartInstance = instantiate(arguments[0], importObject);
 
-var result = dartInstance.exports.main();
-if (result) console.log(result);
+// Initialize async bridge.
+var asyncBridge = WebAssembly.returnPromiseOnSuspend(dartInstance.exports.$asyncBridge);
+
+// Call `main`. If tasks are placed into the event loop (by scheduling tasks
+// explicitly or awaiting Futures), these will automatically keep the script
+// alive even after `main` returns.
+dartInstance.exports.main();
diff --git a/pkg/dart2wasm/dart2wasm.md b/pkg/dart2wasm/dart2wasm.md
index 91f51ae..0f9d65f 100644
--- a/pkg/dart2wasm/dart2wasm.md
+++ b/pkg/dart2wasm/dart2wasm.md
@@ -26,7 +26,7 @@
 
 The resulting `.wasm` file can be run with:
 
-`d8 --experimental-wasm-gc --wasm-gc-js-interop pkg/dart2wasm/bin/run_wasm.js -- `*outfile*`.wasm`
+`d8 --experimental-wasm-gc --wasm-gc-js-interop --experimental-wasm-stack-switching --experimental-wasm-type-reflection pkg/dart2wasm/bin/run_wasm.js -- `*outfile*`.wasm`
 
 ## Imports and exports
 
diff --git a/pkg/dart2wasm/lib/code_generator.dart b/pkg/dart2wasm/lib/code_generator.dart
index 28b61f7..426af54 100644
--- a/pkg/dart2wasm/lib/code_generator.dart
+++ b/pkg/dart2wasm/lib/code_generator.dart
@@ -74,6 +74,8 @@
     b = function.body;
   }
 
+  w.Module get m => translator.m;
+
   Member get member => reference.asMember;
 
   w.ValueType get returnType => translator
@@ -155,6 +157,18 @@
       }
     }
 
+    if (member.function!.asyncMarker == AsyncMarker.Async &&
+        !reference.isAsyncInnerReference) {
+      // Generate the async wrapper function, i.e. the function that gets
+      // called when an async function is called. The inner function, containing
+      // the body of the async function, is marked as an async inner reference
+      // and is generated separately.
+      Procedure procedure = member as Procedure;
+      w.BaseFunction inner =
+          translator.functions.getFunction(procedure.asyncInnerReference);
+      return generateAsyncWrapper(procedure.function, inner);
+    }
+
     return generateBody(member);
   }
 
@@ -226,6 +240,77 @@
     b.end();
   }
 
+  /// Generate the async wrapper for an async function and its associated
+  /// stub function.
+  ///
+  /// The async wrapper is the outer function that gets called when the async
+  /// function is called. It bundles up the arguments to the function into an
+  /// arguments struct along with a reference to the stub function.
+  ///
+  /// This struct is passed to the async helper, which allocates a new stack and
+  /// calls the stub function on that stack.
+  ///
+  /// The stub function unwraps the arguments from the struct and calls the
+  /// inner function, containing the implementation of the async function.
+  void generateAsyncWrapper(FunctionNode functionNode, w.BaseFunction inner) {
+    w.DefinedFunction stub =
+        m.addFunction(translator.functions.asyncStubFunctionType);
+    w.BaseFunction asyncHelper =
+        translator.functions.getFunction(translator.asyncHelper.reference);
+
+    w.Instructions stubBody = stub.body;
+    w.Local stubArguments = stub.locals[0];
+    w.Local stubStack = stub.locals[1];
+
+    // Push the type argument to the async helper, specifying the type argument
+    // of the returned `Future`.
+    DartType returnType = functionNode.returnType;
+    DartType innerType = returnType is InterfaceType &&
+            returnType.classNode == translator.coreTypes.futureClass
+        ? returnType.typeArguments.single
+        : const DynamicType();
+    types.makeType(this, innerType);
+
+    // Create struct for stub reference and arguments
+    w.StructType baseStruct = translator.functions.asyncStubBaseStruct;
+    w.StructType argsStruct = m.addStructType("${function.functionName} (args)",
+        fields: baseStruct.fields, superType: baseStruct);
+
+    // Push stub reference
+    w.Global stubGlobal = translator.makeFunctionRef(stub);
+    b.global_get(stubGlobal);
+
+    // Transfer function arguments to inner
+    w.Local argsLocal =
+        stub.addLocal(w.RefType.def(argsStruct, nullable: false));
+    stubBody.local_get(stubArguments);
+    translator.convertType(stub, stubArguments.type, argsLocal.type);
+    stubBody.local_set(argsLocal);
+    int arity = function.type.inputs.length;
+    for (int i = 0; i < arity; i++) {
+      int fieldIndex = argsStruct.fields.length;
+      w.ValueType type = function.locals[i].type;
+      argsStruct.fields.add(w.FieldType(type, mutable: false));
+      b.local_get(function.locals[i]);
+      stubBody.local_get(argsLocal);
+      stubBody.struct_get(argsStruct, fieldIndex);
+    }
+    b.struct_new(argsStruct);
+
+    // Call async helper
+    b.call(asyncHelper);
+    translator.convertType(function, asyncHelper.type.outputs.single,
+        function.type.outputs.single);
+    b.end();
+
+    // Call inner function from stub
+    stubBody.local_get(stubStack);
+    stubBody.call(inner);
+    translator.convertType(
+        stub, inner.type.outputs.single, stub.type.outputs.single);
+    stubBody.end();
+  }
+
   void generateBody(Member member) {
     ParameterInfo paramInfo = translator.paramInfoFor(reference);
     bool hasThis = member.isInstanceMember || member is Constructor;
@@ -312,7 +397,16 @@
   }
 
   /// Generate code for the body of a lambda.
-  void generateLambda(Lambda lambda, Closures closures) {
+  w.DefinedFunction generateLambda(Lambda lambda, Closures closures) {
+    if (lambda.functionNode.asyncMarker == AsyncMarker.Async &&
+        lambda.function == function) {
+      w.DefinedFunction inner =
+          translator.functions.addAsyncInnerFunctionFor(function);
+      generateAsyncWrapper(lambda.functionNode, inner);
+      return CodeGenerator(translator, inner, reference)
+          .generateLambda(lambda, closures);
+    }
+
     this.closures = closures;
 
     final int implicitParams = 1;
@@ -357,6 +451,8 @@
     lambda.functionNode.body!.accept(this);
     _implicitReturn();
     b.end();
+
+    return function;
   }
 
   void _implicitReturn() {
@@ -1071,6 +1167,23 @@
   void visitYieldStatement(YieldStatement node) => defaultStatement(node);
 
   @override
+  w.ValueType visitAwaitExpression(
+      AwaitExpression node, w.ValueType expectedType) {
+    w.BaseFunction awaitHelper =
+        translator.functions.getFunction(translator.awaitHelper.reference);
+
+    // The stack for the suspension is the last parameter to the function.
+    w.Local stack = function.locals[function.type.inputs.length - 1];
+    assert(stack.type == translator.functions.asyncStackType);
+
+    wrap(node.operand, translator.topInfo.nullableType);
+    b.local_get(stack);
+    b.call(awaitHelper);
+
+    return translator.topInfo.nullableType;
+  }
+
+  @override
   w.ValueType visitBlockExpression(
       BlockExpression node, w.ValueType expectedType) {
     node.body.accept(this);
diff --git a/pkg/dart2wasm/lib/functions.dart b/pkg/dart2wasm/lib/functions.dart
index cd71948..09e3dfa 100644
--- a/pkg/dart2wasm/lib/functions.dart
+++ b/pkg/dart2wasm/lib/functions.dart
@@ -28,6 +28,17 @@
   // allocation of that class is encountered
   final Map<int, List<Reference>> _pendingAllocation = {};
 
+  final w.ValueType asyncStackType = const w.RefType.any(nullable: true);
+
+  late final w.FunctionType asyncStubFunctionType = m.addFunctionType(
+      [const w.RefType.data(nullable: false), asyncStackType],
+      [translator.topInfo.nullableType]);
+
+  late final w.StructType asyncStubBaseStruct = m.addStructType("#AsyncStub",
+      fields: [
+        w.FieldType(w.RefType.def(asyncStubFunctionType, nullable: false))
+      ]);
+
   FunctionCollector(this.translator);
 
   w.Module get m => translator.m;
@@ -106,6 +117,10 @@
   w.BaseFunction getFunction(Reference target) {
     return _functions.putIfAbsent(target, () {
       worklist.add(target);
+      if (target.isAsyncInnerReference) {
+        w.BaseFunction outer = getFunction(target.asProcedure.reference);
+        return addAsyncInnerFunctionFor(outer);
+      }
       w.FunctionType ftype = target.isTearOffReference
           ? translator.dispatchTable.selectorForTarget(target).signature
           : target.asMember.accept1(this, target);
@@ -113,6 +128,13 @@
     });
   }
 
+  w.DefinedFunction addAsyncInnerFunctionFor(w.BaseFunction outer) {
+    w.FunctionType ftype = m.addFunctionType(
+        [...outer.type.inputs, asyncStackType],
+        [translator.topInfo.nullableType]);
+    return m.addFunction(ftype, "${outer.functionName} (inner)");
+  }
+
   void activateSelector(SelectorInfo selector) {
     selector.targets.forEach((classId, target) {
       if (!target.asMember.isAbstract) {
@@ -200,8 +222,9 @@
     // The JS embedder will not accept Wasm struct types as parameter or return
     // types for functions called from JS. We need to use eqref instead.
     w.ValueType adjustExternalType(w.ValueType type) {
-      if (isImportOrExport && type.isSubtypeOf(w.RefType.eq())) {
-        return w.RefType.eq();
+      if (isImportOrExport &&
+          type.isSubtypeOf(translator.topInfo.nullableType)) {
+        return w.RefType.eq(nullable: type.nullable);
       }
       return type;
     }
@@ -217,7 +240,9 @@
     List<w.ValueType> outputs = returnType is VoidType ||
             returnType is NeverType ||
             returnType is NullType
-        ? const []
+        ? member.function?.asyncMarker == AsyncMarker.Async
+            ? [adjustExternalType(translator.topInfo.nullableType)]
+            : const []
         : [adjustExternalType(translator.translateType(returnType))];
 
     return m.addFunctionType(inputs, outputs);
diff --git a/pkg/dart2wasm/lib/intrinsics.dart b/pkg/dart2wasm/lib/intrinsics.dart
index 6a9ec59..7a06523 100644
--- a/pkg/dart2wasm/lib/intrinsics.dart
+++ b/pkg/dart2wasm/lib/intrinsics.dart
@@ -1322,6 +1322,22 @@
       }
     }
 
+    // _asyncBridge2
+    if (member.enclosingLibrary.name == "dart.async" &&
+        name == "_asyncBridge2") {
+      w.Local args = paramLocals[0];
+      w.Local stack = paramLocals[1];
+      const int stubFieldIndex = 0;
+
+      b.local_get(args);
+      b.local_get(stack);
+      b.local_get(args);
+      b.ref_cast(translator.functions.asyncStubBaseStruct);
+      b.struct_get(translator.functions.asyncStubBaseStruct, stubFieldIndex);
+      b.call_ref();
+      return true;
+    }
+
     // int members
     if (member.enclosingClass == translator.boxedIntClass &&
         member.function.body == null) {
diff --git a/pkg/dart2wasm/lib/reference_extensions.dart b/pkg/dart2wasm/lib/reference_extensions.dart
index ba9a68c..9f9fd12 100644
--- a/pkg/dart2wasm/lib/reference_extensions.dart
+++ b/pkg/dart2wasm/lib/reference_extensions.dart
@@ -32,19 +32,32 @@
 // implementation for that procedure. This enables a Reference to refer to any
 // implementation relating to a member, including its tear-off, which it can't
 // do in plain kernel.
+// Also add an asyncInnerReference that refers to the inner, suspendable
+// body of an async function, which returns the value to be put into a future.
+// This can be called directly from other async functions if the result is
+// directly awaited.
 
-extension TearOffReference on Procedure {
-  // Use an Expando to avoid keeping the procedure alive.
-  static final Expando<Reference> _tearOffReference = Expando();
+// Use Expandos to avoid keeping the procedure alive.
+final Expando<Reference> _tearOffReference = Expando();
+final Expando<Reference> _asyncInnerReference = Expando();
 
+extension CustomReference on Procedure {
   Reference get tearOffReference =>
       _tearOffReference[this] ??= Reference()..node = this;
+
+  Reference get asyncInnerReference =>
+      _asyncInnerReference[this] ??= Reference()..node = this;
 }
 
-extension IsTearOffReference on Reference {
+extension IsCustomReference on Reference {
   bool get isTearOffReference {
     Member member = asMember;
-    return member is Procedure && member.tearOffReference == this;
+    return member is Procedure && _tearOffReference[member] == this;
+  }
+
+  bool get isAsyncInnerReference {
+    Member member = asMember;
+    return member is Procedure && _asyncInnerReference[member] == this;
   }
 }
 
diff --git a/pkg/dart2wasm/lib/translator.dart b/pkg/dart2wasm/lib/translator.dart
index 7871290..4647f1c 100644
--- a/pkg/dart2wasm/lib/translator.dart
+++ b/pkg/dart2wasm/lib/translator.dart
@@ -110,6 +110,8 @@
   late final Procedure wasmFunctionCall;
   late final Procedure wasmTableCallIndirect;
   late final Procedure stackTraceCurrent;
+  late final Procedure asyncHelper;
+  late final Procedure awaitHelper;
   late final Procedure stringEquals;
   late final Procedure stringInterpolate;
   late final Procedure throwNullCheckError;
@@ -237,6 +239,14 @@
         .firstWhere((p) => p.name.text == "callIndirect");
     stackTraceCurrent =
         stackTraceClass.procedures.firstWhere((p) => p.name.text == "current");
+    asyncHelper = component.libraries
+        .firstWhere((l) => l.name == "dart.async")
+        .procedures
+        .firstWhere((p) => p.name.text == "_asyncHelper");
+    awaitHelper = component.libraries
+        .firstWhere((l) => l.name == "dart.async")
+        .procedures
+        .firstWhere((p) => p.name.text == "_awaitHelper");
     stringEquals =
         stringBaseClass.procedures.firstWhere((p) => p.name.text == "==");
     stringInterpolate = stringBaseClass.procedures
@@ -404,9 +414,10 @@
       }
 
       for (Lambda lambda in codeGen.closures.lambdas.values) {
-        CodeGenerator(this, lambda.function, reference)
-            .generateLambda(lambda, codeGen.closures);
-        _printFunction(lambda.function, "$canonicalName (closure)");
+        w.DefinedFunction lambdaFunction =
+            CodeGenerator(this, lambda.function, reference)
+                .generateLambda(lambda, codeGen.closures);
+        _printFunction(lambdaFunction, "$canonicalName (closure)");
       }
     }
 
diff --git a/pkg/test_runner/lib/src/compiler_configuration.dart b/pkg/test_runner/lib/src/compiler_configuration.dart
index 2821ff5..ba098e5 100644
--- a/pkg/test_runner/lib/src/compiler_configuration.dart
+++ b/pkg/test_runner/lib/src/compiler_configuration.dart
@@ -574,6 +574,8 @@
     return [
       '--experimental-wasm-gc',
       '--wasm-gc-js-interop',
+      '--experimental-wasm-stack-switching',
+      '--experimental-wasm-type-reflection',
       'pkg/dart2wasm/bin/run_wasm.js',
       '--',
       artifact!.filename,
diff --git a/sdk/lib/_internal/wasm/lib/async_patch.dart b/sdk/lib/_internal/wasm/lib/async_patch.dart
new file mode 100644
index 0000000..8d112bd
--- /dev/null
+++ b/sdk/lib/_internal/wasm/lib/async_patch.dart
@@ -0,0 +1,110 @@
+// Copyright (c) 2022, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+// Machinery for async and await.
+//
+// The implementation is based on two mechanisms in the JS Promise integration:
+//
+// The export wrapper: Allocates a new stack and calls the wrapped export on the
+// new stack, passing a suspender object as an extra first argument that
+// represents the new stack.
+//
+// The import wrapper: Takes a suspender object as an extra first argument and
+// calls the wrapped import. If the wrapped import returns a `Promise`, the
+// current stack is suspended, and the `Promise` is forwarded to the
+// corresponding call of the export wrapper, where execution resumes on the
+// original stack. When the `Promise` is resolved, execution resumes on the
+// suspended stack, with the call to the import wrapper returning the value the
+// `Promise` was resolved with.
+//
+// The call sequence when calling an async function is:
+//
+// Caller
+//  -> Outer (function specific, generated by `generateAsyncWrapper`)
+//  -> `_asyncHelper`
+//  -> `_callAsyncBridge` (imported JS function)
+//  -> `_asyncBridge` (via the Promise integration export wrapper)
+//  -> `_asyncBridge2` (intrinsic function)
+//  -> Stub (function specific, generated by `generateAsyncWrapper`)
+//  -> Inner (contains implementation, generated from async inner reference)
+//
+// The call sequence on await is:
+//
+// Function containing await
+//  -> `_awaitHelper`
+//  -> `_futurePromise` (via the Promise integration import wrapper)
+//  -> `new Promise`
+//  -> `Promise` constructor callback
+//  -> `_awaitCallback`
+//  -> `Future.then`
+// `futurePromise` returns the newly created `Promise`, suspending the
+// current execution.
+//
+// When the `Future` completes:
+//
+// `Future.then` callback
+//  -> `_callResolve` (imported JS function)
+//  -> `Promise` resolve function
+// Resolving the `Promise` causes the suspended execution to resume.
+
+import 'dart:_internal' show patch, scheduleCallback, unsafeCastOpaque;
+
+import 'dart:wasm';
+
+@pragma("wasm:entry-point")
+Future<T> _asyncHelper<T>(WasmEqRef args) {
+  Completer<T> completer = Completer();
+  _callAsyncBridge(args, completer);
+  return completer.future;
+}
+
+@pragma("wasm:import", "dart2wasm.callAsyncBridge")
+external void _callAsyncBridge(WasmEqRef args, Completer<Object?> completer);
+
+@pragma("wasm:export", "\$asyncBridge")
+WasmAnyRef? _asyncBridge(
+    WasmAnyRef? stack, WasmDataRef args, Completer<Object?> completer) {
+  try {
+    Object? result = _asyncBridge2(args, stack);
+    completer.complete(result);
+  } catch (e, s) {
+    completer.completeError(e, s);
+  }
+}
+
+external Object? _asyncBridge2(WasmDataRef args, WasmAnyRef? stack);
+
+class _FutureError {
+  final Object exception;
+  final StackTrace stackTrace;
+
+  _FutureError(this.exception, this.stackTrace);
+}
+
+@pragma("wasm:entry-point")
+Object? _awaitHelper(Object? operand, WasmAnyRef? stack) {
+  if (operand is! Future) return operand;
+  WasmAnyRef futureRef = WasmAnyRef.fromObject(operand);
+  Object? result = unsafeCastOpaque(_futurePromise(stack, futureRef));
+  if (result is _FutureError) {
+    // TODO(askesc): Combine stack traces
+    throw result.exception;
+  }
+  return result;
+}
+
+@pragma("wasm:import", "dart2wasm.futurePromise")
+external WasmAnyRef? _futurePromise(WasmAnyRef? stack, WasmAnyRef? future);
+
+@pragma("wasm:export", "\$awaitCallback")
+void _awaitCallback(Future<Object?> future, WasmAnyRef resolve) {
+  future.then((value) {
+    _callResolve(resolve, value);
+  }, onError: (exception, stackTrace) {
+    _callResolve(resolve, _FutureError(exception, stackTrace));
+  });
+}
+
+@pragma("wasm:import", "dart2wasm.callResolve")
+external void _callResolve(WasmAnyRef resolve, Object? result);
diff --git a/sdk/lib/_internal/wasm/lib/timer_patch.dart b/sdk/lib/_internal/wasm/lib/timer_patch.dart
index 47c7f03..9019c45 100644
--- a/sdk/lib/_internal/wasm/lib/timer_patch.dart
+++ b/sdk/lib/_internal/wasm/lib/timer_patch.dart
@@ -6,8 +6,6 @@
 
 // Implementation of `Timer` and `scheduleMicrotask` via the JS event loop.
 
-import 'dart:_internal' show patch, scheduleCallback;
-
 @patch
 class Timer {
   @patch
diff --git a/sdk/lib/libraries.json b/sdk/lib/libraries.json
index d45b6cb..1248afb 100644
--- a/sdk/lib/libraries.json
+++ b/sdk/lib/libraries.json
@@ -179,7 +179,10 @@
       },
       "async": {
         "uri": "async/async.dart",
-        "patches": "_internal/wasm/lib/timer_patch.dart"
+        "patches": [
+          "_internal/wasm/lib/async_patch.dart",
+          "_internal/wasm/lib/timer_patch.dart"
+        ]
       },
       "collection": {
         "uri": "collection/collection.dart",
diff --git a/sdk/lib/libraries.yaml b/sdk/lib/libraries.yaml
index 4641fbc..1941780 100644
--- a/sdk/lib/libraries.yaml
+++ b/sdk/lib/libraries.yaml
@@ -172,7 +172,9 @@
       uri: _internal/wasm/lib/js_helper.dart
     async:
       uri: async/async.dart
-      patches: _internal/wasm/lib/timer_patch.dart
+      patches:
+      - _internal/wasm/lib/async_patch.dart
+      - _internal/wasm/lib/timer_patch.dart
     collection:
       uri: collection/collection.dart
       patches: