Enhance the replace_with_null_aware fix and apply it in more places

Change-Id: I80945b4814cca0eaeca2129e5365c0d7526a8247
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/209340
Commit-Queue: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Konstantin Shcheglov <scheglov@google.com>
diff --git a/pkg/analysis_server/lib/src/services/correction/dart/replace_with_null_aware.dart b/pkg/analysis_server/lib/src/services/correction/dart/replace_with_null_aware.dart
index c6b7dd1..822442d 100644
--- a/pkg/analysis_server/lib/src/services/correction/dart/replace_with_null_aware.dart
+++ b/pkg/analysis_server/lib/src/services/correction/dart/replace_with_null_aware.dart
@@ -10,11 +10,24 @@
 import 'package:analyzer_plugin/utilities/range_factory.dart';
 
 class ReplaceWithNullAware extends CorrectionProducer {
+  /// The kind of correction to be made.
+  final _CorrectionKind correctionKind;
+
+  ReplaceWithNullAware(this.correctionKind);
+
   @override
   FixKind get fixKind => DartFixKind.REPLACE_WITH_NULL_AWARE;
 
   @override
   Future<void> compute(ChangeBuilder builder) async {
+    if (correctionKind == _CorrectionKind.inChain) {
+      await _computeInChain(builder);
+    } else if (correctionKind == _CorrectionKind.single) {
+      await _computeSingle(builder);
+    }
+  }
+
+  Future<void> _computeInChain(ChangeBuilder builder) async {
     var node = coveredNode;
     if (node is Expression) {
       final node_final = node;
@@ -38,6 +51,37 @@
     }
   }
 
+  Future<void> _computeSingle(ChangeBuilder builder) async {
+    var node = coveredNode?.parent;
+    if (node is MethodInvocation) {
+      var operator = node.operator;
+      if (operator != null) {
+        await builder.addDartFileEdit(file, (builder) {
+          builder.addSimpleReplacement(range.token(operator), '?.');
+        });
+      }
+    } else if (node is PrefixedIdentifier) {
+      await builder.addDartFileEdit(file, (builder) {
+        builder.addSimpleReplacement(range.token(node.period), '?.');
+      });
+    } else if (node is PropertyAccess) {
+      await builder.addDartFileEdit(file, (builder) {
+        builder.addSimpleReplacement(range.token(node.operator), '?.');
+      });
+    }
+  }
+
   /// Return an instance of this class. Used as a tear-off in `FixProcessor`.
-  static ReplaceWithNullAware newInstance() => ReplaceWithNullAware();
+  static ReplaceWithNullAware inChain() =>
+      ReplaceWithNullAware(_CorrectionKind.inChain);
+
+  /// Return an instance of this class. Used as a tear-off in `FixProcessor`.
+  static ReplaceWithNullAware single() =>
+      ReplaceWithNullAware(_CorrectionKind.single);
+}
+
+/// The kinds of corrections supported by [ReplaceWithNullAware].
+enum _CorrectionKind {
+  inChain,
+  single,
 }
diff --git a/pkg/analysis_server/lib/src/services/correction/fix_internal.dart b/pkg/analysis_server/lib/src/services/correction/fix_internal.dart
index c04be05..e4e975c 100644
--- a/pkg/analysis_server/lib/src/services/correction/fix_internal.dart
+++ b/pkg/analysis_server/lib/src/services/correction/fix_internal.dart
@@ -930,12 +930,14 @@
     ],
     CompileTimeErrorCode.UNCHECKED_METHOD_INVOCATION_OF_NULLABLE_VALUE: [
       AddNullCheck.newInstance,
+      ReplaceWithNullAware.single,
     ],
     CompileTimeErrorCode.UNCHECKED_OPERATOR_INVOCATION_OF_NULLABLE_VALUE: [
       AddNullCheck.newInstance,
     ],
     CompileTimeErrorCode.UNCHECKED_PROPERTY_ACCESS_OF_NULLABLE_VALUE: [
       AddNullCheck.newInstance,
+      ReplaceWithNullAware.single,
     ],
     CompileTimeErrorCode.UNCHECKED_USE_OF_NULLABLE_VALUE: [
       AddNullCheck.newInstance,
@@ -1045,7 +1047,7 @@
     ],
 
     HintCode.CAN_BE_NULL_AFTER_NULL_AWARE: [
-      ReplaceWithNullAware.newInstance,
+      ReplaceWithNullAware.inChain,
     ],
     HintCode.DEAD_CODE: [
       RemoveDeadCode.newInstance,
diff --git a/pkg/analysis_server/test/src/services/correction/fix/replace_with_null_aware_test.dart b/pkg/analysis_server/test/src/services/correction/fix/replace_with_null_aware_test.dart
index 282e085..84b84d5 100644
--- a/pkg/analysis_server/test/src/services/correction/fix/replace_with_null_aware_test.dart
+++ b/pkg/analysis_server/test/src/services/correction/fix/replace_with_null_aware_test.dart
@@ -11,6 +11,8 @@
 void main() {
   defineReflectiveSuite(() {
     defineReflectiveTests(ReplaceWithNullAwareTest);
+    defineReflectiveTests(UncheckedMethodInvocationOfNullableValueTest);
+    defineReflectiveTests(UncheckedPropertyAccessOfNullableValueTest);
   });
 }
 
@@ -61,3 +63,90 @@
 ''');
   }
 }
+
+@reflectiveTest
+class UncheckedMethodInvocationOfNullableValueTest extends FixProcessorTest {
+  @override
+  FixKind get kind => DartFixKind.REPLACE_WITH_NULL_AWARE;
+
+  Future<void> test_method() async {
+    await resolveTestCode('''
+class C {
+  List<int>? values;
+
+  void m() {
+    if (values != null) {
+      print(values.toList());
+    }
+  }
+}
+''');
+    await assertHasFix('''
+class C {
+  List<int>? values;
+
+  void m() {
+    if (values != null) {
+      print(values?.toList());
+    }
+  }
+}
+''');
+  }
+}
+
+@reflectiveTest
+class UncheckedPropertyAccessOfNullableValueTest extends FixProcessorTest {
+  @override
+  FixKind get kind => DartFixKind.REPLACE_WITH_NULL_AWARE;
+
+  Future<void> test_prefixedIdentifier() async {
+    await resolveTestCode('''
+class C {
+  List<int>? values;
+
+  void m() {
+    if (values != null) {
+      print(values.length);
+    }
+  }
+}
+''');
+    await assertHasFix('''
+class C {
+  List<int>? values;
+
+  void m() {
+    if (values != null) {
+      print(values?.length);
+    }
+  }
+}
+''');
+  }
+
+  Future<void> test_propertyAccess() async {
+    await resolveTestCode('''
+class C {
+  List<int>? values;
+
+  void m() {
+    if (values != null) {
+      print((values).length);
+    }
+  }
+}
+''');
+    await assertHasFix('''
+class C {
+  List<int>? values;
+
+  void m() {
+    if (values != null) {
+      print((values)?.length);
+    }
+  }
+}
+''');
+  }
+}