[analysis_server] Improve Snippet performance with caching and earlier filtering

This shares a cache for mapping Elements to their public LibraryElements across each snippet producer, and also skips snippet producers that produce snippets that won't match any typed prefix.

Change-Id: I6b64b3c55f1030a5eaa7ca1afdcd6c416e4baa08
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/273962
Commit-Queue: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>

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

@@ -299,13 +299,15 @@ class CompletionHandler extends MessageHandler<CompletionParams, CompletionList>
     required ResolvedUnitResult unit,
     required int offset,
     required LineInfo lineInfo,
+    required bool Function(String input) filter,
   }) async {
     final request = DartSnippetRequest(
       unit: unit,
       offset: offset,
     );
     final snippetManager = DartSnippetManager();
-    final snippets = await snippetManager.computeSnippets(request);
+    final snippets =
+        await snippetManager.computeSnippets(request, filter: filter);
 
     return snippets.map((snippet) => snippetToCompletionItem(
           server,
@@ -527,6 +529,7 @@ class CompletionHandler extends MessageHandler<CompletionParams, CompletionList>
               unit: unit,
               offset: offset,
               lineInfo: unit.lineInfo,
+              filter: fuzzy.stringMatches,
             );
             return snippets.where(fuzzy.completionItemMatches).toList();
           });
@@ -835,8 +838,10 @@ class _FuzzyFilterHelper {
       : _matcher = FuzzyMatcher(prefix, matchStyle: MatchStyle.TEXT);
 
   bool completionItemMatches(CompletionItem item) =>
-      _matcher.score(item.filterText ?? item.label) > 0;
+      stringMatches(item.filterText ?? item.label);
 
   bool completionSuggestionMatches(CompletionSuggestion item) =>
-      _matcher.score(item.displayText ?? item.completion) > 0;
+      stringMatches(item.displayText ?? item.completion);
+
+  bool stringMatches(String input) => _matcher.score(input) > 0;
 }

pkg/analysis_server/lib/src/services/snippets/dart/class_declaration.dart

@@ -11,7 +11,10 @@ class ClassDeclaration extends DartSnippetProducer {
   static const prefix = 'class';
   static const label = 'class';
 
-  ClassDeclaration(super.request);
+  ClassDeclaration(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/do_statement.dart

@@ -11,7 +11,10 @@ class DoStatement extends DartSnippetProducer {
   static const prefix = 'do';
   static const label = 'do while';
 
-  DoStatement(super.request);
+  DoStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/flutter_stateful_widget.dart

@@ -21,7 +21,10 @@ class FlutterStatefulWidget extends FlutterSnippetProducer
   @override
   late ClassElement? classKey;
 
-  FlutterStatefulWidget(super.request);
+  FlutterStatefulWidget(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/flutter_stateful_widget_with_animation.dart

@@ -24,7 +24,11 @@ class FlutterStatefulWidgetWithAnimationController
   late ClassElement? classAnimationController;
   late MixinElement? classSingleTickerProviderStateMixin;
 
-  FlutterStatefulWidgetWithAnimationController(super.request);
+  FlutterStatefulWidgetWithAnimationController(super.request,
+      {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/flutter_stateless_widget.dart

@@ -19,7 +19,10 @@ class FlutterStatelessWidget extends FlutterSnippetProducer
   @override
   late ClassElement? classKey;
 
-  FlutterStatelessWidget(super.request);
+  FlutterStatelessWidget(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/for_in_statement.dart

@@ -11,7 +11,10 @@ class ForInStatement extends DartSnippetProducer {
   static const prefix = 'forin';
   static const label = 'for in';
 
-  ForInStatement(super.request);
+  ForInStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/for_statement.dart

@@ -11,7 +11,10 @@ class ForStatement extends DartSnippetProducer {
   static const prefix = 'for';
   static const label = 'for';
 
-  ForStatement(super.request);
+  ForStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/function_declaration.dart

@@ -11,7 +11,10 @@ class FunctionDeclaration extends DartSnippetProducer {
   static const prefix = 'fun';
   static const label = 'fun';
 
-  FunctionDeclaration(super.request);
+  FunctionDeclaration(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/if_else_statement.dart

@@ -11,7 +11,10 @@ class IfElseStatement extends DartSnippetProducer {
   static const prefix = 'ife';
   static const label = 'ife';
 
-  IfElseStatement(super.request);
+  IfElseStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/if_statement.dart

@@ -11,7 +11,10 @@ class IfStatement extends DartSnippetProducer {
   static const prefix = 'if';
   static const label = 'if';
 
-  IfStatement(super.request);
+  IfStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/main_function.dart

@@ -15,7 +15,10 @@ class MainFunction extends DartSnippetProducer {
   static const prefix = 'main';
   static const label = 'main()';
 
-  MainFunction(super.request);
+  MainFunction(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   /// Whether to insert a `List<String> args` parameter in the generated
   /// function.

pkg/analysis_server/lib/src/services/snippets/dart/switch_statement.dart

@@ -11,7 +11,10 @@ class SwitchStatement extends DartSnippetProducer {
   static const prefix = 'switch';
   static const label = 'switch case';
 
-  SwitchStatement(super.request);
+  SwitchStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/test_definition.dart

@@ -11,7 +11,10 @@ class TestDefinition extends DartSnippetProducer {
   static const prefix = 'test';
   static const label = 'test';
 
-  TestDefinition(super.request);
+  TestDefinition(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/test_group_definition.dart

@@ -11,7 +11,10 @@ class TestGroupDefinition extends DartSnippetProducer {
   static const prefix = 'group';
   static const label = 'group';
 
-  TestGroupDefinition(super.request);
+  TestGroupDefinition(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/try_catch_statement.dart

@@ -11,7 +11,10 @@ class TryCatchStatement extends DartSnippetProducer {
   static const prefix = 'try';
   static const label = 'try';
 
-  TryCatchStatement(super.request);
+  TryCatchStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/dart/while_statement.dart

@@ -11,7 +11,10 @@ class WhileStatement extends DartSnippetProducer {
   static const prefix = 'while';
   static const label = 'while';
 
-  WhileStatement(super.request);
+  WhileStatement(super.request, {required super.elementImportCache});
+
+  @override
+  String get snippetPrefix => prefix;
 
   @override
   Future<Snippet> compute() async {

pkg/analysis_server/lib/src/services/snippets/snippet_manager.dart

@@ -24,9 +24,11 @@ import 'package:analysis_server/src/services/snippets/snippet.dart';
 import 'package:analysis_server/src/services/snippets/snippet_context.dart';
 import 'package:analysis_server/src/services/snippets/snippet_producer.dart';
 import 'package:analyzer/dart/analysis/session.dart';
+import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/src/util/file_paths.dart' as file_paths;
 
-typedef SnippetProducerGenerator = SnippetProducer Function(DartSnippetRequest);
+typedef SnippetProducerGenerator = SnippetProducer Function(DartSnippetRequest,
+    {required Map<Element, LibraryElement?> elementImportCache});
 
 /// [DartSnippetManager] determines if a snippet request is Dart specific
 /// and forwards those requests to all Snippet Producers that return `true` from
@@ -61,8 +63,9 @@ class DartSnippetManager {
   };
 
   Future<List<Snippet>> computeSnippets(
-    DartSnippetRequest request,
-  ) async {
+    DartSnippetRequest request, {
+    bool Function(String input)? filter,
+  }) async {
     var pathContext = request.resourceProvider.pathContext;
     if (!file_paths.isDart(pathContext, request.filePath)) {
       return const [];
@@ -74,9 +77,12 @@ class DartSnippetManager {
       if (generators == null) {
         return snippets;
       }
+      final elementImportCache = <Element, LibraryElement?>{};
       for (final generator in generators) {
-        final producer = generator(request);
-        if (await producer.isValid()) {
+        final producer =
+            generator(request, elementImportCache: elementImportCache);
+        final matchesFilter = filter?.call(producer.snippetPrefix) ?? true;
+        if (matchesFilter && await producer.isValid()) {
           snippets.add(await producer.compute());
         }
       }

pkg/analysis_server/lib/src/services/snippets/snippet_producer.dart

@@ -24,12 +24,26 @@ abstract class DartSnippetProducer extends SnippetProducer {
   final LibraryElement libraryElement;
   final bool useSuperParams;
 
-  DartSnippetProducer(super.request)
+  /// Elements that need to be imported for generated code to be valid.
+  ///
+  /// Calling [addImports] will add any required imports to the supplied
+  /// builder.
+  final Set<Element> requiredElementImports = {};
+
+  /// A cache of mappings from Elements to their public Library Elements.
+  ///
+  /// Callers can share this cache across multiple snippet producers to avoid
+  /// repeated searches where they may add imports for the same elements.
+  final Map<Element, LibraryElement?> _elementImportCache;
+
+  DartSnippetProducer(super.request,
+      {required Map<Element, LibraryElement?> elementImportCache})
       : sessionHelper = AnalysisSessionHelper(request.analysisSession),
         utils = CorrectionUtils(request.unit),
         libraryElement = request.unit.libraryElement,
         useSuperParams = request.unit.libraryElement.featureSet
-            .isEnabled(Feature.super_parameters);
+            .isEnabled(Feature.super_parameters),
+        _elementImportCache = elementImportCache;
 
   CodeStyleOptions get codeStyleOptions =>
       sessionHelper.session.analysisContext.analysisOptions.codeStyleOptions;
@@ -44,6 +58,14 @@ abstract class DartSnippetProducer extends SnippetProducer {
   NullabilitySuffix get nullableSuffix => libraryElement.isNonNullableByDefault
       ? NullabilitySuffix.question
       : NullabilitySuffix.none;
+
+  /// Adds public imports for any elements fetched by [getClass] and [getMixin]
+  /// to [builder].
+  Future<void> addImports(DartFileEditBuilder builder) async {
+    final dartBuilder = builder as DartFileEditBuilderImpl;
+    await Future.wait(requiredElementImports.map((element) => dartBuilder
+        .importElementLibrary(element, resultCache: _elementImportCache)));
+  }
 }
 
 abstract class FlutterSnippetProducer extends DartSnippetProducer {
@@ -52,27 +74,12 @@ abstract class FlutterSnippetProducer extends DartSnippetProducer {
   late ClassElement? classWidget;
   late ClassElement? classPlaceholder;
 
-  /// Elements that need to be imported for generated code to be valid.
-  ///
-  /// Calling [getClass] or [getMixin] records elements in this set.
-  /// Calling [addImports] will add any required imports to the supplied
-  /// builder.
-  final Set<Element> _requiredElementImports = {};
-
-  FlutterSnippetProducer(super.request);
-
-  /// Adds public imports for any elements fetched by [getClass] and [getMixin]
-  /// to [builder].
-  Future<void> addImports(DartFileEditBuilder builder) async {
-    final dartBuilder = builder as DartFileEditBuilderImpl;
-    await Future.wait(
-        _requiredElementImports.map(dartBuilder.importElementLibrary));
-  }
+  FlutterSnippetProducer(super.request, {required super.elementImportCache});
 
   Future<ClassElement?> getClass(String name) async {
     final class_ = await sessionHelper.getClass(flutter.widgetsUri, name);
     if (class_ != null) {
-      _requiredElementImports.add(class_);
+      requiredElementImports.add(class_);
     }
     return class_;
   }
@@ -80,7 +87,7 @@ abstract class FlutterSnippetProducer extends DartSnippetProducer {
   Future<MixinElement?> getMixin(String name) async {
     final mixin = await sessionHelper.getMixin(flutter.widgetsUri, name);
     if (mixin != null) {
-      _requiredElementImports.add(mixin);
+      requiredElementImports.add(mixin);
     }
     return mixin;
   }
@@ -192,6 +199,9 @@ abstract class SnippetProducer {
 
   SnippetProducer(this.request);
 
+  /// The prefix a user types to use this snippet.
+  String get snippetPrefix;
+
   Future<Snippet> compute();
 
   Future<bool> isValid() async {

pkg/analysis_server/test/services/snippets/dart/test_support.dart

@@ -36,7 +36,7 @@ abstract class DartSnippetProducerTest extends AbstractSingleUnitTest {
       offset: offsetFromMarker(code),
     );
 
-    final producer = generator(request);
+    final producer = generator(request, elementImportCache: {});
     expect(await producer.isValid(), isFalse);
   }
 
@@ -47,7 +47,7 @@ abstract class DartSnippetProducerTest extends AbstractSingleUnitTest {
       offset: offsetFromMarker(code),
     );
 
-    final producer = generator(request);
+    final producer = generator(request, elementImportCache: {});
     expect(await producer.isValid(), isTrue);
     return producer.compute();
   }

pkg/analysis_server/test/services/snippets/snippet_manager_test.dart

@@ -7,7 +7,8 @@ import 'package:analysis_server/src/services/snippets/snippet.dart';
 import 'package:analysis_server/src/services/snippets/snippet_context.dart';
 import 'package:analysis_server/src/services/snippets/snippet_manager.dart';
 import 'package:analysis_server/src/services/snippets/snippet_producer.dart';
-import 'package:analyzer_plugin/protocol/protocol_common.dart';
+import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer_plugin/protocol/protocol_common.dart' hide Element;
 import 'package:test/test.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
 
@@ -21,6 +22,36 @@ void main() {
 
 @reflectiveTest
 class SnippetManagerTest extends AbstractSingleUnitTest {
+  Future<void> test_filter_match() async {
+    await resolveTestCode('');
+    final request = DartSnippetRequest(
+      unit: testAnalysisResult,
+      offset: 0,
+    );
+
+    final manager = _TestDartSnippetManager({
+      SnippetContext.atTopLevel: [_ValidSnippetProducer.newInstance],
+    });
+    final results =
+        await manager.computeSnippets(request, filter: (String s) => true);
+    expect(results, isNotEmpty);
+  }
+
+  Future<void> test_filter_noMatch() async {
+    await resolveTestCode('');
+    final request = DartSnippetRequest(
+      unit: testAnalysisResult,
+      offset: 0,
+    );
+
+    final manager = _TestDartSnippetManager({
+      SnippetContext.atTopLevel: [_ValidSnippetProducer.newInstance],
+    });
+    final results =
+        await manager.computeSnippets(request, filter: (String s) => false);
+    expect(results, isEmpty);
+  }
+
   Future<void> test_notValidProducers() async {
     await resolveTestCode('');
     final request = DartSnippetRequest(
@@ -45,7 +76,11 @@ class SnippetManagerTest extends AbstractSingleUnitTest {
     final manager = _TestDartSnippetManager({
       SnippetContext.atTopLevel: [_ValidSnippetProducer.newInstance],
       SnippetContext.inClass: [
-        (context) => throw 'Tried to create producer for wrong context',
+        (
+          context, {
+          required Map<Element, LibraryElement?> elementImportCache,
+        }) =>
+            throw 'Tried to create producer for wrong context',
       ]
     });
     final results = await manager.computeSnippets(request);
@@ -75,6 +110,9 @@ class SnippetManagerTest extends AbstractSingleUnitTest {
 class _NotValidSnippetProducer extends SnippetProducer {
   _NotValidSnippetProducer._(super.request);
 
+  @override
+  String get snippetPrefix => 'invalid';
+
   @override
   Future<Snippet> compute() {
     throw UnsupportedError(
@@ -86,7 +124,8 @@ class _NotValidSnippetProducer extends SnippetProducer {
   @override
   Future<bool> isValid() async => false;
 
-  static _NotValidSnippetProducer newInstance(DartSnippetRequest request) =>
+  static _NotValidSnippetProducer newInstance(DartSnippetRequest request,
+          {required Map<Element, LibraryElement?> elementImportCache}) =>
       _NotValidSnippetProducer._(request);
 }
 
@@ -102,10 +141,13 @@ class _TestDartSnippetManager extends DartSnippetManager {
 class _ValidSnippetProducer extends SnippetProducer {
   _ValidSnippetProducer._(super.request);
 
+  @override
+  String get snippetPrefix => 'mysnip';
+
   @override
   Future<Snippet> compute() async {
     return Snippet(
-      'mysnip',
+      snippetPrefix,
       'My Test Snippet',
       'This is a test snippet',
       SourceChange('message'),
@@ -115,6 +157,7 @@ class _ValidSnippetProducer extends SnippetProducer {
   @override
   Future<bool> isValid() async => true;
 
-  static _ValidSnippetProducer newInstance(DartSnippetRequest request) =>
+  static _ValidSnippetProducer newInstance(DartSnippetRequest request,
+          {required Map<Element, LibraryElement?> elementImportCache}) =>
       _ValidSnippetProducer._(request);
 }

pkg/analyzer/lib/src/services/top_level_declarations.dart

@@ -21,7 +21,12 @@ class TopLevelDeclarations {
 
   /// Return the first public library that that exports (but does not necessary
   /// declare) [element].
-  Future<LibraryElement?> publiclyExporting(Element element) async {
+  Future<LibraryElement?> publiclyExporting(Element element,
+      {Map<Element, LibraryElement?>? resultCache}) async {
+    if (resultCache?.containsKey(element) ?? false) {
+      return resultCache![element];
+    }
+
     var declarationFilePath = element.source?.fullName;
     if (declarationFilePath == null) {
       return null;
@@ -48,6 +53,7 @@ class TopLevelDeclarations {
       }
 
       if (_findElement(elementResult.element, element.displayName) != null) {
+        resultCache?[element] = elementResult.element;
         return elementResult.element;
       }
     }

pkg/analyzer_plugin/lib/src/utilities/change_builder/change_builder_dart.dart

@@ -1534,15 +1534,15 @@ class DartFileEditBuilderImpl extends FileEditBuilderImpl
   ///
   /// If the library [element] is declared in is inside the `src` folder, will
   /// try to locate a public URI to import instead.
-  Future<void> importElementLibrary(Element element) async {
-    // TODO(dantup): Add the ability to pass a cache in to this function so
-    //  multiple callers can avoid looking up the same elements.
+  Future<void> importElementLibrary(Element element,
+      {Map<Element, LibraryElement?>? resultCache}) async {
     if (_isDefinedLocally(element) || _getImportElement(element) != null) {
       return;
     }
 
-    var libraryWithElement =
-        await TopLevelDeclarations(resolvedUnit).publiclyExporting(element);
+    var libraryWithElement = resultCache?[element] ??
+        await TopLevelDeclarations(resolvedUnit)
+            .publiclyExporting(element, resultCache: resultCache);
     if (libraryWithElement != null) {
       _elementLibrariesToImport[element] =
           _importLibrary(libraryWithElement.source.uri);

pkg/analyzer_plugin/test/src/utilities/change_builder/change_builder_dart_test.dart

@@ -4,6 +4,7 @@
 
 // ignore_for_file: camel_case_types
 
+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/nullability_suffix.dart';
@@ -12,9 +13,9 @@ import 'package:analyzer/src/dart/element/inheritance_manager3.dart';
 import 'package:analyzer/src/generated/source.dart';
 import 'package:analyzer/src/test_utilities/find_node.dart';
 import 'package:analyzer/src/test_utilities/package_config_file_builder.dart';
-import 'package:analyzer_plugin/protocol/protocol_common.dart';
+import 'package:analyzer_plugin/protocol/protocol_common.dart' hide Element;
 import 'package:analyzer_plugin/src/utilities/change_builder/change_builder_dart.dart'
-    show DartLinkedEditBuilderImpl;
+    show DartFileEditBuilderImpl, DartLinkedEditBuilderImpl;
 import 'package:linter/src/rules.dart';
 import 'package:test/test.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
@@ -1804,6 +1805,12 @@ _prefix0.A1 a1; _prefix0.A2 a2; _prefix1.B b;''');
 @reflectiveTest
 class DartFileEditBuilderImplTest extends AbstractContextTest
     with DartChangeBuilderMixin {
+  Future<ResolvedUnitResult> resolveContent(String path, String content) async {
+    path = convertPath(path);
+    addSource(path, content);
+    return resolveFile(path);
+  }
+
   Future<void> test_convertFunctionFromSyncToAsync_closure() async {
     var path = convertPath('/home/test/lib/test.dart');
     addSource(path, '''var f = () {}''');
@@ -1934,6 +1941,96 @@ void functionAfter() {
 ''');
   }
 
+  Future<void> test_importElementLibrary_libElement() async {
+    var resolvedUnit = await resolveContent('/home/test/lib/test.dart', '');
+
+    var resolvedLibUnit =
+        await resolveContent('/home/test/lib/a.dart', 'class A {}');
+
+    var findNode = FindNode(resolvedLibUnit.content, resolvedLibUnit.unit);
+    var classElement = findNode.classDeclaration('A').declaredElement!;
+    var cache = <Element, LibraryElement?>{};
+    var builder = await newBuilder();
+    await builder.addDartFileEdit(resolvedUnit.path, (builder) async {
+      var builderImpl = builder as DartFileEditBuilderImpl;
+      await builderImpl.importElementLibrary(classElement, resultCache: cache);
+    });
+
+    var edits = getEdits(builder);
+    expect(edits, hasLength(1));
+    expect(edits[0].replacement,
+        equalsIgnoringWhitespace("import 'package:test/a.dart';\n"));
+    expect(cache[classElement], resolvedLibUnit.libraryElement);
+  }
+
+  Future<void> test_importElementLibrary_sdkElement() async {
+    var resolvedUnit = await resolveContent('/home/test/lib/test.dart', '');
+
+    var futureOrElement = resolvedUnit.typeProvider.futureOrElement;
+    var cache = <Element, LibraryElement?>{};
+    var builder = await newBuilder();
+    await builder.addDartFileEdit(resolvedUnit.path, (builder) async {
+      var builderImpl = builder as DartFileEditBuilderImpl;
+      await builderImpl.importElementLibrary(futureOrElement,
+          resultCache: cache);
+    });
+
+    var edits = getEdits(builder);
+    expect(edits, hasLength(1));
+    expect(
+        edits[0].replacement, equalsIgnoringWhitespace("import 'dart:async';"));
+    expect(cache[futureOrElement], futureOrElement.library);
+  }
+
+  Future<void> test_importElementLibrary_srcElement() async {
+    var resolvedUnit = await resolveContent('/home/test/lib/test.dart', '');
+
+    var resolvedSrcUnit =
+        await resolveContent('/home/test/lib/src/a.dart', 'class A {}');
+
+    var resolvedExportUnit =
+        await resolveContent('/home/test/lib/a.dart', "export 'src/a.dart'");
+
+    var findNode = FindNode(resolvedSrcUnit.content, resolvedSrcUnit.unit);
+    var classElement = findNode.classDeclaration('A').declaredElement!;
+    var cache = <Element, LibraryElement?>{};
+    var builder = await newBuilder();
+    await builder.addDartFileEdit(resolvedUnit.path, (builder) async {
+      var builderImpl = builder as DartFileEditBuilderImpl;
+      await builderImpl.importElementLibrary(classElement, resultCache: cache);
+    });
+
+    var edits = getEdits(builder);
+    expect(edits, hasLength(1));
+    expect(edits[0].replacement,
+        equalsIgnoringWhitespace("import 'package:test/a.dart';\n"));
+    expect(cache[classElement], resolvedExportUnit.libraryElement);
+  }
+
+  Future<void> test_importElementLibrary_usesCache() async {
+    var resolvedUnit = await resolveContent('/home/test/lib/test.dart', '');
+
+    // Create a fake file to put into the cache to verify it's being used.
+    var resolvedFakeUnit = await resolveContent('/home/test/lib/fake.dart', '');
+
+    var futureOrElement = resolvedUnit.typeProvider.futureOrElement;
+    var cache = <Element, LibraryElement?>{
+      futureOrElement: resolvedFakeUnit.libraryElement,
+    };
+    var builder = await newBuilder();
+    await builder.addDartFileEdit(resolvedUnit.path, (builder) async {
+      var builderImpl = builder as DartFileEditBuilderImpl;
+      await builderImpl.importElementLibrary(futureOrElement,
+          resultCache: cache);
+    });
+
+    var edits = getEdits(builder);
+    expect(edits, hasLength(1));
+    expect(edits[0].replacement,
+        equalsIgnoringWhitespace("import 'package:test/fake.dart';"));
+    expect(cache[futureOrElement], resolvedFakeUnit.libraryElement);
+  }
+
   Future<void> test_multipleEdits_concurrently() async {
     var initialCode = '00';
     var path = convertPath('/home/test/lib/test.dart');