Merge MethodElement(s) into augmented.

Change-Id: I4965554d33f4345053c2b5abc05ba463a362b641
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/313301
Reviewed-by: Samuel Rawlins <srawlins@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>

pkg/analyzer/lib/src/dart/analysis/driver.dart

@@ -87,7 +87,7 @@ import 'package:analyzer/src/utilities/uri_cache.dart';
 /// TODO(scheglov) Clean up the list of implicitly analyzed files.
 class AnalysisDriver implements AnalysisDriverGeneric {
   /// The version of data format, should be incremented on every format change.
-  static const int DATA_VERSION = 284;
+  static const int DATA_VERSION = 285;
 
   /// The number of exception contexts allowed to write. Once this field is
   /// zero, we stop writing any new exception contexts in this process.

pkg/analyzer/lib/src/dart/element/element.dart

@@ -124,6 +124,9 @@ class AugmentedInlineClassElementImpl extends AugmentedNamedInstanceElementImpl
 
 abstract class AugmentedInstanceElementImpl
     implements AugmentedInstanceElement {
+  @override
+  List<MethodElement> methods = [];
+
   @override
   // TODO: implement accessors
   List<PropertyAccessorElement> get accessors => throw UnimplementedError();
@@ -136,10 +139,6 @@ abstract class AugmentedInstanceElementImpl
   // TODO: implement metadata
   List<ElementAnnotation> get metadata => throw UnimplementedError();
 
-  @override
-  // TODO: implement methods
-  List<MethodElement> get methods => throw UnimplementedError();
-
   @override
   FieldElement? getField(String name) {
     // TODO: implement getField

pkg/analyzer/lib/src/summary2/bundle_reader.dart

@@ -157,6 +157,7 @@ class ClassElementLinkedData extends ElementLinkedData<ClassElementImpl> {
       element.augmentedInternal = augmented;
       augmented.mixins = reader._readInterfaceTypeList();
       augmented.interfaces = reader._readInterfaceTypeList();
+      augmented.methods = reader.readElementList();
     }
 
     applyConstantOffsets?.perform();
@@ -248,7 +249,7 @@ abstract class ElementLinkedData<E extends ElementImpl> {
     ElementImpl element,
   ) {
     var enclosing = element.enclosingElement2;
-    if (enclosing is InterfaceElement) {
+    if (enclosing is InstanceOrAugmentationElement) {
       reader._addTypeParameters(enclosing.typeParameters);
     } else if (enclosing is CompilationUnitElement) {
       // Nothing.
@@ -1739,6 +1740,7 @@ class MixinElementLinkedData extends ElementLinkedData<MixinElementImpl> {
       element.augmentedInternal = augmented;
       augmented.superclassConstraints = reader._readInterfaceTypeList();
       augmented.interfaces = reader._readInterfaceTypeList();
+      augmented.methods = reader.readElementList();
     }
 
     applyConstantOffsets?.perform();
@@ -1840,6 +1842,10 @@ class ResolutionReader {
     throw UnimplementedError('memberFlags: $memberFlags');
   }
 
+  List<T> readElementList<T extends Element>() {
+    return _reader.readTypedListCast<T>(readElement);
+  }
+
   FunctionType? readOptionalFunctionType() {
     var type = readType();
     return type is FunctionType ? type : null;

pkg/analyzer/lib/src/summary2/bundle_writer.dart

@@ -187,6 +187,7 @@ class BundleWriter {
         (augmented) {
           _resolutionSink._writeTypeList(augmented.mixins);
           _resolutionSink._writeTypeList(augmented.interfaces);
+          _resolutionSink._writeElementList(augmented.methods);
         },
       );
 
@@ -476,6 +477,7 @@ class BundleWriter {
         (augmented) {
           _resolutionSink._writeTypeList(augmented.superclassConstraints);
           _resolutionSink._writeTypeList(augmented.interfaces);
+          _resolutionSink._writeElementList(augmented.methods);
         },
       );
 
@@ -811,6 +813,13 @@ class ResolutionSink extends _SummaryDataWriter {
     writeUInt30(elementIndex);
   }
 
+  void _writeElementList(List<Element> elements) {
+    writeUInt30(elements.length);
+    for (final element in elements) {
+      writeElement(element);
+    }
+  }
+
   void _writeFormalParameters(
     List<ParameterElement> parameters, {
     required bool withAnnotations,
@@ -935,11 +944,7 @@ class ResolutionSink extends _SummaryDataWriter {
     }
 
     var enclosing = declaration.enclosingElement2;
-    if (enclosing is TypeParameterizedElement) {
-      if (enclosing is! InterfaceElement && enclosing is! ExtensionElement) {
-        return const <DartType>[];
-      }
-
+    if (enclosing is InstanceOrAugmentationElement) {
       var typeParameters = enclosing.typeParameters;
       if (typeParameters.isEmpty) {
         return const <DartType>[];

pkg/analyzer/lib/src/summary2/data_reader.dart

@@ -126,6 +126,16 @@ class SummaryDataReader {
     }, growable: false);
   }
 
+  List<T> readTypedListCast<T>(Object? Function() read) {
+    var length = readUInt30();
+    if (length == 0) {
+      return const <Never>[];
+    }
+    return List<T>.generate(length, (_) {
+      return read() as T;
+    }, growable: false);
+  }
+
   int readUInt30() {
     var byte = readByte();
     if (byte & 0x80 == 0) {

pkg/analyzer/lib/src/summary2/reference_resolver.dart

@@ -62,10 +62,12 @@ class ReferenceResolver extends ThrowingAstVisitor<void> {
     node.withClause?.accept(this);
     node.implementsClause?.accept(this);
 
-    // TODO(scheglov) implements
-    // scope = InterfaceScope(scope, element);
-    // LinkingNodeContext(node, scope);
-    // node.members.accept(this);
+    final declaration = element.augmentedDeclaration;
+    if (declaration != null) {
+      scope = InterfaceScope(scope, declaration);
+      LinkingNodeContext(node, scope);
+      node.members.accept(this);
+    }
 
     nodesToBuildType.addDeclaration(node);
     scope = outerScope;
@@ -344,10 +346,12 @@ class ReferenceResolver extends ThrowingAstVisitor<void> {
     node.onClause?.accept(this);
     node.implementsClause?.accept(this);
 
-    // TODO(scheglov) implements
-    // scope = InterfaceScope(scope, element);
-    // LinkingNodeContext(node, scope);
-    // node.members.accept(this);
+    final declaration = element.augmentedDeclaration;
+    if (declaration != null) {
+      scope = InterfaceScope(scope, declaration);
+      LinkingNodeContext(node, scope);
+      node.members.accept(this);
+    }
 
     nodesToBuildType.addDeclaration(node);
     scope = outerScope;

pkg/analyzer/lib/src/summary2/types_builder.dart

@@ -10,6 +10,7 @@ import 'package:analyzer/src/dart/ast/ast.dart';
 import 'package:analyzer/src/dart/ast/extensions.dart';
 import 'package:analyzer/src/dart/element/class_hierarchy.dart';
 import 'package:analyzer/src/dart/element/element.dart';
+import 'package:analyzer/src/dart/element/member.dart';
 import 'package:analyzer/src/dart/element/type.dart';
 import 'package:analyzer/src/dart/element/type_algebra.dart';
 import 'package:analyzer/src/dart/element/type_system.dart';
@@ -128,19 +129,7 @@ class TypesBuilder {
       node.implementsClause?.interfaces,
     );
 
-    _updatedAugmented(
-      element,
-      element.augmentedDeclaration,
-      (declaration, substitution) {
-        final augmented = declaration.augmented;
-        augmented.mixins.addAll(
-          substitution.mapInterfaceTypes(element.mixins),
-        );
-        augmented.interfaces.addAll(
-          substitution.mapInterfaceTypes(element.interfaces),
-        );
-      },
-    );
+    _updatedAugmented(element);
   }
 
   void _classDeclaration(ClassDeclaration node) {
@@ -173,6 +162,7 @@ class TypesBuilder {
       element.augmentedInternal = augmented;
       augmented.mixins.addAll(element.mixins);
       augmented.interfaces.addAll(element.interfaces);
+      augmented.methods.addAll(element.methods);
     }
   }
 
@@ -346,19 +336,7 @@ class TypesBuilder {
       node.implementsClause?.interfaces,
     );
 
-    _updatedAugmented(
-      element,
-      element.augmentedDeclaration,
-      (declaration, substitution) {
-        final augmented = declaration.augmented;
-        augmented.superclassConstraints.addAll(
-          substitution.mapInterfaceTypes(element.superclassConstraints),
-        );
-        augmented.interfaces.addAll(
-          substitution.mapInterfaceTypes(element.interfaces),
-        );
-      },
-    );
+    _updatedAugmented(element);
   }
 
   void _mixinDeclaration(MixinDeclaration node) {
@@ -381,6 +359,7 @@ class TypesBuilder {
       element.augmentedInternal = augmented;
       augmented.superclassConstraints.addAll(element.superclassConstraints);
       augmented.interfaces.addAll(element.interfaces);
+      augmented.methods.addAll(element.methods);
     }
   }
 
@@ -424,9 +403,8 @@ class TypesBuilder {
 
   void _updatedAugmented<Declaration extends InstanceElementImpl>(
     InstanceAugmentationElementImpl element,
-    Declaration? declaration,
-    void Function(Declaration declaration, Substitution substitution) f,
   ) {
+    final declaration = element.augmentedDeclaration;
     if (declaration == null) {
       return;
     }
@@ -446,7 +424,38 @@ class TypesBuilder {
       }).toList(),
     );
 
-    f(declaration, substitution);
+    final typeProvider = element.library.typeProvider;
+    final augmented = declaration.augmented;
+
+    if (element is InterfaceAugmentationElementImpl &&
+        augmented is AugmentedInterfaceElementImpl) {
+      augmented.mixins.addAll(
+        substitution.mapInterfaceTypes(element.mixins),
+      );
+      augmented.interfaces.addAll(
+        substitution.mapInterfaceTypes(element.interfaces),
+      );
+    }
+
+    if (element is MixinAugmentationElementImpl &&
+        augmented is AugmentedMixinElementImpl) {
+      augmented.superclassConstraints.addAll(
+        substitution.mapInterfaceTypes(element.superclassConstraints),
+      );
+    }
+
+    if (augmented is AugmentedInstanceElementImpl) {
+      MethodElement mapMethodElement(MethodElement element) {
+        if (substitution.map.isEmpty) {
+          return element;
+        }
+        return MethodMember(typeProvider, element, substitution, false);
+      }
+
+      augmented.methods.addAll(
+        element.methods.map(mapMethodElement),
+      );
+    }
   }
 
   /// The [FunctionType] to use when a function type is expected for a type

pkg/analyzer/test/src/summary/element_text.dart

@@ -207,19 +207,33 @@ class _ElementWriter {
       }
     }
 
+    final augmented = e.augmented;
+
+    void writeMethods() {
+      _writeElements('methods', augmented.methods, (element) {
+        if (element.enclosingElement2 == e) {
+          _writeMethodElement(element);
+        } else {
+          _sink.writeIndent();
+          _elementPrinter.writeElement(element);
+        }
+      });
+    }
+
     _sink.writelnWithIndent('augmented');
     _sink.withIndent(() {
-      final augmented = e.augmented;
       switch (augmented) {
         case AugmentedClassElement():
           _elementPrinter.writeTypeList('mixins', augmented.mixins);
           _elementPrinter.writeTypeList('interfaces', augmented.interfaces);
+          writeMethods();
         case AugmentedMixinElement():
           _elementPrinter.writeTypeList(
             'superclassConstraints',
             augmented.superclassConstraints,
           );
           _elementPrinter.writeTypeList('interfaces', augmented.interfaces);
+          writeMethods();
       }
       // TODO(scheglov) Add other types and properties
     });
@@ -466,7 +480,7 @@ class _ElementWriter {
     _sink.withIndent(() {
       _writeElements('fields', e.fields, _writePropertyInducingElement);
       _writeElements('accessors', e.accessors, _writePropertyAccessorElement);
-      _writeElements('methods', e.methods, _writeMethodElement);
+      _writeMethods(e.methods);
     });
 
     _assertNonSyntheticElementSelf(e);
@@ -612,7 +626,7 @@ class _ElementWriter {
       }
 
       _writeElements('accessors', e.accessors, _writePropertyAccessorElement);
-      _writeElements('methods', e.methods, _writeMethodElement);
+      _writeMethods(e.methods);
 
       if (e is InterfaceElementImpl) {
         _writeAugmented(e);
@@ -711,6 +725,10 @@ class _ElementWriter {
     }
   }
 
+  void _writeMethods(List<MethodElement> elements) {
+    _writeElements('methods', elements, _writeMethodElement);
+  }
+
   void _writeName(Element e) {
     final String name;
     switch (e) {

pkg/analyzer/test/src/summary/elements_test.dart

@@ -347,6 +347,102 @@ library
 ''');
   }
 
+  test_augmented_methods() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+augment class A {
+  void bar() {}
+}
+''');
+
+    var library = await buildLibrary(r'''
+import augment 'a.dart';
+class A {
+  void foo() {}
+}
+''');
+
+    checkElementText(library, r'''
+library
+  definingUnit
+    classes
+      class A @31
+        augmentation: self::@augmentation::package:test/a.dart::@classAugmentation::A
+        constructors
+          synthetic @-1
+        methods
+          foo @42
+            returnType: void
+        augmented
+          methods
+            foo @42
+              returnType: void
+            self::@augmentation::package:test/a.dart::@classAugmentation::A::@method::bar
+  augmentationImports
+    package:test/a.dart
+      definingUnit
+        classAugmentations
+          augment class A @43
+            augmentationTarget: self::@class::A
+            augmentedDeclaration: self::@class::A
+            methods
+              bar @54
+                returnType: void
+''');
+  }
+
+  test_augmented_methods_generic() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+augment class A<T2> {
+  T2 bar() => throw 0;
+}
+''');
+
+    var library = await buildLibrary(r'''
+import augment 'a.dart';
+class A<T> {
+  T foo() => throw 0;
+}
+''');
+
+    checkElementText(library, r'''
+library
+  definingUnit
+    classes
+      class A @31
+        typeParameters
+          covariant T @33
+            defaultType: dynamic
+        augmentation: self::@augmentation::package:test/a.dart::@classAugmentation::A
+        constructors
+          synthetic @-1
+        methods
+          foo @42
+            returnType: T
+        augmented
+          methods
+            foo @42
+              returnType: T
+            MethodMember
+              base: self::@augmentation::package:test/a.dart::@classAugmentation::A::@method::bar
+              substitution: {T2: T}
+  augmentationImports
+    package:test/a.dart
+      definingUnit
+        classAugmentations
+          augment class A @43
+            typeParameters
+              covariant T2 @45
+                defaultType: dynamic
+            augmentationTarget: self::@class::A
+            augmentedDeclaration: self::@class::A
+            methods
+              bar @56
+                returnType: T2
+''');
+  }
+
   test_augmented_mixins() async {
     newFile('$testPackageLibPath/a.dart', r'''
 library augment 'test.dart';
@@ -45462,6 +45558,105 @@ library
 ''');
   }
 
+  test_augmented_methods() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+augment mixin A {
+  void bar() {}
+}
+''');
+
+    var library = await buildLibrary(r'''
+import augment 'a.dart';
+mixin A {
+  void foo() {}
+}
+''');
+
+    checkElementText(library, r'''
+library
+  definingUnit
+    mixins
+      mixin A @31
+        augmentation: self::@augmentation::package:test/a.dart::@mixinAugmentation::A
+        superclassConstraints
+          Object
+        methods
+          foo @42
+            returnType: void
+        augmented
+          superclassConstraints
+            Object
+          methods
+            foo @42
+              returnType: void
+            self::@augmentation::package:test/a.dart::@mixinAugmentation::A::@method::bar
+  augmentationImports
+    package:test/a.dart
+      definingUnit
+        mixinAugmentations
+          augment mixin A @43
+            augmentationTarget: self::@mixin::A
+            augmentedDeclaration: self::@mixin::A
+            methods
+              bar @54
+                returnType: void
+''');
+  }
+
+  test_augmented_methods_generic() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+library augment 'test.dart';
+augment mixin A<T2> {
+  T2 bar() => throw 0;
+}
+''');
+
+    var library = await buildLibrary(r'''
+import augment 'a.dart';
+mixin A<T> {
+  T foo() => throw 0;
+}
+''');
+
+    checkElementText(library, r'''
+library
+  definingUnit
+    mixins
+      mixin A @31
+        typeParameters
+          covariant T @33
+            defaultType: dynamic
+        augmentation: self::@augmentation::package:test/a.dart::@mixinAugmentation::A
+        superclassConstraints
+          Object
+        methods
+          foo @42
+            returnType: T
+        augmented
+          superclassConstraints
+            Object
+          methods
+            foo @42
+              returnType: T
+            MethodMember
+              base: self::@augmentation::package:test/a.dart::@mixinAugmentation::A::@method::bar
+              substitution: {T2: T}
+  augmentationImports
+    package:test/a.dart
+      definingUnit
+        mixinAugmentations
+          augment mixin A @43
+            typeParameters
+              covariant T2 @45
+            augmentationTarget: self::@mixin::A
+            augmentedDeclaration: self::@mixin::A
+            methods
+              bar @56
+                returnType: T2
+''');
+  }
+
   test_augmented_superclassConstraints() async {
     newFile('$testPackageLibPath/a.dart', r'''
 library augment 'test.dart';