pkg/compiler/lib/src/kernel/transformations/modular/async_lowering.dart

// 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.

import 'package:kernel/ast.dart';
import 'package:kernel/class_hierarchy.dart';
import 'package:kernel/core_types.dart';
import 'package:kernel/kernel.dart';
import 'package:kernel/type_environment.dart';

class _FunctionData {
  final List<AwaitExpression> awaits = [];
  final Set<ReturnStatement> returnStatements = {};
  // If we find certain control flow statements in this function, we choose to
  // not lower it.
  bool shouldLower = true;

  _FunctionData();

  /// Returns true if all [AwaitExpression]s are children of [ReturnStatement]s.
  bool allAwaitsDirectReturn() {
    return awaits.every(
      (awaitExpression) => returnStatements.contains(awaitExpression.parent),
    );
  }
}

/// Handles simplification of basic 'async' functions into [Future]s.
///
/// The JS expansion of an async/await function is a complex state machine.
/// In many cases 'async' functions either do not use 'await' or have very
/// simple 'await' logic which is easily captured by [Future]. By making this
/// transformation we avoid substantial over head from the state machine.
class AsyncLowering {
  final List<_FunctionData> _functions = [];
  final CoreTypes _coreTypes;
  final TypeEnvironment _typeEnvironment;
  late final StatefulStaticTypeContext _typeContext =
      StatefulStaticTypeContext.flat(_typeEnvironment);

  AsyncLowering(this._coreTypes, ClassHierarchy hierarchy)
    : _typeEnvironment = TypeEnvironment(_coreTypes, hierarchy);

  void enterMember(Member node) {
    _typeContext.enterMember(node);
  }

  void exitMember(Member node) {
    _typeContext.leaveMember(node);
  }

  bool _shouldTryAsyncLowering(FunctionNode node) =>
      node.asyncMarker == AsyncMarker.Async &&
      node.emittedValueType != null &&
      _functions.last.shouldLower;

  void enterFunction(FunctionNode node) {
    _functions.add(_FunctionData());
  }

  void _exitFunction() {
    _functions.removeLast();
  }

  void _updateFunctionBody(FunctionNode node, Statement? newBody) {
    node.body = newBody;
    newBody?.parent = node;
  }

  void _wrapBodySync(FunctionNode node) {
    node.asyncMarker = AsyncMarker.Sync;
    final futureValueType = node.emittedValueType!;
    _updateFunctionBody(
      node,
      ReturnStatement(
        StaticInvocation(
          _coreTypes.futureSyncFactory,
          Arguments(
            ExpressionList(
              FunctionExpression(
                FunctionNode(
                  node.body,
                  returnType: FutureOrType(
                    futureValueType,
                    futureValueType.nullability,
                  ),
                ),
              ),
            ),
            types: DartTypeList(futureValueType),
          ),
        ),
      ),
    );
  }

  void _transformDirectReturnAwaits(
    FunctionNode node,
    _FunctionData functionData,
  ) {
    // If every await is the direct child of a return statement then we can
    // do the following transformation:
    // return await e; --> return Future<T>.value(e);
    final futureValueType = node.emittedValueType!;
    for (final awaitExpression in functionData.awaits) {
      awaitExpression.replaceWith(
        StaticInvocation(
          _coreTypes.futureValueFactory,
          Arguments(
            ExpressionList(awaitExpression.operand),
            types: DartTypeList(futureValueType),
          ),
        ),
      );
    }
  }

  void _transformAsyncFunctionNode(FunctionNode node) {
    assert(_functions.isNotEmpty, 'Must be within a function scope.');
    final functionData = _functions.last;
    if (!_canLowerReturns(node, functionData)) return;
    var isLowered = false;
    if (functionData.awaits.isEmpty) {
      // There are no awaits within this function so convert to a simple
      // Future.sync call with the function's returned expressions. We use
      // this over Future.value because the expression can throw and async
      // functions defer exceptions as errors on returned Future. Future.sync
      // has the same deferred behavior.
      //
      // Before:
      // Future<int> foo() async {
      //   doSomething();
      //   return 3;
      // }
      //
      // After:
      // Future<int> foo() {
      //   return Future.sync(() {
      //     doSomething();
      //     return 3;
      //   });
      // }
      //
      // Edge cases to consider:
      // 1) Function doesn't include a return expression. (e.g. Future<void>)
      //    In this case we call Future.value(null).
      // 2) The returned expression might itself be a future. Future.sync will
      //    handle the unpacking of the returned future in that case.
      // 3) The return type of the function is not specified. In this case we
      //    instantiate Future.value with 'dynamic'.
      isLowered = true;
    } else if (functionData.allAwaitsDirectReturn()) {
      _transformDirectReturnAwaits(node, functionData);
      isLowered = true;
    }
    if (isLowered) {
      _wrapBodySync(node);
    }
  }

  bool _canLowerReturns(FunctionNode node, _FunctionData functionData) {
    DartType? futureType;
    for (final returnStatement in functionData.returnStatements) {
      final expression = returnStatement.expression;
      if (expression == null || expression is AwaitExpression) continue;
      futureType ??= InterfaceType(
        _coreTypes.futureClass,
        Nullability.nonNullable,
        DartTypeList(node.emittedValueType!),
      );
      if (!_isDefinitelyNonFutureValue(expression, futureType)) {
        return false;
      }
    }
    return true;
  }

  bool _isDefinitelyNonFutureValue(Expression expression, DartType futureType) {
    final exprType = expression.getStaticType(_typeContext);
    if (exprType is NullType) return true;
    if (exprType is TypeParameterType) return false;
    return !_typeEnvironment.isSubtypeOf(
          exprType,
          _coreTypes.futureNullableRawType,
        ) &&
        !_typeEnvironment.isSubtypeOf(futureType, exprType);
  }

  void transformFunctionNodeAndExit(FunctionNode node) {
    if (_shouldTryAsyncLowering(node)) _transformAsyncFunctionNode(node);
    _exitFunction();
  }

  void visitAwaitExpression(AwaitExpression expression) {
    assert(
      _functions.isNotEmpty,
      'Awaits must be within the scope of a function.',
    );
    _functions.last.awaits.add(expression);
  }

  void visitReturnStatement(ReturnStatement statement) {
    _functions.last.returnStatements.add(statement);
  }

  void visitForInStatement(ForInStatement statement) {
    if (statement.isAsync && _functions.isNotEmpty) {
      _functions.last.shouldLower = false;
    }
  }

  void visitTry() {
    _functions.last.shouldLower = false;
  }
}