Migration: add a data structure for tracking class hierarchy information.

The new data structure, called `DecoratedClassHierarchy`, computes and
caches the variable substitutions necessary to relate a class to one
of its superclasses.  In future CLs, we'll use this information to
create the appropriate contraints for override relationships, and to
compute the correct types when a method dispatch resolves to a method
declared in a superclass.

Note that the analyzer uses `InhertanceManager2` both to determine
what a method call resolves to and to record the appropriate
substitutions.  We can take advantage of its determination of what a
method call resolves to (that information is stored in the resolved
AST), but we need to cmopute the substitutions separately because the
substitutions need to be decorated with nullability nodes.

Change-Id: Ifb334972e509b77151d41f7e0700cd590d7b5417
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/106800
Commit-Queue: Paul Berry <paulberry@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Konstantin Shcheglov <scheglov@google.com>
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
new file mode 100644
index 0000000..52b3511
--- /dev/null
+++ b/pkg/nnbd_migration/lib/src/decorated_class_hierarchy.dart
@@ -0,0 +1,91 @@
+// Copyright (c) 2019, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/src/dart/element/handle.dart';
+import 'package:nnbd_migration/src/decorated_type.dart';
+import 'package:nnbd_migration/src/node_builder.dart';
+import 'package:nnbd_migration/src/nullability_node.dart';
+
+/// Responsible for building and maintaining information about nullability
+/// decorations related to the class hierarchy.
+///
+/// For instance, if one class is a subclass of the other, we record the
+/// nullabilities of all the types involved in the subclass relationship.
+class DecoratedClassHierarchy {
+  final VariableRepository _variables;
+
+  final NullabilityGraph _graph;
+
+  /// Cache for speeding up the computation of
+  /// [_getGenericSupertypeDecorations].
+  final Map<ClassElement, Map<ClassElement, DecoratedType>>
+      _genericSupertypeDecorations = {};
+
+  DecoratedClassHierarchy(this._variables, this._graph);
+
+  /// Retrieves a [DecoratedType] describing how [class_] implements
+  /// [superclass].
+  ///
+  /// If [class_] does not implement [superclass], raises an exception.
+  ///
+  /// Note that the returned [DecoratedType] will have a node of `never`,
+  /// because the relationship between a class and its superclass is not
+  /// nullable.
+  DecoratedType getDecoratedSupertype(
+      ClassElement class_, ClassElement superclass) {
+    assert(class_ is! ClassElementHandle);
+    assert(superclass is! ClassElementHandle);
+    if (superclass.typeParameters.isEmpty) {
+      return DecoratedType(superclass.type, _graph.never);
+    }
+    return _getGenericSupertypeDecorations(class_)[superclass] ??
+        (throw StateError('Unrelated types'));
+  }
+
+  /// Computes a map whose keys are all the superclasses of [class_], and whose
+  /// values indicate how [class_] implements each superclass.
+  Map<ClassElement, DecoratedType> _getGenericSupertypeDecorations(
+      ClassElement class_) {
+    var decorations = _genericSupertypeDecorations[class_];
+    if (decorations == null) {
+      // Call ourselves recursively to compute how each of [class_]'s direct
+      // superclasses relates to all of its transitive superclasses.
+      decorations = _genericSupertypeDecorations[class_] = {};
+      var decoratedDirectSupertypes =
+          _variables.decoratedDirectSupertypes(class_);
+      for (var entry in decoratedDirectSupertypes.entries) {
+        var superclass = entry.key;
+        var decoratedSupertype = entry.value;
+        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++) {
+          substitution[supertype.typeParameters[i]] =
+              decoratedSupertype.typeArguments[i];
+        }
+        // Apply that substitution to the relation between [superclass] and
+        // each of its transitive superclasses, to determine the relation
+        // between [class_] and the transitive superclass.
+        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);
+        }
+        // Also record the relation between [class_] and its direct
+        // superclass.
+        decorations[superclass] ??= decoratedSupertype;
+      }
+    }
+    return decorations;
+  }
+}
diff --git a/pkg/nnbd_migration/lib/src/decorated_type.dart b/pkg/nnbd_migration/lib/src/decorated_type.dart
index a0a070f..57bae17 100644
--- a/pkg/nnbd_migration/lib/src/decorated_type.dart
+++ b/pkg/nnbd_migration/lib/src/decorated_type.dart
@@ -162,10 +162,22 @@
           returnType: returnType._substitute(
               substitution, undecoratedResult.returnType),
           positionalParameters: newPositionalParameters);
+    } else if (type is InterfaceType && undecoratedResult is InterfaceType) {
+      List<DecoratedType> newTypeArguments = [];
+      for (int i = 0; i < typeArguments.length; i++) {
+        newTypeArguments.add(typeArguments[i]
+            .substitute(substitution, undecoratedResult.typeArguments[i]));
+      }
+      return DecoratedType(undecoratedResult, node,
+          typeArguments: newTypeArguments);
     } else if (type is TypeParameterType) {
       var inner = substitution[type.element];
-      return DecoratedType(undecoratedResult,
-          NullabilityNode.forSubstitution(inner?.node, node));
+      if (inner == null) {
+        return this;
+      } else {
+        return inner
+            .withNode(NullabilityNode.forSubstitution(inner.node, node));
+      }
     } else if (type is VoidType) {
       return this;
     }
diff --git a/pkg/nnbd_migration/lib/src/node_builder.dart b/pkg/nnbd_migration/lib/src/node_builder.dart
index e472ffd..38d89ef 100644
--- a/pkg/nnbd_migration/lib/src/node_builder.dart
+++ b/pkg/nnbd_migration/lib/src/node_builder.dart
@@ -6,6 +6,7 @@
 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/handle.dart';
 import 'package:analyzer/src/dart/element/type.dart';
 import 'package:analyzer/src/generated/resolver.dart';
 import 'package:analyzer/src/generated/source.dart';
@@ -48,8 +49,13 @@
 
   final TypeProvider _typeProvider;
 
+  /// For convenience, a [DecoratedType] representing non-nullable `Object`.
+  final DecoratedType _nonNullableObjectType;
+
   NodeBuilder(this._variables, this._source, this.listener, this._graph,
-      this._typeProvider);
+      this._typeProvider)
+      : _nonNullableObjectType =
+            DecoratedType(_typeProvider.objectType, _graph.never);
 
   @override
   DecoratedType visitCompilationUnit(CompilationUnit node) {
@@ -72,6 +78,32 @@
   }
 
   @override
+  DecoratedType visitClassDeclaration(ClassDeclaration node) {
+    node.metadata.accept(this);
+    node.name.accept(this);
+    node.typeParameters?.accept(this);
+    node.nativeClause?.accept(this);
+    node.members.accept(this);
+    _handleSupertypeClauses(
+        node.declaredElement,
+        node.extendsClause?.superclass,
+        node.withClause,
+        node.implementsClause,
+        null);
+    return null;
+  }
+
+  @override
+  visitClassTypeAlias(ClassTypeAlias node) {
+    node.metadata.accept(this);
+    node.name.accept(this);
+    node.typeParameters?.accept(this);
+    _handleSupertypeClauses(node.declaredElement, node.superclass,
+        node.withClause, node.implementsClause, null);
+    return null;
+  }
+
+  @override
   DecoratedType visitConstructorDeclaration(ConstructorDeclaration node) {
     if (node.factoryKeyword != null) {
       // Factory constructors can return null, but we don't want to propagate a
@@ -145,6 +177,17 @@
   }
 
   @override
+  visitMixinDeclaration(MixinDeclaration node) {
+    node.metadata.accept(this);
+    node.name?.accept(this);
+    node.typeParameters?.accept(this);
+    node.members.accept(this);
+    _handleSupertypeClauses(
+        node.declaredElement, null, null, node.implementsClause, node.onClause);
+    return null;
+  }
+
+  @override
   DecoratedType visitNode(AstNode node) {
     if (listener != null) {
       try {
@@ -226,8 +269,18 @@
         _namedParameters = previousNamedParameters;
       }
     }
-    var decoratedType = DecoratedTypeAnnotation(
-        type, NullabilityNode.forTypeAnnotation(node.end), node.end,
+    NullabilityNode nullabilityNode;
+    var parent = node.parent;
+    if (parent is ExtendsClause ||
+        parent is ImplementsClause ||
+        parent is WithClause ||
+        parent is OnClause ||
+        parent is ClassTypeAlias) {
+      nullabilityNode = _graph.never;
+    } else {
+      nullabilityNode = NullabilityNode.forTypeAnnotation(node.end);
+    }
+    var decoratedType = DecoratedTypeAnnotation(type, nullabilityNode, node.end,
         typeArguments: typeArguments,
         returnType: returnType,
         positionalParameters: positionalParameters,
@@ -346,6 +399,41 @@
     _variables.recordDecoratedElementType(declaredElement, functionType);
   }
 
+  void _handleSupertypeClauses(
+      ClassElement declaredElement,
+      TypeName superclass,
+      WithClause withClause,
+      ImplementsClause implementsClause,
+      OnClause onClause) {
+    var supertypes = <TypeName>[];
+    supertypes.add(superclass);
+    if (withClause != null) {
+      supertypes.addAll(withClause.mixinTypes);
+    }
+    if (implementsClause != null) {
+      supertypes.addAll(implementsClause.interfaces);
+    }
+    if (onClause != null) {
+      supertypes.addAll(onClause.superclassConstraints);
+    }
+    var decoratedSupertypes = <ClassElement, DecoratedType>{};
+    for (var supertype in supertypes) {
+      DecoratedType decoratedSupertype;
+      if (supertype == null) {
+        decoratedSupertype = _nonNullableObjectType;
+      } else {
+        decoratedSupertype = supertype.accept(this);
+      }
+      var class_ = (decoratedSupertype.type as InterfaceType).element;
+      if (class_ is ClassElementHandle) {
+        class_ = (class_ as ClassElementHandle).actualElement;
+      }
+      decoratedSupertypes[class_] = decoratedSupertype;
+    }
+    _variables.recordDecoratedDirectSupertypes(
+        declaredElement, decoratedSupertypes);
+  }
+
   @alwaysThrows
   void _unimplemented(AstNode node, String message) {
     CompilationUnit unit = node.root as CompilationUnit;
@@ -369,6 +457,11 @@
 /// ([NodeBuilder], which finds all the variables that need to be
 /// constrained).
 abstract class VariableRecorder {
+  /// Associates a [class_] with decorated type information for the superclasses
+  /// it directly implements/extends/etc.
+  void recordDecoratedDirectSupertypes(ClassElement class_,
+      Map<ClassElement, DecoratedType> decoratedDirectSupertypes);
+
   /// Associates decorated type information with the given [element].
   void recordDecoratedElementType(Element element, DecoratedType type);
 
@@ -392,6 +485,11 @@
 /// results of the first ([NodeBuilder], which finds all the
 /// variables that need to be constrained).
 abstract class VariableRepository {
+  /// Given a [class_], gets the decorated type information for the superclasses
+  /// it directly implements/extends/etc.
+  Map<ClassElement, DecoratedType> decoratedDirectSupertypes(
+      ClassElement class_);
+
   /// Retrieves the [DecoratedType] associated with the static type of the given
   /// [element].
   ///
diff --git a/pkg/nnbd_migration/lib/src/variables.dart b/pkg/nnbd_migration/lib/src/variables.dart
index 03a2055..390a135 100644
--- a/pkg/nnbd_migration/lib/src/variables.dart
+++ b/pkg/nnbd_migration/lib/src/variables.dart
@@ -4,6 +4,7 @@
 
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer/src/dart/element/handle.dart';
 import 'package:analyzer/src/generated/source.dart';
 import 'package:nnbd_migration/src/conditional_discard.dart';
 import 'package:nnbd_migration/src/decorated_type.dart';
@@ -17,6 +18,9 @@
 
   final _decoratedElementTypes = <Element, DecoratedType>{};
 
+  final _decoratedDirectSupertypes =
+      <ClassElement, Map<ClassElement, DecoratedType>>{};
+
   final _decoratedTypeAnnotations =
       <Source, Map<int, DecoratedTypeAnnotation>>{};
 
@@ -25,6 +29,14 @@
   Variables(this._graph);
 
   @override
+  Map<ClassElement, DecoratedType> decoratedDirectSupertypes(
+      ClassElement class_) {
+    assert(class_ is! ClassElementHandle);
+    return _decoratedDirectSupertypes[class_] ??
+        _decorateDirectSupertypes(class_);
+  }
+
+  @override
   DecoratedType decoratedElementType(Element element, {bool create: false}) =>
       _decoratedElementTypes[element] ??= create
           ? DecoratedType.forElement(element, _graph)
@@ -51,6 +63,19 @@
         source, ConditionalModification(node, conditionalDiscard));
   }
 
+  @override
+  void recordDecoratedDirectSupertypes(ClassElement class_,
+      Map<ClassElement, DecoratedType> decoratedDirectSupertypes) {
+    assert(() {
+      assert(class_ is! ClassElementHandle);
+      for (var key in decoratedDirectSupertypes.keys) {
+        assert(key is! ClassElementHandle);
+      }
+      return true;
+    }());
+    _decoratedDirectSupertypes[class_] = decoratedDirectSupertypes;
+  }
+
   void recordDecoratedElementType(Element element, DecoratedType type) {
     _decoratedElementTypes[element] = type;
   }
@@ -130,6 +155,20 @@
     (_potentialModifications[source] ??= []).add(potentialModification);
   }
 
+  /// Creates an entry [_decoratedDirectSupertypes] for an already-migrated
+  /// class.
+  Map<ClassElement, DecoratedType> _decorateDirectSupertypes(
+      ClassElement class_) {
+    if (class_.type.isObject) {
+      // TODO(paulberry): this special case is just to get the basic
+      // infrastructure working (necessary since all classes derive from
+      // Object).  Once we have the full implementation this case shouldn't be
+      // needed.
+      return const {};
+    }
+    throw UnimplementedError('TODO(paulberry)');
+  }
+
   int _uniqueOffsetForTypeAnnotation(TypeAnnotation typeAnnotation) =>
       typeAnnotation is GenericFunctionType
           ? typeAnnotation.functionKeyword.offset
diff --git a/pkg/nnbd_migration/test/abstract_single_unit.dart b/pkg/nnbd_migration/test/abstract_single_unit.dart
index 25c5f87..2023c7e 100644
--- a/pkg/nnbd_migration/test/abstract_single_unit.dart
+++ b/pkg/nnbd_migration/test/abstract_single_unit.dart
@@ -8,6 +8,7 @@
 import 'package:analyzer/error/error.dart';
 import 'package:analyzer/src/error/codes.dart';
 import 'package:analyzer/src/generated/source.dart';
+import 'package:analyzer/src/test_utilities/find_element.dart';
 import 'package:analyzer/src/test_utilities/find_node.dart';
 import 'package:test/test.dart';
 
@@ -26,6 +27,7 @@
   CompilationUnitElement testUnitElement;
   LibraryElement testLibraryElement;
   FindNode findNode;
+  FindElement findElement;
 
   void addTestSource(String code, [Uri uri]) {
     testCode = code;
@@ -50,6 +52,7 @@
     testUnitElement = testUnit.declaredElement;
     testLibraryElement = testUnitElement.library;
     findNode = FindNode(code, testUnit);
+    findElement = FindElement(testUnit);
   }
 
   @override
diff --git a/pkg/nnbd_migration/test/decorated_class_hierarchy_test.dart b/pkg/nnbd_migration/test/decorated_class_hierarchy_test.dart
new file mode 100644
index 0000000..9e09dca
--- /dev/null
+++ b/pkg/nnbd_migration/test/decorated_class_hierarchy_test.dart
@@ -0,0 +1,157 @@
+// Copyright (c) 2019, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:analyzer/dart/ast/ast.dart';
+import 'package:nnbd_migration/src/decorated_class_hierarchy.dart';
+import 'package:nnbd_migration/src/nullability_node.dart';
+import 'package:test/test.dart';
+import 'package:test_reflective_loader/test_reflective_loader.dart';
+
+import 'migration_visitor_test_base.dart';
+
+main() {
+  defineReflectiveSuite(() {
+    defineReflectiveTests(DecoratedClassHierarchyTest);
+  });
+}
+
+@reflectiveTest
+class DecoratedClassHierarchyTest extends MigrationVisitorTestBase {
+  DecoratedClassHierarchy _hierarchy;
+
+  @override
+  Future<CompilationUnit> analyze(String code) async {
+    var unit = await super.analyze(code);
+    _hierarchy = DecoratedClassHierarchy(variables, graph);
+    return unit;
+  }
+
+  test_getDecoratedSupertype_complex() async {
+    await analyze('''
+class Base<T> {}
+class Intermediate<U> extends Base<List<U>> {}
+class Derived<V> extends Intermediate<Map<int, V>> {}
+''');
+    var decoratedSupertype = _hierarchy.getDecoratedSupertype(
+        findElement.class_('Derived'), findElement.class_('Base'));
+    var listRef = decoratedTypeAnnotation('List');
+    var uRef = decoratedTypeAnnotation('U>>');
+    var mapRef = decoratedTypeAnnotation('Map');
+    var intRef = decoratedTypeAnnotation('int');
+    var vRef = decoratedTypeAnnotation('V>>');
+    expect(decoratedSupertype.type.toString(), 'Base<List<Map<int, V>>>');
+    expect(decoratedSupertype.node, same(never));
+    var baseArgs = decoratedSupertype.typeArguments;
+    expect(baseArgs, hasLength(1));
+    expect(baseArgs[0].type.toString(), 'List<Map<int, V>>');
+    expect(baseArgs[0].node, same(listRef.node));
+    var listArgs = baseArgs[0].typeArguments;
+    expect(listArgs, hasLength(1));
+    expect(listArgs[0].type.toString(), 'Map<int, V>');
+    var mapNode = listArgs[0].node as NullabilityNodeForSubstitution;
+    expect(mapNode.innerNode, same(mapRef.node));
+    expect(mapNode.outerNode, same(uRef.node));
+    var mapArgs = listArgs[0].typeArguments;
+    expect(mapArgs, hasLength(2));
+    expect(mapArgs[0].type.toString(), 'int');
+    expect(mapArgs[0].node, same(intRef.node));
+    expect(mapArgs[1].type.toString(), 'V');
+    expect(mapArgs[1].node, same(vRef.node));
+  }
+
+  test_getDecoratedSupertype_extends_simple() async {
+    await analyze('''
+class Base<T, U> {}
+class Derived<V, W> extends Base<V, W> {}
+''');
+    var decoratedSupertype = _hierarchy.getDecoratedSupertype(
+        findElement.class_('Derived'), findElement.class_('Base'));
+    var vRef = decoratedTypeAnnotation('V, W> {');
+    var wRef = decoratedTypeAnnotation('W> {');
+    expect(decoratedSupertype.type.toString(), 'Base<V, W>');
+    expect(decoratedSupertype.node, same(never));
+    expect(decoratedSupertype.typeArguments, hasLength(2));
+    expect(decoratedSupertype.typeArguments[0].type.toString(), 'V');
+    expect(decoratedSupertype.typeArguments[0].node, same(vRef.node));
+    expect(decoratedSupertype.typeArguments[1].type.toString(), 'W');
+    expect(decoratedSupertype.typeArguments[1].node, same(wRef.node));
+  }
+
+  test_getDecoratedSupertype_implements_simple() async {
+    await analyze('''
+class Base<T, U> {}
+class Derived<V, W> implements Base<V, W> {}
+''');
+    var decoratedSupertype = _hierarchy.getDecoratedSupertype(
+        findElement.class_('Derived'), findElement.class_('Base'));
+    var vRef = decoratedTypeAnnotation('V, W> {');
+    var wRef = decoratedTypeAnnotation('W> {');
+    expect(decoratedSupertype.type.toString(), 'Base<V, W>');
+    expect(decoratedSupertype.node, same(never));
+    expect(decoratedSupertype.typeArguments, hasLength(2));
+    expect(decoratedSupertype.typeArguments[0].type.toString(), 'V');
+    expect(decoratedSupertype.typeArguments[0].node, same(vRef.node));
+    expect(decoratedSupertype.typeArguments[1].type.toString(), 'W');
+    expect(decoratedSupertype.typeArguments[1].node, same(wRef.node));
+  }
+
+  test_getDecoratedSupertype_not_generic() async {
+    await analyze('''
+class Base {}
+class Derived<T> extends Base {}
+''');
+    var decoratedSupertype = _hierarchy.getDecoratedSupertype(
+        findElement.class_('Derived'), findElement.class_('Base'));
+    expect(decoratedSupertype.type.toString(), 'Base');
+    expect(decoratedSupertype.node, same(never));
+    expect(decoratedSupertype.typeArguments, isEmpty);
+  }
+
+  test_getDecoratedSupertype_on_simple() async {
+    await analyze('''
+class Base<T, U> {}
+mixin Derived<V, W> on Base<V, W> {}
+''');
+    var decoratedSupertype = _hierarchy.getDecoratedSupertype(
+        findElement.mixin('Derived'), findElement.class_('Base'));
+    var vRef = decoratedTypeAnnotation('V, W> {');
+    var wRef = decoratedTypeAnnotation('W> {');
+    expect(decoratedSupertype.type.toString(), 'Base<V, W>');
+    expect(decoratedSupertype.node, same(never));
+    expect(decoratedSupertype.typeArguments, hasLength(2));
+    expect(decoratedSupertype.typeArguments[0].type.toString(), 'V');
+    expect(decoratedSupertype.typeArguments[0].node, same(vRef.node));
+    expect(decoratedSupertype.typeArguments[1].type.toString(), 'W');
+    expect(decoratedSupertype.typeArguments[1].node, same(wRef.node));
+  }
+
+  test_getDecoratedSupertype_unrelated_type() async {
+    await analyze('''
+class A<T> {}
+class B<T> {}
+''');
+    expect(
+        () => _hierarchy.getDecoratedSupertype(
+            findElement.class_('A'), findElement.class_('B')),
+        throwsA(TypeMatcher<StateError>()));
+  }
+
+  test_getDecoratedSupertype_with_simple() async {
+    await analyze('''
+class Base<T, U> {}
+class Derived<V, W> extends Object with Base<V, W> {}
+''');
+    var decoratedSupertype = _hierarchy.getDecoratedSupertype(
+        findElement.class_('Derived'), findElement.class_('Base'));
+    var vRef = decoratedTypeAnnotation('V, W> {');
+    var wRef = decoratedTypeAnnotation('W> {');
+    expect(decoratedSupertype.type.toString(), 'Base<V, W>');
+    expect(decoratedSupertype.node, same(never));
+    expect(decoratedSupertype.typeArguments, hasLength(2));
+    expect(decoratedSupertype.typeArguments[0].type.toString(), 'V');
+    expect(decoratedSupertype.typeArguments[0].node, same(vRef.node));
+    expect(decoratedSupertype.typeArguments[1].type.toString(), 'W');
+    expect(decoratedSupertype.typeArguments[1].node, same(wRef.node));
+  }
+}
diff --git a/pkg/nnbd_migration/test/migration_visitor_test_base.dart b/pkg/nnbd_migration/test/migration_visitor_test_base.dart
index 2e2e9d1..d7f26f6 100644
--- a/pkg/nnbd_migration/test/migration_visitor_test_base.dart
+++ b/pkg/nnbd_migration/test/migration_visitor_test_base.dart
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analyzer/dart/ast/ast.dart';
+import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/src/generated/resolver.dart';
 import 'package:analyzer/src/generated/source.dart';
 import 'package:meta/meta.dart';
@@ -137,6 +138,10 @@
     fail('Expected union between $x and $y, not found');
   }
 
+  Map<ClassElement, DecoratedType> decoratedDirectSupertypes(String name) {
+    return variables.decoratedDirectSupertypes(findElement.classOrMixin(name));
+  }
+
   /// Gets the [DecoratedType] associated with the generic function type
   /// annotation whose text is [text].
   DecoratedType decoratedGenericFunctionTypeAnnotation(String text) {
diff --git a/pkg/nnbd_migration/test/node_builder_test.dart b/pkg/nnbd_migration/test/node_builder_test.dart
index ef4476f..6ed3113 100644
--- a/pkg/nnbd_migration/test/node_builder_test.dart
+++ b/pkg/nnbd_migration/test/node_builder_test.dart
@@ -41,6 +41,158 @@
     expect(decoratedType.node, same(never));
   }
 
+  test_directSupertypes_class_extends() async {
+    await analyze('''
+class C<T, U> {}
+class D<V> extends C<int, V> {}
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V> {').node));
+  }
+
+  test_directSupertypes_class_extends_default() async {
+    await analyze('''
+class C<T, U> {}
+''');
+    var types = decoratedDirectSupertypes('C');
+    var decorated = types[typeProvider.objectType.element];
+    expect(decorated.type.toString(), 'Object');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, isEmpty);
+  }
+
+  test_directSupertypes_class_implements() async {
+    await analyze('''
+class C<T, U> {}
+class D<V> implements C<int, V> {}
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V> {').node));
+  }
+
+  test_directSupertypes_class_with() async {
+    await analyze('''
+class C<T, U> {}
+class D<V> extends Object with C<int, V> {}
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V> {').node));
+  }
+
+  test_directSupertypes_classAlias_extends() async {
+    await analyze('''
+class M {}
+class C<T, U> {}
+class D<V> = C<int, V> with M;
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V> w').node));
+  }
+
+  test_directSupertypes_classAlias_implements() async {
+    await analyze('''
+class M {}
+class C<T, U> {}
+class D<V> = Object with M implements C<int, V>;
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V>;').node));
+  }
+
+  test_directSupertypes_classAlias_with() async {
+    await analyze('''
+class C<T, U> {}
+class D<V> = Object with C<int, V>;
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V>;').node));
+  }
+
+  test_directSupertypes_mixin_extends_default() async {
+    await analyze('''
+mixin C<T, U> {}
+''');
+    var types = decoratedDirectSupertypes('C');
+    var decorated = types[typeProvider.objectType.element];
+    expect(decorated.type.toString(), 'Object');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, isEmpty);
+  }
+
+  test_directSupertypes_mixin_implements() async {
+    await analyze('''
+class C<T, U> {}
+mixin D<V> implements C<int, V> {}
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V> {').node));
+  }
+
+  test_directSupertypes_mixin_on() async {
+    await analyze('''
+class C<T, U> {}
+mixin D<V> on C<int, V> {}
+''');
+    var types = decoratedDirectSupertypes('D');
+    var decorated = types[findElement.class_('C')];
+    expect(decorated.type.toString(), 'C<int, V>');
+    expect(decorated.node, same(never));
+    expect(decorated.typeArguments, hasLength(2));
+    expect(decorated.typeArguments[0].node,
+        same(decoratedTypeAnnotation('int').node));
+    expect(decorated.typeArguments[1].node,
+        same(decoratedTypeAnnotation('V> {').node));
+  }
+
   test_dynamic_type() async {
     await analyze('''
 dynamic f() {}
diff --git a/pkg/nnbd_migration/test/test_all.dart b/pkg/nnbd_migration/test/test_all.dart
index 4241cf2..a496a02 100644
--- a/pkg/nnbd_migration/test/test_all.dart
+++ b/pkg/nnbd_migration/test/test_all.dart
@@ -5,6 +5,7 @@
 import 'package:test_reflective_loader/test_reflective_loader.dart';
 
 import 'api_test.dart' as api_test;
+import 'decorated_class_hierarchy_test.dart' as decorated_class_hierarchy_test;
 import 'graph_builder_test.dart' as graph_builder_test;
 import 'node_builder_test.dart' as node_builder_test;
 import 'nullability_node_test.dart' as nullability_node_test;
@@ -12,6 +13,7 @@
 main() {
   defineReflectiveSuite(() {
     api_test.main();
+    decorated_class_hierarchy_test.main();
     graph_builder_test.main();
     node_builder_test.main();
     nullability_node_test.main();