Migrator: edge between extension 'on' type and invocation targets; #40023

Change-Id: Ia2f30acec0da2449567d93f99b4c0b0a583c07ac
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/130960
Commit-Queue: Samuel Rawlins <srawlins@google.com>
Reviewed-by: Mike Fairhurst <mfairhurst@google.com>

pkg/nnbd_migration/lib/src/edge_builder.dart

@@ -210,11 +210,28 @@ class EdgeBuilder extends GeneralizingAstVisitor<DecoratedType>
     Map<TypeParameterElement, DecoratedType> substitution;
     Element baseElement = element.declaration;
     if (targetType != null) {
-      var classElement = baseElement.enclosingElement as ClassElement;
-      if (classElement.typeParameters.isNotEmpty) {
-        substitution = _decoratedClassHierarchy
-            .asInstanceOf(targetType, classElement)
-            .asSubstitution;
+      var enclosingElement = baseElement.enclosingElement;
+      if (enclosingElement is ClassElement) {
+        if (enclosingElement.typeParameters.isNotEmpty) {
+          substitution = _decoratedClassHierarchy
+              .asInstanceOf(targetType, enclosingElement)
+              .asSubstitution;
+        }
+      } else {
+        assert(enclosingElement is ExtensionElement);
+        final extensionElement = enclosingElement as ExtensionElement;
+        final extendedType = extensionElement.extendedType;
+        if (extendedType is InterfaceType) {
+          if (extensionElement.typeParameters.isNotEmpty) {
+            substitution = _decoratedClassHierarchy
+                .asInstanceOf(targetType, extendedType.element)
+                .asSubstitution;
+          }
+        } else {
+          // TODO(srawlins): Handle generic typedef. Others?
+          _unimplemented(
+              null, 'Extension on $extendedType (${extendedType.runtimeType}');
+        }
       }
     }
     DecoratedType decoratedBaseType;
@@ -977,7 +994,7 @@ class EdgeBuilder extends GeneralizingAstVisitor<DecoratedType>
       } else if (isNullAware) {
         targetType = target.accept(this);
       } else {
-        targetType = _handleTarget(target, node.methodName.name);
+        targetType = _handleTarget(target, node.methodName.name, callee);
       }
     } else if (target == null && callee.enclosingElement is ClassElement) {
       targetType = _thisOrSuper(node);
@@ -2191,7 +2208,7 @@ class EdgeBuilder extends GeneralizingAstVisitor<DecoratedType>
     } else if (isNullAware) {
       targetType = target.accept(this);
     } else {
-      targetType = _handleTarget(target, propertyName.name);
+      targetType = _handleTarget(target, propertyName.name, callee);
     }
     if (callee == null) {
       // Dynamic dispatch.
@@ -2217,9 +2234,17 @@ class EdgeBuilder extends GeneralizingAstVisitor<DecoratedType>
     }
   }
 
-  DecoratedType _handleTarget(Expression target, String name) {
+  DecoratedType _handleTarget(Expression target, String name, Element method) {
     if (isDeclaredOnObject(name)) {
       return target.accept(this);
+    } else if (method is MethodElement &&
+        method.enclosingElement is ExtensionElement) {
+      // Extension methods can be called on a `null` target, when the `on` type
+      // of the extension is nullable.
+      _handleAssignment(target,
+          destinationType:
+              _variables.decoratedElementType(method.enclosingElement));
+      return target.accept(this);
     } else {
       return _checkExpressionNotNull(target);
     }

pkg/nnbd_migration/lib/src/node_builder.dart

@@ -201,7 +201,7 @@ class NodeBuilder extends GeneralizingAstVisitor<DecoratedType>
   }
 
   @override
-  visitEnumDeclaration(EnumDeclaration node) {
+  DecoratedType visitEnumDeclaration(EnumDeclaration node) {
     node.metadata.accept(this);
     node.name.accept(this);
     var classElement = node.declaredElement;
@@ -241,6 +241,16 @@ class NodeBuilder extends GeneralizingAstVisitor<DecoratedType>
     return null;
   }
 
+  @override
+  DecoratedType visitExtensionDeclaration(ExtensionDeclaration node) {
+    node.metadata.accept(this);
+    node.typeParameters?.accept(this);
+    var type = node.extendedType.accept(this);
+    _variables.recordDecoratedElementType(node.declaredElement, type);
+    node.members.accept(this);
+    return null;
+  }
+
   @override
   DecoratedType visitFieldFormalParameter(FieldFormalParameter node) {
     return _handleFormalParameter(

pkg/nnbd_migration/test/api_test.dart

@@ -1310,6 +1310,43 @@ extension E on String? {
     await _checkSingleFileChanges(content, expected);
   }
 
+  Future<void> test_extension_nullableOnType_typeArgument() async {
+    var content = '''
+extension E on List<String> {
+  void m() {}
+}
+void f(List<String> list) => list.m();
+void g() => f([null]);
+''';
+    var expected = '''
+extension E on List<String?> {
+  void m() {}
+}
+void f(List<String?> list) => list.m();
+void g() => f([null]);
+''';
+    await _checkSingleFileChanges(content, expected);
+  }
+
+  @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/40023')
+  Future<void> test_extension_nullableOnType_typeVariable() async {
+    var content = '''
+extension E<T> on List<T> {
+  void m() {}
+}
+void f<U>(List<U> list) => list.m();
+void g() => f([null]);
+''';
+    var expected = '''
+extension E<T> on List<T?> {
+  void m() {}
+}
+void f<U>(List<U?> list) => list.m();
+void g() => f([null]);
+''';
+    await _checkSingleFileChanges(content, expected);
+  }
+
   @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/40023')
   Future<void> test_extension_nullableOnType_viaExplicitInvocation() async {
     var content = '''
@@ -1329,7 +1366,6 @@ void f() => E(null).m();
     await _checkSingleFileChanges(content, expected);
   }
 
-  @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/40023')
   Future<void> test_extension_nullableOnType_viaImplicitInvocation() async {
     var content = '''
 class C {}
@@ -3834,6 +3870,16 @@ class _ProvisionalApiTestWithFixBuilder extends _ProvisionalApiTestBase
   @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/38472')
   Future<void> test_enum() => super.test_enum();
 
+  @override
+  @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/40023')
+  Future<void> test_extension_nullableOnType_typeArgument() =>
+      super.test_extension_nullableOnType_typeArgument();
+
+  @override
+  @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/40023')
+  Future<void> test_extension_nullableOnType_viaImplicitInvocation() =>
+      super.test_extension_nullableOnType_viaImplicitInvocation();
+
   @override
   @FailingTest(issue: 'https://github.com/dart-lang/sdk/issues/38472')
   Future<void> test_field_initializer_typed_list_literal() =>

pkg/nnbd_migration/test/edge_builder_test.dart

@@ -3662,6 +3662,22 @@ extension E on C {
         hard: true);
   }
 
+  Future<void> test_methodInvocation_extension_nullTarget() async {
+    await analyze('''
+class C {}
+extension on C /*1*/ {
+  void m() {}
+}
+void f() {
+  C c = null;
+  c.m();
+}
+''');
+    assertEdge(decoratedTypeAnnotation('C c').node,
+        decoratedTypeAnnotation('C /*1*/').node,
+        hard: true);
+  }
+
   Future<void> test_methodInvocation_extension_unnamed() async {
     await analyze('''
 class C {