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

// @dart=2.12

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

class _FunctionData {
  final List<AwaitExpression> awaits = [];
  final List<ReturnStatement> returnStatements = [];

  _FunctionData();
}

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

  AsyncLowering(this._coreTypes);

  bool _shouldTryAsyncLowering(FunctionNode node) =>
      node.asyncMarker == AsyncMarker.Async;

  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.futureValueType!;
    _updateFunctionBody(
        node,
        ExpressionStatement(StaticInvocation(
            _coreTypes.futureSyncFactory,
            Arguments([
              FunctionExpression(FunctionNode(node.body,
                  returnType: FutureOrType(
                      futureValueType, futureValueType.nullability)))
            ], types: [
              futureValueType
            ]))));
  }

  void _transformAsyncFunctionNode(FunctionNode node) {
    assert(_functions.isNotEmpty, 'Must be within a function scope.');
    final functionData = _functions.last;
    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'.
      _wrapBodySync(node);
    }
  }

  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);
  }
}