[DAS] Fixes create method/function return type

Fixes: #60306
Change-Id: Iba6222bbb5d0b63ca94920cad6d6abfd571ed830
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/417780
Auto-Submit: Felipe Morschel <[email protected]>
Commit-Queue: Samuel Rawlins <[email protected]>
Reviewed-by: Konstantin Shcheglov <[email protected]>
Reviewed-by: Samuel Rawlins <[email protected]>

Author: FMorschel

pkg/analysis_server/test/src/services/correction/fix/create_function_test.dart

@@ -726,6 +726,42 @@ void bar(int i) {
     await assert_returnType_bool('while ( test() ) {}');
   }
 
+  Future<void> test_returnType_closure_expression() async {
+    await resolveTestCode('''
+void f(List<int> list) {
+  list.where((i) => myMethod(i));
+}
+''');
+    await assertHasFix('''
+void f(List<int> list) {
+  list.where((i) => myMethod(i));
+}
+
+bool myMethod(int i) {
+}
+''');
+  }
+
+  Future<void> test_returnType_closure_return() async {
+    await resolveTestCode('''
+void f(List<int> list) {
+  list.where((i) {
+    return myMethod(i);
+  });
+}
+''');
+    await assertHasFix('''
+void f(List<int> list) {
+  list.where((i) {
+    return myMethod(i);
+  });
+}
+
+bool myMethod(int i) {
+}
+''');
+  }
+
   Future<void> test_returnType_fromAssignment_eq() async {
     await resolveTestCode('''
 void f() {

pkg/analysis_server/test/src/services/correction/fix/create_method_test.dart

@@ -1253,6 +1253,48 @@ class E {}
 ''', target: '$testPackageLibPath/test2.dart');
   }
 
+  Future<void> test_returnType_closure_expression() async {
+    await resolveTestCode('''
+class A {
+  void m(List<A> list) {
+    list.where((a) => a.myMethod());
+  }
+}
+''');
+    await assertHasFix('''
+class A {
+  void m(List<A> list) {
+    list.where((a) => a.myMethod());
+  }
+
+  bool myMethod() {}
+}
+''');
+  }
+
+  Future<void> test_returnType_closure_return() async {
+    await resolveTestCode('''
+class A {
+  void m(List<A> list) {
+    list.where((a) {
+      return a.myMethod();
+    });
+  }
+}
+''');
+    await assertHasFix('''
+class A {
+  void m(List<A> list) {
+    list.where((a) {
+      return a.myMethod();
+    });
+  }
+
+  bool myMethod() {}
+}
+''');
+  }
+
   Future<void> test_static() async {
     await resolveTestCode('''
 extension E on String {}

pkg/analysis_server_plugin/lib/edit/dart/correction_producer.dart

@@ -511,11 +511,17 @@ abstract class ResolvedCorrectionProducer
     }
     // `=> myFunction();`.
     if (parent is ExpressionFunctionBody) {
+      if (_closureReturnType(expression) case var returnType?) {
+        return returnType;
+      }
       var executable = expression.enclosingExecutableElement2;
       return executable?.returnType;
     }
     // `return myFunction();`.
     if (parent is ReturnStatement) {
+      if (_closureReturnType(expression) case var returnType?) {
+        return returnType;
+      }
       var executable = expression.enclosingExecutableElement2;
       return executable?.returnType;
     }
@@ -639,6 +645,19 @@ abstract class ResolvedCorrectionProducer
 
   bool isEnabled(Feature feature) =>
       libraryElement2.featureSet.isEnabled(feature);
+
+  /// Looks if the [expression] is directly inside a closure and returns the
+  /// return type of the closure.
+  DartType? _closureReturnType(Expression expression) {
+    if (expression.enclosingClosure
+        case FunctionExpression(:var correspondingParameter, :var staticType)) {
+      if (correspondingParameter?.type ?? staticType
+          case FunctionType(:var returnType)) {
+        return returnType;
+      }
+    }
+    return null;
+  }
 }
 
 final class StubCorrectionProducerContext implements CorrectionProducerContext {

pkg/analyzer/lib/src/utilities/extensions/ast.dart

@@ -26,6 +26,23 @@ extension AstNodeExtension on AstNode {
     return result;
   }
 
+  /// The [FunctionExpression] that encloses this node directly or `null` if
+  /// there is another enclosing executable element.
+  FunctionExpression? get enclosingClosure {
+    for (var node in withParents) {
+      switch (node) {
+        case FunctionExpression(:var parent)
+            when parent is! FunctionDeclaration:
+          return node;
+        case FunctionDeclaration() ||
+              ConstructorDeclaration() ||
+              MethodDeclaration():
+          break;
+      }
+    }
+    return null;
+  }
+
   /// The [ExecutableElement2] of the enclosing executable [AstNode].
   ExecutableElement2? get enclosingExecutableElement2 {
     for (var node in withParents) {