Elements. Add reference for MixinElement2.

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

Author: scheglov

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

@@ -96,7 +96,7 @@ import 'package:meta/meta.dart';
 // TODO(scheglov): Clean up the list of implicitly analyzed files.
 class AnalysisDriver {
   /// The version of data format, should be incremented on every format change.
-  static const int DATA_VERSION = 403;
+  static const int DATA_VERSION = 404;
 
   /// 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

@@ -175,13 +175,19 @@ abstract class AugmentedInterfaceElementImpl
 
 class AugmentedMixinElementImpl extends AugmentedInterfaceElementImpl
     with MaybeAugmentedMixinElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final MixinElementImpl declaration;
 
   @override
   List<InterfaceType> superclassConstraints = [];
 
-  AugmentedMixinElementImpl(this.declaration);
+  AugmentedMixinElementImpl(this.reference, this.declaration) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 }
 
 class BindPatternVariableElementImpl extends PatternVariableElementImpl
@@ -7202,6 +7208,9 @@ mixin MaybeAugmentedMixinElementMixin on MaybeAugmentedInterfaceElementMixin
   @override
   bool get isBase => declaration.isBase;
 
+  /// See [ElementImpl2.reference].
+  Reference get reference;
+
   @override
   T? accept2<T>(ElementVisitor2<T> visitor) {
     return visitor.visitMixinElement(this);
@@ -7737,21 +7746,14 @@ class MixinElementImpl extends ClassOrMixinElementImpl
   /// The list will be empty if this class is not a mixin declaration.
   late List<String> superInvokedNames;
 
-  late MaybeAugmentedMixinElementMixin augmentedInternal =
-      NotAugmentedMixinElementImpl(this);
+  late MaybeAugmentedMixinElementMixin augmentedInternal;
 
   /// Initialize a newly created class element to have the given [name] at the
   /// given [offset] in the file that contains the declaration of this element.
   MixinElementImpl(super.name, super.offset);
 
   @override
   MaybeAugmentedMixinElementMixin get augmented {
-    if (isAugmentation) {
-      if (augmentationTarget case var augmentationTarget?) {
-        return augmentationTarget.augmented;
-      }
-    }
-
     linkedData?.read(this);
     return augmentedInternal;
   }
@@ -7771,7 +7773,7 @@ class MixinElementImpl extends ClassOrMixinElementImpl
   }
 
   @override
-  MixinElement2 get element => super.element as MixinElement2;
+  MaybeAugmentedMixinElementMixin get element => augmented;
 
   @override
   bool get isBase {
@@ -8468,10 +8470,16 @@ abstract class NotAugmentedInterfaceElementImpl
 
 class NotAugmentedMixinElementImpl extends NotAugmentedInterfaceElementImpl
     with MaybeAugmentedMixinElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final MixinElementImpl element;
 
-  NotAugmentedMixinElementImpl(this.element);
+  NotAugmentedMixinElementImpl(this.reference, this.element) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 
   @override
   MixinElementImpl get declaration => element;
@@ -8483,7 +8491,7 @@ class NotAugmentedMixinElementImpl extends NotAugmentedInterfaceElementImpl
 
   @override
   AugmentedMixinElementImpl toAugmented() {
-    var augmented = AugmentedMixinElementImpl(declaration);
+    var augmented = AugmentedMixinElementImpl(reference, declaration);
     declaration.augmentedInternal = augmented;
     return augmented;
   }

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

@@ -344,13 +344,18 @@ abstract class InstanceElementBuilder<F extends InstanceElementImpl>
 class MixinElementBuilder extends InstanceElementBuilder<MixinElementImpl> {
   MixinElementBuilder({
     required super.firstFragment,
+    required MaybeAugmentedMixinElementMixin element,
   }) {
+    firstFragment.augmentedInternal = element;
     addFields(firstFragment.fields);
     addAccessors(firstFragment.accessors);
     addMethods(firstFragment.methods);
   }
 
   void augment(MixinElementImpl fragment) {
+    lastFragment.augmentation = fragment;
+    lastFragment = fragment;
+
     var element = _ensureAugmented();
     fragment.augmentedInternal = firstFragment.augmentedInternal;
 

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

@@ -1369,12 +1369,24 @@ class LibraryReader {
   ) {
     var resolutionOffset = _baseResolutionOffset + _reader.readUInt30();
     var reference = _readReference();
+
+    var reference2 = _readReference();
+    var isAugmented = _reader.readBool();
+
     var fragmentName = _readFragmentName();
     var name = fragmentName?.name ?? '';
 
     var element = MixinElementImpl(name, -1);
     element.name2 = fragmentName;
 
+    if (reference2.element2 case MaybeAugmentedMixinElementMixin element2?) {
+      element.augmentedInternal = element2;
+    } else if (isAugmented) {
+      AugmentedMixinElementImpl(reference2, element);
+    } else {
+      NotAugmentedMixinElementImpl(reference2, element);
+    }
+
     var linkedData = MixinElementLinkedData(
       reference: reference,
       libraryReader: this,
@@ -1931,7 +1943,7 @@ class MixinElementLinkedData extends ElementLinkedData<MixinElementImpl> {
 
     if (element.augmentationTarget == null) {
       if (reader.readBool()) {
-        var augmented = AugmentedMixinElementImpl(element);
+        var augmented = element.augmentedInternal as AugmentedMixinElementImpl;
         element.augmentedInternal = augmented;
         augmented.superclassConstraints = reader._readInterfaceTypeList();
         augmented.interfaces = reader._readInterfaceTypeList();

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

@@ -542,6 +542,8 @@ class BundleWriter {
     _sink.writeUInt30(_resolutionSink.offset);
 
     _writeReference(element);
+    _writeReference2(element.augmented.reference);
+    _sink.writeBool(element.augmentedInternal is AugmentedMixinElementImpl);
     _writeFragmentName(element.name2);
     MixinElementFlags.write(_sink, element);
     _resolutionSink._writeAnnotationList(element.metadata);

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

@@ -1107,51 +1107,77 @@ class ElementBuilder extends ThrowingAstVisitor<void> {
     var name = nameToken.lexeme;
     var fragmentName = _buildFragmentName(nameToken);
 
-    var element = MixinElementImpl(name, nameToken.offset);
-    element.name2 = fragmentName;
-    element.isAugmentation = node.augmentKeyword != null;
-    element.isBase = node.baseKeyword != null;
-    element.metadata = _buildAnnotations(node.metadata);
-    _setCodeRange(element, node);
-    _setDocumentation(element, node);
+    var fragment = MixinElementImpl(name, nameToken.offset);
+    fragment.name2 = fragmentName;
+    fragment.isAugmentation = node.augmentKeyword != null;
+    fragment.isBase = node.baseKeyword != null;
+    fragment.metadata = _buildAnnotations(node.metadata);
+    _setCodeRange(fragment, node);
+    _setDocumentation(fragment, node);
 
-    node.declaredElement = element;
-    _linker.elementNodes[element] = node;
+    node.declaredElement = fragment;
+    _linker.elementNodes[fragment] = node;
 
     var refName = fragmentName?.name ?? '${_nextUnnamedId++}';
-    var reference = _enclosingContext.addMixin(refName, element);
-    if (!element.isAugmentation) {
+    var reference = _enclosingContext.addMixin(refName, fragment);
+    if (!fragment.isAugmentation) {
       _libraryBuilder.declare(name, reference);
     }
 
-    var holder = _EnclosingContext(reference, element);
+    var holder = _EnclosingContext(reference, fragment);
     _withEnclosing(holder, () {
       node.typeParameters?.accept(this);
       _visitPropertyFirst<FieldDeclaration>(node.members);
     });
-    element.typeParameters = holder.typeParameters;
-    element.accessors = holder.propertyAccessors;
-    element.fields = holder.fields;
-    element.methods = holder.methods;
+    fragment.typeParameters = holder.typeParameters;
+    fragment.accessors = holder.propertyAccessors;
+    fragment.fields = holder.fields;
+    fragment.methods = holder.methods;
 
     node.onClause?.accept(this);
     node.implementsClause?.accept(this);
 
-    _libraryBuilder.updateAugmentationTarget(name, element);
+    // TODO(scheglov): remove it eventually
+    _libraryBuilder.updateAugmentationTarget0(name, fragment);
 
-    if (element.augmentationTarget != null) {
-      switch (_libraryBuilder.getElementBuilder(name)) {
-        case MixinElementBuilder builder:
-          builder.augment(element);
-      }
+    var elementBuilder = _libraryBuilder.getElementBuilder(name);
+    elementBuilder?.setPreviousFor(fragment);
+
+    // If the fragment is an augmentation, and the corresponding builder
+    // has correct type, add the fragment to the builder. Otherwise, create
+    // a new builder.
+    if (fragment.isAugmentation && elementBuilder is MixinElementBuilder) {
+      elementBuilder.augment(fragment);
     } else {
+      var libraryRef = _libraryBuilder.reference;
+      var containerRef = libraryRef.getChild('@mixin');
+      var elementReference = containerRef.addChild(refName);
+      var element = NotAugmentedMixinElementImpl(elementReference, fragment);
+
       _libraryBuilder.putElementBuilder(
         name,
         MixinElementBuilder(
-          firstFragment: element,
+          firstFragment: fragment,
+          element: element,
         ),
       );
     }
+
+    // _libraryBuilder.updateAugmentationTarget(name, element);
+    //
+    // if (element.augmentationTarget != null) {
+    //   switch (_libraryBuilder.getElementBuilder(name)) {
+    //     case MixinElementBuilder builder:
+    //       builder.augment(element);
+    //   }
+    // } else {
+    //   _libraryBuilder.putElementBuilder(
+    //     name,
+    //     MixinElementBuilder(
+    //       firstFragment: element,
+    //     ),
+    //   );
+    // }
   }
 
   @override

pkg/analyzer/test/generated/elements_types_mixin.dart

@@ -414,6 +414,7 @@ mixin ElementsTypesMixin {
     List<InterfaceType> interfaces = const [],
   }) {
     var element = MixinElementImpl(name, 0);
+    NotAugmentedMixinElementImpl(Reference.root(), element);
     element.isAugmentation = isAugmentation;
     element.enclosingElement3 = testLibrary.definingCompilationUnit;
     element.typeParameters = typeParameters;
@@ -663,7 +664,7 @@ extension MixinElementImplExtension on MixinElementImpl {
   void addAugmentations(List<MixinElementImpl> augmentations) {
     expect(this.augmented, TypeMatcher<NotAugmentedMixinElementImpl>());
 
-    var augmented = AugmentedMixinElementImpl(this);
+    var augmented = AugmentedMixinElementImpl(Reference.root(), this);
     augmentedInternal = augmented;
 
     var augmentationTarget = this;

pkg/analyzer/test/src/dart/resolution/enum_test.dart

@@ -538,7 +538,7 @@ WithClause
     NamedType
       name: M
       element: <testLibraryFragment>::@mixin::M
-      element2: <testLibraryFragment>::@mixin::M#element
+      element2: <testLibrary>::@mixin::M
       type: M
 ''');
   }
@@ -569,12 +569,12 @@ WithClause
             type: int
         rightBracket: >
       element: <testLibraryFragment>::@mixin::M1
-      element2: <testLibraryFragment>::@mixin::M1#element
+      element2: <testLibrary>::@mixin::M1
       type: M1<int>
     NamedType
       name: M2
       element: <testLibraryFragment>::@mixin::M2
-      element2: <testLibraryFragment>::@mixin::M2#element
+      element2: <testLibrary>::@mixin::M2
       type: M2<int>
 ''');
   }

pkg/analyzer/test/src/dart/resolution/extension_method_test.dart

@@ -535,7 +535,7 @@ extension E on M {}
 NamedType
   name: M
   element: <testLibraryFragment>::@mixin::M
-  element2: <testLibraryFragment>::@mixin::M#element
+  element2: <testLibrary>::@mixin::M
   type: M
 ''');
   }
@@ -654,7 +654,7 @@ extension on M {}
 NamedType
   name: M
   element: <testLibraryFragment>::@mixin::M
-  element2: <testLibraryFragment>::@mixin::M#element
+  element2: <testLibrary>::@mixin::M
   type: M
 ''');
   }

pkg/analyzer/test/src/dart/resolution/mixin_test.dart

@@ -30,7 +30,7 @@ WithClause
     NamedType
       name: M
       element: <testLibraryFragment>::@mixin::M
-      element2: <testLibraryFragment>::@mixin::M#element
+      element2: <testLibrary>::@mixin::M
       type: M
 ''');
   }
@@ -49,7 +49,7 @@ WithClause
     NamedType
       name: M
       element: <testLibraryFragment>::@mixin::M
-      element2: <testLibraryFragment>::@mixin::M#element
+      element2: <testLibrary>::@mixin::M
       type: M
 ''');
   }