Support for 'Wrap with Xyz' for expressions in ExpressionFunctionBody.

R=brianwilkerson@google.com

Bug: https://github.com/flutter/flutter-intellij/issues/3500
Change-Id: I9e14d6cd1a79070fd489740a881199707be83e90
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/103563
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>
diff --git a/pkg/analysis_server/lib/src/utilities/flutter.dart b/pkg/analysis_server/lib/src/utilities/flutter.dart
index b803440..867a30b 100644
--- a/pkg/analysis_server/lib/src/utilities/flutter.dart
+++ b/pkg/analysis_server/lib/src/utilities/flutter.dart
@@ -284,6 +284,7 @@
       if (isWidgetExpression(node)) {
         var parent = node.parent;
         if (parent is ArgumentList ||
+            parent is ExpressionFunctionBody && parent.expression == node ||
             parent is ListLiteral ||
             parent is NamedExpression && parent.expression == node ||
             parent is Statement) {
diff --git a/pkg/analysis_server/test/src/services/correction/assist/flutter_wrap_center_test.dart b/pkg/analysis_server/test/src/services/correction/assist/flutter_wrap_center_test.dart
index f51ffa7..4e1af86 100644
--- a/pkg/analysis_server/test/src/services/correction/assist/flutter_wrap_center_test.dart
+++ b/pkg/analysis_server/test/src/services/correction/assist/flutter_wrap_center_test.dart
@@ -81,4 +81,20 @@
 }
 ''');
   }
+
+  test_expressionFunctionBody() async {
+    addFlutterPackage();
+    await resolveTestUnit('''
+import 'package:flutter/widgets.dart';
+class FakeFlutter {
+  main() => /*caret*/Container();
+}
+''');
+    await assertHasAssist('''
+import 'package:flutter/widgets.dart';
+class FakeFlutter {
+  main() => Center(child: Container());
+}
+''');
+  }
 }
diff --git a/pkg/analysis_server/test/src/utilities/flutter_test.dart b/pkg/analysis_server/test/src/utilities/flutter_test.dart
index 6bb17e3..5a558ac 100644
--- a/pkg/analysis_server/test/src/utilities/flutter_test.dart
+++ b/pkg/analysis_server/test/src/utilities/flutter_test.dart
@@ -279,6 +279,16 @@
     expect(flutter.identifyWidgetExpression(expression), expression);
   }
 
+  test_identifyWidgetExpression_parent_expressionFunctionBody() async {
+    await resolveTestUnit('''
+import 'package:flutter/widgets.dart';
+
+main(Widget widget) => widget; // ref
+''');
+    Expression expression = findNodeAtString("widget; // ref");
+    expect(flutter.identifyWidgetExpression(expression), expression);
+  }
+
   test_identifyWidgetExpression_parent_expressionStatement() async {
     await resolveTestUnit('''
 import 'package:flutter/widgets.dart';