Add implementations for LibraryAugmentationElementImpl and friends.

Change-Id: I575332526eebd7cd2d87054f3b1395fabbd4d574
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/244724
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>
diff --git a/pkg/analysis_server/lib/src/utilities/extensions/element.dart b/pkg/analysis_server/lib/src/utilities/extensions/element.dart
index 1bc51b1..bfe5606 100644
--- a/pkg/analysis_server/lib/src/utilities/extensions/element.dart
+++ b/pkg/analysis_server/lib/src/utilities/extensions/element.dart
@@ -38,7 +38,7 @@
       ancestor = ancestor.enclosingElement;
     }
     return ancestor is CompilationUnitElement &&
-        ancestor.enclosingElement.hasDeprecated;
+        ancestor.enclosingElement2.hasDeprecated;
   }
 
   /// Return this element and all its enclosing elements.
diff --git a/pkg/analysis_server/test/protocol_server_test.dart b/pkg/analysis_server/test/protocol_server_test.dart
index 02e6971..603ea68 100644
--- a/pkg/analysis_server/test/protocol_server_test.dart
+++ b/pkg/analysis_server/test/protocol_server_test.dart
@@ -240,11 +240,13 @@
     EnumTester<engine.ElementKind, ElementKind>()
         .run(convertElementKind, exceptions: {
       // TODO(paulberry): do any of the exceptions below constitute bugs?
+      engine.ElementKind.AUGMENTATION_IMPORT: ElementKind.UNKNOWN,
       engine.ElementKind.DYNAMIC: ElementKind.UNKNOWN,
       engine.ElementKind.ERROR: ElementKind.UNKNOWN,
       engine.ElementKind.EXPORT: ElementKind.UNKNOWN,
       engine.ElementKind.GENERIC_FUNCTION_TYPE: ElementKind.FUNCTION_TYPE_ALIAS,
       engine.ElementKind.IMPORT: ElementKind.UNKNOWN,
+      engine.ElementKind.LIBRARY_AUGMENTATION: ElementKind.UNKNOWN,
       engine.ElementKind.NAME: ElementKind.UNKNOWN,
       engine.ElementKind.NEVER: ElementKind.UNKNOWN,
       engine.ElementKind.UNIVERSE: ElementKind.UNKNOWN
diff --git a/pkg/analyzer/lib/dart/element/element.dart b/pkg/analyzer/lib/dart/element/element.dart
index 024d39b..1ac9cdd 100644
--- a/pkg/analyzer/lib/dart/element/element.dart
+++ b/pkg/analyzer/lib/dart/element/element.dart
@@ -409,6 +409,10 @@
   @override
   LibraryElement get enclosingElement;
 
+  /// Return the library, or library augmentation that encloses this unit.
+  @experimental
+  LibraryOrAugmentationElement get enclosingElement2;
+
   /// Return a list containing all of the enums contained in this compilation
   /// unit.
   List<ClassElement> get enums;
@@ -867,70 +871,76 @@
 ///
 /// Clients may not extend, implement or mix-in this class.
 class ElementKind implements Comparable<ElementKind> {
-  static const ElementKind CLASS = ElementKind('CLASS', 0, "class");
+  static const ElementKind AUGMENTATION_IMPORT =
+      ElementKind('AUGMENTATION_IMPORT', 0, "augmentation import");
+
+  static const ElementKind CLASS = ElementKind('CLASS', 1, "class");
 
   static const ElementKind COMPILATION_UNIT =
-      ElementKind('COMPILATION_UNIT', 1, "compilation unit");
+      ElementKind('COMPILATION_UNIT', 2, "compilation unit");
 
   static const ElementKind CONSTRUCTOR =
-      ElementKind('CONSTRUCTOR', 2, "constructor");
+      ElementKind('CONSTRUCTOR', 3, "constructor");
 
-  static const ElementKind DYNAMIC = ElementKind('DYNAMIC', 3, "<dynamic>");
+  static const ElementKind DYNAMIC = ElementKind('DYNAMIC', 4, "<dynamic>");
 
-  static const ElementKind ENUM = ElementKind('ENUM', 4, "enum");
+  static const ElementKind ENUM = ElementKind('ENUM', 5, "enum");
 
-  static const ElementKind ERROR = ElementKind('ERROR', 5, "<error>");
+  static const ElementKind ERROR = ElementKind('ERROR', 6, "<error>");
 
   static const ElementKind EXPORT =
-      ElementKind('EXPORT', 6, "export directive");
+      ElementKind('EXPORT', 7, "export directive");
 
-  static const ElementKind EXTENSION = ElementKind('EXTENSION', 7, "extension");
+  static const ElementKind EXTENSION = ElementKind('EXTENSION', 8, "extension");
 
-  static const ElementKind FIELD = ElementKind('FIELD', 8, "field");
+  static const ElementKind FIELD = ElementKind('FIELD', 9, "field");
 
-  static const ElementKind FUNCTION = ElementKind('FUNCTION', 9, "function");
+  static const ElementKind FUNCTION = ElementKind('FUNCTION', 10, "function");
 
   static const ElementKind GENERIC_FUNCTION_TYPE =
-      ElementKind('GENERIC_FUNCTION_TYPE', 10, 'generic function type');
+      ElementKind('GENERIC_FUNCTION_TYPE', 11, 'generic function type');
 
-  static const ElementKind GETTER = ElementKind('GETTER', 11, "getter");
+  static const ElementKind GETTER = ElementKind('GETTER', 12, "getter");
 
   static const ElementKind IMPORT =
-      ElementKind('IMPORT', 12, "import directive");
+      ElementKind('IMPORT', 13, "import directive");
 
-  static const ElementKind LABEL = ElementKind('LABEL', 13, "label");
+  static const ElementKind LABEL = ElementKind('LABEL', 14, "label");
 
-  static const ElementKind LIBRARY = ElementKind('LIBRARY', 14, "library");
+  static const ElementKind LIBRARY = ElementKind('LIBRARY', 15, "library");
+
+  static const ElementKind LIBRARY_AUGMENTATION =
+      ElementKind('LIBRARY_AUGMENTATION', 16, "library augmentation");
 
   static const ElementKind LOCAL_VARIABLE =
-      ElementKind('LOCAL_VARIABLE', 15, "local variable");
+      ElementKind('LOCAL_VARIABLE', 17, "local variable");
 
-  static const ElementKind METHOD = ElementKind('METHOD', 16, "method");
+  static const ElementKind METHOD = ElementKind('METHOD', 18, "method");
 
-  static const ElementKind NAME = ElementKind('NAME', 17, "<name>");
+  static const ElementKind NAME = ElementKind('NAME', 19, "<name>");
 
-  static const ElementKind NEVER = ElementKind('NEVER', 18, "<never>");
+  static const ElementKind NEVER = ElementKind('NEVER', 20, "<never>");
 
   static const ElementKind PARAMETER =
-      ElementKind('PARAMETER', 19, "parameter");
+      ElementKind('PARAMETER', 21, "parameter");
 
-  static const ElementKind PREFIX = ElementKind('PREFIX', 20, "import prefix");
+  static const ElementKind PREFIX = ElementKind('PREFIX', 22, "import prefix");
 
-  static const ElementKind SETTER = ElementKind('SETTER', 21, "setter");
+  static const ElementKind SETTER = ElementKind('SETTER', 23, "setter");
 
   static const ElementKind TOP_LEVEL_VARIABLE =
-      ElementKind('TOP_LEVEL_VARIABLE', 22, "top level variable");
+      ElementKind('TOP_LEVEL_VARIABLE', 24, "top level variable");
 
   static const ElementKind FUNCTION_TYPE_ALIAS =
-      ElementKind('FUNCTION_TYPE_ALIAS', 23, "function type alias");
+      ElementKind('FUNCTION_TYPE_ALIAS', 25, "function type alias");
 
   static const ElementKind TYPE_PARAMETER =
-      ElementKind('TYPE_PARAMETER', 24, "type parameter");
+      ElementKind('TYPE_PARAMETER', 26, "type parameter");
 
   static const ElementKind TYPE_ALIAS =
-      ElementKind('TYPE_ALIAS', 25, "type alias");
+      ElementKind('TYPE_ALIAS', 27, "type alias");
 
-  static const ElementKind UNIVERSE = ElementKind('UNIVERSE', 26, "<universe>");
+  static const ElementKind UNIVERSE = ElementKind('UNIVERSE', 28, "<universe>");
 
   static const List<ElementKind> values = [
     CLASS,
@@ -1017,6 +1027,8 @@
 /// * ThrowingElementVisitor which implements every visit method by throwing an
 ///   exception.
 abstract class ElementVisitor<R> {
+  R? visitAugmentationImportElement(AugmentationImportElement element);
+
   R? visitClassElement(ClassElement element);
 
   R? visitCompilationUnitElement(CompilationUnitElement element);
@@ -1039,6 +1051,8 @@
 
   R? visitLabelElement(LabelElement element);
 
+  R? visitLibraryAugmentationElement(LibraryAugmentationElement element);
+
   R? visitLibraryElement(LibraryElement element);
 
   R? visitLocalVariableElement(LocalVariableElement element);
@@ -1640,6 +1654,10 @@
   @override
   LibraryElement get enclosingElement;
 
+  /// Return the library, or library augmentation that encloses this element.
+  @experimental
+  LibraryOrAugmentationElement get enclosingElement2;
+
   @override
   String get name;
 
diff --git a/pkg/analyzer/lib/dart/element/visitor.dart b/pkg/analyzer/lib/dart/element/visitor.dart
index aede2b6..ace3ca2 100644
--- a/pkg/analyzer/lib/dart/element/visitor.dart
+++ b/pkg/analyzer/lib/dart/element/visitor.dart
@@ -83,6 +83,10 @@
   const GeneralizingElementVisitor();
 
   @override
+  R? visitAugmentationImportElement(AugmentationImportElement element) =>
+      visitElement(element);
+
+  @override
   R? visitClassElement(ClassElement element) => visitElement(element);
 
   @override
@@ -129,6 +133,10 @@
   R? visitLabelElement(LabelElement element) => visitElement(element);
 
   @override
+  R? visitLibraryAugmentationElement(LibraryAugmentationElement element) =>
+      visitElement(element);
+
+  @override
   R? visitLibraryElement(LibraryElement element) => visitElement(element);
 
   R? visitLocalElement(LocalElement element) {
@@ -202,6 +210,12 @@
   const RecursiveElementVisitor();
 
   @override
+  R? visitAugmentationImportElement(AugmentationImportElement element) {
+    element.visitChildren(this);
+    return null;
+  }
+
+  @override
   R? visitClassElement(ClassElement element) {
     element.visitChildren(this);
     return null;
@@ -268,6 +282,12 @@
   }
 
   @override
+  R? visitLibraryAugmentationElement(LibraryAugmentationElement element) {
+    element.visitChildren(this);
+    return null;
+  }
+
+  @override
   R? visitLibraryElement(LibraryElement element) {
     element.visitChildren(this);
     return null;
@@ -345,6 +365,9 @@
   const SimpleElementVisitor();
 
   @override
+  R? visitAugmentationImportElement(AugmentationImportElement element) => null;
+
+  @override
   R? visitClassElement(ClassElement element) => null;
 
   @override
@@ -380,6 +403,10 @@
   R? visitLabelElement(LabelElement element) => null;
 
   @override
+  R? visitLibraryAugmentationElement(LibraryAugmentationElement element) =>
+      null;
+
+  @override
   R? visitLibraryElement(LibraryElement element) => null;
 
   @override
@@ -426,6 +453,10 @@
   const ThrowingElementVisitor();
 
   @override
+  R? visitAugmentationImportElement(AugmentationImportElement element) =>
+      _throw(element);
+
+  @override
   R? visitClassElement(ClassElement element) => _throw(element);
 
   @override
@@ -462,6 +493,10 @@
   R? visitLabelElement(LabelElement element) => _throw(element);
 
   @override
+  R? visitLibraryAugmentationElement(LibraryAugmentationElement element) =>
+      _throw(element);
+
+  @override
   R? visitLibraryElement(LibraryElement element) => _throw(element);
 
   @override
diff --git a/pkg/analyzer/lib/src/dart/element/element.dart b/pkg/analyzer/lib/src/dart/element/element.dart
index 392afd4..94e99c7 100644
--- a/pkg/analyzer/lib/src/dart/element/element.dart
+++ b/pkg/analyzer/lib/src/dart/element/element.dart
@@ -13,6 +13,8 @@
 import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:analyzer/dart/element/scope.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/dart/element/type_provider.dart';
+import 'package:analyzer/dart/element/type_system.dart';
 import 'package:analyzer/error/error.dart';
 import 'package:analyzer/src/context/source.dart';
 import 'package:analyzer/src/dart/analysis/experiments.dart';
@@ -518,6 +520,30 @@
   }
 }
 
+/// A concrete implementation of a [AugmentationImportElement].
+class AugmentationImportElementImpl extends UriReferencedElementImpl
+    implements AugmentationImportElement {
+  @override
+  final LibraryAugmentationElement augmentation;
+
+  AugmentationImportElementImpl({
+    required this.augmentation,
+    required int nameOffset,
+  }) : super(null, nameOffset);
+
+  @override
+  LibraryOrAugmentationElementImpl get enclosingElement {
+    return super.enclosingElement as LibraryOrAugmentationElementImpl;
+  }
+
+  @override
+  ElementKind get kind => ElementKind.AUGMENTATION_IMPORT;
+
+  @override
+  T? accept<T>(ElementVisitor<T> visitor) =>
+      visitor.visitAugmentationImportElement(this);
+}
+
 /// An [AbstractClassElementImpl] which is a class.
 class ClassElementImpl extends AbstractClassElementImpl {
   /// For classes which are not mixin applications, a list containing all of the
@@ -1077,6 +1103,10 @@
       super.enclosingElement as LibraryElement;
 
   @override
+  LibraryOrAugmentationElement get enclosingElement2 =>
+      super.enclosingElement as LibraryOrAugmentationElement;
+
+  @override
   CompilationUnitElementImpl get enclosingUnit {
     return this;
   }
@@ -3626,9 +3656,55 @@
   T? accept<T>(ElementVisitor<T> visitor) => visitor.visitLabelElement(this);
 }
 
+class LibraryAugmentationElementImpl extends LibraryOrAugmentationElementImpl
+    implements LibraryAugmentationElement {
+  @override
+  final LibraryOrAugmentationElement augmented;
+
+  LibraryAugmentationElementImpl({
+    required this.augmented,
+    required int nameOffset,
+  }) : super(name: null, nameOffset: nameOffset);
+
+  @override
+  // TODO: implement accessibleExtensions
+  List<ExtensionElement> get accessibleExtensions => throw UnimplementedError();
+
+  @override
+  FeatureSet get featureSet => augmented.featureSet;
+
+  @override
+  bool get isNonNullableByDefault => augmented.isNonNullableByDefault;
+
+  @override
+  ElementKind get kind => ElementKind.LIBRARY_AUGMENTATION;
+
+  @override
+  LibraryLanguageVersion get languageVersion => augmented.languageVersion;
+
+  @override
+  // TODO: implement scope
+  Scope get scope => throw UnimplementedError();
+
+  @override
+  AnalysisSession get session => augmented.session;
+
+  @override
+  TypeProvider get typeProvider => augmented.typeProvider;
+
+  @override
+  TypeSystem get typeSystem => augmented.typeSystem;
+
+  @override
+  T? accept<T>(ElementVisitor<T> visitor) {
+    return visitor.visitLibraryAugmentationElement(this);
+  }
+}
+
 /// A concrete implementation of a [LibraryElement].
-class LibraryElementImpl extends _ExistingElementImpl
-    implements LibraryElement, MacroTargetElementContainer {
+class LibraryElementImpl extends LibraryOrAugmentationElementImpl
+    with _HasLibraryMixin
+    implements LibraryElement {
   /// The analysis context in which this library is defined.
   @override
   final AnalysisContext context;
@@ -3654,21 +3730,10 @@
   @override
   final FeatureSet featureSet;
 
-  /// The compilation unit that defines this library.
-  late CompilationUnitElementImpl _definingCompilationUnit;
-
   /// The entry point for this library, or `null` if this library does not have
   /// an entry point.
   FunctionElement? _entryPoint;
 
-  /// A list containing specifications of all of the imports defined in this
-  /// library.
-  List<ImportElement> _imports = _Sentinel.importElement;
-
-  /// A list containing specifications of all of the exports defined in this
-  /// library.
-  List<ExportElement> _exports = _Sentinel.exportElement;
-
   /// A list containing all of the compilation units that are included in this
   /// library using a `part` directive.
   List<CompilationUnitElement> _parts = const <CompilationUnitElement>[];
@@ -3689,9 +3754,6 @@
   /// computed yet.
   Namespace? _publicNamespace;
 
-  /// The cached list of prefixes.
-  List<PrefixElement>? _prefixes;
-
   /// The scope of this library, `null` if it has not been created yet.
   LibraryScope? _scope;
 
@@ -3703,30 +3765,12 @@
   LibraryElementImpl(this.context, this.session, String name, int offset,
       this.nameLength, this.featureSet)
       : linkedData = null,
-        super(name, offset);
+        super(name: name, nameOffset: offset);
 
   @override
   List<ExtensionElement> get accessibleExtensions => scope.extensions;
 
   @override
-  List<AugmentationImportElement> get augmentationImports {
-    // TODO(scheglov): implement augmentationImports
-    throw UnimplementedError();
-  }
-
-  @override
-  CompilationUnitElementImpl get definingCompilationUnit =>
-      _definingCompilationUnit;
-
-  /// Set the compilation unit that defines this library to the given
-  ///  compilation[unit].
-  set definingCompilationUnit(CompilationUnitElement unit) {
-    assert((unit as CompilationUnitElementImpl).librarySource == unit.source);
-    (unit as CompilationUnitElementImpl).enclosingElement = this;
-    _definingCompilationUnit = unit;
-  }
-
-  @override
   CompilationUnitElementImpl get enclosingUnit {
     return _definingCompilationUnit;
   }
@@ -3779,19 +3823,6 @@
     return _exports;
   }
 
-  /// Set the specifications of all of the exports defined in this library to
-  /// the given list of [exports].
-  set exports(List<ExportElement> exports) {
-    for (ExportElement exportElement in exports) {
-      (exportElement as ExportElementImpl).enclosingElement = this;
-    }
-    _exports = exports;
-  }
-
-  List<ExportElement> get exports_unresolved {
-    return _exports;
-  }
-
   bool get hasPartOfDirective {
     return hasModifier(Modifier.HAS_PART_OF_DIRECTIVE);
   }
@@ -3821,24 +3852,6 @@
     return _imports;
   }
 
-  /// Set the specifications of all of the imports defined in this library to
-  /// the given list of [imports].
-  set imports(List<ImportElement> imports) {
-    for (ImportElement importElement in imports) {
-      (importElement as ImportElementImpl).enclosingElement = this;
-      var prefix = importElement.prefix as PrefixElementImpl?;
-      if (prefix != null) {
-        prefix.enclosingElement = this;
-      }
-    }
-    _imports = imports;
-    _prefixes = null;
-  }
-
-  List<ImportElement> get imports_unresolved {
-    return _imports;
-  }
-
   @override
   bool get isBrowserApplication =>
       entryPoint != null && isOrImportsBrowserLibrary;
@@ -4091,9 +4104,6 @@
   @override
   void visitChildren(ElementVisitor visitor) {
     super.visitChildren(visitor);
-    _definingCompilationUnit.accept(visitor);
-    safelyVisitChildren(exports, visitor);
-    safelyVisitChildren(imports, visitor);
     safelyVisitChildren(_parts, visitor);
   }
 
@@ -4139,6 +4149,134 @@
   }
 }
 
+/// A concrete implementation of a [LibraryOrAugmentationElement].
+abstract class LibraryOrAugmentationElementImpl extends ElementImpl
+    implements LibraryOrAugmentationElement, MacroTargetElementContainer {
+  /// The compilation unit that defines this library.
+  late CompilationUnitElementImpl _definingCompilationUnit;
+
+  List<AugmentationImportElement> _augmentationImports =
+      _Sentinel.augmentationImportElement;
+
+  /// A list containing specifications of all of the imports defined in this
+  /// library.
+  List<ImportElement> _imports = _Sentinel.importElement;
+
+  /// A list containing specifications of all of the exports defined in this
+  /// library.
+  List<ExportElement> _exports = _Sentinel.exportElement;
+
+  /// The cached list of prefixes.
+  List<PrefixElement>? _prefixes;
+
+  LibraryOrAugmentationElementImpl({
+    required String? name,
+    required int nameOffset,
+  }) : super(name, nameOffset);
+
+  @override
+  List<AugmentationImportElement> get augmentationImports {
+    return _augmentationImports;
+  }
+
+  set augmentationImports(List<AugmentationImportElement> imports) {
+    for (final importElement in imports) {
+      (importElement as AugmentationImportElementImpl).enclosingElement = this;
+    }
+    _augmentationImports = imports;
+  }
+
+  @override
+  CompilationUnitElementImpl get definingCompilationUnit =>
+      _definingCompilationUnit;
+
+  /// Set the compilation unit that defines this library to the given
+  ///  compilation[unit].
+  set definingCompilationUnit(CompilationUnitElement unit) {
+    assert((unit as CompilationUnitElementImpl).librarySource == unit.source);
+    (unit as CompilationUnitElementImpl).enclosingElement = this;
+    _definingCompilationUnit = unit;
+  }
+
+  @override
+  CompilationUnitElementImpl get enclosingUnit {
+    return _definingCompilationUnit;
+  }
+
+  @override
+  List<ExportElement> get exports {
+    return _exports;
+  }
+
+  /// Set the specifications of all of the exports defined in this library to
+  /// the given list of [exports].
+  set exports(List<ExportElement> exports) {
+    for (ExportElement exportElement in exports) {
+      (exportElement as ExportElementImpl).enclosingElement = this;
+    }
+    _exports = exports;
+  }
+
+  List<ExportElement> get exports_unresolved {
+    return _exports;
+  }
+
+  @override
+  String get identifier => '${_definingCompilationUnit.source.uri}';
+
+  @override
+  List<ImportElement> get imports {
+    return _imports;
+  }
+
+  /// Set the specifications of all of the imports defined in this library to
+  /// the given list of [imports].
+  set imports(List<ImportElement> imports) {
+    for (ImportElement importElement in imports) {
+      (importElement as ImportElementImpl).enclosingElement = this;
+      var prefix = importElement.prefix as PrefixElementImpl?;
+      if (prefix != null) {
+        prefix.enclosingElement = this;
+      }
+    }
+    _imports = imports;
+    _prefixes = null;
+  }
+
+  List<ImportElement> get imports_unresolved {
+    return _imports;
+  }
+
+  @override
+  List<PrefixElement> get prefixes =>
+      _prefixes ??= buildPrefixesFromImports(imports);
+
+  @override
+  Source get source {
+    return _definingCompilationUnit.source;
+  }
+
+  @override
+  void visitChildren(ElementVisitor visitor) {
+    super.visitChildren(visitor);
+    _definingCompilationUnit.accept(visitor);
+    safelyVisitChildren(exports, visitor);
+    safelyVisitChildren(imports, visitor);
+  }
+
+  static List<PrefixElement> buildPrefixesFromImports(
+      List<ImportElement> imports) {
+    HashSet<PrefixElement> prefixes = HashSet<PrefixElement>();
+    for (ImportElement element in imports) {
+      var prefix = element.prefix;
+      if (prefix != null) {
+        prefixes.add(prefix);
+      }
+    }
+    return prefixes.toList(growable: false);
+  }
+}
+
 /// A concrete implementation of a [LocalVariableElement].
 class LocalVariableElementImpl extends NonParameterVariableElementImpl
     implements LocalVariableElement {
@@ -4971,6 +5109,10 @@
       super.enclosingElement as LibraryElement;
 
   @override
+  LibraryOrAugmentationElementImpl get enclosingElement2 =>
+      super.enclosingElement as LibraryOrAugmentationElementImpl;
+
+  @override
   ElementKind get kind => ElementKind.PREFIX;
 
   @override
@@ -5994,6 +6136,8 @@
       List.unmodifiable([]);
   static final List<ExportElement> exportElement = List.unmodifiable([]);
   static final List<FieldElement> fieldElement = List.unmodifiable([]);
+  static final List<AugmentationImportElement> augmentationImportElement =
+      List.unmodifiable([]);
   static final List<ImportElement> importElement = List.unmodifiable([]);
   static final List<MethodElement> methodElement = List.unmodifiable([]);
   static final List<PropertyAccessorElement> propertyAccessorElement =
diff --git a/pkg/analyzer/lib/src/dart/element/extensions.dart b/pkg/analyzer/lib/src/dart/element/extensions.dart
index decbd59..f5c453d 100644
--- a/pkg/analyzer/lib/src/dart/element/extensions.dart
+++ b/pkg/analyzer/lib/src/dart/element/extensions.dart
@@ -82,7 +82,7 @@
     }
 
     return ancestor is CompilationUnitElement &&
-        ancestor.enclosingElement.hasDoNotStore;
+        ancestor.enclosingElement2.hasDoNotStore;
   }
 
   /// Return `true` if this element is an instance member of a class or mixin.
diff --git a/pkg/analyzer/test/src/dart/analysis/results/get_element_declaration_test.dart b/pkg/analyzer/test/src/dart/analysis/results/get_element_declaration_test.dart
index c34e500..3de528c 100644
--- a/pkg/analyzer/test/src/dart/analysis/results/get_element_declaration_test.dart
+++ b/pkg/analyzer/test/src/dart/analysis/results/get_element_declaration_test.dart
@@ -300,7 +300,7 @@
     await resolveTestCode(r'''
 library foo;
 ''');
-    var element = findElement.unitElement.enclosingElement;
+    var element = this.result.libraryElement;
     var result = await getElementDeclaration(element);
     expect(result, isNull);
   }