Elements. Add reference for enum, extension, extension type elements.

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

Author: scheglov

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

@@ -98,7 +98,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 = 404;
+  static const int DATA_VERSION = 405;
 
   /// 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

@@ -122,26 +122,44 @@ class AugmentedClassElementImpl extends AugmentedInterfaceElementImpl
 
 class AugmentedEnumElementImpl extends AugmentedInterfaceElementImpl
     with MaybeAugmentedEnumElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final EnumElementImpl declaration;
 
-  AugmentedEnumElementImpl(this.declaration);
+  AugmentedEnumElementImpl(this.reference, this.declaration) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 }
 
 class AugmentedExtensionElementImpl extends AugmentedInstanceElementImpl
     with MaybeAugmentedExtensionElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final ExtensionElementImpl declaration;
 
-  AugmentedExtensionElementImpl(this.declaration);
+  AugmentedExtensionElementImpl(this.reference, this.declaration) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 }
 
 class AugmentedExtensionTypeElementImpl extends AugmentedInterfaceElementImpl
     with MaybeAugmentedExtensionTypeElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final ExtensionTypeElementImpl declaration;
 
-  AugmentedExtensionTypeElementImpl(this.declaration);
+  AugmentedExtensionTypeElementImpl(this.reference, this.declaration) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 }
 
 abstract class AugmentedInstanceElementImpl
@@ -286,8 +304,7 @@ class ClassElementImpl extends ClassOrMixinElementImpl
 
   @override
   MaybeAugmentedClassElementMixin get augmented {
-    linkedData?.read(this);
-    return augmentedInternal;
+    return element;
   }
 
   AugmentedClassElementImpl? get augmentedIfReally {
@@ -312,7 +329,8 @@ class ClassElementImpl extends ClassOrMixinElementImpl
 
   @override
   MaybeAugmentedClassElementMixin get element {
-    return augmented;
+    linkedData?.read(this);
+    return augmentedInternal;
   }
 
   @override
@@ -3368,23 +3386,15 @@ class ElementLocationImpl implements ElementLocation {
 class EnumElementImpl extends InterfaceElementImpl
     with AugmentableElement<EnumElementImpl>
     implements EnumElement, EnumFragment {
-  late MaybeAugmentedEnumElementMixin augmentedInternal =
-      NotAugmentedEnumElementImpl(this);
+  late MaybeAugmentedEnumElementMixin 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.
   EnumElementImpl(super.name, super.offset);
 
   @override
   MaybeAugmentedEnumElementMixin get augmented {
-    if (isAugmentation) {
-      if (augmentationTarget case var augmentationTarget?) {
-        return augmentationTarget.augmented;
-      }
-    }
-
-    linkedData?.read(this);
-    return augmentedInternal;
+    return element;
   }
 
   AugmentedEnumElementImpl? get augmentedIfReally {
@@ -3409,7 +3419,10 @@ class EnumElementImpl extends InterfaceElementImpl
   List<FieldElement2> get constants2 => constants.cast<FieldElement2>();
 
   @override
-  EnumElement2 get element => super.element as EnumElement2;
+  MaybeAugmentedEnumElementMixin get element {
+    linkedData?.read(this);
+    return augmentedInternal;
+  }
 
   @override
   ElementKind get kind => ElementKind.ENUM;
@@ -3640,8 +3653,7 @@ abstract class ExecutableElementImpl2 extends FunctionTypedElementImpl2
 class ExtensionElementImpl extends InstanceElementImpl
     with AugmentableElement<ExtensionElementImpl>
     implements ExtensionElement, ExtensionFragment {
-  late MaybeAugmentedExtensionElementMixin augmentedInternal =
-      NotAugmentedExtensionElementImpl(this);
+  late MaybeAugmentedExtensionElementMixin augmentedInternal;
 
   /// Initialize a newly created extension element to have the given [name] at
   /// the given [offset] in the file that contains the declaration of this
@@ -3650,14 +3662,7 @@ class ExtensionElementImpl extends InstanceElementImpl
 
   @override
   MaybeAugmentedExtensionElementMixin get augmented {
-    if (isAugmentation) {
-      if (augmentationTarget case var augmentationTarget?) {
-        return augmentationTarget.augmented;
-      }
-    }
-
-    linkedData?.read(this);
-    return augmentedInternal;
+    return element;
   }
 
   AugmentedExtensionElementImpl? get augmentedIfReally {
@@ -3687,7 +3692,10 @@ class ExtensionElementImpl extends InstanceElementImpl
   String get displayName => name ?? '';
 
   @override
-  ExtensionElement2 get element => super.element as ExtensionElement2;
+  MaybeAugmentedExtensionElementMixin get element {
+    linkedData?.read(this);
+    return augmentedInternal;
+  }
 
   @override
   DartType get extendedType {
@@ -3769,8 +3777,7 @@ class ExtensionElementImpl extends InstanceElementImpl
 class ExtensionTypeElementImpl extends InterfaceElementImpl
     with AugmentableElement<ExtensionTypeElementImpl>
     implements ExtensionTypeElement, ExtensionTypeFragment {
-  late MaybeAugmentedExtensionTypeElementMixin augmentedInternal =
-      NotAugmentedExtensionTypeElementImpl(this);
+  late MaybeAugmentedExtensionTypeElementMixin augmentedInternal;
 
   /// Whether the element has direct or indirect reference to itself,
   /// in representation.
@@ -3784,14 +3791,7 @@ class ExtensionTypeElementImpl extends InterfaceElementImpl
 
   @override
   MaybeAugmentedExtensionTypeElementMixin get augmented {
-    if (isAugmentation) {
-      if (augmentationTarget case var augmentationTarget?) {
-        return augmentationTarget.augmented;
-      }
-    }
-
-    linkedData?.read(this);
-    return augmentedInternal;
+    return element;
   }
 
   AugmentedExtensionTypeElementImpl? get augmentedIfReally {
@@ -3809,7 +3809,10 @@ class ExtensionTypeElementImpl extends InterfaceElementImpl
   }
 
   @override
-  ExtensionTypeElement2 get element => super.element as ExtensionTypeElement2;
+  MaybeAugmentedExtensionTypeElementMixin get element {
+    linkedData?.read(this);
+    return augmentedInternal;
+  }
 
   @override
   ElementKind get kind {
@@ -6748,6 +6751,9 @@ mixin MaybeAugmentedEnumElementMixin on MaybeAugmentedInterfaceElementMixin
   @override
   EnumFragment get firstFragment => declaration;
 
+  /// See [ElementImpl2.reference].
+  Reference get reference;
+
   @override
   T? accept2<T>(ElementVisitor2<T> visitor) {
     return visitor.visitEnumElement(this);
@@ -6765,6 +6771,9 @@ mixin MaybeAugmentedExtensionElementMixin on MaybeAugmentedInstanceElementMixin
   @override
   ExtensionElementImpl get firstFragment => declaration;
 
+  /// See [ElementImpl2.reference].
+  Reference get reference;
+
   @override
   DartType get thisType => extendedType;
 
@@ -6795,6 +6804,9 @@ mixin MaybeAugmentedExtensionTypeElementMixin
   @override
   ConstructorElement2 get primaryConstructor2 => primaryConstructor.element;
 
+  /// See [ElementImpl2.reference].
+  Reference get reference;
+
   @override
   FieldElement2 get representation2 =>
       representation.asElement2 as FieldElement2;
@@ -7754,8 +7766,7 @@ class MixinElementImpl extends ClassOrMixinElementImpl
 
   @override
   MaybeAugmentedMixinElementMixin get augmented {
-    linkedData?.read(this);
-    return augmentedInternal;
+    return element;
   }
 
   AugmentedMixinElementImpl? get augmentedIfReally {
@@ -7773,7 +7784,10 @@ class MixinElementImpl extends ClassOrMixinElementImpl
   }
 
   @override
-  MaybeAugmentedMixinElementMixin get element => augmented;
+  MaybeAugmentedMixinElementMixin get element {
+    linkedData?.read(this);
+    return augmentedInternal;
+  }
 
   @override
   bool get isBase {
@@ -8349,35 +8363,47 @@ class NotAugmentedClassElementImpl extends NotAugmentedInterfaceElementImpl
 
 class NotAugmentedEnumElementImpl extends NotAugmentedInterfaceElementImpl
     with MaybeAugmentedEnumElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final EnumElementImpl element;
 
-  NotAugmentedEnumElementImpl(this.element);
+  NotAugmentedEnumElementImpl(this.reference, this.element) {
+    reference.element2 = this;
+    element.augmentedInternal = this;
+  }
 
   @override
   EnumElementImpl get declaration => element;
 
   @override
   AugmentedEnumElementImpl toAugmented() {
-    var augmented = AugmentedEnumElementImpl(declaration);
+    var augmented = AugmentedEnumElementImpl(reference, declaration);
     declaration.augmentedInternal = augmented;
     return augmented;
   }
 }
 
 class NotAugmentedExtensionElementImpl extends NotAugmentedInstanceElementImpl
     with MaybeAugmentedExtensionElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final ExtensionElementImpl element;
 
-  NotAugmentedExtensionElementImpl(this.element);
+  NotAugmentedExtensionElementImpl(this.reference, this.element) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 
   @override
   ExtensionElementImpl get declaration => element;
 
   @override
   AugmentedExtensionElementImpl toAugmented() {
-    var augmented = AugmentedExtensionElementImpl(declaration);
+    var augmented = AugmentedExtensionElementImpl(reference, declaration);
     augmented.extendedType = extendedType;
     declaration.augmentedInternal = augmented;
     return augmented;
@@ -8387,17 +8413,23 @@ class NotAugmentedExtensionElementImpl extends NotAugmentedInstanceElementImpl
 class NotAugmentedExtensionTypeElementImpl
     extends NotAugmentedInterfaceElementImpl
     with MaybeAugmentedExtensionTypeElementMixin {
+  @override
+  final Reference reference;
+
   @override
   final ExtensionTypeElementImpl element;
 
-  NotAugmentedExtensionTypeElementImpl(this.element);
+  NotAugmentedExtensionTypeElementImpl(this.reference, this.element) {
+    reference.element2 = this;
+    declaration.augmentedInternal = this;
+  }
 
   @override
   ExtensionTypeElementImpl get declaration => element;
 
   @override
   AugmentedExtensionTypeElementImpl toAugmented() {
-    var augmented = AugmentedExtensionTypeElementImpl(declaration);
+    var augmented = AugmentedExtensionTypeElementImpl(reference, declaration);
     augmented.primaryConstructor = primaryConstructor;
     augmented.representation = representation;
     declaration.augmentedInternal = augmented;