Enhance add_null_check to cover more cases

Change-Id: I9817a300bc4600b72396fb1cd35463e6fa2f3953
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/177420
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/add_null_check.dart b/pkg/analysis_server/lib/src/services/correction/dart/add_null_check.dart
index 38fa836..4c04dd2 100644
--- a/pkg/analysis_server/lib/src/services/correction/dart/add_null_check.dart
+++ b/pkg/analysis_server/lib/src/services/correction/dart/add_null_check.dart
@@ -43,11 +43,18 @@
       toType = target.staticParameterElement.type;
     } else if (parent is IndexExpression) {
       toType = parent.realTarget.staticType;
+    } else if ((parent is PrefixedIdentifier && target == parent.prefix) ||
+        (parent is PropertyAccess && target == parent.target) ||
+        (parent is MethodInvocation && target == parent.target) ||
+        (parent is FunctionExpressionInvocation && target == parent.function)) {
+      // No need to set the `toType` because there isn't any need for a type
+      // check.
     } else {
       return;
     }
-    if (!typeSystem.isAssignableTo(
-        toType, typeSystem.promoteToNonNull(fromType))) {
+    if (toType != null &&
+        !typeSystem.isAssignableTo(
+            toType, typeSystem.promoteToNonNull(fromType))) {
       // The reason that `fromType` can't be assigned to `toType` is more than
       // just because it's nullable, in which case a null check won't fix the
       // problem.
diff --git a/pkg/analysis_server/test/src/services/correction/fix/add_null_check_test.dart b/pkg/analysis_server/test/src/services/correction/fix/add_null_check_test.dart
index 4dd8288..b445c7e 100644
--- a/pkg/analysis_server/test/src/services/correction/fix/add_null_check_test.dart
+++ b/pkg/analysis_server/test/src/services/correction/fix/add_null_check_test.dart
@@ -97,6 +97,21 @@
 ''');
   }
 
+  Future<void> test_functionExpressionInvocation() async {
+    await resolveTestCode('''
+int f(C c) => c.func();
+class C {
+  int Function()? get func => null;
+}
+''');
+    await assertHasFix('''
+int f(C c) => c.func!();
+class C {
+  int Function()? get func => null;
+}
+''');
+  }
+
   Future<void> test_indexExpression() async {
     await resolveTestCode('''
 void f (List<String>? args) {
@@ -134,4 +149,31 @@
 ''');
     await assertNoFix();
   }
+
+  Future<void> test_methodInvocation() async {
+    await resolveTestCode('''
+String f(String? s) => s.substring(0);
+''');
+    await assertHasFix('''
+String f(String? s) => s!.substring(0);
+''');
+  }
+
+  Future<void> test_prefixedIdentifier() async {
+    await resolveTestCode('''
+int f(String? s) => s.length;
+''');
+    await assertHasFix('''
+int f(String? s) => s!.length;
+''');
+  }
+
+  Future<void> test_propertyAccess() async {
+    await resolveTestCode('''
+int f(String? s) => (s).length;
+''');
+    await assertHasFix('''
+int f(String? s) => (s)!.length;
+''');
+  }
 }