Elements. Add 'ElementImpl2.reference', implement for ClassElement2.

Change-Id: Ie555b0032bba93f9c3d0dcd4cdef819eec6e1187
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/391841
Reviewed-by: Phil Quitslund <pquitslund@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>

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 = 402;
+  static const int DATA_VERSION = 403;
 
   /// 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

@@ -108,10 +108,16 @@ mixin AugmentableElement<T extends ElementImpl> on ElementImpl {
 
 class AugmentedClassElementImpl extends AugmentedInterfaceElementImpl
     with MaybeAugmentedClassElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final ClassElementImpl declaration;
 
-  AugmentedClassElementImpl(this.declaration);
+  AugmentedClassElementImpl(this.reference, this.declaration) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 }
 
 class AugmentedEnumElementImpl extends AugmentedInterfaceElementImpl
@@ -209,8 +215,7 @@ class BindPatternVariableElementImpl2 extends PatternVariableElementImpl2
 class ClassElementImpl extends ClassOrMixinElementImpl
     with AugmentableElement<ClassElementImpl>
     implements ClassElement, ClassFragment {
-  late MaybeAugmentedClassElementMixin augmentedInternal =
-      NotAugmentedClassElementImpl(this);
+  late MaybeAugmentedClassElementMixin 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.
@@ -275,12 +280,6 @@ class ClassElementImpl extends ClassOrMixinElementImpl
 
   @override
   MaybeAugmentedClassElementMixin get augmented {
-    if (isAugmentation) {
-      if (augmentationTarget case var augmentationTarget?) {
-        return augmentationTarget.augmented;
-      }
-    }
-
     linkedData?.read(this);
     return augmentedInternal;
   }
@@ -306,7 +305,9 @@ class ClassElementImpl extends ClassOrMixinElementImpl
   }
 
   @override
-  ClassElement2 get element => super.element as ClassElement2;
+  MaybeAugmentedClassElementMixin get element {
+    return augmented;
+  }
 
   @override
   set fields(List<FieldElementImpl> fields) {
@@ -3176,6 +3177,13 @@ abstract class ElementImpl2 implements Element2 {
   @override
   Element2 get nonSynthetic2 => this;
 
+  /// The reference of this element, used during reading summaries.
+  ///
+  /// Can be `null` if this element cannot be referenced from outside,
+  /// for example a [LocalFunctionElement], a [TypeParameterElement],
+  /// a positional [FormalParameterElement], etc.
+  Reference? get reference => null;
+
   @override
   AnalysisSession? get session {
     return enclosingElement2?.session;
@@ -6680,6 +6688,9 @@ mixin MaybeAugmentedClassElementMixin on MaybeAugmentedInterfaceElementMixin
   @override
   bool get isValidMixin => declaration.isValidMixin;
 
+  /// See [ElementImpl2.reference].
+  Reference get reference;
+
   @override
   T? accept2<T>(ElementVisitor2<T> visitor) {
     return visitor.visitClassElement(this);
@@ -8295,17 +8306,23 @@ abstract class NonParameterVariableElementImpl extends VariableElementImpl
 
 class NotAugmentedClassElementImpl extends NotAugmentedInterfaceElementImpl
     with MaybeAugmentedClassElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final ClassElementImpl element;
 
-  NotAugmentedClassElementImpl(this.element);
+  NotAugmentedClassElementImpl(this.reference, this.element) {
+    reference.element2 = this;
+    element.augmentedInternal = this;
+  }
 
   @override
   ClassElementImpl get declaration => element;
 
   @override
   AugmentedClassElementImpl toAugmented() {
-    var augmented = AugmentedClassElementImpl(declaration);
+    var augmented = AugmentedClassElementImpl(reference, declaration);
     declaration.augmentedInternal = augmented;
     return augmented;
   }
@@ -8889,6 +8906,7 @@ class PrefixElementImpl extends _ExistingElementImpl implements PrefixElement {
 }
 
 class PrefixElementImpl2 extends ElementImpl2 implements PrefixElement2 {
+  @override
   final Reference reference;
 
   @override

pkg/analyzer/lib/src/generated/testing/element_factory.dart

@@ -15,6 +15,7 @@ import 'package:analyzer/src/dart/element/type.dart';
 import 'package:analyzer/src/generated/engine.dart';
 import 'package:analyzer/src/generated/source.dart' show NonExistingSource;
 import 'package:analyzer/src/generated/utilities_dart.dart';
+import 'package:analyzer/src/summary2/reference.dart';
 import 'package:collection/collection.dart';
 import 'package:meta/meta.dart';
 import 'package:path/path.dart';
@@ -44,6 +45,7 @@ class ElementFactory {
       String typeName, InterfaceType? superclassType,
       [List<String>? parameterNames]) {
     ClassElementImpl element = ClassElementImpl(typeName, 0);
+    NotAugmentedClassElementImpl(Reference.root(), element);
     element.constructors = const <ConstructorElementImpl>[];
     element.supertype = superclassType;
     if (parameterNames != null) {
@@ -68,6 +70,7 @@ class ElementFactory {
     supertype ??= objectType;
 
     var element = ClassElementImpl(name, 0);
+    NotAugmentedClassElementImpl(Reference.root(), element);
     element.typeParameters = typeParameters;
     element.supertype = supertype;
     element.mixins = mixins;

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

@@ -13,23 +13,28 @@ import 'package:analyzer/src/utilities/extensions/string.dart';
 
 class AugmentedClassDeclarationBuilder
     extends AugmentedInstanceDeclarationBuilder {
-  final ClassElementImpl declaration;
+  final ClassElementImpl firstFragment;
 
   AugmentedClassDeclarationBuilder({
-    required this.declaration,
+    required MaybeAugmentedClassElementMixin element,
+    required this.firstFragment,
   }) {
-    addFields(declaration.fields);
-    addConstructors(declaration.constructors);
-    addAccessors(declaration.accessors);
-    addMethods(declaration.methods);
+    firstFragment.augmentedInternal = element;
+    addFields(firstFragment.fields);
+    addConstructors(firstFragment.constructors);
+    addAccessors(firstFragment.accessors);
+    addMethods(firstFragment.methods);
   }
 
-  void augment(ClassElementImpl element) {
-    addFields(element.fields);
-    addConstructors(element.constructors);
-    addAccessors(element.accessors);
-    addMethods(element.methods);
-    _updatedAugmented(element);
+  void augment(ClassElementImpl fragment) {
+    var element = _ensureAugmented(firstFragment);
+    fragment.augmentedInternal = firstFragment.augmentedInternal;
+
+    addFields(fragment.fields);
+    addConstructors(fragment.constructors);
+    addAccessors(fragment.accessors);
+    addMethods(fragment.methods);
+    _updatedAugmented(element, fragment);
   }
 }
 
@@ -47,11 +52,14 @@ class AugmentedEnumDeclarationBuilder
   }
 
   void augment(EnumElementImpl element) {
+    var augmented = _ensureAugmented(declaration);
+    element.augmentedInternal = declaration.augmentedInternal;
+
     addFields(element.fields);
     addConstructors(element.constructors);
     addAccessors(element.accessors);
     addMethods(element.methods);
-    _updatedAugmented(element);
+    _updatedAugmented(augmented, element);
   }
 }
 
@@ -68,10 +76,13 @@ class AugmentedExtensionDeclarationBuilder
   }
 
   void augment(ExtensionElementImpl element) {
+    var augmented = _ensureAugmented(declaration);
+    element.augmentedInternal = declaration.augmentedInternal;
+
     addFields(element.fields);
     addAccessors(element.accessors);
     addMethods(element.methods);
-    _updatedAugmented(element);
+    _updatedAugmented(augmented, element);
   }
 }
 
@@ -89,11 +100,14 @@ class AugmentedExtensionTypeDeclarationBuilder
   }
 
   void augment(ExtensionTypeElementImpl element) {
+    var augmented = _ensureAugmented(declaration);
+    element.augmentedInternal = declaration.augmentedInternal;
+
     addFields(element.fields);
     addConstructors(element.constructors);
     addAccessors(element.accessors);
     addMethods(element.methods);
-    _updatedAugmented(element);
+    _updatedAugmented(augmented, element);
   }
 }
 
@@ -183,7 +197,7 @@ abstract class AugmentedInstanceDeclarationBuilder {
     }
   }
 
-  AugmentedInstanceElementImpl? _ensureAugmented(
+  AugmentedInstanceElementImpl _ensureAugmented(
     InstanceElementImpl augmentation,
   ) {
     var maybeAugmented = augmentation.augmented;
@@ -229,13 +243,10 @@ abstract class AugmentedInstanceDeclarationBuilder {
     return target;
   }
 
-  void _updatedAugmented(InstanceElementImpl augmentation) {
-    assert(augmentation.augmentationTarget != null);
-    var augmented = _ensureAugmented(augmentation);
-    if (augmented == null) {
-      return;
-    }
-
+  void _updatedAugmented(
+    AugmentedInstanceElementImpl augmented,
+    InstanceElementImpl augmentation,
+  ) {
     var declaration = augmented.declaration;
     var declarationTypeParameters = declaration.typeParameters;
 
@@ -320,10 +331,13 @@ class AugmentedMixinDeclarationBuilder
   }
 
   void augment(MixinElementImpl element) {
+    var augmented = _ensureAugmented(declaration);
+    element.augmentedInternal = declaration.augmentedInternal;
+
     addFields(element.fields);
     addAccessors(element.accessors);
     addMethods(element.methods);
-    _updatedAugmented(element);
+    _updatedAugmented(augmented, element);
   }
 }
 

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

@@ -154,7 +154,7 @@ class ClassElementLinkedData extends ElementLinkedData<ClassElementImpl> {
 
     if (element.augmentationTarget == null) {
       if (reader.readBool()) {
-        var augmented = AugmentedClassElementImpl(element);
+        var augmented = element.augmentedInternal as AugmentedClassElementImpl;
         element.augmentedInternal = augmented;
         augmented.mixins = reader._readInterfaceTypeList();
         augmented.interfaces = reader._readInterfaceTypeList();
@@ -754,12 +754,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 = ClassElementImpl(name, -1);
     element.name2 = fragmentName;
 
+    if (reference2.element2 case MaybeAugmentedClassElementMixin element2?) {
+      element.augmentedInternal = element2;
+    } else if (isAugmented) {
+      AugmentedClassElementImpl(reference2, element);
+    } else {
+      NotAugmentedClassElementImpl(reference2, element);
+    }
+
     var linkedData = ClassElementLinkedData(
       reference: reference,
       libraryReader: this,
@@ -1809,13 +1821,6 @@ class LibraryReader {
       );
     });
 
-    unitElement.parts = _reader.readTypedList(() {
-      return _readPartElement(
-        libraryElement: libraryElement,
-        containerUnit: unitElement,
-      );
-    });
-
     _readClasses(unitElement, unitReference);
     _readEnums(unitElement, unitReference);
     _readExtensions(unitElement, unitReference);
@@ -1839,6 +1844,13 @@ class LibraryReader {
       );
     });
 
+    unitElement.parts = _reader.readTypedList(() {
+      return _readPartElement(
+        libraryElement: libraryElement,
+        containerUnit: unitElement,
+      );
+    });
+
     return unitElement;
   }
 

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

@@ -156,6 +156,8 @@ class BundleWriter {
     _sink.writeUInt30(_resolutionSink.offset);
 
     _writeReference(element);
+    _writeReference2(element.augmented.reference);
+    _sink.writeBool(element.augmentedInternal is AugmentedClassElementImpl);
     _writeFragmentName(element.name2);
     ClassElementFlags.write(_sink, element);
 
@@ -755,8 +757,9 @@ class BundleWriter {
     _writeList(unitElement.libraryImports, _writeImportElement);
     _writeList(unitElement.libraryExports, _writeExportElement);
 
+    // Write the metadata for parts here, even though we write parts below.
+    // The reason is that resolution data must be in a single chunk.
     _writePartElementsMetadata(unitElement);
-    _writeList(unitElement.parts, _writePartElement);
 
     _writeList(unitElement.classes, _writeClassElement);
     _writeList(unitElement.enums, _writeEnumElement);
@@ -781,6 +784,10 @@ class BundleWriter {
       _sink.writeStringUtf8(macroGenerated.code);
       _sink.writeUint8List(macroGenerated.informativeBytes);
     });
+
+    // Write parts after this library fragment, so that when we read, we
+    // process fragments of declarations in the same order as we build them.
+    _writeList(unitElement.parts, _writePartElement);
   }
 
   static TypeParameterVarianceTag _encodeVariance(