Avoid using `findLibraryByName`

drift_dev/lib/src/analysis/backend.dart

@@ -44,6 +44,10 @@ abstract class DriftBackend {
   /// resolved.
   Future<Expression> resolveExpression(
       Uri context, String dartExpression, Iterable<String> imports);
+
+  /// Resolves the Dart element named [reference] in the [imports] of [context].
+  Future<Element?> resolveTopLevelElement(
+      Uri context, String reference, Iterable<Uri> imports);
 }
 
 /// Thrown when attempting to read a Dart library from a file that's not a

drift_dev/lib/src/analysis/preprocess_drift.dart

@@ -81,15 +81,13 @@ class DriftPreprocessor {
         .map((token) => token.dartCode)
         .toList();
 
-    var dartHelperFile = '';
     final codeToField = <String, String>{};
+    final seenFiles = <Uri>{uri};
+    final queue = [...directImports];
 
     if (dartLexemes.isNotEmpty) {
       // Imports in drift files are transitive, so we need to find all
       // transitive Dart sources to import into the generated helper file.
-      final seenFiles = <Uri>{uri};
-      final queue = [...directImports];
-
       while (queue.isNotEmpty) {
         final foundImport = queue.removeLast();
 
@@ -120,27 +118,25 @@ class DriftPreprocessor {
           }
         }
       }
+    }
 
-      final importedDartFiles =
-          seenFiles.where((uri) => url.extension(uri.path) == '.dart');
+    final importedDartFiles =
+        seenFiles.where((uri) => url.extension(uri.path) == '.dart');
 
-      // to analyze the expressions, generate a fake Dart file that declares each
-      // expression in a `var`, we can then read the static type when resolving
-      // file later.
+    // to analyze the expressions, generate a fake Dart file that declares each
+    // expression in a `var`, we can then read the static type when resolving
+    // file later.
 
-      final dartBuffer = StringBuffer();
-      for (final import in importedDartFiles) {
-        final importUri = import.toString();
-        dartBuffer.writeln('import ${asDartLiteral(importUri)};');
-      }
+    final dartBuffer = StringBuffer();
+    for (final import in importedDartFiles) {
+      final importUri = import.toString();
+      dartBuffer.writeln('import ${asDartLiteral(importUri)};');
+    }
 
-      for (var i = 0; i < dartLexemes.length; i++) {
-        final name = 'expr_$i';
-        dartBuffer.writeln('var $name = ${dartLexemes[i]};');
-        codeToField[dartLexemes[i]] = name;
-      }
-
-      dartHelperFile = dartBuffer.toString();
+    for (var i = 0; i < dartLexemes.length; i++) {
+      final name = 'expr_$i';
+      dartBuffer.writeln('var $name = ${dartLexemes[i]};');
+      codeToField[dartLexemes[i]] = name;
     }
 
     final declaredTablesAndViews = <String>[];
@@ -158,6 +154,6 @@ class DriftPreprocessor {
       directImports.toList(),
     );
 
-    return DriftPreprocessor._(result, dartHelperFile);
+    return DriftPreprocessor._(result, dartBuffer.toString());
   }
 }

drift_dev/lib/src/analysis/resolver/dart/helper.dart

@@ -89,11 +89,12 @@ class KnownDriftTypes {
   }
 
   static Future<KnownDriftTypes> resolve(DriftAnalysisDriver driver) async {
-    final library = await driver.backend
-        .readDart(Uri.parse('package:drift/src/drift_dev_helper.dart'));
+    final library = await driver.backend.readDart(uri);
 
     return KnownDriftTypes._fromLibrary(library);
   }
+
+  static final Uri uri = Uri.parse('package:drift/src/drift_dev_helper.dart');
 }
 
 Expression? returnExpressionOfMethod(MethodDeclaration method) {

drift_dev/lib/src/analysis/resolver/drift/element_resolver.dart

@@ -67,19 +67,8 @@ abstract class DriftElementResolver<T extends DiscoveredElement>
         // are available.
         .followedBy([AnnotatedDartCode.dartCore]);
 
-    for (final import in dartImports) {
-      LibraryElement library;
-      try {
-        library = await resolver.driver.backend.readDart(import);
-      } on NotALibraryException {
-        continue;
-      }
-
-      final foundElement = library.exportNamespace.get(identifier);
-      if (foundElement != null) return foundElement;
-    }
-
-    return null;
+    return await resolver.driver.backend
+        .resolveTopLevelElement(file.ownUri, identifier, dartImports);
   }
 
   /// Resolves [identifier] to a Dart element declaring a type, or reports an

drift_dev/lib/src/backends/analyzer_context_backend.dart

@@ -151,6 +151,40 @@ class AnalysisContextBackend extends DriftBackend {
     }
   }
 
+  @override
+  Future<Element?> resolveTopLevelElement(
+      Uri context, String reference, Iterable<Uri> imports) async {
+    // Create a fake file next to the content
+    final path = _pathOfUri(context)!;
+    final pathContext = provider.pathContext;
+    final pathForTemp = pathContext.join(
+        pathContext.dirname(path), 'moor_temp_${imports.hashCode}.dart');
+
+    final content = StringBuffer();
+    for (final import in imports) {
+      content.writeln('import "$import";');
+    }
+
+    provider.setOverlay(
+      pathForTemp,
+      content: content.toString(),
+      modificationStamp: DateTime.now().millisecondsSinceEpoch,
+    );
+
+    try {
+      final result =
+          await this.context.currentSession.getResolvedLibrary(pathForTemp);
+
+      if (result is ResolvedLibraryResult) {
+        return result.element.scope.lookup(reference).getter;
+      }
+    } finally {
+      provider.removeOverlay(path);
+    }
+
+    return null;
+  }
+
   @override
   Uri resolveUri(Uri base, String uriString) {
     final resolved = base.resolve(uriString);

drift_dev/lib/src/backends/build/backend.dart

@@ -11,6 +11,7 @@ import 'package:build/build.dart' as build;
 import '../../analysis/backend.dart';
 import '../../analysis/driver/driver.dart';
 import '../../analysis/preprocess_drift.dart';
+import '../../analysis/resolver/dart/helper.dart';
 
 class DriftBuildBackend extends DriftBackend {
   final BuildStep _buildStep;
@@ -97,6 +98,30 @@ class DriftBuildBackend extends DriftBackend {
     }
     return initializer;
   }
+
+  @override
+  Future<Element?> resolveTopLevelElement(
+      Uri context, String reference, Iterable<Uri> imports) async {
+    final original = AssetId.resolve(context);
+    final tempDart = original.changeExtension('.expr.temp.dart');
+
+    if (await _buildStep.canRead(tempDart)) {
+      final library = await _buildStep.resolver.libraryFor(tempDart);
+
+      return library.scope.lookup(reference).getter;
+    } else {
+      // If there's no temporary file whose imports we can use, then that means
+      // that there aren't any Dart imports in [context] at all. So we just need
+      // to look it up in `dart:core`.
+      // For that, resolve a library we know exists and likely has been resolved
+      // already.
+      final libraryWeKnowExists = await _buildStep.resolver
+          .libraryFor(AssetId.resolve(KnownDriftTypes.uri));
+      final dartCore = libraryWeKnowExists.typeProvider.objectElement.library;
+
+      return dartCore.exportNamespace.get(reference);
+    }
+  }
 }
 
 class BuildCacheReader implements AnalysisResultCacheReader {

drift_dev/lib/src/services/schema/sqlite_to_drift.dart

@@ -102,6 +102,12 @@ class _SingleFileNoAnalyzerBackend extends DriftBackend {
     _noAnalyzer();
   }
 
+  @override
+  Future<Element?> resolveTopLevelElement(
+      Uri context, String reference, Iterable<Uri> imports) {
+    _noAnalyzer();
+  }
+
   @override
   Uri resolveUri(Uri base, String uriString) {
     return uri;

drift_dev/test/analysis/test_utils.dart

@@ -184,6 +184,39 @@ class TestBackend extends DriftBackend {
     }
   }
 
+  @override
+  Future<Element?> resolveTopLevelElement(
+      Uri context, String reference, Iterable<Uri> imports) async {
+    final fileContents = StringBuffer();
+    for (final import in imports) {
+      fileContents.writeln("import '$import';");
+    }
+
+    final path = '${_pathFor(context)}.imports.dart';
+
+    await _setupDartAnalyzer();
+
+    final resourceProvider = _resourceProvider!;
+    final analysisContext = _dartContext!;
+
+    resourceProvider.setOverlay(path,
+        content: fileContents.toString(), modificationStamp: 1);
+
+    try {
+      final result =
+          await analysisContext.currentSession.getResolvedLibrary(path);
+
+      if (result is ResolvedLibraryResult) {
+        final lookup = result.element.scope.lookup(reference);
+        return lookup.getter;
+      }
+    } finally {
+      resourceProvider.removeOverlay(path);
+    }
+
+    return null;
+  }
+
   @override
   Future<LibraryElement> readDart(Uri uri) async {
     await ensureHasDartAnalyzer();