Migration: propagate nullability inference through override relationships: return types.

Parameter types will be handled in a follow-up CL.

Change-Id: I71862fa6cf27aeede24bef60c096f631f6d8f60d
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/107100
Reviewed-by: Dan Rubel <danrubel@google.com>
diff --git a/pkg/nnbd_migration/lib/src/decorated_class_hierarchy.dart b/pkg/nnbd_migration/lib/src/decorated_class_hierarchy.dart
index 52b3511..368bfbd 100644
--- a/pkg/nnbd_migration/lib/src/decorated_class_hierarchy.dart
+++ b/pkg/nnbd_migration/lib/src/decorated_class_hierarchy.dart
@@ -62,11 +62,8 @@
         var supertype = decoratedSupertype.type as InterfaceType;
         // Compute a type substitution to determine how [class_] relates to
         // this specific [superclass].
-        var argumentTypes = supertype.typeArguments;
-        var parameterTypes =
-            supertype.typeParameters.map((p) => p.type).toList();
         Map<TypeParameterElement, DecoratedType> substitution = {};
-        for (int i = 0; i < argumentTypes.length; i++) {
+        for (int i = 0; i < supertype.typeArguments.length; i++) {
           substitution[supertype.typeParameters[i]] =
               decoratedSupertype.typeArguments[i];
         }
@@ -76,10 +73,7 @@
         var recursiveSupertypeDecorations =
             _getGenericSupertypeDecorations(superclass);
         for (var entry in recursiveSupertypeDecorations.entries) {
-          var undecoratedResult =
-              entry.value.type.substitute2(argumentTypes, parameterTypes);
-          decorations[entry.key] ??=
-              entry.value.substitute(substitution, undecoratedResult);
+          decorations[entry.key] ??= entry.value.substitute(substitution);
         }
         // Also record the relation between [class_] and its direct
         // superclass.
diff --git a/pkg/nnbd_migration/lib/src/decorated_type.dart b/pkg/nnbd_migration/lib/src/decorated_type.dart
index 57bae17..1fca000 100644
--- a/pkg/nnbd_migration/lib/src/decorated_type.dart
+++ b/pkg/nnbd_migration/lib/src/decorated_type.dart
@@ -96,14 +96,41 @@
     return decoratedType;
   }
 
+  /// If `this` represents an interface type, returns the substitution necessary
+  /// to produce this type using the class's type as a starting point.
+  /// Otherwise throws an exception.
+  ///
+  /// For instance, if `this` represents `List<int?1>`, returns the substitution
+  /// `{T: int?1}`, where `T` is the [TypeParameterElement] for `List`'s type
+  /// parameter.
+  Map<TypeParameterElement, DecoratedType> get asSubstitution {
+    var type = this.type;
+    if (type is InterfaceType) {
+      return Map<TypeParameterElement, DecoratedType>.fromIterables(
+          type.element.typeParameters, typeArguments);
+    } else {
+      throw StateError(
+          'Tried to convert a non-interface type to a substitution');
+    }
+  }
+
   /// Apply the given [substitution] to this type.
   ///
   /// [undecoratedResult] is the result of the substitution, as determined by
-  /// the normal type system.
+  /// the normal type system.  If not supplied, it is inferred.
   DecoratedType substitute(
       Map<TypeParameterElement, DecoratedType> substitution,
-      DartType undecoratedResult) {
+      [DartType undecoratedResult]) {
     if (substitution.isEmpty) return this;
+    if (undecoratedResult == null) {
+      List<DartType> argumentTypes = [];
+      List<DartType> parameterTypes = [];
+      for (var entry in substitution.entries) {
+        argumentTypes.add(entry.value.type);
+        parameterTypes.add(entry.key.type);
+      }
+      undecoratedResult = type.substitute2(argumentTypes, parameterTypes);
+    }
     return _substitute(substitution, undecoratedResult);
   }
 
diff --git a/pkg/nnbd_migration/lib/src/edge_builder.dart b/pkg/nnbd_migration/lib/src/edge_builder.dart
index 66d84c2..6688d6c 100644
--- a/pkg/nnbd_migration/lib/src/edge_builder.dart
+++ b/pkg/nnbd_migration/lib/src/edge_builder.dart
@@ -7,12 +7,14 @@
 import 'package:analyzer/dart/ast/visitor.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/src/dart/element/inheritance_manager2.dart';
 import 'package:analyzer/src/dart/element/member.dart';
 import 'package:analyzer/src/generated/resolver.dart';
 import 'package:analyzer/src/generated/source.dart';
 import 'package:meta/meta.dart';
 import 'package:nnbd_migration/nnbd_migration.dart';
 import 'package:nnbd_migration/src/conditional_discard.dart';
+import 'package:nnbd_migration/src/decorated_class_hierarchy.dart';
 import 'package:nnbd_migration/src/decorated_type.dart';
 import 'package:nnbd_migration/src/edge_origin.dart';
 import 'package:nnbd_migration/src/expression_checks.dart';
@@ -27,6 +29,8 @@
 /// variables that will determine its nullability.  For `visit...` methods that
 /// don't visit expressions, `null` will be returned.
 class EdgeBuilder extends GeneralizingAstVisitor<DecoratedType> {
+  final InheritanceManager2 _inheritanceManager;
+
   /// The repository of constraint variables and decorated types (from a
   /// previous pass over the source code).
   final VariableRepository _variables;
@@ -38,6 +42,8 @@
   /// The file being analyzed.
   final Source _source;
 
+  final DecoratedClassHierarchy _decoratedClassHierarchy;
+
   /// For convenience, a [DecoratedType] representing non-nullable `Object`.
   final DecoratedType _notNullType;
 
@@ -77,9 +83,11 @@
 
   NullabilityNode _lastConditionalNode;
 
-  EdgeBuilder(TypeProvider typeProvider, this._variables, this._graph,
-      this._source, this.listener)
-      : _notNullType = DecoratedType(typeProvider.objectType, _graph.never),
+  EdgeBuilder(TypeProvider typeProvider, TypeSystem typeSystem, this._variables,
+      this._graph, this._source, this.listener)
+      : _decoratedClassHierarchy = DecoratedClassHierarchy(_variables, _graph),
+        _inheritanceManager = InheritanceManager2(typeSystem),
+        _notNullType = DecoratedType(typeProvider.objectType, _graph.never),
         _nonNullableBoolType =
             DecoratedType(typeProvider.boolType, _graph.never),
         _nonNullableTypeType =
@@ -490,11 +498,31 @@
   DecoratedType visitMethodDeclaration(MethodDeclaration node) {
     node.parameters?.accept(this);
     assert(_currentFunctionType == null);
-    _currentFunctionType =
-        _variables.decoratedElementType(node.declaredElement);
+    var declaredElement = node.declaredElement;
+    _currentFunctionType = _variables.decoratedElementType(declaredElement);
     _inConditionalControlFlow = false;
     try {
       node.body.accept(this);
+      var classElement = declaredElement.enclosingElement as ClassElement;
+      for (var overridden in _inheritanceManager.getOverridden(
+              classElement.type,
+              Name(classElement.library.source.uri, declaredElement.name)) ??
+          const []) {
+        var overriddenElement = overridden.element as ExecutableElement;
+        assert(overriddenElement is! ExecutableMember);
+        var overriddenClass =
+            overriddenElement.enclosingElement as ClassElement;
+        var decoratedOverriddenFunctionType =
+            _variables.decoratedElementType(overriddenElement);
+        var decoratedSupertype = _decoratedClassHierarchy.getDecoratedSupertype(
+            classElement, overriddenClass);
+        var substitution = decoratedSupertype.asSubstitution;
+        _checkAssignment(null,
+            source: _currentFunctionType,
+            destination:
+                decoratedOverriddenFunctionType.substitute(substitution),
+            hard: true);
+      }
     } finally {
       _currentFunctionType = null;
     }
diff --git a/pkg/nnbd_migration/lib/src/nullability_migration_impl.dart b/pkg/nnbd_migration/lib/src/nullability_migration_impl.dart
index a37dfb1..6fede39 100644
--- a/pkg/nnbd_migration/lib/src/nullability_migration_impl.dart
+++ b/pkg/nnbd_migration/lib/src/nullability_migration_impl.dart
@@ -59,8 +59,8 @@
 
   void processInput(ResolvedUnitResult result) {
     var unit = result.unit;
-    unit.accept(EdgeBuilder(result.typeProvider, _variables, _graph,
-        unit.declaredElement.source, _permissive ? listener : null));
+    unit.accept(EdgeBuilder(result.typeProvider, result.typeSystem, _variables,
+        _graph, unit.declaredElement.source, _permissive ? listener : null));
   }
 }
 
diff --git a/pkg/nnbd_migration/test/api_test.dart b/pkg/nnbd_migration/test/api_test.dart
index 41dd341..90f1633 100644
--- a/pkg/nnbd_migration/test/api_test.dart
+++ b/pkg/nnbd_migration/test/api_test.dart
@@ -1043,6 +1043,88 @@
     await _checkSingleFileChanges(content, expected);
   }
 
+  test_override_return_type_non_nullable() async {
+    var content = '''
+abstract class Base {
+  int/*!*/ f();
+}
+class Derived extends Base {
+  int f() => g();
+}
+int g() => null;
+''';
+    var expected = '''
+abstract class Base {
+  int/*!*/ f();
+}
+class Derived extends Base {
+  int f() => g()!;
+}
+int? g() => null;
+''';
+    await _checkSingleFileChanges(content, expected);
+  }
+
+  test_override_return_type_nullable() async {
+    var content = '''
+abstract class Base {
+  int f();
+}
+class Derived extends Base {
+  int f() => null;
+}
+''';
+    var expected = '''
+abstract class Base {
+  int? f();
+}
+class Derived extends Base {
+  int? f() => null;
+}
+''';
+    await _checkSingleFileChanges(content, expected);
+  }
+
+  test_override_return_type_nullable_substitution_complex() async {
+    var content = '''
+abstract class Base<T> {
+  T f();
+}
+class Derived extends Base<List<int>> {
+  List<int> f() => <int>[null];
+}
+''';
+    var expected = '''
+abstract class Base<T> {
+  T f();
+}
+class Derived extends Base<List<int?>> {
+  List<int?> f() => <int?>[null];
+}
+''';
+    await _checkSingleFileChanges(content, expected);
+  }
+
+  test_override_return_type_nullable_substitution_simple() async {
+    var content = '''
+abstract class Base<T> {
+  T f();
+}
+class Derived extends Base<int> {
+  int f() => null;
+}
+''';
+    var expected = '''
+abstract class Base<T> {
+  T f();
+}
+class Derived extends Base<int?> {
+  int? f() => null;
+}
+''';
+    await _checkSingleFileChanges(content, expected);
+  }
+
   test_parameter_genericFunctionType() async {
     var content = '''
 int f(int x, int Function(int i) g) {
diff --git a/pkg/nnbd_migration/test/edge_builder_test.dart b/pkg/nnbd_migration/test/edge_builder_test.dart
index 18da745..6580d6d 100644
--- a/pkg/nnbd_migration/test/edge_builder_test.dart
+++ b/pkg/nnbd_migration/test/edge_builder_test.dart
@@ -25,7 +25,8 @@
   @override
   Future<CompilationUnit> analyze(String code) async {
     var unit = await super.analyze(code);
-    unit.accept(EdgeBuilder(typeProvider, variables, graph, testSource, null));
+    unit.accept(EdgeBuilder(
+        typeProvider, typeSystem, variables, graph, testSource, null));
     return unit;
   }
 
@@ -1273,6 +1274,48 @@
     expect(never.isNullable, isFalse);
   }
 
+  test_override_return_type_getter() async {
+    await analyze('''
+abstract class Base {
+  int/*1*/ get x;
+}
+class Derived extends Base {
+  int/*2*/ get x => null;
+}
+''');
+    var int1 = decoratedTypeAnnotation('int/*1*/');
+    var int2 = decoratedTypeAnnotation('int/*2*/');
+    assertEdge(int2.node, int1.node, hard: true);
+  }
+
+  test_override_return_type_method() async {
+    await analyze('''
+abstract class Base {
+  int/*1*/ f();
+}
+class Derived extends Base {
+  int/*2*/ f() => null;
+}
+''');
+    var int1 = decoratedTypeAnnotation('int/*1*/');
+    var int2 = decoratedTypeAnnotation('int/*2*/');
+    assertEdge(int2.node, int1.node, hard: true);
+  }
+
+  test_override_return_type_operator() async {
+    await analyze('''
+abstract class Base {
+  Base/*1*/ operator-();
+}
+class Derived extends Base {
+  Derived/*2*/ operator-() => null;
+}
+''');
+    var base1 = decoratedTypeAnnotation('Base/*1*/');
+    var derived2 = decoratedTypeAnnotation('Derived/*2*/');
+    assertEdge(derived2.node, base1.node, hard: true);
+  }
+
   test_parenthesizedExpression() async {
     await analyze('''
 int f() {
diff --git a/pkg/nnbd_migration/test/migration_visitor_test_base.dart b/pkg/nnbd_migration/test/migration_visitor_test_base.dart
index d7f26f6..7fdc709 100644
--- a/pkg/nnbd_migration/test/migration_visitor_test_base.dart
+++ b/pkg/nnbd_migration/test/migration_visitor_test_base.dart
@@ -97,6 +97,8 @@
 
   TypeProvider get typeProvider => testAnalysisResult.typeProvider;
 
+  TypeSystem get typeSystem => testAnalysisResult.typeSystem;
+
   Future<CompilationUnit> analyze(String code) async {
     await resolveTestUnit(code);
     testUnit