Elements. Migrate src/dart/analysis/search.dart

Replace index / search of `CompilationUnitElement` with special
`referencesLibraryFragment(LibraryFragment)`, because they are not
elements.


Change-Id: I4de4ce42270bdad92996bedbd06f108a7c71dffe
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/412681
Reviewed-by: Paul Berry <paulberry@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>
diff --git a/pkg/analysis_server/lib/src/services/search/search_engine_internal.dart b/pkg/analysis_server/lib/src/services/search/search_engine_internal.dart
index cab52f0..9e7398f 100644
--- a/pkg/analysis_server/lib/src/services/search/search_engine_internal.dart
+++ b/pkg/analysis_server/lib/src/services/search/search_engine_internal.dart
@@ -35,7 +35,8 @@
             _searchDirectSubtypes(type, searchEngineCache, performance),
       );
       for (var directResult in directResults) {
-        var directSubtype = directResult.enclosingElement2 as InterfaceElement2;
+        var directSubtype =
+            directResult.enclosingFragment.element as InterfaceElement2;
         if (allSubtypes.add(directSubtype)) {
           await addSubtypes(directSubtype);
         }
@@ -124,7 +125,7 @@
     var drivers = _drivers.toList();
     var searchedFiles = _createSearchedFiles(drivers);
     for (var driver in drivers) {
-      var elements = await driver.search.classMembers2(name, searchedFiles);
+      var elements = await driver.search.classMembers(name, searchedFiles);
       allDeclarations.addAll(elements.map(SearchMatchImpl.forElement));
     }
     return allDeclarations;
@@ -161,7 +162,7 @@
     var drivers = _drivers.toList();
     var searchedFiles = _createSearchedFiles(drivers);
     for (var driver in drivers) {
-      var results = await driver.search.references2(element, searchedFiles);
+      var results = await driver.search.references(element, searchedFiles);
       allResults.addAll(results);
     }
     return allResults.map(SearchMatchImpl.forSearchResult).toList();
@@ -194,7 +195,7 @@
     var regExp = RegExp(pattern);
     var drivers = _drivers.toList();
     for (var driver in drivers) {
-      var elements = await driver.search.topLevelElements2(regExp);
+      var elements = await driver.search.topLevelElements(regExp);
       allElements.addAll(elements);
     }
     return allElements.map(SearchMatchImpl.forElement).toList();
@@ -336,7 +337,7 @@
       libraryFragment.element.firstFragment.source,
       libraryFragment.source,
       libraryFragment.element,
-      result.enclosingElement2,
+      result.enclosingFragment.element,
       result.isResolved,
       result.isQualified,
       toMatchKind(result.kind),
diff --git a/pkg/analyzer/lib/src/dart/analysis/index.dart b/pkg/analyzer/lib/src/dart/analysis/index.dart
index 7d435bb..54e9e51 100644
--- a/pkg/analyzer/lib/src/dart/analysis/index.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/index.dart
@@ -830,9 +830,17 @@
   }
 
   @override
-  void visitExportDirective(ExportDirective node) {
-    var element = node.libraryExport;
-    recordUriReference(element?.exportedLibrary2, node.uri);
+  void visitExportDirective(covariant ExportDirectiveImpl node) {
+    if (node.libraryExport case var libraryExport?) {
+      if (libraryExport.exportedLibrary2 case var exportedLibrary?) {
+        assembler.addLibraryFragmentReference(
+          target: exportedLibrary.firstFragment,
+          uriOffset: node.uri.offset,
+          uriLength: node.uri.length,
+        );
+      }
+    }
+
     super.visitExportDirective(node);
   }
 
@@ -898,9 +906,17 @@
   }
 
   @override
-  void visitImportDirective(ImportDirective node) {
-    var element = node.libraryImport;
-    recordUriReference(element?.importedLibrary2, node.uri);
+  void visitImportDirective(covariant ImportDirectiveImpl node) {
+    if (node.libraryImport case var libraryImport?) {
+      if (libraryImport.importedLibrary2 case var importedLibrary?) {
+        assembler.addLibraryFragmentReference(
+          target: importedLibrary.firstFragment,
+          uriOffset: node.uri.offset,
+          uriLength: node.uri.length,
+        );
+      }
+    }
+
     super.visitImportDirective(node);
   }
 
diff --git a/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart b/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
index 0380eec..904402b 100644
--- a/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/library_analyzer.dart
@@ -1111,7 +1111,7 @@
 /// Analysis result for single file.
 class UnitAnalysisResult {
   final FileState file;
-  final CompilationUnit unit;
+  final CompilationUnitImpl unit;
   final List<AnalysisError> errors;
 
   UnitAnalysisResult(this.file, this.unit, this.errors);
diff --git a/pkg/analyzer/lib/src/dart/analysis/results.dart b/pkg/analyzer/lib/src/dart/analysis/results.dart
index 7816bf9..172448b 100644
--- a/pkg/analyzer/lib/src/dart/analysis/results.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/results.dart
@@ -6,7 +6,6 @@
 
 import 'package:analyzer/dart/analysis/results.dart';
 import 'package:analyzer/dart/analysis/session.dart';
-import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/token.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
 import 'package:analyzer/dart/element/element.dart';
@@ -16,6 +15,7 @@
 import 'package:analyzer/file_system/file_system.dart';
 import 'package:analyzer/source/line_info.dart';
 import 'package:analyzer/src/dart/analysis/file_state.dart';
+import 'package:analyzer/src/dart/ast/ast.dart';
 import 'package:analyzer/src/dart/element/element.dart';
 import 'package:analyzer/src/dart/element/type_provider.dart';
 import 'package:analyzer/src/dart/element/type_system.dart';
@@ -490,7 +490,7 @@
 class ResolvedUnitResultImpl extends FileResultImpl
     implements ResolvedUnitResult {
   @override
-  final CompilationUnit unit;
+  final CompilationUnitImpl unit;
 
   @override
   final List<AnalysisError> errors;
@@ -538,5 +538,5 @@
   });
 
   @override
-  LibraryFragment get fragment => element;
+  CompilationUnitElementImpl get fragment => element;
 }
diff --git a/pkg/analyzer/lib/src/dart/analysis/search.dart b/pkg/analyzer/lib/src/dart/analysis/search.dart
index 8cc7036..90695e1 100644
--- a/pkg/analyzer/lib/src/dart/analysis/search.dart
+++ b/pkg/analyzer/lib/src/dart/analysis/search.dart
@@ -2,13 +2,9 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
-// ignore_for_file: analyzer_use_new_elements
-
 import 'package:analyzer/dart/analysis/results.dart';
-import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/syntactic_entity.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
-import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/element2.dart';
 import 'package:analyzer/dart/element/visitor.dart';
 import 'package:analyzer/source/line_info.dart';
@@ -17,78 +13,85 @@
 import 'package:analyzer/src/dart/analysis/file_state.dart';
 import 'package:analyzer/src/dart/analysis/index.dart';
 import 'package:analyzer/src/dart/analysis/results.dart';
+import 'package:analyzer/src/dart/ast/ast.dart';
 import 'package:analyzer/src/dart/ast/extensions.dart';
 import 'package:analyzer/src/dart/ast/utilities.dart';
 import 'package:analyzer/src/dart/element/element.dart';
 import 'package:analyzer/src/summary/idl.dart';
 import 'package:analyzer/src/util/performance/operation_performance.dart';
 import 'package:analyzer/src/utilities/cancellation.dart';
-import 'package:analyzer/src/utilities/extensions/element.dart';
 import 'package:analyzer/src/utilities/fuzzy_matcher.dart';
 import 'package:collection/collection.dart';
 
-Element _getEnclosingElement(CompilationUnitElement unitElement, int offset) {
-  var finder = _ContainingElementFinder(offset);
-  unitElement.accept(finder);
-  var element = finder.containingElement;
-  if (element == null) {
+Fragment _getEnclosingFragment(
+  CompilationUnitElementImpl libraryFragment,
+  int offset,
+) {
+  var finder = _ContainingFragmentFinder(offset);
+  libraryFragment.accept(finder);
+  var result = finder.containingFragment;
+  if (result == null) {
     throw StateError(
-      'No containing element in ${unitElement.source.fullName} at $offset',
+      'No containing fragment in ${libraryFragment.source.fullName} at $offset',
     );
   }
-  return element;
+  return result;
 }
 
-DeclarationKind? _getSearchElementKind(Element element) {
-  if (element is EnumElement) {
+DeclarationKind? _getSearchElementKind(Element2 element) {
+  if (element is EnumElement2) {
     return DeclarationKind.ENUM;
   }
 
-  if (element is ExtensionTypeElement) {
+  if (element is ExtensionTypeElement2) {
     return DeclarationKind.EXTENSION_TYPE;
   }
 
-  if (element is MixinElement) {
+  if (element is MixinElement2) {
     return DeclarationKind.MIXIN;
   }
 
-  if (element is ClassElement) {
+  if (element is ClassElement2) {
     if (element.isMixinApplication) {
       return DeclarationKind.CLASS_TYPE_ALIAS;
     }
     return DeclarationKind.CLASS;
   }
 
-  if (element is ConstructorElement) {
+  if (element is ConstructorElement2) {
     return DeclarationKind.CONSTRUCTOR;
   }
 
-  if (element is ExtensionElement) {
+  if (element is ExtensionElement2) {
     return DeclarationKind.EXTENSION;
   }
 
-  if (element is FieldElement) {
+  if (element is FieldElement2) {
     if (element.isEnumConstant) return DeclarationKind.ENUM_CONSTANT;
     return DeclarationKind.FIELD;
   }
 
-  if (element is FunctionElement) {
+  if (element is LocalFunctionElement || element is TopLevelFunctionElement) {
     return DeclarationKind.FUNCTION;
   }
 
-  if (element is MethodElement) {
+  if (element is MethodElement2) {
     return DeclarationKind.METHOD;
   }
 
-  if (element is PropertyAccessorElement) {
-    return element.isGetter ? DeclarationKind.GETTER : DeclarationKind.SETTER;
+  if (element is GetterElement) {
+    return DeclarationKind.GETTER;
   }
 
-  if (element is TypeAliasElement) {
+  if (element is SetterElement) {
+    return DeclarationKind.SETTER;
+  }
+
+  if (element is TypeAliasElement2) {
     return DeclarationKind.TYPE_ALIAS;
   }
 
-  if (element is VariableElement) {
+  if (element is VariableElement2) {
     return DeclarationKind.VARIABLE;
   }
 
@@ -200,19 +203,19 @@
   }
 }
 
-/// Visitor that adds [SearchResult]s for references to the [importElement].
+/// Visitor that adds [SearchResult]s for references to the [import].
 class ImportElementReferencesVisitor extends RecursiveAstVisitor<void> {
   final List<SearchResult> results = <SearchResult>[];
 
-  final LibraryImportElement importElement;
-  final CompilationUnitElement enclosingUnitElement;
+  final LibraryImport import;
+  final CompilationUnitElementImpl enclosingLibraryFragment;
 
-  late final Set<Element> importedElements;
+  late final Set<Element2> importedElements;
 
   ImportElementReferencesVisitor(
-      LibraryImportElement element, this.enclosingUnitElement)
-      : importElement = element {
-    importedElements = element.namespace.definedNames.values.toSet();
+      LibraryImport element, this.enclosingLibraryFragment)
+      : import = element {
+    importedElements = element.namespace.definedNames2.values.toSet();
   }
 
   @override
@@ -223,16 +226,16 @@
 
   @override
   void visitNamedType(NamedType node) {
-    if (importedElements.contains(node.element)) {
-      var importElementPrefix = importElement.prefix;
+    if (importedElements.contains(node.element2)) {
+      var prefixFragment = import.prefix2;
       var importPrefix = node.importPrefix;
-      if (importElementPrefix == null) {
+      if (prefixFragment == null) {
         if (importPrefix == null) {
           _addResult(node.offset, 0);
         }
       } else {
         if (importPrefix != null &&
-            importPrefix.element == importElementPrefix.element) {
+            importPrefix.element2 == prefixFragment.element) {
           var offset = importPrefix.offset;
           var end = importPrefix.period.end;
           _addResult(offset, end - offset);
@@ -249,24 +252,24 @@
     if (node.inDeclarationContext()) {
       return;
     }
-    if (importElement.prefix != null) {
-      if (node.staticElement == importElement.prefix?.element) {
+    if (import.prefix2 != null) {
+      if (node.element == import.prefix2?.element) {
         var parent = node.parent;
         if (parent is PrefixedIdentifier && parent.prefix == node) {
-          var element = parent.writeOrReadElement?.declaration;
+          var element = parent.writeOrReadElement2?.baseElement;
           if (importedElements.contains(element)) {
             _addResultForPrefix(node, parent.identifier);
           }
         }
         if (parent is MethodInvocation && parent.target == node) {
-          var element = parent.methodName.staticElement?.declaration;
+          var element = parent.methodName.element?.baseElement;
           if (importedElements.contains(element)) {
             _addResultForPrefix(node, parent.methodName);
           }
         }
       }
     } else {
-      var element = node.writeOrReadElement?.declaration;
+      var element = node.writeOrReadElement2?.baseElement;
       if (importedElements.contains(element)) {
         _addResult(node.offset, 0);
       }
@@ -274,9 +277,9 @@
   }
 
   void _addResult(int offset, int length) {
-    Element enclosingElement =
-        _getEnclosingElement(enclosingUnitElement, offset);
-    results.add(SearchResult._(enclosingElement, SearchResultKind.REFERENCE,
+    var enclosingFragment =
+        _getEnclosingFragment(enclosingLibraryFragment, offset);
+    results.add(SearchResult._(enclosingFragment, SearchResultKind.REFERENCE,
         offset, length, true, false));
   }
 
@@ -308,44 +311,39 @@
   Search(this._driver);
 
   /// Returns class or mixin members with the given [name].
-  Future<List<Element>> classMembers(
+  Future<List<Element2>> classMembers(
       String name, SearchedFiles searchedFiles) async {
-    List<Element> elements = <Element>[];
+    var elements = <Element2>[];
 
-    void addElement(Element element) {
+    void addElement(Element2 element) {
       if (!element.isSynthetic && element.displayName == name) {
         elements.add(element);
       }
     }
 
-    void addElements(InterfaceElement element) {
-      element.accessors.forEach(addElement);
-      element.fields.forEach(addElement);
-      element.methods.forEach(addElement);
+    void addElements(InterfaceElement2 element) {
+      element.getters2.forEach(addElement);
+      element.setters2.forEach(addElement);
+      element.fields2.forEach(addElement);
+      element.methods2.forEach(addElement);
     }
 
     var files = await _driver.getFilesDefiningClassMemberName(name);
     for (var file in files) {
       if (searchedFiles.add(file.path, this)) {
-        var unitResult = await _driver.getUnitElement(file.path);
-        if (unitResult is UnitElementResult) {
-          unitResult.element.classes.forEach(addElements);
-          unitResult.element.enums.forEach(addElements);
-          unitResult.element.extensionTypes.forEach(addElements);
-          unitResult.element.mixins.forEach(addElements);
+        var libraryResult = await _driver.getLibraryByUri(file.uriStr);
+        if (libraryResult is LibraryElementResultImpl) {
+          var element = libraryResult.element;
+          element.classes.forEach(addElements);
+          element.enums.forEach(addElements);
+          element.extensionTypes.forEach(addElements);
+          element.mixins.forEach(addElements);
         }
       }
     }
     return elements;
   }
 
-  /// Returns class or mixin members with the given [name].
-  Future<List<Element2>> classMembers2(
-      String name, SearchedFiles searchedFiles) async {
-    var elements = await classMembers(name, searchedFiles);
-    return elements.map((e) => e.asElement2!).toList();
-  }
-
   /// Return the prefixes used to reference the [element] in any of the
   /// compilation units in the [library]. The returned set will include an empty
   /// string if the element is referenced without a prefix.
@@ -356,7 +354,7 @@
       var index = await _driver.getIndex(unit.source.fullName);
       if (index != null) {
         _IndexRequest request = _IndexRequest(index);
-        int elementId = request.findElementId(element.asElement!);
+        int elementId = request.findElementId(element);
         if (elementId != -1) {
           var prefixList = index.elementImportPrefixes[elementId].split(',');
           prefixes.addAll(prefixList);
@@ -368,34 +366,28 @@
 
   /// Returns references to the [element].
   Future<List<SearchResult>> references(
-      Element? element, SearchedFiles searchedFiles) async {
+      Element2? element, SearchedFiles searchedFiles) async {
     if (element == null) {
       return const <SearchResult>[];
     }
 
     ElementKind kind = element.kind;
-    if (element is ExtensionElement ||
-        element is InterfaceElement ||
-        element is PropertyAccessorElement && element.isSetter ||
-        element is TypeAliasElement) {
+    if (element is ExtensionElement2 ||
+        element is InterfaceElement2 ||
+        element is SetterElement ||
+        element is TypeAliasElement2) {
       return _searchReferences(element, searchedFiles);
-    } else if (element is ConstructorElement) {
+    } else if (element is ConstructorElement2) {
       return await _searchReferences_Constructor(element, searchedFiles);
-    } else if (element is CompilationUnitElementImpl) {
-      return _searchReferences_CompilationUnit(element);
-    } else if (element is PropertyAccessorElement && element.isGetter) {
+    } else if (element is GetterElement) {
       return _searchReferences_Getter(element, searchedFiles);
-    } else if (element is PropertyInducingElement) {
+    } else if (element is PropertyInducingElement2) {
       return _searchReferences_Field(element, searchedFiles);
-    } else if (kind == ElementKind.FUNCTION || kind == ElementKind.METHOD) {
-      if (element.enclosingElement3 is ExecutableElement) {
-        return _searchReferences_Local(
-            element, (n) => n is Block, searchedFiles);
-      }
+    } else if (element is LocalFunctionElement) {
+      return _searchReferences_Local(element, (n) => n is Block, searchedFiles);
+    } else if (element is ExecutableElement2) {
       return _searchReferences_Function(element, searchedFiles);
-    } else if (element is LibraryImportElement) {
-      return _searchReferences_Import(element, searchedFiles);
-    } else if (element is PatternVariableElementImpl) {
+    } else if (element is PatternVariableElementImpl2) {
       return _searchReferences_PatternVariable(element, searchedFiles);
     } else if (kind == ElementKind.LABEL ||
         kind == ElementKind.LOCAL_VARIABLE) {
@@ -409,25 +401,19 @@
               n is SwitchExpression ||
               n.parent is CompilationUnit,
           searchedFiles);
-    } else if (element is LibraryElement) {
+    } else if (element is LibraryElementImpl) {
       return _searchReferences_Library(element, searchedFiles);
-    } else if (element is ParameterElement) {
+    } else if (element is FormalParameterElement) {
       return _searchReferences_Parameter(element, searchedFiles);
-    } else if (element is PrefixElement) {
+    } else if (element is PrefixElementImpl2) {
       return _searchReferences_Prefix(element, searchedFiles);
-    } else if (element is TypeParameterElement) {
+    } else if (element is TypeParameterElement2) {
       return _searchReferences_Local(
           element, (n) => n.parent is CompilationUnit, searchedFiles);
     }
     return const <SearchResult>[];
   }
 
-  /// Returns references to the [element].
-  Future<List<SearchResult>> references2(
-      Element2? element, SearchedFiles searchedFiles) {
-    return references(element.asElement, searchedFiles);
-  }
-
   Future<List<LibraryFragmentSearchMatch>> referencesLibraryFragment(
     LibraryFragment libraryFragment,
   ) async {
@@ -474,7 +460,7 @@
     List<SearchResult> results = <SearchResult>[];
     await _addResults(
       results,
-      type.asElement,
+      type,
       searchedFiles,
       const {
         IndexRelationKind.IS_EXTENDED_BY:
@@ -493,12 +479,14 @@
   /// Return direct [SubtypeResult]s for either the [type] or [subtype].
   Future<List<SubtypeResult>> subtypes(SearchedFiles searchedFiles,
       {InterfaceElement2? type, SubtypeResult? subtype}) async {
-    var type1 = type?.asElement;
+    var type1 = type;
     String name;
     String id;
     if (type1 != null) {
-      name = type1.name;
-      id = '${type1.librarySource.uri};${type1.source.uri};$name';
+      name = type1.name3!;
+      var librarySource = type1.library2.firstFragment.source;
+      var source = type1.firstFragment.libraryFragment.source;
+      id = '${librarySource.uri};${source.uri};$name';
     } else {
       name = subtype!.name;
       id = subtype.id;
@@ -527,10 +515,10 @@
   }
 
   /// Returns top-level elements with names matching the given [regExp].
-  Future<List<Element>> topLevelElements(RegExp regExp) async {
-    List<Element> elements = <Element>[];
+  Future<List<Element2>> topLevelElements(RegExp regExp) async {
+    var elements = <Element2>[];
 
-    void addElement(Element element) {
+    void addElement(Element2 element) {
       if (!element.isSynthetic && regExp.hasMatch(element.displayName)) {
         elements.add(element);
       }
@@ -538,29 +526,24 @@
 
     List<FileState> knownFiles = _driver.fsState.knownFiles.toList();
     for (FileState file in knownFiles) {
-      var unitResult = await _driver.getUnitElement(file.path);
-      if (unitResult is UnitElementResult) {
-        CompilationUnitElement unitElement = unitResult.element;
-        unitElement.accessors.forEach(addElement);
-        unitElement.classes.forEach(addElement);
-        unitElement.enums.forEach(addElement);
-        unitElement.extensions.forEach(addElement);
-        unitElement.extensionTypes.forEach(addElement);
-        unitElement.functions.forEach(addElement);
-        unitElement.mixins.forEach(addElement);
-        unitElement.topLevelVariables.forEach(addElement);
-        unitElement.typeAliases.forEach(addElement);
+      var libraryResult = await _driver.getLibraryByUri(file.uriStr);
+      if (libraryResult is LibraryElementResult) {
+        var element = libraryResult.element2;
+        element.getters.forEach(addElement);
+        element.classes.forEach(addElement);
+        element.enums.forEach(addElement);
+        element.extensions.forEach(addElement);
+        element.extensionTypes.forEach(addElement);
+        element.topLevelFunctions.forEach(addElement);
+        element.mixins.forEach(addElement);
+        element.setters.forEach(addElement);
+        element.topLevelVariables.forEach(addElement);
+        element.typeAliases.forEach(addElement);
       }
     }
     return elements;
   }
 
-  /// Returns top-level elements with names matching the given [regExp].
-  Future<List<Element2>> topLevelElements2(RegExp regExp) async {
-    var results = await topLevelElements(regExp);
-    return results.map((e) => e.asElement2!).toList();
-  }
-
   /// Returns unresolved references to the given [name].
   Future<List<SearchResult>> unresolvedMemberReferences(
       String? name, SearchedFiles searchedFiles) async {
@@ -598,17 +581,17 @@
 
   Future<void> _addResults(
       List<SearchResult> results,
-      Element element,
+      Element2 element,
       SearchedFiles searchedFiles,
       Map<IndexRelationKind, SearchResultKind> relationToResultKind,
       {List<FileState>? filesToCheck}) async {
     // Prepare the element name.
     String name = element.displayName;
-    if (element is ConstructorElement) {
-      name = element.enclosingElement3.displayName;
+    if (element is ConstructorElement2) {
+      name = element.enclosingElement2.displayName;
     }
 
-    var elementPath = element.source!.fullName;
+    var elementPath = element.firstFragment.libraryFragment!.source.fullName;
     var elementFile = _driver.fsState.getExistingFromPath(elementPath);
     if (elementFile == null) {
       return;
@@ -617,7 +600,7 @@
     // Prepare the list of files that reference the element name.
     var files = <FileState>[];
     if (name.startsWith('_')) {
-      String libraryPath = element.library!.source.fullName;
+      String libraryPath = element.library2!.firstFragment.source.fullName;
       if (searchedFiles.add(libraryPath, this)) {
         var libraryFile = _driver.fsState.getFileForPath(libraryPath);
         var libraryKind = libraryFile.kind;
@@ -663,7 +646,7 @@
   /// Add results for [element] usage in the given [file].
   Future<void> _addResultsInFile(
       List<SearchResult> results,
-      Element element,
+      Element2 element,
       Map<IndexRelationKind, SearchResultKind> relationToResultKind,
       String file) async {
     var index = await _driver.getIndex(file);
@@ -678,13 +661,13 @@
     }
   }
 
-  Future<CompilationUnitElement?> _getUnitElement(String file) async {
+  Future<CompilationUnitElementImpl?> _getUnitElement(String file) async {
     var result = await _driver.getUnitElement(file);
-    return result is UnitElementResult ? result.element : null;
+    return result is UnitElementResultImpl ? result.fragment : null;
   }
 
   Future<List<SearchResult>> _searchReferences(
-      Element element, SearchedFiles searchedFiles) async {
+      Element2 element, SearchedFiles searchedFiles) async {
     List<SearchResult> results = <SearchResult>[];
     await _addResults(results, element, searchedFiles,
         const {IndexRelationKind.IS_REFERENCED_BY: SearchResultKind.REFERENCE});
@@ -725,21 +708,12 @@
           }
         }
       }
-
-      await _addResultsInFile(
-        results,
-        element,
-        const {
-          IndexRelationKind.IS_REFERENCED_BY: SearchResultKind.REFERENCE,
-        },
-        reference.path,
-      );
     }
     return results;
   }
 
   Future<List<SearchResult>> _searchReferences_Constructor(
-      ConstructorElement element, SearchedFiles searchedFiles) async {
+      ConstructorElement2 element, SearchedFiles searchedFiles) async {
     List<SearchResult> results = <SearchResult>[];
     await _addResults(results, element, searchedFiles, const {
       IndexRelationKind.IS_INVOKED_BY: SearchResultKind.INVOCATION,
@@ -753,10 +727,10 @@
   }
 
   Future<List<SearchResult>> _searchReferences_Field(
-      PropertyInducingElement field, SearchedFiles searchedFiles) async {
+      PropertyInducingElement2 field, SearchedFiles searchedFiles) async {
     List<SearchResult> results = <SearchResult>[];
-    PropertyAccessorElement? getter = field.getter;
-    PropertyAccessorElement? setter = field.setter;
+    var getter = field.getter2;
+    var setter = field.setter2;
     if (!field.isSynthetic) {
       await _addResults(results, field, searchedFiles, const {
         IndexRelationKind.IS_WRITTEN_BY: SearchResultKind.WRITE,
@@ -777,9 +751,9 @@
   }
 
   Future<List<SearchResult>> _searchReferences_Function(
-      Element element, SearchedFiles searchedFiles) async {
+      Element2 element, SearchedFiles searchedFiles) async {
     List<SearchResult> results = <SearchResult>[];
-    await _addResults(results, element.declaration!, searchedFiles, const {
+    await _addResults(results, element.baseElement, searchedFiles, const {
       IndexRelationKind.IS_REFERENCED_BY: SearchResultKind.REFERENCE,
       IndexRelationKind.IS_INVOKED_BY: SearchResultKind.INVOCATION
     });
@@ -787,7 +761,7 @@
   }
 
   Future<List<SearchResult>> _searchReferences_Getter(
-      PropertyAccessorElement getter, SearchedFiles searchedFiles) async {
+      GetterElement getter, SearchedFiles searchedFiles) async {
     List<SearchResult> results = <SearchResult>[];
     await _addResults(results, getter, searchedFiles, const {
       IndexRelationKind.IS_REFERENCED_BY: SearchResultKind.REFERENCE,
@@ -797,15 +771,15 @@
   }
 
   Future<List<SearchResult>> _searchReferences_Import(
-      LibraryImportElement element, SearchedFiles searchedFiles) async {
+      LibraryImportElementImpl element, SearchedFiles searchedFiles) async {
     String path = element.source.fullName;
     if (!searchedFiles.add(path, this)) {
       return const <SearchResult>[];
     }
 
     List<SearchResult> results = <SearchResult>[];
-    LibraryElement libraryElement = element.library;
-    for (CompilationUnitElement unitElement in libraryElement.units) {
+    LibraryElementImpl libraryElement = element.library;
+    for (var unitElement in libraryElement.units) {
       String unitPath = unitElement.source.fullName;
       var unitResult = await _driver.getResolvedUnit(unitPath);
       if (unitResult is ResolvedUnitResult) {
@@ -818,24 +792,25 @@
   }
 
   Future<List<SearchResult>> _searchReferences_Library(
-      LibraryElement element, SearchedFiles searchedFiles) async {
+      LibraryElementImpl element, SearchedFiles searchedFiles) async {
     String path = element.source.fullName;
     if (!searchedFiles.add(path, this)) {
       return const <SearchResult>[];
     }
 
     List<SearchResult> results = <SearchResult>[];
-    for (CompilationUnitElement unitElement in element.units) {
+    for (var unitElement in element.units) {
       String unitPath = unitElement.source.fullName;
       var unitResult = await _driver.getResolvedUnit(unitPath);
-      if (unitResult is ResolvedUnitResult) {
-        CompilationUnit unit = unitResult.unit;
-        for (Directive directive in unit.directives) {
-          if (directive is PartOfDirective && directive.element == element) {
+      if (unitResult is ResolvedUnitResultImpl) {
+        var unit = unitResult.unit;
+        for (var directive in unit.directives) {
+          if (directive is PartOfDirectiveImpl &&
+              directive.element == element) {
             var targetEntity = directive.libraryName ?? directive.uri;
             results.add(
               SearchResult._(
-                unit.declaredElement!,
+                unit.declaredFragment!,
                 SearchResultKind.REFERENCE,
                 targetEntity!.offset,
                 targetEntity.length,
@@ -850,22 +825,23 @@
     return results;
   }
 
-  Future<List<SearchResult>> _searchReferences_Local(Element element,
+  Future<List<SearchResult>> _searchReferences_Local(Element2 element,
       bool Function(AstNode n) isRootNode, SearchedFiles searchedFiles) async {
-    String path = element.source!.fullName;
+    String path = element.firstFragment.libraryFragment!.source.fullName;
     if (!searchedFiles.add(path, this)) {
       return const <SearchResult>[];
     }
 
     // Prepare the unit.
     var unitResult = await _driver.getResolvedUnit(path);
-    if (unitResult is! ResolvedUnitResult) {
+    if (unitResult is! ResolvedUnitResultImpl) {
       return const <SearchResult>[];
     }
     var unit = unitResult.unit;
 
     // Prepare the node.
-    var node = NodeLocator(element.nameOffset).searchWithin(unit);
+    var node =
+        NodeLocator(element.firstFragment.nameOffset2!).searchWithin(unit);
     if (node == null) {
       return const <SearchResult>[];
     }
@@ -875,7 +851,7 @@
         isRootNode(node) || node is ClassMember || node is CompilationUnit);
     assert(
       enclosingNode != null && enclosingNode is! CompilationUnit,
-      'Did not find enclosing node for local "${element.name}". '
+      'Did not find enclosing node for local "${element.name3}". '
       'Perhaps the isRootNode function is missing a condition to locate the '
       'outermost node where this element is in scope?',
     );
@@ -884,13 +860,13 @@
     }
 
     // Find the matches.
-    var visitor = _LocalReferencesVisitor({element}, unit.declaredElement!);
+    var visitor = _LocalReferencesVisitor({element}, unit.declaredFragment!);
     enclosingNode.accept(visitor);
     return visitor.results;
   }
 
   Future<List<SearchResult>> _searchReferences_Parameter(
-      ParameterElement parameter, SearchedFiles searchedFiles) async {
+      FormalParameterElement parameter, SearchedFiles searchedFiles) async {
     List<SearchResult> results = <SearchResult>[];
     results.addAll(await _searchReferences_Local(
       parameter,
@@ -902,29 +878,29 @@
     ));
     if (parameter.isNamed ||
         parameter.isOptionalPositional ||
-        parameter.enclosingElement3 is ConstructorElement) {
+        parameter.enclosingElement2 is ConstructorElement2) {
       results.addAll(await _searchReferences(parameter, searchedFiles));
     }
     return results;
   }
 
   Future<List<SearchResult>> _searchReferences_PatternVariable(
-    PatternVariableElementImpl element,
+    PatternVariableElementImpl2 element,
     SearchedFiles searchedFiles,
   ) async {
-    String path = element.source.fullName;
+    String path = element.firstFragment.libraryFragment.source.fullName;
     if (!searchedFiles.add(path, this)) {
       return const <SearchResult>[];
     }
 
     var rootVariable = element.rootVariable;
-    var transitiveVariables = rootVariable is JoinPatternVariableElementImpl
+    var transitiveVariables = rootVariable is JoinPatternVariableElementImpl2
         ? rootVariable.transitiveVariables
         : [rootVariable];
 
     // Prepare a binding element for the variable.
     var bindElement = transitiveVariables
-        .whereType<BindPatternVariableElementImpl>()
+        .whereType<BindPatternVariableElementImpl2>()
         .firstOrNull;
     if (bindElement == null) {
       return const <SearchResult>[];
@@ -941,22 +917,22 @@
     // Find the matches.
     var visitor = _LocalReferencesVisitor(
       transitiveVariables.toSet(),
-      bindElement.enclosingUnit,
+      bindElement.firstFragment.libraryFragment,
     );
     rootNode.accept(visitor);
     return visitor.results;
   }
 
   Future<List<SearchResult>> _searchReferences_Prefix(
-      PrefixElement element, SearchedFiles searchedFiles) async {
-    String path = element.source.fullName;
+      PrefixElementImpl2 element, SearchedFiles searchedFiles) async {
+    String path = element.firstFragment.libraryFragment.source.fullName;
     if (!searchedFiles.add(path, this)) {
       return const <SearchResult>[];
     }
 
     List<SearchResult> results = <SearchResult>[];
-    LibraryElement libraryElement = element.library;
-    for (CompilationUnitElement unitElement in libraryElement.units) {
+    var libraryElement = element.library2;
+    for (var unitElement in libraryElement.units) {
       String unitPath = unitElement.source.fullName;
       var unitResult = await _driver.getResolvedUnit(unitPath);
       if (unitResult is ResolvedUnitResult) {
@@ -1011,8 +987,8 @@
 
 /// A single search result.
 class SearchResult {
-  /// The deep most element that contains this result.
-  final Element enclosingElement;
+  /// The deep most fragment that contains this result.
+  final Fragment enclosingFragment;
 
   /// The kind of the element usage.
   final SearchResultKind kind;
@@ -1029,25 +1005,16 @@
   /// Whether the result is a resolved reference to the element.
   final bool isQualified;
 
-  SearchResult._(this.enclosingElement, this.kind, this.offset, this.length,
+  SearchResult._(this.enclosingFragment, this.kind, this.offset, this.length,
       this.isResolved, this.isQualified);
 
-  Element2 get enclosingElement2 => enclosingElement.asElement2!;
-
-  Fragment get enclosingFragment {
-    if (enclosingElement case CompilationUnitElementImpl libraryFragment) {
-      return libraryFragment;
-    }
-    return enclosingElement as Fragment;
-  }
-
   @override
   String toString() {
     StringBuffer buffer = StringBuffer();
     buffer.write("SearchResult(kind=");
     buffer.write(kind);
-    buffer.write(", enclosingElement=");
-    buffer.write(enclosingElement);
+    buffer.write(", enclosingFragment=");
+    buffer.write(enclosingFragment);
     buffer.write(", offset=");
     buffer.write(offset);
     buffer.write(", length=");
@@ -1119,195 +1086,20 @@
   }
 }
 
-/// A visitor that finds the deep-most [Element] that contains the [offset].
-class _ContainingElementFinder extends GeneralizingElementVisitor<void> {
+/// A visitor that finds the deep-most [ElementImpl] that contains the [offset].
+class _ContainingFragmentFinder extends GeneralizingElementVisitor<void> {
   final int offset;
-  Element? containingElement;
+  Fragment? containingFragment;
 
-  _ContainingElementFinder(this.offset);
+  _ContainingFragmentFinder(this.offset);
 
   @override
-  void visitElement(Element element) {
-    if (element is ElementImpl) {
-      if (element.codeOffset != null &&
-          element.codeOffset! <= offset &&
-          offset <= element.codeOffset! + element.codeLength!) {
-        containingElement = element;
-        super.visitElement(element);
-      }
-    }
-  }
-}
-
-class _FindCompilationUnitDeclarations {
-  final CompilationUnitElement unit;
-  final String filePath;
-  final LineInfo lineInfo;
-  final WorkspaceSymbols result;
-  final int? maxResults;
-  final FuzzyMatcher matcher;
-  final void Function(Declaration) collect;
-
-  _FindCompilationUnitDeclarations(
-    this.unit,
-    this.filePath,
-    this.result,
-    this.maxResults,
-    this.matcher,
-    this.collect,
-  ) : lineInfo = unit.lineInfo;
-
-  void compute(CancellationToken? cancellationToken) {
-    if (result.hasMoreDeclarationsThan(maxResults)) {
-      return;
-    }
-
-    _addAccessors(unit.accessors);
-    _addClasses(unit.classes);
-    _addClasses(unit.enums);
-    _addClasses(unit.mixins);
-    _addExtensions(unit.extensions);
-    _addClasses(unit.extensionTypes);
-    _addFunctions(unit.functions);
-    _addTypeAliases(unit.typeAliases);
-    _addVariables(unit.topLevelVariables);
-  }
-
-  void _addAccessors(List<PropertyAccessorElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      if (!element.isSynthetic) {
-        _addDeclaration(element, element.displayName);
-      }
-    }
-  }
-
-  void _addClasses(List<InterfaceElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      _addDeclaration(element, element.name);
-      _addAccessors(element.accessors);
-      _addConstructors(element.constructors);
-      _addFields(element.fields);
-      _addMethods(element.methods);
-    }
-  }
-
-  void _addConstructors(List<ConstructorElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      if (!element.isSynthetic) {
-        _addDeclaration(element, element.name);
-      }
-    }
-  }
-
-  void _addDeclaration(Element element, String name) {
-    if (result.hasMoreDeclarationsThan(maxResults)) {
-      throw const _MaxNumberOfDeclarationsError();
-    }
-
-    if (matcher.score(name) < 0) {
-      return;
-    }
-
-    var enclosing = element.enclosingElement3;
-
-    String? className;
-    String? mixinName;
-    if (enclosing is EnumElement) {
-      // skip
-    } else if (enclosing is MixinElement) {
-      mixinName = enclosing.name;
-    } else if (enclosing is InterfaceElement) {
-      className = enclosing.name;
-    }
-
-    var kind = _getSearchElementKind(element);
-    if (kind == null) {
-      return;
-    }
-
-    String? parameters;
-    if (element is ExecutableElement) {
-      var displayString = element.getDisplayString();
-      var parameterIndex = displayString.indexOf('(');
-      if (parameterIndex > 0) {
-        parameters = displayString.substring(parameterIndex);
-      }
-    }
-
-    element as ElementImpl; // to access codeOffset/codeLength
-    var locationOffset = element.nameOffset;
-    var locationStart = lineInfo.getLocation(locationOffset);
-
-    collect(
-      Declaration(
-        result._getPathIndex(filePath),
-        lineInfo,
-        name,
-        kind,
-        locationOffset,
-        locationStart.lineNumber,
-        locationStart.columnNumber,
-        element.codeOffset ?? 0,
-        element.codeLength ?? 0,
-        className,
-        mixinName,
-        parameters,
-      ),
-    );
-  }
-
-  void _addExtensions(List<ExtensionElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      var name = element.name;
-      if (name != null) {
-        _addDeclaration(element, name);
-      }
-      _addAccessors(element.accessors);
-      _addFields(element.fields);
-      _addMethods(element.methods);
-    }
-  }
-
-  void _addFields(List<FieldElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      if (!element.isSynthetic) {
-        _addDeclaration(element, element.name);
-      }
-    }
-  }
-
-  void _addFunctions(List<FunctionElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      _addDeclaration(element, element.name);
-    }
-  }
-
-  void _addMethods(List<MethodElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      _addDeclaration(element, element.name);
-    }
-  }
-
-  void _addTypeAliases(List<TypeAliasElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      _addDeclaration(element, element.name);
-    }
-  }
-
-  void _addVariables(List<TopLevelVariableElement> elements) {
-    for (var i = 0; i < elements.length; i++) {
-      var element = elements[i];
-      if (!element.isSynthetic) {
-        _addDeclaration(element, element.name);
-      }
+  void visitElement(covariant ElementImpl element) {
+    if (element.codeOffset != null &&
+        element.codeOffset! <= offset &&
+        offset <= element.codeOffset! + element.codeLength!) {
+      containingFragment = element as Fragment;
+      super.visitElement(element);
     }
   }
 }
@@ -1368,25 +1160,21 @@
             continue;
           }
 
-          var units = libraryElement.units;
-          for (var i = 0; i < units.length; i++) {
-            var unit = units[i];
-            var filePath = unit.source.fullName;
-            if (onlyForFile != null && filePath != onlyForFile) {
-              continue;
-            }
-            performance.run('unitDeclarations', (performance) {
-              var finder = _FindCompilationUnitDeclarations(
-                unit,
-                filePath,
-                result,
-                maxResults,
-                matcher,
-                result.declarations.add,
-              );
-              finder.compute(cancellationToken);
-            });
+          var filePath = libraryElement.firstFragment.source.fullName;
+          if (onlyForFile != null && filePath != onlyForFile) {
+            continue;
           }
+
+          performance.run('libraryDeclarations', (performance) {
+            var finder = _FindLibraryDeclarations(
+              libraryElement.firstFragment,
+              result,
+              maxResults,
+              matcher,
+              result.declarations.add,
+            );
+            finder.compute(cancellationToken);
+          });
         }
 
         // Periodically yield and check cancellation token.
@@ -1404,6 +1192,208 @@
   }
 }
 
+class _FindLibraryDeclarations {
+  final LibraryElementImpl library;
+  final WorkspaceSymbols result;
+  final int? maxResults;
+  final FuzzyMatcher matcher;
+  final void Function(Declaration) collect;
+
+  _FindLibraryDeclarations(
+    CompilationUnitElementImpl unit,
+    this.result,
+    this.maxResults,
+    this.matcher,
+    this.collect,
+  ) : library = unit.element;
+
+  void compute(CancellationToken? cancellationToken) {
+    if (result.hasMoreDeclarationsThan(maxResults)) {
+      return;
+    }
+
+    _addClasses(library.classes);
+    _addGetters(library.getters);
+    _addClasses(library.enums);
+    _addClasses(library.mixins);
+    _addExtensions(library.extensions);
+    _addClasses(library.extensionTypes);
+    _addSetters(library.setters);
+    _addFunctions(library.topLevelFunctions);
+    _addVariables(library.topLevelVariables);
+    _addTypeAliases(library.typeAliases);
+  }
+
+  void _addClasses(List<InterfaceElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      _addDeclaration(element, element.name3!);
+      _addGetters(element.getters2);
+      _addConstructors(element.constructors2);
+      _addFields(element.fields2);
+      _addMethods(element.methods2);
+      _addSetters(element.setters2);
+    }
+  }
+
+  void _addConstructors(List<ConstructorElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      if (!element.isSynthetic) {
+        _addDeclaration(element, element.name3!);
+      }
+    }
+  }
+
+  void _addDeclaration(Element2 element, String name) {
+    if (result.hasMoreDeclarationsThan(maxResults)) {
+      throw const _MaxNumberOfDeclarationsError();
+    }
+
+    if (matcher.score(name) < 0) {
+      return;
+    }
+
+    var enclosing = element.enclosingElement2;
+
+    String? className;
+    String? mixinName;
+    if (enclosing is EnumElement2) {
+      // skip
+    } else if (enclosing is MixinElement2) {
+      mixinName = enclosing.name3;
+    } else if (enclosing is InterfaceElement2) {
+      className = enclosing.name3;
+    }
+
+    var kind = _getSearchElementKind(element);
+    if (kind == null) {
+      return;
+    }
+
+    String? parameters;
+    if (element is ExecutableElement2) {
+      var displayString = element.displayString2();
+      var parameterIndex = displayString.indexOf('(');
+      if (parameterIndex > 0) {
+        parameters = displayString.substring(parameterIndex);
+      }
+    }
+
+    var firstFragment = element.firstFragment;
+    var libraryFragment = firstFragment.libraryFragment;
+    if (libraryFragment == null) {
+      return;
+    }
+
+    var filePath = libraryFragment.source.fullName;
+
+    var locationOffset = firstFragment.nameOffset2;
+    if (locationOffset == null) {
+      if (firstFragment is ConstructorFragment) {
+        locationOffset = firstFragment.typeNameOffset;
+      }
+    }
+
+    if (locationOffset == null) {
+      return;
+    }
+    var lineInfo = libraryFragment.lineInfo;
+    var locationStart = lineInfo.getLocation(locationOffset);
+
+    var fragmentImpl =
+        firstFragment as ElementImpl; // to access codeOffset/codeLength
+
+    collect(
+      Declaration(
+        result._getPathIndex(filePath),
+        lineInfo,
+        name,
+        kind,
+        locationOffset,
+        locationStart.lineNumber,
+        locationStart.columnNumber,
+        fragmentImpl.codeOffset ?? 0,
+        fragmentImpl.codeLength ?? 0,
+        className,
+        mixinName,
+        parameters,
+      ),
+    );
+  }
+
+  void _addExtensions(List<ExtensionElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      var name = element.name3;
+      if (name != null) {
+        _addDeclaration(element, name);
+      }
+      _addFields(element.fields2);
+      _addGetters(element.getters2);
+      _addMethods(element.methods2);
+      _addSetters(element.setters2);
+    }
+  }
+
+  void _addFields(List<FieldElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      if (!element.isSynthetic) {
+        _addDeclaration(element, element.name3!);
+      }
+    }
+  }
+
+  void _addFunctions(List<TopLevelFunctionElement> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      _addDeclaration(element, element.name3!);
+    }
+  }
+
+  void _addGetters(List<GetterElement> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      if (!element.isSynthetic) {
+        _addDeclaration(element, element.displayName);
+      }
+    }
+  }
+
+  void _addMethods(List<MethodElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      _addDeclaration(element, element.name3!);
+    }
+  }
+
+  void _addSetters(List<SetterElement> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      if (!element.isSynthetic) {
+        _addDeclaration(element, element.displayName);
+      }
+    }
+  }
+
+  void _addTypeAliases(List<TypeAliasElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      _addDeclaration(element, element.name3!);
+    }
+  }
+
+  void _addVariables(List<TopLevelVariableElement2> elements) {
+    for (var i = 0; i < elements.length; i++) {
+      var element = elements[i];
+      if (!element.isSynthetic) {
+        _addDeclaration(element, element.name3!);
+      }
+    }
+  }
+}
+
 class _IndexRequest {
   final AnalysisDriverUnitIndex index;
 
@@ -1444,16 +1434,16 @@
 
   /// Return the [element]'s identifier in the [index] or `-1` if the
   /// [element] is not referenced in the [index].
-  int findElementId(Element element) {
-    IndexElementInfo info = IndexElementInfo(element.asElement2!);
-    element = info.element.asElement!;
+  int findElementId(Element2 element) {
+    IndexElementInfo info = IndexElementInfo(element);
+    element = info.element;
     // Find the id of the element's unit.
     int unitId = getUnitId(element);
     if (unitId == -1) {
       return -1;
     }
     // Prepare information about the element.
-    var components = ElementNameComponents(element.asElement2!);
+    var components = ElementNameComponents(element);
     int unitMemberId = index.getStringId(components.unitMemberName);
     if (unitMemberId == -1) {
       return -1;
@@ -1491,12 +1481,12 @@
   /// a relation with the kind from [relationToResultKind].
   ///
   /// The function [getEnclosingUnitElement] is used to lazily compute the
-  /// enclosing [CompilationUnitElement] if there is a relation of an
+  /// enclosing [CompilationUnitElementImpl] if there is a relation of an
   /// interesting kind.
   Future<List<SearchResult>> getRelations(
       int elementId,
       Map<IndexRelationKind, SearchResultKind> relationToResultKind,
-      Future<CompilationUnitElement?> Function()
+      Future<CompilationUnitElementImpl?> Function()
           getEnclosingUnitElement) async {
     // Find the first usage of the element.
     int i = _findFirstOccurrence(index.usedElements, elementId);
@@ -1505,7 +1495,7 @@
     }
     // Create locations for every usage of the element.
     List<SearchResult> results = <SearchResult>[];
-    CompilationUnitElement? enclosingUnitElement;
+    CompilationUnitElementImpl? enclosingUnitElement;
     for (;
         i < index.usedElements.length && index.usedElements[i] == elementId;
         i++) {
@@ -1515,10 +1505,10 @@
         int offset = index.usedElementOffsets[i];
         enclosingUnitElement ??= await getEnclosingUnitElement();
         if (enclosingUnitElement != null) {
-          Element enclosingElement =
-              _getEnclosingElement(enclosingUnitElement, offset);
+          var enclosingFragment =
+              _getEnclosingFragment(enclosingUnitElement, offset);
           results.add(SearchResult._(
-            enclosingElement,
+            enclosingFragment,
             resultKind,
             offset,
             index.usedElementLengths[i],
@@ -1531,10 +1521,10 @@
     return results;
   }
 
-  /// Return the identifier of the [CompilationUnitElement] containing the
+  /// Return the identifier of the [CompilationUnitElementIml] containing the
   /// [element] in the [index] or `-1` if not found.
-  int getUnitId(Element element) {
-    var unitElement = getUnitElement(element.asElement2!);
+  int getUnitId(Element2 element) {
+    var unitElement = getUnitElement(element);
     return index.getLibraryFragmentId(unitElement);
   }
 
@@ -1543,7 +1533,7 @@
   Future<List<SearchResult>> getUnresolvedMemberReferences(
       String name,
       Map<IndexRelationKind, SearchResultKind> relationToResultKind,
-      Future<CompilationUnitElement?> Function()
+      Future<CompilationUnitElementImpl?> Function()
           getEnclosingUnitElement) async {
     // Find the name identifier.
     int nameId = index.getStringId(name);
@@ -1559,7 +1549,7 @@
 
     // Create results for every usage of the name.
     List<SearchResult> results = <SearchResult>[];
-    CompilationUnitElement? enclosingUnitElement;
+    CompilationUnitElementImpl? enclosingUnitElement;
     for (; i < index.usedNames.length && index.usedNames[i] == nameId; i++) {
       IndexRelationKind relationKind = index.usedNameKinds[i];
       SearchResultKind? resultKind = relationToResultKind[relationKind];
@@ -1567,9 +1557,9 @@
         int offset = index.usedNameOffsets[i];
         enclosingUnitElement ??= await getEnclosingUnitElement();
         if (enclosingUnitElement != null) {
-          Element enclosingElement =
-              _getEnclosingElement(enclosingUnitElement, offset);
-          results.add(SearchResult._(enclosingElement, resultKind, offset,
+          var enclosingFragment =
+              _getEnclosingFragment(enclosingUnitElement, offset);
+          results.add(SearchResult._(enclosingFragment, resultKind, offset,
               name.length, false, index.usedNameIsQualifiedFlags[i]));
         }
       }
@@ -1600,14 +1590,14 @@
 class _LocalReferencesVisitor extends RecursiveAstVisitor<void> {
   final List<SearchResult> results = <SearchResult>[];
 
-  final Set<Element> elements;
-  final CompilationUnitElement enclosingUnitElement;
+  final Set<Element2> elements;
+  final CompilationUnitElementImpl enclosingLibraryFragment;
 
-  _LocalReferencesVisitor(this.elements, this.enclosingUnitElement);
+  _LocalReferencesVisitor(this.elements, this.enclosingLibraryFragment);
 
   @override
   void visitAssignedVariablePattern(AssignedVariablePattern node) {
-    if (elements.contains(node.element)) {
+    if (elements.contains(node.element2)) {
       _addResult(node, SearchResultKind.WRITE);
     }
 
@@ -1623,7 +1613,7 @@
 
   @override
   void visitImportPrefixReference(ImportPrefixReference node) {
-    var element = node.element;
+    var element = node.element2;
     if (elements.contains(element)) {
       _addResult(node.name, SearchResultKind.REFERENCE);
     }
@@ -1631,7 +1621,7 @@
 
   @override
   void visitNamedType(NamedType node) {
-    var element = node.element;
+    var element = node.element2;
     if (elements.contains(element)) {
       _addResult(node.name2, SearchResultKind.REFERENCE);
     }
@@ -1645,15 +1635,15 @@
     if (node.inDeclarationContext()) {
       return;
     }
-    var element = node.staticElement;
+    var element = node.element;
     if (elements.contains(element)) {
       var parent = node.parent;
       SearchResultKind kind = SearchResultKind.REFERENCE;
-      if (element is FunctionElement) {
+      if (element is LocalFunctionElement) {
         if (parent is MethodInvocation && parent.methodName == node) {
           kind = SearchResultKind.INVOCATION;
         }
-      } else if (element is VariableElement) {
+      } else if (element is VariableElement2) {
         bool isGet = node.inGetterContext();
         bool isSet = node.inSetterContext();
         if (isGet && isSet) {
@@ -1674,9 +1664,9 @@
 
   void _addResult(SyntacticEntity entity, SearchResultKind kind) {
     bool isQualified = entity is AstNode ? entity.parent is Label : false;
-    Element enclosingElement =
-        _getEnclosingElement(enclosingUnitElement, entity.offset);
-    results.add(SearchResult._(enclosingElement, kind, entity.offset,
+    var enclosingFragment =
+        _getEnclosingFragment(enclosingLibraryFragment, entity.offset);
+    results.add(SearchResult._(enclosingFragment, kind, entity.offset,
         entity.length, true, isQualified));
   }
 }
diff --git a/pkg/analyzer/lib/src/dart/element/element.dart b/pkg/analyzer/lib/src/dart/element/element.dart
index 36d1372..bf9e9f0 100644
--- a/pkg/analyzer/lib/src/dart/element/element.dart
+++ b/pkg/analyzer/lib/src/dart/element/element.dart
@@ -3157,7 +3157,7 @@
     return enclosingElement3?.source;
   }
 
-  /// Whether to include the [nameOffset] in [identifier] to disambiguiate
+  /// Whether to include the [nameOffset] in [identifier] to disambiguate
   /// elements that might otherwise have the same identifier.
   bool get _includeNameOffsetInIdentifier {
     var element = this;
@@ -9940,8 +9940,14 @@
       _wrappedElement.isVisitingWhenClause = value;
 
   @override
-  JoinPatternVariableElement2? get join2 =>
-      JoinPatternVariableElementImpl2(_wrappedElement.join!);
+  JoinPatternVariableElementImpl2? get join2 {
+    return _wrappedElement.join?.asElement2;
+  }
+
+  /// Return the root [join2], or self.
+  PatternVariableElementImpl2 get rootVariable {
+    return join2?.rootVariable ?? this;
+  }
 
   @override
   PatternVariableElementImpl get _wrappedElement =>
diff --git a/pkg/analyzer/lib/src/dart/micro/resolve_file.dart b/pkg/analyzer/lib/src/dart/micro/resolve_file.dart
index 1e4ee56..51012c7 100644
--- a/pkg/analyzer/lib/src/dart/micro/resolve_file.dart
+++ b/pkg/analyzer/lib/src/dart/micro/resolve_file.dart
@@ -298,7 +298,7 @@
       if (element is LocalVariableElement ||
           (element is ParameterElement && !element.isNamed)) {
         await collectReferences2(element.source!.fullName, performance!);
-      } else if (element is LibraryImportElement) {
+      } else if (element is LibraryImportElementImpl) {
         return await _searchReferences_Import(element);
       } else {
         var result = performance!.run('getFilesContaining', (performance) {
@@ -892,10 +892,10 @@
   }
 
   Future<List<CiderSearchMatch>> _searchReferences_Import(
-      LibraryImportElement element) async {
+      LibraryImportElementImpl element) async {
     var results = <CiderSearchMatch>[];
-    LibraryElement libraryElement = element.library;
-    for (CompilationUnitElement unitElement in libraryElement.units) {
+    var libraryElement = element.library;
+    for (var unitElement in libraryElement.units) {
       String unitPath = unitElement.source.fullName;
       var unitResult = await resolve(path: unitPath);
       var visitor = ImportElementReferencesVisitor(element, unitElement);
diff --git a/pkg/analyzer/lib/src/lint/constants.dart b/pkg/analyzer/lib/src/lint/constants.dart
index e0c08d0..1d79beb 100644
--- a/pkg/analyzer/lib/src/lint/constants.dart
+++ b/pkg/analyzer/lib/src/lint/constants.dart
@@ -15,7 +15,6 @@
 import 'package:analyzer/src/dart/constant/potentially_constant.dart';
 import 'package:analyzer/src/dart/constant/utilities.dart';
 import 'package:analyzer/src/dart/constant/value.dart';
-import 'package:analyzer/src/dart/element/element.dart';
 import 'package:analyzer/src/error/codes.dart';
 
 /// The result of attempting to evaluate an expression as a constant.
@@ -88,11 +87,11 @@
   /// Whether [ConstantVerifier] reports an error when computing the value of
   /// `this` as a constant.
   bool get hasConstantVerifierError {
-    var unitNode = thisOrAncestorOfType<CompilationUnit>();
+    var unitNode = thisOrAncestorOfType<CompilationUnitImpl>();
     var unitFragment = unitNode?.declaredFragment;
     if (unitFragment == null) return false;
 
-    var libraryElement = unitFragment.element as LibraryElementImpl;
+    var libraryElement = unitFragment.element;
     var declaredVariables = libraryElement.session.declaredVariables;
 
     var dependenciesFinder = ConstantExpressionsDependenciesFinder();
@@ -159,13 +158,13 @@
   /// Returns a [LinterConstantEvaluationResult], containing both the computed
   /// constant value, and a list of errors that occurred during the computation.
   LinterConstantEvaluationResult computeConstantValue() {
-    var unitNode = thisOrAncestorOfType<CompilationUnit>();
+    var unitNode = thisOrAncestorOfType<CompilationUnitImpl>();
     var unitFragment = unitNode?.declaredFragment;
     if (unitFragment == null) {
       return LinterConstantEvaluationResult._(null, []);
     }
 
-    var libraryElement = unitFragment.element as LibraryElementImpl;
+    var libraryElement = unitFragment.element;
     var declaredVariables = libraryElement.session.declaredVariables;
 
     var evaluationEngine = ConstantEvaluationEngine(
diff --git a/pkg/analyzer/lib/src/test_utilities/find_element.dart b/pkg/analyzer/lib/src/test_utilities/find_element.dart
index 55d182e..0e1b27d 100644
--- a/pkg/analyzer/lib/src/test_utilities/find_element.dart
+++ b/pkg/analyzer/lib/src/test_utilities/find_element.dart
@@ -21,7 +21,7 @@
   @override
   CompilationUnitElement get unitElement => unit.declaredElement!;
 
-  LibraryExportElement export(String targetUri) {
+  LibraryExportElementImpl export(String targetUri) {
     LibraryExportElement? result;
 
     for (var export in unitElement.libraryExports) {
@@ -35,7 +35,7 @@
     }
 
     if (result != null) {
-      return result;
+      return result as LibraryExportElementImpl;
     }
     throw StateError('Not found: $targetUri');
   }
@@ -53,7 +53,8 @@
     throw StateError('Not found: $name');
   }
 
-  LibraryImportElement import(String targetUri, {bool mustBeUnique = true}) {
+  LibraryImportElementImpl import(String targetUri,
+      {bool mustBeUnique = true}) {
     LibraryImportElement? importElement;
 
     for (var libraryFragment in unitElement.withEnclosing) {
@@ -70,7 +71,7 @@
     }
 
     if (importElement != null) {
-      return importElement;
+      return importElement as LibraryImportElementImpl;
     }
     throw StateError('Not found: $targetUri');
   }
@@ -288,7 +289,9 @@
 
   ImportFindElement(this.import);
 
-  LibraryElement get importedLibrary => import.importedLibrary!;
+  LibraryElementImpl get importedLibrary {
+    return import.importedLibrary as LibraryElementImpl;
+  }
 
   PrefixElement? get prefix => import.prefix?.element;
 
diff --git a/pkg/analyzer/lib/src/utilities/extensions/element.dart b/pkg/analyzer/lib/src/utilities/extensions/element.dart
index 4df71ba..c3208f7 100644
--- a/pkg/analyzer/lib/src/utilities/extensions/element.dart
+++ b/pkg/analyzer/lib/src/utilities/extensions/element.dart
@@ -544,6 +544,13 @@
   }
 }
 
+extension JoinPatternVariableElementImplExtension
+    on JoinPatternVariableElementImpl {
+  JoinPatternVariableElementImpl2 get asElement2 {
+    return element;
+  }
+}
+
 extension LabelElement2Extension on LabelElement2 {
   LabelElement get asElement {
     return firstFragment as LabelElement;
@@ -692,6 +699,12 @@
   }
 }
 
+extension PatternVariableElementImplExtension on PatternVariableElementImpl {
+  PatternVariableElementImpl2 get asElement2 {
+    return element;
+  }
+}
+
 extension PrefixElement2Extension on PrefixElement2 {
   PrefixElement get asElement {
     return (this as PrefixElementImpl2).asElement;
diff --git a/pkg/analyzer/test/error/error_reporter_test.dart b/pkg/analyzer/test/error/error_reporter_test.dart
index e31475e..49d97dd 100644
--- a/pkg/analyzer/test/error/error_reporter_test.dart
+++ b/pkg/analyzer/test/error/error_reporter_test.dart
@@ -145,7 +145,7 @@
     var fa = findNode.topLevelVariableDeclaration('fa');
     var fb = findNode.topLevelVariableDeclaration('fb');
 
-    var source = result.unit.declaredFragment!.libraryFragment!.source;
+    var source = result.unit.declaredFragment!.source;
     var reporter = ErrorReporter(listener, source);
     reporter.atNode(
       findNode.simple('x'),
@@ -176,7 +176,7 @@
     var ba = findNode.topLevelVariableDeclaration('ba');
     var bb = findNode.topLevelVariableDeclaration('bb');
 
-    var source = result.unit.declaredFragment!.libraryFragment!.source;
+    var source = result.unit.declaredFragment!.source;
     var reporter = ErrorReporter(listener, source);
     reporter.atNode(
       findNode.simple('x'),
diff --git a/pkg/analyzer/test/src/dart/analysis/index_test.dart b/pkg/analyzer/test/src/dart/analysis/index_test.dart
index a4cd66b..b32a781 100644
--- a/pkg/analyzer/test/src/dart/analysis/index_test.dart
+++ b/pkg/analyzer/test/src/dart/analysis/index_test.dart
@@ -1054,9 +1054,10 @@
     await _indexTestUnit('''
 export 'lib.dart';
 ''');
-    var element = findElement.export('package:test/lib.dart').exportedLibrary!;
-    assertElementIndexText(element, r'''
-7 1:8 |'lib.dart'| IS_REFERENCED_BY qualified
+    var export = findElement.export('package:test/lib.dart');
+    var fragment = export.exportedLibrary2!.firstFragment;
+    assertLibraryFragmentIndexText(fragment, r'''
+7 1:8 |'lib.dart'|
 ''');
   }
 
@@ -1067,9 +1068,10 @@
     await _indexTestUnit('''
 import 'lib.dart';
 ''');
-    var element = findElement.import('package:test/lib.dart').importedLibrary!;
-    assertElementIndexText(element, r'''
-7 1:8 |'lib.dart'| IS_REFERENCED_BY qualified
+    var import = findElement.import('package:test/lib.dart');
+    var fragment = import.importedLibrary2!.firstFragment;
+    assertLibraryFragmentIndexText(fragment, r'''
+7 1:8 |'lib.dart'|
 ''');
   }
 
@@ -1872,8 +1874,9 @@
 ''');
 
     var importFind = findElement.importFind('package:test/foo.dart');
-    assertElementIndexText(importFind.importedLibrary, r'''
-7 1:8 |"foo.dart"| IS_REFERENCED_BY qualified
+    var importedFragment = importFind.importedLibrary.firstFragment;
+    assertLibraryFragmentIndexText(importedFragment, r'''
+7 1:8 |"foo.dart"|
 ''');
 
     FunctionElement bar = importFind.topFunction('bar');
diff --git a/pkg/analyzer/test/src/dart/analysis/search_test.dart b/pkg/analyzer/test/src/dart/analysis/search_test.dart
index 5ce1061..4ba0391 100644
--- a/pkg/analyzer/test/src/dart/analysis/search_test.dart
+++ b/pkg/analyzer/test/src/dart/analysis/search_test.dart
@@ -87,7 +87,7 @@
     String expected,
   ) async {
     var searchedFiles = SearchedFiles();
-    var results = await driver.search.references2(element, searchedFiles);
+    var results = await driver.search.references(element, searchedFiles);
     var actual = _getSearchResultsText(results);
     if (actual != expected) {
       print(actual);
@@ -266,7 +266,7 @@
     offset: 16 2:7
     codeOffset: 12 + 5
     className: C
-  CONSTRUCTOR <unnamed>
+  CONSTRUCTOR new
     offset: 21 3:3
     codeOffset: 21 + 4
     className: C
@@ -439,7 +439,7 @@
   EXTENSION_TYPE E
     offset: 15 1:16
     codeOffset: 0 + 79
-  CONSTRUCTOR <unnamed>
+  CONSTRUCTOR new
     offset: 15 1:16
     codeOffset: 16 + 8
     className: E
@@ -760,7 +760,7 @@
     var testDriver = driverFor(testFile);
 
     // No references, but this is not the most important.
-    var references = await testDriver.search.references2(A, searchedFiles);
+    var references = await testDriver.search.references(A, searchedFiles);
     expect(references, isEmpty);
 
     // We should not add the file to known files. It is not in the
@@ -1831,7 +1831,6 @@
 ''');
   }
 
-
   test_searchReferences_MethodElement_class() async {
     await resolveTestCode('''
 class A {
@@ -3022,7 +3021,7 @@
 
     void assertHasResult(String uriStr, String name, {bool not = false}) {
       var matcher = contains(predicate((SearchResult r) {
-        var element = r.enclosingElement2;
+        var element = r.enclosingFragment.element;
         return element.library2!.uri.toString() == uriStr &&
             element.name3 == name;
       }));
@@ -3064,7 +3063,6 @@
     expect(c.id, endsWith('c.dart;C'));
   }
 
-
   test_subtypes_enum() async {
     await resolveTestCode('''
 class A {}
@@ -3209,13 +3207,13 @@
     var f = findElement2.function('f');
     var g = findElement2.topVar('g');
     RegExp regExp = RegExp(r'^[ABCDfg]$');
-    expect(await driver.search.topLevelElements2(regExp),
+    expect(await driver.search.topLevelElements(regExp),
         unorderedEquals([a, b, c, d, f, g]));
   }
 
   Future<List<Element2>> _findClassMembers(String name) {
     var searchedFiles = SearchedFiles();
-    return driver.search.classMembers2(name, searchedFiles);
+    return driver.search.classMembers(name, searchedFiles);
   }
 
   String _getDeclarationsText(
diff --git a/pkg/analyzer/test/src/diagnostics/invalid_language_override_greater_test.dart b/pkg/analyzer/test/src/diagnostics/invalid_language_override_greater_test.dart
index 08b838c..25752f5 100644
--- a/pkg/analyzer/test/src/diagnostics/invalid_language_override_greater_test.dart
+++ b/pkg/analyzer/test/src/diagnostics/invalid_language_override_greater_test.dart
@@ -3,7 +3,6 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analyzer/src/dart/analysis/experiments.dart';
-import 'package:analyzer/src/dart/ast/ast.dart';
 import 'package:analyzer/src/error/codes.g.dart';
 import 'package:pub_semver/pub_semver.dart';
 import 'package:test/test.dart';
@@ -64,8 +63,7 @@
     required Version package,
     required Version? override,
   }) {
-    var unitImpl = result.unit as CompilationUnitImpl;
-    var languageVersion = unitImpl.languageVersion;
+    var languageVersion = result.unit.languageVersion;
     expect(languageVersion.package, package);
     expect(languageVersion.override, override);
   }