Elements. Start moving toward element builders.

Change-Id: Iad5abe08363a658f48a24d7ba8103c5102314508
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/391497
Reviewed-by: Brian Wilkerson <[email protected]>
Reviewed-by: Phil Quitslund <[email protected]>
Commit-Queue: Konstantin Shcheglov <[email protected]>

Author: scheglov

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

@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer/dart/element/element2.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';
@@ -12,12 +13,10 @@ import 'package:analyzer/src/utilities/extensions/element.dart';
 import 'package:analyzer/src/utilities/extensions/string.dart';
 
 class AugmentedClassDeclarationBuilder
-    extends AugmentedInstanceDeclarationBuilder {
-  final ClassElementImpl firstFragment;
-
+    extends AugmentedInstanceDeclarationBuilder<ClassElementImpl> {
   AugmentedClassDeclarationBuilder({
+    required super.firstFragment,
     required MaybeAugmentedClassElementMixin element,
-    required this.firstFragment,
   }) {
     firstFragment.augmentedInternal = element;
     addFields(firstFragment.fields);
@@ -27,6 +26,9 @@ class AugmentedClassDeclarationBuilder
   }
 
   void augment(ClassElementImpl fragment) {
+    lastFragment.augmentation = fragment;
+    lastFragment = fragment;
+
     var element = _ensureAugmented(firstFragment);
     fragment.augmentedInternal = firstFragment.augmentedInternal;
 
@@ -39,21 +41,19 @@ class AugmentedClassDeclarationBuilder
 }
 
 class AugmentedEnumDeclarationBuilder
-    extends AugmentedInstanceDeclarationBuilder {
-  final EnumElementImpl declaration;
-
+    extends AugmentedInstanceDeclarationBuilder<EnumElementImpl> {
   AugmentedEnumDeclarationBuilder({
-    required this.declaration,
+    required super.firstFragment,
   }) {
-    addFields(declaration.fields);
-    addConstructors(declaration.constructors);
-    addAccessors(declaration.accessors);
-    addMethods(declaration.methods);
+    addFields(firstFragment.fields);
+    addConstructors(firstFragment.constructors);
+    addAccessors(firstFragment.accessors);
+    addMethods(firstFragment.methods);
   }
 
   void augment(EnumElementImpl element) {
-    var augmented = _ensureAugmented(declaration);
-    element.augmentedInternal = declaration.augmentedInternal;
+    var augmented = _ensureAugmented(firstFragment);
+    element.augmentedInternal = firstFragment.augmentedInternal;
 
     addFields(element.fields);
     addConstructors(element.constructors);
@@ -64,20 +64,18 @@ class AugmentedEnumDeclarationBuilder
 }
 
 class AugmentedExtensionDeclarationBuilder
-    extends AugmentedInstanceDeclarationBuilder {
-  final ExtensionElementImpl declaration;
-
+    extends AugmentedInstanceDeclarationBuilder<ExtensionElementImpl> {
   AugmentedExtensionDeclarationBuilder({
-    required this.declaration,
+    required super.firstFragment,
   }) {
-    addFields(declaration.fields);
-    addAccessors(declaration.accessors);
-    addMethods(declaration.methods);
+    addFields(firstFragment.fields);
+    addAccessors(firstFragment.accessors);
+    addMethods(firstFragment.methods);
   }
 
   void augment(ExtensionElementImpl element) {
-    var augmented = _ensureAugmented(declaration);
-    element.augmentedInternal = declaration.augmentedInternal;
+    var augmented = _ensureAugmented(firstFragment);
+    element.augmentedInternal = firstFragment.augmentedInternal;
 
     addFields(element.fields);
     addAccessors(element.accessors);
@@ -87,21 +85,19 @@ class AugmentedExtensionDeclarationBuilder
 }
 
 class AugmentedExtensionTypeDeclarationBuilder
-    extends AugmentedInstanceDeclarationBuilder {
-  final ExtensionTypeElementImpl declaration;
-
+    extends AugmentedInstanceDeclarationBuilder<ExtensionTypeElementImpl> {
   AugmentedExtensionTypeDeclarationBuilder({
-    required this.declaration,
+    required super.firstFragment,
   }) {
-    addFields(declaration.fields);
-    addConstructors(declaration.constructors);
-    addAccessors(declaration.accessors);
-    addMethods(declaration.methods);
+    addFields(firstFragment.fields);
+    addConstructors(firstFragment.constructors);
+    addAccessors(firstFragment.accessors);
+    addMethods(firstFragment.methods);
   }
 
   void augment(ExtensionTypeElementImpl element) {
-    var augmented = _ensureAugmented(declaration);
-    element.augmentedInternal = declaration.augmentedInternal;
+    var augmented = _ensureAugmented(firstFragment);
+    element.augmentedInternal = firstFragment.augmentedInternal;
 
     addFields(element.fields);
     addConstructors(element.constructors);
@@ -111,13 +107,18 @@ class AugmentedExtensionTypeDeclarationBuilder
   }
 }
 
-abstract class AugmentedInstanceDeclarationBuilder {
+abstract class AugmentedInstanceDeclarationBuilder<
+    F extends InstanceElementImpl> extends FragmentedElementBuilder<F> {
   final Map<String, FieldElementImpl> fields = {};
   final Map<String, ConstructorElementImpl> constructors = {};
   final Map<String, PropertyAccessorElementImpl> getters = {};
   final Map<String, PropertyAccessorElementImpl> setters = {};
   final Map<String, MethodElementImpl> methods = {};
 
+  AugmentedInstanceDeclarationBuilder({
+    required super.firstFragment,
+  });
+
   void addAccessors(List<PropertyAccessorElementImpl> elements) {
     for (var element in elements) {
       var name = element.name;
@@ -319,20 +320,18 @@ abstract class AugmentedInstanceDeclarationBuilder {
 }
 
 class AugmentedMixinDeclarationBuilder
-    extends AugmentedInstanceDeclarationBuilder {
-  final MixinElementImpl declaration;
-
+    extends AugmentedInstanceDeclarationBuilder<MixinElementImpl> {
   AugmentedMixinDeclarationBuilder({
-    required this.declaration,
+    required super.firstFragment,
   }) {
-    addFields(declaration.fields);
-    addAccessors(declaration.accessors);
-    addMethods(declaration.methods);
+    addFields(firstFragment.fields);
+    addAccessors(firstFragment.accessors);
+    addMethods(firstFragment.methods);
   }
 
   void augment(MixinElementImpl element) {
-    var augmented = _ensureAugmented(declaration);
-    element.augmentedInternal = declaration.augmentedInternal;
+    var augmented = _ensureAugmented(firstFragment);
+    element.augmentedInternal = firstFragment.augmentedInternal;
 
     addFields(element.fields);
     addAccessors(element.accessors);
@@ -402,6 +401,32 @@ class AugmentedTopVariablesBuilder {
   }
 }
 
+/// A builder for top-level fragmented elements, e.g. classes.
+class FragmentedElementBuilder<F extends Fragment> {
+  final F firstFragment;
+  F lastFragment;
+
+  FragmentedElementBuilder({
+    required this.firstFragment,
+  }) : lastFragment = firstFragment;
+
+  /// If [fragment] is an augmentation, set its previous fragment to
+  /// [lastFragment].
+  ///
+  /// We invoke this method on any [FragmentedElementBuilder] associated with
+  /// the name of [fragment], even if it is not a correct builder for this
+  /// [fragment]. So, the [lastFragment] might have a wrong type, but we still
+  /// want to remember it for generating the corresponding diagnostic.
+  void setPreviousFor(AugmentableElement fragment) {
+    if (fragment.isAugmentation) {
+      // TODO(scheglov): hopefully the type check can be removed in the future.
+      if (lastFragment case ElementImpl lastFragment) {
+        fragment.augmentationTargetAny = lastFragment;
+      }
+    }
+  }
+}
+
 extension<T extends ExecutableElement> on List<T> {
   Iterable<T> get notAugmented {
     return where((e) => e.augmentation == null);

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

@@ -140,25 +140,29 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
     node.withClause?.accept(this);
     node.implementsClause?.accept(this);
 
-    _libraryBuilder.updateAugmentationTarget(name, element);
+    // TODO(scheglov): remove it eventually
+    _libraryBuilder.updateAugmentationTarget0(name, element);
+
+    var augmentedBuilder = _libraryBuilder.getAugmentedBuilder(name);
+    augmentedBuilder?.setPreviousFor(element);
 
     // If the fragment is an augmentation, and the corresponding builder
     // has correct type, add the fragment to the builder. Otherwise, create
     // a new builder.
-    var augmentedBuilder = _libraryBuilder.getAugmentedBuilder(name);
-    if (element.augmentationTarget != null &&
+    if (element.isAugmentation &&
         augmentedBuilder is AugmentedClassDeclarationBuilder) {
       augmentedBuilder.augment(element);
     } else {
-      var containerRef = _libraryBuilder.reference.getChild('@class');
+      var libraryRef = _libraryBuilder.reference;
+      var containerRef = libraryRef.getChild('@class');
       var elementReference = containerRef.addChild(refName);
       var element2 = NotAugmentedClassElementImpl(elementReference, element);
 
       _libraryBuilder.putAugmentedBuilder(
         name,
         AugmentedClassDeclarationBuilder(
-          element: element2,
           firstFragment: element,
+          element: element2,
         ),
       );
     }
@@ -479,7 +483,7 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       _libraryBuilder.putAugmentedBuilder(
         name,
         AugmentedEnumDeclarationBuilder(
-          declaration: element,
+          firstFragment: element,
         ),
       );
     }
@@ -548,7 +552,7 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
         _libraryBuilder.putAugmentedBuilder(
           name,
           AugmentedExtensionDeclarationBuilder(
-            declaration: element,
+            firstFragment: element,
           ),
         );
       }
@@ -622,7 +626,7 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       _libraryBuilder.putAugmentedBuilder(
         name,
         AugmentedExtensionTypeDeclarationBuilder(
-          declaration: element,
+          firstFragment: element,
         ),
       );
     }
@@ -770,6 +774,10 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       }
 
       _libraryBuilder.topVariables.addAccessor(element);
+      _libraryBuilder.putAugmentedBuilder(
+        name,
+        FragmentedElementBuilder(firstFragment: element),
+      );
     } else if (node.isSetter) {
       var element = PropertyAccessorElementImpl(name, nameOffset);
       element.name2 = fragmentName;
@@ -785,6 +793,10 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       }
 
       _libraryBuilder.topVariables.addAccessor(element);
+      _libraryBuilder.putAugmentedBuilder(
+        '$name=',
+        FragmentedElementBuilder(firstFragment: element),
+      );
     } else {
       var element = FunctionElementImpl(name, nameOffset);
       element.name2 = fragmentName;
@@ -794,6 +806,10 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       executableElement = element;
 
       _libraryBuilder.updateAugmentationTarget(name, element);
+      _libraryBuilder.putAugmentedBuilder(
+        name,
+        FragmentedElementBuilder(firstFragment: element),
+      );
     }
 
     executableElement.hasImplicitReturnType = node.returnType == null;
@@ -1133,7 +1149,7 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
       _libraryBuilder.putAugmentedBuilder(
         name,
         AugmentedMixinDeclarationBuilder(
-          declaration: element,
+          firstFragment: element,
         ),
       );
     }

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

@@ -82,8 +82,7 @@ class LibraryBuilder with MacroApplicationsContainer {
       Map.identity();
 
   /// The top-level elements that can be augmented.
-  final Map<String, AugmentedInstanceDeclarationBuilder> _augmentedBuilders =
-      {};
+  final Map<String, FragmentedElementBuilder> _augmentedBuilders = {};
 
   /// The top-level variables and accessors that can be augmented.
   late final AugmentedTopVariablesBuilder topVariables =
@@ -472,7 +471,7 @@ class LibraryBuilder with MacroApplicationsContainer {
     }
   }
 
-  AugmentedInstanceDeclarationBuilder? getAugmentedBuilder(String name) {
+  FragmentedElementBuilder? getAugmentedBuilder(String name) {
     return _augmentedBuilders[name];
   }
 
@@ -654,9 +653,9 @@ class LibraryBuilder with MacroApplicationsContainer {
 
   void putAugmentedBuilder(
     String name,
-    AugmentedInstanceDeclarationBuilder element,
+    FragmentedElementBuilder builder,
   ) {
-    _augmentedBuilders[name] = element;
+    _augmentedBuilders[name] = builder;
   }
 
   void replaceConstFieldsIfNoConstConstructor() {
@@ -797,6 +796,14 @@ class LibraryBuilder with MacroApplicationsContainer {
     _augmentationTargets[name] = augmentation;
   }
 
+  // TODO(scheglov): Remove this method when we use only builders.
+  void updateAugmentationTarget0<T extends ElementImpl>(
+    String name,
+    AugmentableElement<T> augmentation,
+  ) {
+    _augmentationTargets[name] = augmentation;
+  }
+
   /// Updates the element of the macro augmentation.
   void updateInputMacroAugmentation() {
     if (inputMacroPartInclude case var import?) {

pkg/analyzer/test/src/diagnostics/augmentation_of_different_declaration_kind_test.dart

@@ -652,8 +652,7 @@ part of 'a.dart';
 
 augment class foo {}
 ''', [
-      error(CompileTimeErrorCode.AUGMENTATION_OF_DIFFERENT_DECLARATION_KIND, 19,
-          7),
+      error(CompileTimeErrorCode.AUGMENTATION_WITHOUT_DECLARATION, 19, 7),
     ]);
   }
 }