[dart2wasm] Reduce size of trivial async function bodies

[mkustermann]

In ACX code there's a lot of static async functions that simply return a constant (e.g. grep for ulanHotShi).

For this example

main() async {
  print(await load1());
  print(await load2());
  print(await load3());
  print(await load4());
}

Future<Foo> load1() async => const Foo(1);
Future<Foo> load2() async => const Foo(2);
Future<Foo> load3() async => const Foo(3);
Future<Foo> load4() async => const Foo(4);

class Foo {
  final int i;
  const Foo(this.i);
}

dart2wasm currently compiles the inner bodies to

  (func $load1 inner (;193;) (param $var0 (ref $_AsyncSuspendState_35)) (param $var1 (ref null $#Top)) (param $var2 (ref null $#Top)) (param $var3 (ref null $_StringStackTrace_5)) (result (ref null $#Top))
    (local $var4 (ref $#Top))
    (local $var5 (ref $_StringStackTrace_5))
    try
      block $label1
        block $label0
          block $label2
            local.get $var0
            struct.get $_AsyncSuspendState_35 $field3
            br_table $label0 $label1 $label2
          end $label2
          unreachable
        end $label0
        local.get $var0
        struct.get $_AsyncSuspendState_35 $field4
        global.get $global208
        call $_AsyncCompleter.complete (body)
        ref.null none
        return
      end $label1
      local.get $var0
      struct.get $_AsyncSuspendState_35 $field4
      ref.null none
      call $_AsyncCompleter.complete (body)
      ref.null none
      return
    catch $tag0
      local.set $var5
      local.set $var4
      local.get $var0
      struct.get $_AsyncSuspendState_35 $field4
      local.get $var4
      local.get $var5
      call $_Completer.completeError
      ref.null none
      return
    catch_all
      local.get $var0
      struct.get $_AsyncSuspendState_35 $field4
      call $_JavaScriptError
      call $StackTrace.current
      call $_Completer.completeError
      ref.null none
      return
    end
    unreachable
  )

There's a few things here:

• The body of the try/catch block has two cases where it calls _AsyncCompleter.complete even though the original Dart source code just has a single return statement
  => It seems this body should only have a single call to complete the async completer.
• The original dart code function bodies are just returning constants. They are effect free and cannot cause exceptions to be thrown
  => We could avoid the try/catch in this case

/cc @osa1

[osa1]

I started looking into this.

Ideally load1 inner above could just be:

local.get $var0
struct.get $_AsyncSuspendState_35 $field4
global.get $global208
call $_AsyncCompleter.complete (body)
ref.null none
return

But we generate two unused states:

• One state for the default case in br_table. br_table requires a default and currently we generate a default state with just unreachable (indicating a bug in dart2wasm). We could just use the last state as the default.
  Update: With the comment referenced below we now use the last state as the default.

• One state for the fall-through case when the function body doesn't return.

  This handles the cases where in kernel not all code paths explicitly return.

  This is the same case that _implicitReturn in the non-async code generator handles:

  sdk/pkg/dart2wasm/lib/code_generator.dart

  Lines 564 to 575 in 93df675

  	void _implicitReturn() {
  	if (outputs.isNotEmpty) {
  	w.ValueType returnType = outputs.single;
  	if (returnType is w.RefType && returnType.nullable) {
  	// Dart body may have an implicit return null.
  	b.ref_null(returnType.heapType.bottomType);
  	} else {
  	b.comment("Unreachable implicit return");
  	b.unreachable();
  	}
  	}
  	}

  We could analyze the kernel AST to see if there's an implicit return in the function, and eliminate this state if there isn't. We could also get rid of _implcitReturn with this.

And then we generate more code for handling exceptions returned from awaited futures.

When the function body is just return const ... we know there are no awaited futures and no exceptions can be thrown, so in this case this state can be eliminated.

Anything other than return const ... needs to be analyzed, but I suspect the analysis would have to be very conservative to be feasible and return "can throw" in majority of the cases. It might make more sense to handle known special cases instead, like return const ... and perhaps also with function bodies with just an await: foo() async => await bar(). Update: This is also done, with the second referenced commit below.

[osa1]

We could analyze the kernel AST to see if there's an implicit return in the function, and eliminate this state if there isn't. We could also get rid of _implcitReturn with this.

With #55939 kernel will have a return statement at the end of a function if needed, so we can just wait for that issue to be fixed and then remove the implicit returns.

[mkustermann]

It's not uncommon that one has to implement an interface where methods are async, so one implements them with async functions but the bodies are actually synchronous. e.g.

Future foo() async {
  print('a');
  print('b');
  return bar();
}

Currently we generate the outer and inner function for this, allocate a AsyncSuspendState, ...

But actually we could turn this into basically the wasm equivalent of this:

Future foo()  {
  final c = Completer();
  try {
    print('a');
    print('b');
    c.complete(null);
  } catch (e, s) {
    c.completeError(e, s);
  }
  return c.future;
}

[biggs0125]

It might be worth sharing this with dart2js: https://github.com/dart-lang/sdk/blob/main/pkg/compiler/lib/src/kernel/transformations/modular/async_lowering.dart

We built this out a while ago to handle similar cases where async/awaits could be lowered into simpler Future-based code.

There were some more complex cases we were considering covering but the error handling starts getting more convoluted.

[lrhn]

Indeed, if an async function is just returning a constant, which doesn't implement Future, it can be reduced to => Future.value(constant);.

Anything more is unnecessary overhead.

[osa1]

I started running the await lowering pass linked above in dart2wasm: https://dart-review.googlesource.com/c/sdk/+/420140.

With the recent improvements (see commits referenced above), Future.sync actually makes the code sizes for load1 larger, because Future.sync is quite large compared to the code we generate currently, and to call it we need a closure which means more code.

Dart code:

Future<Foo> load1() async => const Foo(1);

Current compilation (-O4):

(func $load1 inner (;180;) (param $var0 (ref $_AsyncSuspendState_32)) (param $var1 (ref null $#Top)) (param $var2 (ref null $#Top)) (param $var3 (ref null $_StringStackTrace_5)) (result (ref null $#Top))
  local.get $var0
  struct.get $_AsyncSuspendState_32 $_completer
  global.get $global184
  call $_AsyncCompleter.complete
  ref.null none
)

;; Call site: (the "outer" function, inlined)
ref.func $load1 inner
ref.null none
global.get $global88
call $_AsyncCompleter
call $_newAsyncSuspendState
local.tee $var5
ref.null none
ref.null none
ref.null none
call $load1 inner

With the code converted to Future.sync with the pass linked above:

;; Kernel (simplified)
static method load1() → asy::Future<test::Foo> /* emittedValueType= test::Foo */ /* originally async */ 
    return Future::sync<test::Foo>(() → FutureOr<test::Foo> => asy::Future::value<test::Foo>(#C2));

;; Call site (smaller than previous version)
local.get $var0
i32.const 37
global.get $global91
global.get $global84
global.get $global189
struct.new $#Closure-0-0_112
call $new Future.sync

This is smaller, but it needs Future.sync, which looks like this (-O4):

  (func $new Future.sync (;176;) (param $var0 (ref $#Closure-0-0_112)) (result (ref $_Future))
    (local $var1 (ref $#Top))
    (local $var2 (ref $#Top))
    (local $var3 (ref $_StringStackTrace_7))
    (local $var4 (ref $_StringStackTrace_7))
    (local $var5 (ref null $#Top))
    (local $var6 (ref $_Future))
    (local $var7 (ref $_Future_93))
    block $label0
      block $label1 (result (ref $#Top)) (result (ref $_StringStackTrace_7))
        try
          local.get $var0
          struct.get $#Closure-0-0_112 $field2
          local.get $var0
          struct.get $#Closure-0-0_112 $field3
          struct.get $#Vtable-0-0 $field0
          call_ref $type9
          local.set $var5
          br $label0
        catch $tag0
          local.set $var3
          local.set $var1
          local.get $var1
          local.get $var3
          br $label1
        catch_all
          call $StackTrace.current
          local.set $var4
          call $_JavaScriptError
          local.get $var4
          br $label1
        end
        unreachable
      end $label1
      local.set $var3
      local.tee $var1
      local.set $var2
      local.get $var3
      local.set $var4
      global.get $global87
      call $_Future
      local.tee $var6
      local.get $var2
      local.get $var4
      call $AsyncError
      call $_Future._asyncCompleteErrorObject
      local.get $var6
      return
    end $label0
    local.get $var5
    global.get $global87
    call $<obj> is Future<T0>_451
    if (result (ref $_Future_93))
      local.get $var5
      ref.cast $_Future_93
    else
      i32.const 70
      global.get $global87
      i64.const 0
      global.get $Zone._current
      ref.null none
      struct.new $_Future_93
      local.tee $var7
      local.get $var5
      call $_Future._setValue
      drop
      local.get $var7
    end
  )

And it brings a bunch of other functions as well, for example $<obj> is Future<T0>_451 also doesn't exist in the current version.

Future.sync argument is also a new closure that we don't have in the current version.

One would think that if we have a dozen or so of these functions then the size of Future.sync should be smaller than each of the functions compiled, but actually the ACX demo also gets larger, from 19,133,733 bytes to 21,132,529 bytes (1,998,796 bytes, +10%).

So it seems like the current compilation of simple async bodies is actually better than Future.sync.

[mkustermann]

@osa1 Why Future.sync(() => x)? Wouldn't we want to use Future.value(x) as suggested by Lasse?

I guess the dart2js transformation doesn't handle the special case where the synchronous body is just returning a constant? See code

[osa1]

Good question -- the pass is supposed to generate Future.value calls as well (it has comments about it, and code that generates Future.value calls), but at least in the original code it doesn't convert ... async => const Blah(); to ... => Future.value(const Blah()). I'll investigate.

[osa1]

It looks like it always generates a function body with Future.sync, to handle exceptions.

It doesn't have special cases we added in 68fd76f to avoid Future.sync and only generate a Future.value.

I'll try to add those special cases to this pass and see if it improves code size.

I suspect the final code size will still not be better than the current one, because we generate very small code for these cases already.

[osa1]

I decided to not update the pass, because it will introduce Future.sync in some cases, which we never want to do.

Instead I convert simple async function bodies to Future.values and remove the async marker in an existing kernel-to-kernel pass: https://dart-review.googlesource.com/c/sdk/+/422120.

It reduces ACX demo size from 19,132,496 to 18,751,175 (381,321 bytes, -1.9%).

Here "simple" means the same thing as in 68fd76f: a const or basic literal.

With this we can also revert 68fd76f as the cases it handles today will be transformed to Future.value.

We should also look into improving the standard library Future and _Future code. Currently they use Zone, which has no relevance to Wasm as far as I know. These Zone stuff cause a lot of redundant code generation. For example, _Future._completeAsync goes through this just to call scheduleMicrotask:

(func $_Future._asyncComplete (;172;) (param $var0 (ref $_Future)) (param $var1 (ref null $#Top))
  (local $var2 (ref $<context org-dartlang-sdk:///lib/async/future_impl.dart:804:31>_87))
  block $label0
    local.get $var1
    br_on_null $label0
    drop
  end $label0
  local.get $var0
  local.get $var1
  struct.new $<context org-dartlang-sdk:///lib/async/future_impl.dart:804:31>_87
  local.set $var2
  local.get $var0
  call $_Future._setPendingComplete
  local.get $var0
  struct.get $_Future $field3
  i32.const 37
  global.get $global82
  local.get $var2
  global.get $global177
  struct.new $#Closure-0-0_88
  call $_RootZone.scheduleMicrotask
)

(func $_RootZone.scheduleMicrotask (;108;) (param $var0 (ref $#Top)) (param $var1 (ref $#Closure-0-0))
  global.get $global69
  local.get $var0
  ref.eq
  i32.eqz
  if
    global.get $global69
    local.get $var0
    call $_Zone.inSameErrorZone
    if (result (ref $#Closure-0-0))
      i32.const 37
      i32.const 12
      i32.const 0
      global.get $global60
      global.get $global60
      global.get $global81
      global.get $global60
      i64.const 0
      global.get $global64
      struct.new $_FunctionType_100
      local.get $var0
      local.get $var1
      struct.new $<context org-dartlang-sdk:///lib/async/zone.dart:1799:34>_42
      global.get $global150
      struct.new $#Closure-0-0_43
    else
      i32.const 37
      global.get $global82
      local.get $var0
      local.get $var1
      struct.new $<context org-dartlang-sdk:///lib/async/zone.dart:1799:34>_42
      global.get $global151
      struct.new $#Closure-0-0_45
    end
    local.set $var1
  end
  local.get $var1
  call $_scheduleAsyncCallback
)

(func $_scheduleAsyncCallback (;110;) (param $var0 (ref $#Closure-0-0))
  (local $var1 (ref $_AsyncCallbackEntry_46))
  (local $var2 (ref null $_AsyncCallbackEntry_46))
  local.get $var0
  call $_AsyncCallbackEntry
  local.set $var1
  global.get $_lastCallback
  local.tee $var2
  ref.is_null
  if
    local.get $var1
    global.set $_lastCallback
    local.get $var1
    global.set $_nextCallback
    global.get $_isInCallbackLoop
    i32.eqz
    if
      global.get $global87
      call $dart2wasm._190 (import)
    end
  else
    local.get $var2
    local.get $var1
    struct.set $_AsyncCallbackEntry_46 $field2
    local.get $var1
    global.set $_lastCallback
  end
)

(Where the import _190 is scheduleMicrotask)

It could call scheduleMicrotask directly. That's what we do in the async functions compiled to the state machines.