[DAS] Fixes wrong return type - add async

[email protected]

Fixes #50460

Change-Id: I0761364ec0292696ddcdba81213ab4d65a2a5e14
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/394361
Commit-Queue: Brian Wilkerson <[email protected]>
Reviewed-by: Samuel Rawlins <[email protected]>
Reviewed-by: Brian Wilkerson <[email protected]>
Auto-Submit: Felipe Morschel <[email protected]>

pkg/analysis_server/lib/src/services/correction/dart/add_async.dart

@@ -7,20 +7,23 @@ import 'package:analysis_server_plugin/edit/dart/correction_producer.dart';
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/dart/element/type_system.dart';
 import 'package:analyzer_plugin/utilities/change_builder/change_builder_core.dart';
 import 'package:analyzer_plugin/utilities/fixes/fixes.dart';
 
 class AddAsync extends ResolvedCorrectionProducer {
   // TODO(pq): consider adding a variation that adds an `await` as well.
 
-  /// A flag indicating whether this producer is producing a fix in the case
-  /// where a function is missing a return at the end.
-  final bool isForMissingReturn;
+  final _Type _type;
 
   /// Initialize a newly created producer.
-  AddAsync({required super.context}) : isForMissingReturn = false;
+  AddAsync({required super.context}) : _type = _Type.others;
+
+  AddAsync.missingReturn({required super.context})
+    : _type = _Type.missingReturn;
 
-  AddAsync.missingReturn({required super.context}) : isForMissingReturn = true;
+  AddAsync.wrongReturnType({required super.context})
+    : _type = _Type.wrongReturnType;
 
   @override
   CorrectionApplicability get applicability =>
@@ -33,43 +36,100 @@ class AddAsync extends ResolvedCorrectionProducer {
 
   @override
   Future<void> compute(ChangeBuilder builder) async {
-    if (isForMissingReturn) {
-      var node = this.node;
-      FunctionBody? body;
-      DartType? returnType;
-      switch (node) {
-        case FunctionDeclaration():
-          body = node.functionExpression.body;
-          if (node.declaredFragment?.element case var declaredElement?) {
-            returnType = declaredElement.returnType;
-          } else if (node.declaredFragment case var declaredFragment?) {
-            returnType = declaredFragment.element.returnType;
-          }
-        case MethodDeclaration():
-          body = node.body;
-          returnType = node.declaredFragment!.element.returnType;
-      }
-      if (body == null || returnType == null) {
-        return;
-      }
-      if (_isFutureVoid(returnType) && _hasNoReturns(body)) {
-        var final_body = body;
-        await builder.addDartFileEdit(file, (builder) {
-          builder.addSimpleInsertion(final_body.offset, 'async ');
-        });
-      }
-    } else {
-      var body = node.thisOrAncestorOfType<FunctionBody>();
-      if (body != null && body.keyword == null) {
-        await builder.addDartFileEdit(file, (builder) {
-          builder.convertFunctionFromSyncToAsync(
-            body: body,
-            typeSystem: typeSystem,
-            typeProvider: typeProvider,
-          );
-        });
-      }
+    switch (_type) {
+      case _Type.missingReturn:
+        var node = this.node;
+        FunctionBody? body;
+        DartType? returnType;
+        switch (node) {
+          case FunctionDeclaration():
+            body = node.functionExpression.body;
+            if (node.declaredFragment?.element case var declaredElement?) {
+              returnType = declaredElement.returnType;
+            } else if (node.declaredFragment case var declaredFragment?) {
+              returnType = declaredFragment.element.returnType;
+            }
+          case MethodDeclaration():
+            body = node.body;
+            returnType = node.declaredFragment!.element.returnType;
+        }
+        if (body == null || returnType == null) {
+          return;
+        }
+        if (_isFutureVoid(returnType) && _hasNoReturns(body)) {
+          var final_body = body;
+          await builder.addDartFileEdit(file, (builder) {
+            builder.addSimpleInsertion(final_body.offset, 'async ');
+          });
+        }
+      case _Type.wrongReturnType:
+        var body = node.thisOrAncestorOfType<FunctionBody>();
+        if (body == null) {
+          return;
+        }
+        if (body.keyword != null) {
+          return;
+        }
+        var expectedReturnType = _getStaticFunctionType(body);
+        if (expectedReturnType == null) {
+          return;
+        }
+
+        if (expectedReturnType is! InterfaceType ||
+            !expectedReturnType.isDartAsyncFuture) {
+          return;
+        }
+
+        var visitor = _ReturnTypeTester(
+          typeSystem,
+          expectedReturnType.typeArguments.first,
+        );
+        if (visitor.returnsAreAssignable(body)) {
+          await builder.addDartFileEdit(file, (builder) {
+            builder.convertFunctionFromSyncToAsync(
+              body: body,
+              typeSystem: typeSystem,
+              typeProvider: typeProvider,
+            );
+          });
+        }
+      case _Type.others:
+        var body = node.thisOrAncestorOfType<FunctionBody>();
+        if (body != null && body.keyword == null) {
+          await builder.addDartFileEdit(file, (builder) {
+            builder.convertFunctionFromSyncToAsync(
+              body: body,
+              typeSystem: typeSystem,
+              typeProvider: typeProvider,
+            );
+          });
+        }
+    }
+  }
+
+  DartType? _getStaticFunctionType(FunctionBody body) {
+    // Gets return type for expression functions.
+    if (body case ExpressionFunctionBody(:FunctionExpression parent)) {
+      return parent.declaredFragment?.element.returnType;
+    }
+
+    if (body is! BlockFunctionBody) {
+      return null;
+    }
+
+    // Gets return type for methods.
+    if (node.thisOrAncestorOfType<MethodDeclaration>() case var method?
+        when method.body == body) {
+      return method.declaredFragment?.element.returnType;
+    }
+
+    // Gets return type for functions.
+    if (node.thisOrAncestorOfType<FunctionDeclaration>() case var function?
+        when function.functionExpression.body == body) {
+      return function.declaredFragment?.element.returnType;
     }
+
+    return null;
   }
 
   /// Return `true` if there are no return statements in the given function
@@ -107,3 +167,83 @@ class _ReturnFinder extends RecursiveAstVisitor<void> {
     foundReturn = true;
   }
 }
+
+/// An AST visitor used to test if all return statements in a function body
+/// are assignable to a given type.
+class _ReturnTypeTester extends RecursiveAstVisitor<void> {
+  /// A flag indicating whether a return statement was visited.
+  bool _foundOneReturn = false;
+
+  /// A flag indicating whether the return type is assignable considering
+  /// [isAssignable].
+  bool _returnsAreAssignable = true;
+
+  /// The type system used to check assignability.
+  final TypeSystem typeSystem;
+
+  /// The type that the return statements should be assignable to.
+  final DartType futureOf;
+
+  /// Initialize a newly created visitor.
+  _ReturnTypeTester(this.typeSystem, this.futureOf);
+
+  /// Tests whether a type is assignable to the [futureOf] type.
+  bool isAssignable(DartType type) {
+    if (typeSystem.isAssignableTo(type, futureOf)) {
+      return true;
+    }
+    if (type is InterfaceType && type.isDartAsyncFuture) {
+      return typeSystem.isAssignableTo(type.typeArguments.first, futureOf);
+    }
+    return false;
+  }
+
+  /// Returns `true` if all return statements in the given [node] are
+  /// assignable to the [futureOf] type.
+  bool returnsAreAssignable(AstNode node) {
+    _returnsAreAssignable = true;
+    _foundOneReturn = false;
+    node.accept(this);
+    return _returnsAreAssignable && _foundOneReturn;
+  }
+
+  @override
+  void visitExpressionFunctionBody(ExpressionFunctionBody node) {
+    _foundOneReturn = true;
+    if (node.expression.staticType case var type?) {
+      _returnsAreAssignable &= isAssignable(type);
+    } else {
+      _returnsAreAssignable = false;
+    }
+  }
+
+  @override
+  void visitFunctionExpression(FunctionExpression node) {
+    // Return statements within closures aren't counted.
+  }
+
+  @override
+  void visitReturnStatement(ReturnStatement node) {
+    _foundOneReturn = true;
+    if (node.expression?.staticType case var type?) {
+      _returnsAreAssignable &= isAssignable(type);
+    } else {
+      _returnsAreAssignable = false;
+    }
+  }
+}
+
+enum _Type {
+  /// Indicates whether the producer is producing a fix in the case
+  /// where a function has a return statement with the wrong return type.
+  wrongReturnType,
+
+  /// Indicates whether the producer is producing a fix in the case
+  /// where a function is missing a return at the end.
+  missingReturn,
+
+  /// Indicates whether the producer is producing a fix that adds `async`
+  /// to a function that is missing it. In cases where the error/lint is
+  /// good enough to suggest adding `async` to a function, this is valid.
+  others,
+}

pkg/analysis_server/lib/src/services/correction/error_fix_status.yaml

@@ -1427,7 +1427,7 @@ CompileTimeErrorCode.RETURN_IN_GENERATOR:
   notes: |-
     Replace with a yield.
 CompileTimeErrorCode.RETURN_OF_INVALID_TYPE_FROM_CLOSURE:
-  status: noFix
+  status: hasFix
 CompileTimeErrorCode.RETURN_OF_INVALID_TYPE_FROM_CONSTRUCTOR:
   status: noFix
 CompileTimeErrorCode.RETURN_OF_INVALID_TYPE_FROM_FUNCTION:

pkg/analysis_server/lib/src/services/correction/fix_internal.dart

@@ -947,11 +947,16 @@ final _builtInNonLintProducers = <ErrorCode, List<ProducerGenerator>>{
   CompileTimeErrorCode.RECORD_LITERAL_ONE_POSITIONAL_NO_TRAILING_COMMA: [
     AddTrailingComma.new,
   ],
+  CompileTimeErrorCode.RETURN_OF_INVALID_TYPE_FROM_CLOSURE: [
+    AddAsync.wrongReturnType,
+  ],
   CompileTimeErrorCode.RETURN_OF_INVALID_TYPE_FROM_FUNCTION: [
+    AddAsync.wrongReturnType,
     MakeReturnTypeNullable.new,
     ReplaceReturnType.new,
   ],
   CompileTimeErrorCode.RETURN_OF_INVALID_TYPE_FROM_METHOD: [
+    AddAsync.wrongReturnType,
     MakeReturnTypeNullable.new,
     ReplaceReturnType.new,
   ],