Use LibraryElement(s) to find extensions for 'Import Library' quick fix.

So, we don't need ExtensionCache anymore.

This fixes the issue with importing a library from lib/src sometimes.

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

Author: scheglov

pkg/analysis_server/lib/plugin/edit/fix/fix_dart.dart

@@ -3,7 +3,6 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analysis_server/plugin/edit/fix/fix_core.dart';
-import 'package:analysis_server/src/services/completion/dart/extension_cache.dart';
 import 'package:analyzer/dart/analysis/results.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/instrumentation/service.dart';
@@ -13,9 +12,6 @@ import 'package:analyzer_plugin/utilities/change_builder/change_workspace.dart';
 ///
 /// Clients may not extend, implement or mix-in this class.
 abstract class DartFixContext implements FixContext {
-  /// Return the extension cache used to find available extensions.
-  ExtensionCache get extensionCache;
-
   /// Return the instrumentation service used to report errors that prevent a
   /// fix from being composed.
   InstrumentationService get instrumentationService;
@@ -31,4 +27,8 @@ abstract class DartFixContext implements FixContext {
   Future<Map<LibraryElement, List<Element>>> getTopLevelDeclarations(
     String name,
   );
+
+  /// Return libraries with extensions that declare non-static public
+  /// extension members with the [memberName].
+  Stream<LibraryElement> librariesWithExtensions(String memberName);
 }

pkg/analysis_server/lib/src/analysis_server.dart

@@ -783,7 +783,6 @@ class ServerContextManagerCallbacks extends ContextManagerCallbacks {
     var path = result.path;
 
     analysisServer.getDocumentationCacheFor(result)?.cacheFromResult(result);
-    analysisServer.getExtensionCacheFor(result)?.cacheFromResult(result);
 
     var unit = result.unit;
     if (analysisServer._hasAnalysisServiceSubscription(

pkg/analysis_server/lib/src/analysis_server_abstract.dart

@@ -15,7 +15,6 @@ import 'package:analysis_server/src/plugin/plugin_watcher.dart';
 import 'package:analysis_server/src/server/crash_reporting_attachments.dart';
 import 'package:analysis_server/src/server/diagnostic_server.dart';
 import 'package:analysis_server/src/services/completion/dart/documentation_cache.dart';
-import 'package:analysis_server/src/services/completion/dart/extension_cache.dart';
 import 'package:analysis_server/src/services/correction/namespace.dart';
 import 'package:analysis_server/src/services/pub/pub_api.dart';
 import 'package:analysis_server/src/services/pub/pub_command.dart';
@@ -95,10 +94,6 @@ abstract class AbstractAnalysisServer {
   /// each context.
   Map<AnalysisContext, DocumentationCache> documentationForContext = {};
 
-  /// A map from analysis contexts to the extension cache associated with
-  /// each context.
-  Map<AnalysisContext, ExtensionCache> extensionForContext = {};
-
   /// The DiagnosticServer for this AnalysisServer. If available, it can be used
   /// to start an http diagnostics server or return the port for an existing
   /// server.
@@ -252,7 +247,6 @@ abstract class AbstractAnalysisServer {
   void addContextsToDeclarationsTracker() {
     declarationsTracker?.discardContexts();
     documentationForContext.clear();
-    extensionForContext.clear();
     for (var driver in driverMap.values) {
       declarationsTracker?.addContext(driver.analysisContext!);
     }
@@ -383,14 +377,6 @@ abstract class AbstractAnalysisServer {
     return element;
   }
 
-  /// Return the object used to cache information about extensions in the
-  /// context that produced the [result], or `null` if there is no cache for the
-  /// context.
-  ExtensionCache? getExtensionCacheFor(ResolvedUnitResult result) {
-    var context = result.session.analysisContext;
-    return extensionForContext.putIfAbsent(context, () => ExtensionCache());
-  }
-
   /// Return a [Future] that completes with the resolved [AstNode] at the
   /// given [offset] of the given [file], or with `null` if there is no node as
   /// the [offset].

pkg/analysis_server/lib/src/cider/fixes.dart

@@ -97,4 +97,7 @@ class _CiderDartFixContextImpl extends DartFixContextImpl {
     }
     return result;
   }
+
+  @override
+  Stream<LibraryElement> librariesWithExtensions(String memberName) async* {}
 }

pkg/analysis_server/lib/src/edit/edit_domain.dart

@@ -603,8 +603,7 @@ class EditDomainHandler extends AbstractRequestHandler {
         if (errorLine == requestLine) {
           var workspace = DartChangeWorkspace(server.currentSessions);
           var context = DartFixContextImpl(
-              server.instrumentationService, workspace, result, error,
-              extensionCache: server.getExtensionCacheFor(result));
+              server.instrumentationService, workspace, result, error);
 
           List<Fix> fixes;
           try {

pkg/analysis_server/lib/src/lsp/handlers/handler_code_actions.dart

@@ -346,8 +346,7 @@ class CodeActionHandler extends MessageHandler<CodeActionParams,
         }
         var workspace = DartChangeWorkspace(server.currentSessions);
         var context = DartFixContextImpl(
-            server.instrumentationService, workspace, unit, error,
-            extensionCache: server.getExtensionCacheFor(unit));
+            server.instrumentationService, workspace, unit, error);
         final fixes = await fixContributor.computeFixes(context);
         if (fixes.isNotEmpty) {
           final diagnostic = toDiagnostic(

pkg/analysis_server/lib/src/lsp/lsp_analysis_server.dart

@@ -896,7 +896,6 @@ class LspServerContextManagerCallbacks extends ContextManagerCallbacks {
     var path = result.path;
 
     analysisServer.getDocumentationCacheFor(result)?.cacheFromResult(result);
-    analysisServer.getExtensionCacheFor(result)?.cacheFromResult(result);
 
     final unit = result.unit;
     if (analysisServer.shouldSendClosingLabelsFor(path)) {

pkg/analysis_server/lib/src/services/completion/dart/extension_cache.dart

@@ -6,7 +6,6 @@ import 'dart:math' as math;
 
 import 'package:_fe_analyzer_shared/src/scanner/token.dart';
 import 'package:analysis_server/plugin/edit/fix/fix_dart.dart';
-import 'package:analysis_server/src/services/completion/dart/extension_cache.dart';
 import 'package:analysis_server/src/services/correction/fix/data_driven/transform_override_set.dart';
 import 'package:analysis_server/src/services/correction/util.dart';
 import 'package:analysis_server/src/utilities/flutter.dart';
@@ -412,9 +411,6 @@ abstract class _AbstractCorrectionProducer {
   /// Returns the EOL to use for this [CompilationUnit].
   String get eol => utils.endOfLine;
 
-  /// Return the extension cache used to find available extensions.
-  ExtensionCache get extensionCache => _context.dartFixContext!.extensionCache;
-
   String get file => _context.file;
 
   Flutter get flutter => Flutter.instance;
@@ -525,6 +521,12 @@ abstract class _AbstractCorrectionProducer {
     return false;
   }
 
+  /// Return libraries with extensions that declare non-static public
+  /// extension members with the [memberName].
+  Stream<LibraryElement> librariesWithExtensions(String memberName) {
+    return _context.dartFixContext!.librariesWithExtensions(memberName);
+  }
+
   /// Return `true` if the given [node] is in a location where an implicit
   /// constructor invocation would be allowed.
   bool mightBeImplicitConstructor(AstNode node) {

pkg/analysis_server/lib/src/services/correction/dart/abstract_producer.dart

@@ -9,7 +9,6 @@ import 'package:analysis_server/src/services/correction/fix.dart';
 import 'package:analysis_server/src/services/correction/namespace.dart';
 import 'package:analysis_server/src/services/linter/lint_names.dart';
 import 'package:analysis_server/src/utilities/extensions/element.dart';
-import 'package:analyzer/dart/analysis/results.dart';
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
@@ -45,16 +44,9 @@ class ImportLibrary extends MultiCorrectionProducer {
         if (targetType == null) {
           return;
         }
-        var definingLibraries =
-            extensionCache.membersByName[memberName]?.toList();
-        if (definingLibraries != null) {
-          for (var definingLibrary in definingLibraries) {
-            var libraryPath = definingLibrary.libraryPath;
-            var uri = sessionHelper.session.uriConverter.pathToUri(libraryPath);
-            if (uri != null) {
-              yield* _importExtensionInLibrary(uri, targetType, memberName);
-            }
-          }
+        await for (var libraryToImport in librariesWithExtensions(memberName)) {
+          yield* _importExtensionInLibrary(
+              libraryToImport, targetType, memberName);
         }
       }
 
@@ -146,15 +138,18 @@ class ImportLibrary extends MultiCorrectionProducer {
   }
 
   Stream<CorrectionProducer> _importExtensionInLibrary(
-      Uri uri, DartType targetType, String memberName) async* {
+    LibraryElement libraryToImport,
+    DartType targetType,
+    String memberName,
+  ) async* {
     // Look to see whether the library at the [uri] is already imported. If it
     // is, then we can check the extension elements without needing to perform
     // additional analysis.
     var foundImport = false;
     for (var imp in libraryElement.imports) {
       // prepare element
       var importedLibrary = imp.importedLibrary;
-      if (importedLibrary == null || importedLibrary.source.uri != uri) {
+      if (importedLibrary == null || importedLibrary != libraryToImport) {
         continue;
       }
       foundImport = true;
@@ -169,8 +164,8 @@ class ImportLibrary extends MultiCorrectionProducer {
             // TODO(brianwilkerson) Support removing the extension name from a
             //  hide combinator.
           } else if (combinator is ShowElementCombinator) {
-            yield _ImportLibraryShow(
-                uri.toString(), combinator, extension.name!);
+            yield _ImportLibraryShow(libraryToImport.source.uri.toString(),
+                combinator, extension.name!);
           }
         }
       }
@@ -180,7 +175,8 @@ class ImportLibrary extends MultiCorrectionProducer {
     // correction producer that will either add an import or not based on the
     // result of analyzing the library.
     if (!foundImport) {
-      yield _ImportLibraryContainingExtension(uri, targetType, memberName);
+      yield _ImportLibraryContainingExtension(
+          libraryToImport, targetType, memberName);
     }
   }
 
@@ -429,8 +425,8 @@ enum _ImportKind {
 /// A correction processor that can add an import of a library containing an
 /// extension, but which does so only if the extension applies to a given type.
 class _ImportLibraryContainingExtension extends CorrectionProducer {
-  /// The URI of the library defining the extension.
-  Uri uri;
+  /// The library defining the extension.
+  LibraryElement library;
 
   /// The type of the target that the extension must apply to.
   DartType targetType;
@@ -441,7 +437,11 @@ class _ImportLibraryContainingExtension extends CorrectionProducer {
   /// The URI that is being proposed for the import directive.
   String _uriText = '';
 
-  _ImportLibraryContainingExtension(this.uri, this.targetType, this.memberName);
+  _ImportLibraryContainingExtension(
+    this.library,
+    this.targetType,
+    this.memberName,
+  );
 
   @override
   List<Object> get fixArguments => [_uriText];
@@ -451,16 +451,12 @@ class _ImportLibraryContainingExtension extends CorrectionProducer {
 
   @override
   Future<void> compute(ChangeBuilder builder) async {
-    var result = await sessionHelper.session.getLibraryByUri(uri.toString());
-    if (result is LibraryElementResult) {
-      var library = result.element;
-      if (library
-          .matchingExtensionsWithMember(libraryElement, targetType, memberName)
-          .isNotEmpty) {
-        await builder.addDartFileEdit(file, (builder) {
-          _uriText = builder.importLibrary(uri);
-        });
-      }
+    if (library
+        .matchingExtensionsWithMember(libraryElement, targetType, memberName)
+        .isNotEmpty) {
+      await builder.addDartFileEdit(file, (builder) {
+        _uriText = builder.importLibrary(library.source.uri);
+      });
     }
   }
 }