Avoid resolving all transitive imports in build

drift_dev/build.yaml

@@ -23,7 +23,7 @@ builders:
   # Regular build flow, emitting a shared part file for source_gen to pick up
   drift_dev:
     import: "package:drift_dev/integrations/build.dart"
-    builder_factories: ["analyzer", "driftBuilder"]
+    builder_factories: ["discover", "analyzer", "driftBuilder"]
     build_extensions:
       ".dart": [".drift.g.part", ".dart.drift_module.json"]
       ".drift": [".drift.drift_module.json"]
@@ -54,7 +54,7 @@ builders:
     applies_builders: [":analyzer"]
   analyzer:
     import: "package:drift_dev/integrations/build.dart"
-    builder_factories: ["analyzer"]
+    builder_factories: ["discover", "analyzer"]
     build_extensions:
       ".dart": [".dart.drift_module.json"]
       ".drift": [".drift.drift_module.json"]

drift_dev/lib/integrations/build.dart

@@ -5,6 +5,8 @@ import 'package:drift_dev/src/backends/build/preprocess_builder.dart';
 
 Builder preparingBuilder(BuilderOptions options) => PreprocessBuilder();
 
+Builder discover(BuilderOptions options) => DriftDiscover(options);
+
 Builder analyzer(BuilderOptions options) => DriftAnalyzer(options);
 
 Builder driftBuilder(BuilderOptions options) =>

drift_dev/lib/src/analysis/driver/cache.dart

@@ -7,7 +7,7 @@ import 'state.dart';
 class DriftAnalysisCache {
   final Map<Uri, CachedSerializationResult> serializationCache = {};
   final Map<Uri, FileState> knownFiles = {};
-  final Map<DriftElementId, DiscoveredElement> discoveredElements = {};
+  final Map<DriftElementId, DriftElementKind> discoveredElements = {};
 
   FileState stateForUri(Uri uri) {
     return knownFiles[uri] ?? notifyFileChanged(uri);
@@ -26,15 +26,11 @@ class DriftAnalysisCache {
 
   void notifyFileDeleted(Uri uri) {}
 
-  void postFileDiscoveryResults(FileState state) {
+  void knowsLocalElements(FileState state) {
     discoveredElements.removeWhere((key, _) => key.libraryUri == state.ownUri);
 
-    final discovery = state.discovery;
+    for (final (id, kind) in state.definedElements) {
-    if (discovery != null) {
+      discoveredElements[id] = kind;
-      discoveredElements.addAll({
-        for (final definedHere in discovery.locallyDefinedElements)
-          definedHere.ownId: definedHere,
-      });
     }
   }
 

drift_dev/lib/src/analysis/driver/driver.dart

@@ -107,7 +107,7 @@ class DriftAnalysisDriver {
     final reader = cacheReader;
     if (reader == null) return null;
 
-    final found = await reader.readCacheFor(uri);
+    final found = await reader.readElementCacheFor(uri);
     if (found == null) return null;
 
     final parsed = json.decode(found) as Map<String, Object?>;
@@ -140,7 +140,8 @@ class DriftAnalysisDriver {
 
     final cachedImports = cache.serializationCache[state.ownUri]?.cachedImports;
     if (cachedImports != null && state.discovery == null) {
-      state.cachedImports = cachedImports;
+      state.cachedDiscovery ??= CachedDiscoveryResults(true, cachedImports, {});
+
       for (final import in cachedImports) {
         final found = cache.stateForUri(import);
 
@@ -156,53 +157,93 @@ class DriftAnalysisDriver {
     return allRecovered;
   }
 
-  /// Runs the first step (element discovery) on a file with the given [uri].
+  Future<void> discoverIfNecessary(
-  Future<FileState> prepareFileForAnalysis(
+    FileState file, {
-    Uri uri, {
-    bool needsDiscovery = true,
     bool warnIfFileDoesntExist = true,
   }) async {
-    var known = cache.knownFiles[uri] ?? cache.notifyFileChanged(uri);
+    if (file.discovery == null) {
-
+      await DiscoverStep(this, file)
-    if (known.discovery == null && needsDiscovery) {
-      await DiscoverStep(this, known)
           .discover(warnIfFileDoesntExist: warnIfFileDoesntExist);
-      cache.postFileDiscoveryResults(known);
+      cache.knowsLocalElements(file);
+    }
+  }
 
-      // todo: Mark elements that need to be analyzed again
+  /// Runs the first step (discovering local elements) on a file with the given
+  /// [uri].
+  Future<FileState> findLocalElements(
+    Uri uri, {
+    bool warnIfFileDoesntExist = true,
+  }) async {
+    final known = cache.knownFiles[uri] ?? cache.notifyFileChanged(uri);
 
-      // To analyze a drift file, we also need to be able to analyze imports.
+    if (known.cachedDiscovery != null || known.discovery != null) {
-      final state = known.discovery;
+      // We already know local elements.
-      if (state is DiscoveredDriftFile) {
+      return known;
-        for (final import in state.imports) {
+    }
-          // todo: We shouldn't unconditionally crawl files like this. The build
-          // backend should emit prepared file results in a previous step which
-          // should be used here.
-          final file = await prepareFileForAnalysis(import.importedUri);
 
-          if (file.discovery?.isValidImport != true) {
+    // First, try to read cached results.
-            known.errorsDuringDiscovery.add(
+    final reader = cacheReader;
-              DriftAnalysisError.inDriftFile(
+    CachedDiscoveryResults? cached;
-                import.ast,
+
-                'The imported file, `${import.importedUri}`, does not exist or '
+    if (reader != null) {
-                "can't be imported.",
+      cached = await reader.readDiscovery(uri);
-              ),
+
-            );
+      if (cached == null && reader.findsLocalElementsReliably) {
-          }
+        // There are no locally defined elements, since otherwise the reader
-        }
+        // would have found them.
-      } else if (state is DiscoveredDartLibrary) {
+        cached = CachedDiscoveryResults(false, const [], const {});
-        for (final import in state.importDependencies) {
-          // We might import a generated file that doesn't exist yet, that
-          // should not be a user-visible error. Users will notice because the
-          // import is reported as an error by the analyzer either way.
-          await prepareFileForAnalysis(import, warnIfFileDoesntExist: false);
-        }
       }
     }
 
+    if (cached != null) {
+      known.cachedDiscovery = cached;
+      cache.knowsLocalElements(known);
+    } else {
+      await discoverIfNecessary(
+        known,
+        warnIfFileDoesntExist: warnIfFileDoesntExist,
+      );
+    }
+
     return known;
   }
 
+  Future<void> _findLocalElementsInAllImports(FileState known) async {
+    // To analyze references in elements, we also need to know locally defined
+    // elements in all imports.
+    final state = known.discovery;
+    if (state is DiscoveredDriftFile) {
+      for (final import in state.imports) {
+        // todo: We shouldn't unconditionally crawl files like this. The build
+        // backend should emit prepared file results in a previous step which
+        // should be used here.
+        final file = await findLocalElements(import.importedUri);
+
+        if (file.isValidImport != true) {
+          known.errorsDuringDiscovery.add(
+            DriftAnalysisError.inDriftFile(
+              import.ast,
+              'The imported file, `${import.importedUri}`, does not exist or '
+              "can't be imported.",
+            ),
+          );
+        } else {
+          await _findLocalElementsInAllImports(file);
+        }
+      }
+    } else {
+      for (final import in known.imports ?? const <Uri>[]) {
+        await findLocalElements(
+          import,
+          // We might import a generated file that doesn't exist yet, that
+          // should not be a user-visible error. Users will notice because the
+          // import is reported as an error by the analyzer either way.
+          warnIfFileDoesntExist: false,
+        );
+      }
+    }
+  }
+
   /// Runs the second analysis step (element analysis) on a file.
   ///
   /// The file, as well as all imports, should have undergone the first analysis
@@ -232,8 +273,6 @@ class DriftAnalysisDriver {
   /// necessary work up until that point.
   Future<FileState> resolveElements(Uri uri) async {
     var known = cache.stateForUri(uri);
-    await prepareFileForAnalysis(uri, needsDiscovery: false);
-
     if (known.isFullyAnalyzed) {
       // Well, there's nothing to do now.
       return known;
@@ -247,8 +286,10 @@ class DriftAnalysisDriver {
     }
 
     // We couldn't recover all analyzed elements. Let's run an analysis run
-    // now then.
+    // then.
-    await prepareFileForAnalysis(uri, needsDiscovery: true);
+    await discoverIfNecessary(known);
+    await _findLocalElementsInAllImports(known);
+
     await _analyzePrepared(known);
     return known;
   }
@@ -299,8 +340,19 @@ class DriftAnalysisDriver {
 /// This class is responsible for recovering both assets in a subsequent build-
 /// step.
 abstract class AnalysisResultCacheReader {
+  /// Whether [readDiscovery] only returns `null` when the file under the URI
+  /// is not relevant to drift.
+  bool get findsLocalElementsReliably;
+
+  /// Whether [readElementCacheFor] is guaranteed to return all elements defined
+  /// in the supplied `uri`, or whether it could be that we just didn't analyze
+  /// that file yet.
+  bool get findsResolvedElementsReliably;
+
+  Future<CachedDiscoveryResults?> readDiscovery(Uri uri);
+
   Future<LibraryElement?> readTypeHelperFor(Uri uri);
-  Future<String?> readCacheFor(Uri uri);
+  Future<String?> readElementCacheFor(Uri uri);
 }
 
 /// Thrown by a local element resolver when an element could not be resolved and

drift_dev/lib/src/analysis/driver/state.dart

@@ -15,7 +15,7 @@ class FileState {
   final Uri ownUri;
 
   DiscoveredFileState? discovery;
-  List<Uri>? cachedImports;
+  CachedDiscoveryResults? cachedDiscovery;
 
   final List<DriftAnalysisError> errorsDuringDiscovery = [];
 
@@ -26,7 +26,12 @@ class FileState {
 
   FileState(this.ownUri);
 
-  Iterable<Uri>? get imports => discovery?.importDependencies ?? cachedImports;
+  bool get isValidImport {
+    return (cachedDiscovery?.isValidImport ?? discovery?.isValidImport) == true;
+  }
+
+  Iterable<Uri>? get imports =>
+      discovery?.importDependencies ?? cachedDiscovery?.imports;
 
   String get extension => url.extension(ownUri.path);
 
@@ -35,6 +40,22 @@ class FileState {
     return analyzedElements.any((e) => e is BaseDriftAccessor);
   }
 
+  Iterable<(DriftElementId, DriftElementKind)> get definedElements sync* {
+    final discovery = this.discovery;
+    final cached = cachedDiscovery;
+
+    if (discovery != null) {
+      for (final element in discovery.locallyDefinedElements) {
+        yield (element.ownId, element.kind);
+      }
+    } else if (cached != null) {
+      for (final MapEntry(:key, :value)
+          in cached.locallyDefinedElements.entries) {
+        yield (id(key), value);
+      }
+    }
+  }
+
   /// All analyzed [DriftElement]s found in this library.
   @visibleForTesting
   Iterable<DriftElement> get analyzedElements {
@@ -90,6 +111,18 @@ class FileState {
   }
 }
 
+class CachedDiscoveryResults {
+  final bool isValidImport;
+  final List<Uri> imports;
+  final Map<String, DriftElementKind> locallyDefinedElements;
+
+  CachedDiscoveryResults(
+    this.isValidImport,
+    this.imports,
+    this.locallyDefinedElements,
+  );
+}
+
 abstract class DiscoveredFileState {
   final List<DiscoveredElement> locallyDefinedElements;
 
@@ -156,6 +189,7 @@ class UnknownFile extends DiscoveredFileState {
 
 abstract class DiscoveredElement {
   final DriftElementId ownId;
+  DriftElementKind get kind;
 
   DiscoveredElement(this.ownId);
 

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

@@ -91,9 +91,9 @@ class DartAccessorResolver
       } else {
         includes.add(import);
 
-        final resolved = await resolver.driver.prepareFileForAnalysis(
+        final resolved = await resolver.driver
-            discovered.ownId.libraryUri.resolveUri(import));
+            .findLocalElements(discovered.ownId.libraryUri.resolveUri(import));
-        if (resolved.discovery?.isValidImport != true) {
+        if (!resolved.isValidImport) {
           reportError(
             DriftAnalysisError.forDartElement(
                 element, '`$value` could not be imported'),

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

@@ -116,17 +116,17 @@ class DiscoverStep {
 
             imports.add(DriftFileImport(node, uri));
           } else if (node is TableInducingStatement) {
-            pendingElements
+            pendingElements.add(DiscoveredDriftTable(
-                .add(DiscoveredDriftTable(_id(node.createdName), node));
+                _id(node.createdName), DriftElementKind.table, node));
           } else if (node is CreateViewStatement) {
-            pendingElements
+            pendingElements.add(DiscoveredDriftView(
-                .add(DiscoveredDriftView(_id(node.createdName), node));
+                _id(node.createdName), DriftElementKind.view, node));
           } else if (node is CreateIndexStatement) {
-            pendingElements
+            pendingElements.add(DiscoveredDriftIndex(
-                .add(DiscoveredDriftIndex(_id(node.indexName), node));
+                _id(node.indexName), DriftElementKind.dbIndex, node));
           } else if (node is CreateTriggerStatement) {
-            pendingElements
+            pendingElements.add(DiscoveredDriftTrigger(
-                .add(DiscoveredDriftTrigger(_id(node.triggerName), node));
+                _id(node.triggerName), DriftElementKind.trigger, node));
           } else if (node is DeclaredStatement) {
             String name;
 
@@ -137,7 +137,8 @@ class DiscoverStep {
               name = '\$drift_${specialQueryNameCount++}';
             }
 
-            pendingElements.add(DiscoveredDriftStatement(_id(name), node));
+            pendingElements.add(DiscoveredDriftStatement(
+                _id(name), DriftElementKind.definedQuery, node));
           }
         }
 

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

@@ -4,11 +4,15 @@ import 'package:drift/drift.dart' show DriftView;
 import 'package:sqlparser/sqlparser.dart';
 
 import '../driver/state.dart';
+import '../results/element.dart';
 
 class DiscoveredDriftElement<AST extends AstNode> extends DiscoveredElement {
   final AST sqlNode;
 
-  DiscoveredDriftElement(super.ownId, this.sqlNode);
+  @override
+  final DriftElementKind kind;
+
+  DiscoveredDriftElement(super.ownId, this.kind, this.sqlNode);
 }
 
 typedef DiscoveredDriftTable = DiscoveredDriftElement<TableInducingStatement>;
@@ -25,6 +29,9 @@ abstract class DiscoveredDartElement<DE extends Element>
 }
 
 class DiscoveredDartTable extends DiscoveredDartElement<ClassElement> {
+  @override
+  DriftElementKind get kind => DriftElementKind.table;
+
   DiscoveredDartTable(super.ownId, super.dartElement);
 }
 
@@ -32,6 +39,9 @@ class DiscoveredDartView extends DiscoveredDartElement<ClassElement> {
   /// The [DriftView] annotation on this class, if there is any.
   DartObject? viewAnnotation;
 
+  @override
+  DriftElementKind get kind => DriftElementKind.view;
+
   DiscoveredDartView(super.ownId, super.dartElement, this.viewAnnotation);
 }
 
@@ -39,6 +49,11 @@ class DiscoveredBaseAccessor extends DiscoveredDartElement<ClassElement> {
   final bool isDatabase;
   final DartObject annotation;
 
+  @override
+  DriftElementKind get kind => isAccessor
+      ? DriftElementKind.databaseAccessor
+      : DriftElementKind.database;
+
   bool get isAccessor => !isDatabase;
 
   DiscoveredBaseAccessor(

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

@@ -114,7 +114,12 @@ class DriftResolver {
       _currentDependencyPath.add(reference);
 
       try {
-        final resolved = await resolveDiscovered(pending);
+        final owningFile = driver.cache.stateForUri(reference.libraryUri);
+        await driver.discoverIfNecessary(owningFile);
+        final discovered = owningFile.discovery!.locallyDefinedElements
+            .firstWhere((e) => e.ownId == reference);
+
+        final resolved = await resolveDiscovered(discovered);
         return ResolvedReferenceFound(resolved);
       } catch (e, s) {
         driver.backend.log.warning('Could not analze $reference', e, s);
@@ -134,7 +139,8 @@ class DriftResolver {
   Future<ResolveReferencedElementResult> resolveDartReference(
       DriftElementId owner, Element element) async {
     final uri = await driver.backend.uriOfDart(element.library!);
-    final state = await driver.prepareFileForAnalysis(uri);
+    final state = driver.cache.stateForUri(uri);
+    await driver.discoverIfNecessary(driver.cache.stateForUri(uri));
 
     final discovered = state.discovery?.locallyDefinedElements
         .whereType<DiscoveredDartElement>()
@@ -164,7 +170,7 @@ class DriftResolver {
     for (final available in driver.cache.crawl(file)) {
       final localElementIds = {
         ...available.analysis.keys,
-        ...?available.discovery?.locallyDefinedElements.map((e) => e.ownId),
+        ...available.definedElements.map((e) => e.$1),
       };
 
       for (final definedLocally in localElementIds) {

drift_dev/lib/src/analysis/results/database.dart

@@ -66,6 +66,9 @@ class DriftDatabase extends BaseDriftAccessor {
     this.schemaVersion,
     this.accessors = const [],
   });
+
+  @override
+  DriftElementKind get kind => DriftElementKind.database;
 }
 
 /// A Dart class with a similar API to a database, providing a view over a
@@ -86,6 +89,9 @@ class DatabaseAccessor extends BaseDriftAccessor {
     required this.databaseClass,
     required this.ownType,
   });
+
+  @override
+  DriftElementKind get kind => DriftElementKind.databaseAccessor;
 }
 
 /// A query defined on a [BaseDriftAccessor].

drift_dev/lib/src/analysis/results/element.dart

@@ -78,6 +78,8 @@ abstract class DriftElement {
   final DriftElementId id;
   final DriftDeclaration declaration;
 
+  DriftElementKind get kind;
+
   /// All elements referenced by this element.
   ///
   /// References include the following:
@@ -104,6 +106,18 @@ abstract class DriftElement {
   DriftElement(this.id, this.declaration);
 }
 
+enum DriftElementKind {
+  table,
+  view,
+  dbIndex,
+  trigger,
+  database,
+  databaseAccessor,
+  definedQuery;
+
+  static Map<String, DriftElementKind> byName = values.asNameMap();
+}
+
 abstract class DriftSchemaElement extends DriftElement {
   DriftSchemaElement(super.id, super.declaration);
 

drift_dev/lib/src/analysis/results/index.dart

@@ -27,6 +27,9 @@ class DriftIndex extends DriftSchemaElement {
     required this.createStmt,
   });
 
+  @override
+  DriftElementKind get kind => DriftElementKind.dbIndex;
+
   @override
   String get dbGetterName => DriftSchemaElement.dbFieldName(id.name);
 

drift_dev/lib/src/analysis/results/query.dart

@@ -66,6 +66,9 @@ class DefinedSqlQuery extends DriftElement implements DriftQueryDeclaration {
     }
   }
 
+  @override
+  DriftElementKind get kind => DriftElementKind.definedQuery;
+
   /// All in-line Dart source code literals embedded into the query.
   final List<String> dartTokens;
 

drift_dev/lib/src/analysis/results/table.dart

@@ -84,6 +84,9 @@ class DriftTable extends DriftElementWithResultSet {
 
   late final DriftColumn? rowid = _findRowId();
 
+  @override
+  DriftElementKind get kind => DriftElementKind.table;
+
   /// Whether this is a virtual table, created with a `CREATE VIRTUAL TABLE`
   /// statement in SQL.
   bool get isVirtual => virtualTableData != null;

drift_dev/lib/src/analysis/results/trigger.dart

@@ -31,6 +31,9 @@ class DriftTrigger extends DriftSchemaElement {
     required this.writes,
   });
 
+  @override
+  DriftElementKind get kind => DriftElementKind.trigger;
+
   @override
   String get dbGetterName => DriftSchemaElement.dbFieldName(id.name);
 

drift_dev/lib/src/analysis/results/view.dart

@@ -43,6 +43,9 @@ class DriftView extends DriftElementWithResultSet {
   @override
   String get dbGetterName => DriftSchemaElement.dbFieldName(id.name);
 
+  @override
+  DriftElementKind get kind => DriftElementKind.view;
+
   /// Obtains all tables transitively referenced by the declaration of this
   /// view.
   ///

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

@@ -11,6 +11,52 @@ import '../../writer/writer.dart';
 import 'backend.dart';
 import 'exception.dart';
 
+class DriftDiscover extends Builder {
+  final DriftOptions options;
+
+  DriftDiscover(BuilderOptions options)
+      : options = DriftOptions.fromJson(options.config);
+
+  @override
+  Map<String, List<String>> get buildExtensions => const {
+        '.drift': [
+          '.drift.drift_elements.json',
+        ],
+        '.dart': [
+          '.dart.drift_elements.json',
+        ],
+      };
+
+  @override
+  Future<void> build(BuildStep buildStep) async {
+    final backend = DriftBuildBackend(buildStep);
+    final driver = DriftAnalysisDriver(backend, options);
+
+    final prepared = await driver.findLocalElements(buildStep.inputId.uri);
+    final discovery = prepared.discovery;
+
+    if (discovery != null) {
+      await buildStep.writeAsString(
+        buildStep.allowedOutputs.single,
+        json.encode({
+          'valid_import': discovery.isValidImport,
+          'imports': [
+            for (final import in discovery.importDependencies)
+              import.toString(),
+          ],
+          'elements': [
+            for (final entry in discovery.locallyDefinedElements)
+              {
+                'kind': entry.kind.name,
+                'name': entry.ownId.name,
+              }
+          ]
+        }),
+      );
+    }
+  }
+}
+
 class DriftAnalyzer extends Builder {
   final DriftOptions options;
 
@@ -32,7 +78,9 @@ class DriftAnalyzer extends Builder {
   @override
   Future<void> build(BuildStep buildStep) async {
     final backend = DriftBuildBackend(buildStep);
-    final driver = DriftAnalysisDriver(backend, options);
+    final driver = DriftAnalysisDriver(backend, options)
+      ..cacheReader =
+          BuildCacheReader(buildStep, findsLocalElementsReliably: true);
 
     final results = await driver.resolveElements(buildStep.inputId.uri);
     var hadWarnings = false;

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

@@ -2,6 +2,8 @@ import 'dart:convert';
 
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element.dart';
+import 'package:drift_dev/src/analysis/driver/state.dart';
+import 'package:drift_dev/src/analysis/results/element.dart';
 import 'package:logging/logging.dart';
 import 'package:build/build.dart';
 import 'package:build/build.dart' as build;
@@ -100,10 +102,43 @@ class DriftBuildBackend extends DriftBackend {
 class BuildCacheReader implements AnalysisResultCacheReader {
   final BuildStep _buildStep;
 
-  BuildCacheReader(this._buildStep);
+  @override
+  final bool findsLocalElementsReliably;
+  @override
+  final bool findsResolvedElementsReliably;
+
+  BuildCacheReader(
+    this._buildStep, {
+    this.findsLocalElementsReliably = false,
+    this.findsResolvedElementsReliably = false,
+  });
 
   @override
-  Future<String?> readCacheFor(Uri uri) async {
+  Future<CachedDiscoveryResults?> readDiscovery(Uri uri) async {
+    // For the format, see the `DriftDiscover` builder in `analyzer.dart`.
+    final assetId = AssetId.resolve(uri).addExtension('.drift_elements.json');
+
+    if (await _buildStep.canRead(assetId)) {
+      final results = json.decode(await _buildStep.readAsString(assetId));
+      final rawImports = results['imports'] as List;
+      final rawElements = (results['elements'] as List).cast<Map>();
+
+      return CachedDiscoveryResults(
+        results['valid_import'] as bool,
+        [for (final import in rawImports) Uri.parse(import as String)],
+        {
+          for (final element in rawElements)
+            (element['name'] as String):
+                DriftElementKind.byName[element['kind']]!,
+        },
+      );
+    }
+
+    return null;
+  }
+
+  @override
+  Future<String?> readElementCacheFor(Uri uri) async {
     // These files are generated by the `DriftAnalyzer` builder
     final assetId = AssetId.resolve(uri).addExtension('.drift_module.json');
     if (await _buildStep.canRead(assetId)) {

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

@@ -135,7 +135,14 @@ class _DriftBuildRun {
 
   _DriftBuildRun(this.options, this.mode, this.buildStep)
       : driver = DriftAnalysisDriver(DriftBuildBackend(buildStep), options)
-          ..cacheReader = BuildCacheReader(buildStep);
+          ..cacheReader = BuildCacheReader(
+            buildStep,
+            // The discovery and analyzer builders will have emitted IR for
+            // every relevant file in a previous build step that this builder
+            // has a dependency on.
+            findsResolvedElementsReliably: !mode.embeddedAnalyzer,
+            findsLocalElementsReliably: !mode.embeddedAnalyzer,
+          );
 
   Future<void> run() async {
     await _warnAboutDeprecatedOptions();
@@ -214,12 +221,10 @@ class _DriftBuildRun {
   /// Checks if the input file contains elements drift should generate code for.
   Future<bool> _checkForElementsToBuild() async {
     if (mode.embeddedAnalyzer) {
-      // Run the discovery step, which we'll have to run either way, to check if
+      // Check if there are any elements defined locally that would need code
-      // there are any elements to generate code for.
+      // to be generated for this file.
-      final state = await driver.prepareFileForAnalysis(buildStep.inputId.uri,
+      final state = await driver.findLocalElements(buildStep.inputId.uri);
-          needsDiscovery: true);
+      return state.definedElements.isNotEmpty;
-
-      return state.discovery?.locallyDefinedElements.isNotEmpty == true;
     } else {
       // An analysis step should have already run for this asset. If we can't
       // pick up results from that, there is no code for drift to generate.
@@ -237,8 +242,8 @@ class _DriftBuildRun {
         buildStep.inputId.extension != '.dart') {
       // For modular drift file generation, we need to know about imports which
       // are only available when discovery ran.
-      await driver.prepareFileForAnalysis(buildStep.inputId.uri,
+      final state = driver.cache.stateForUri(buildStep.inputId.uri);
-          needsDiscovery: true);
+      await driver.discoverIfNecessary(state);
     }
 
     return true;

drift_dev/test/analysis/resolver/discover_test.dart

@@ -17,7 +17,7 @@ CREATE VIEW my_view AS SELECT whatever FROM unknown_table;
       });
 
       final uri = Uri.parse('package:a/main.drift');
-      final state = await backend.driver.prepareFileForAnalysis(uri);
+      final state = await backend.driver.findLocalElements(uri);
       final discovered = state.discovery;
 
       DriftElementId id(String name) => DriftElementId(uri, name);
@@ -51,8 +51,7 @@ CREATE TABLE valid_2 (bar INTEGER);
 ''',
       });
 
-      final state = await backend.driver
+      final state = await backend.discoverLocalElements('package:a/main.drift');
-          .prepareFileForAnalysis(Uri.parse('package:a/main.drift'));
       expect(state.errorsDuringDiscovery, [
         isDriftError(contains('Expected a table name')),
       ]);
@@ -72,8 +71,7 @@ CREATE VIEW a AS VALUES(1,2,3);
 ''',
       });
 
-      final state = await backend.driver
+      final state = await backend.discoverLocalElements('package:a/main.drift');
-          .prepareFileForAnalysis(Uri.parse('package:a/main.drift'));
       expect(state.errorsDuringDiscovery, [
         isDriftError(contains('already defines an element named `a`')),
       ]);
@@ -89,8 +87,7 @@ CREATE VIEW a AS VALUES(1,2,3);
           'a|lib/b.drift': "CREATE TABLE foo (bar INTEGER);",
         });
 
-        final state = await backend.driver
+        final state = await backend.discoverLocalElements('package:a/a.drift');
-            .prepareFileForAnalysis(Uri.parse('package:a/a.drift'));
         expect(state, hasNoErrors);
         expect(
           state.discovery,
@@ -100,10 +97,11 @@ CREATE VIEW a AS VALUES(1,2,3);
             [Uri.parse('package:a/b.drift')],
           ),
         );
+
         expect(
           backend.driver.cache.knownFiles[Uri.parse('package:a/b.drift')],
-          isNotNull,
+          isNull,
-          reason: 'Import should have been prepared as well',
+          reason: 'Discovering local elements should not prepare other files',
         );
       });
 
@@ -113,23 +111,9 @@ CREATE VIEW a AS VALUES(1,2,3);
           'a|lib/b.drift': "import 'a.drift';",
         });
 
-        final state = await backend.driver
+        final state = await backend.discoverLocalElements('package:a/a.drift');
-            .prepareFileForAnalysis(Uri.parse('package:a/a.drift'));
         expect(state, hasNoErrors);
       });
-
-      test('emits warning on invalid import', () async {
-        final backend = TestBackend.inTest({
-          'a|lib/a.drift': "import 'b.drift';",
-        });
-
-        final state = await backend.driver
-            .prepareFileForAnalysis(Uri.parse('package:a/a.drift'));
-        expect(state.errorsDuringDiscovery, [
-          isDriftError(contains(
-              'The imported file, `package:a/b.drift`, does not exist'))
-        ]);
-      });
     });
   });
 
@@ -145,7 +129,7 @@ part 'a.dart';
       });
 
       final uri = Uri.parse('package:a/a.dart');
-      final state = await backend.driver.prepareFileForAnalysis(uri);
+      final state = await backend.driver.findLocalElements(uri);
 
       expect(state, hasNoErrors);
       expect(state.discovery, isA<NotADartLibrary>());
@@ -169,7 +153,7 @@ class Groups extends Table {
       });
 
       final uri = Uri.parse('package:a/a.dart');
-      final state = await backend.driver.prepareFileForAnalysis(uri);
+      final state = await backend.driver.findLocalElements(uri);
 
       expect(state, hasNoErrors);
       expect(
@@ -207,7 +191,7 @@ abstract class BaseRelationTable extends Table {
       });
 
       final uri = Uri.parse('package:a/a.dart');
-      final state = await backend.driver.prepareFileForAnalysis(uri);
+      final state = await backend.driver.findLocalElements(uri);
 
       expect(state, hasNoErrors);
       expect(
@@ -244,8 +228,7 @@ class InvalidGetter extends Table {
       });
 
       for (final source in backend.sourceContents.keys) {
-        final state =
+        final state = await backend.driver.findLocalElements(Uri.parse(source));
-            await backend.driver.prepareFileForAnalysis(Uri.parse(source));
 
         expect(
           state.errorsDuringDiscovery,
@@ -274,8 +257,8 @@ class B extends Table {
 ''',
       });
 
-      final state = await backend.driver
+      final state =
-          .prepareFileForAnalysis(Uri.parse('package:a/a.dart'));
+          await backend.driver.findLocalElements(Uri.parse('package:a/a.dart'));
 
       expect(state.errorsDuringDiscovery, [
         isDriftError(contains('already defines an element named `tbl`')),

drift_dev/test/analysis/resolver/resolver_test.dart

@@ -170,4 +170,16 @@ END;
       });
     });
   });
+
+  test('emits warning on invalid import', () async {
+    final backend = TestBackend.inTest({
+      'a|lib/a.drift': "import 'b.drift';",
+    });
+
+    final state = await backend.analyze('package:a/a.drift');
+    expect(state.errorsDuringDiscovery, [
+      isDriftError(
+          contains('The imported file, `package:a/b.drift`, does not exist'))
+    ]);
+  });
 }

drift_dev/test/analysis/test_utils.dart

@@ -208,6 +208,10 @@ class TestBackend extends DriftBackend {
 
   Future<void> dispose() async {}
 
+  Future<FileState> discoverLocalElements(String uriString) {
+    return driver.findLocalElements(Uri.parse(uriString));
+  }
+
   Future<FileState> analyze(String uriString) {
     return driver.fullyAnalyze(Uri.parse(uriString));
   }

drift_dev/test/utils.dart

@@ -44,6 +44,7 @@ Future<RecordingAssetWriter> emulateDriftBuild({
   ]);
 
   final stages = [
+    discover(options),
     preparingBuilder(options),
     analyzer(options),
     modularBuild ? modular(options) : driftBuilderNotShared(options),