Also optimize looking up Dart references

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

@@ -29,8 +29,8 @@ class DriftAnalysisCache {
   void knowsLocalElements(FileState state) {
     discoveredElements.removeWhere((key, _) => key.libraryUri == state.ownUri);
 
-    for (final (id, kind) in state.definedElements) {
-      discoveredElements[id] = kind;
+    for (final defined in state.definedElements) {
+      discoveredElements[defined.ownId] = defined.kind;
     }
   }
 

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

@@ -140,7 +140,7 @@ class DriftAnalysisDriver {
 
     final cachedImports = cache.serializationCache[state.ownUri]?.cachedImports;
     if (cachedImports != null && state.discovery == null) {
-      state.cachedDiscovery ??= CachedDiscoveryResults(true, cachedImports, {});
+      state.cachedDiscovery ??= CachedDiscoveryResults(true, cachedImports, []);
 
       for (final import in cachedImports) {
         final found = cache.stateForUri(import);
@@ -191,7 +191,7 @@ class DriftAnalysisDriver {
       if (cached == null && reader.findsLocalElementsReliably) {
         // There are no locally defined elements, since otherwise the reader
         // would have found them.
-        cached = CachedDiscoveryResults(false, const [], const {});
+        cached = CachedDiscoveryResults(false, const [], const []);
       }
     }
 

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

@@ -40,19 +40,16 @@ class FileState {
     return analyzedElements.any((e) => e is BaseDriftAccessor);
   }
 
-  Iterable<(DriftElementId, DriftElementKind)> get definedElements sync* {
+  Iterable<ExistingDriftElement> get definedElements {
     final discovery = this.discovery;
     final cached = cachedDiscovery;
 
     if (discovery != null) {
-      for (final element in discovery.locallyDefinedElements) {
-        yield (element.ownId, element.kind);
-      }
+      return discovery.locallyDefinedElements;
     } else if (cached != null) {
-      for (final MapEntry(:key, :value)
-          in cached.locallyDefinedElements.entries) {
-        yield (id(key), value);
-      }
+      return cached.locallyDefinedElements;
+    } else {
+      return const Iterable.empty();
     }
   }
 
@@ -114,7 +111,7 @@ class FileState {
 class CachedDiscoveryResults {
   final bool isValidImport;
   final List<Uri> imports;
-  final Map<String, DriftElementKind> locallyDefinedElements;
+  final List<ExistingDriftElement> locallyDefinedElements;
 
   CachedDiscoveryResults(
     this.isValidImport,
@@ -187,8 +184,30 @@ class UnknownFile extends DiscoveredFileState {
   UnknownFile() : super(const []);
 }
 
-abstract class DiscoveredElement {
+class ExistingDriftElement {
   final DriftElementId ownId;
+  final DriftElementKind kind;
+
+  /// When this element was defined via a Dart class, the name of that class.
+  ///
+  /// Together with the [DriftElementId.libraryUri] of the [id], this can be
+  /// used to find the Drift element for an [Element] without resolving all
+  /// available imports multiple times.
+  final String? dartElementName;
+
+  ExistingDriftElement({
+    required this.ownId,
+    required this.kind,
+    this.dartElementName,
+  });
+}
+
+/// Information about the syntax of a known drift element which can be used to
+/// fully resolve it.
+abstract class DiscoveredElement implements ExistingDriftElement {
+  @override
+  final DriftElementId ownId;
+  @override
   DriftElementKind get kind;
 
   DiscoveredElement(this.ownId);

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

@@ -12,6 +12,9 @@ class DiscoveredDriftElement<AST extends AstNode> extends DiscoveredElement {
   @override
   final DriftElementKind kind;
 
+  @override
+  String? get dartElementName => null;
+
   DiscoveredDriftElement(super.ownId, this.kind, this.sqlNode);
 }
 
@@ -25,6 +28,9 @@ abstract class DiscoveredDartElement<DE extends Element>
     extends DiscoveredElement {
   final DE dartElement;
 
+  @override
+  String? get dartElementName => dartElement.name;
+
   DiscoveredDartElement(super.ownId, this.dartElement);
 }
 

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

@@ -140,14 +140,12 @@ class DriftResolver {
       DriftElementId owner, Element element) async {
     final uri = await driver.backend.uriOfDart(element.library!);
     final state = driver.cache.stateForUri(uri);
-    await driver.discoverIfNecessary(driver.cache.stateForUri(uri));
 
-    final discovered = state.discovery?.locallyDefinedElements
-        .whereType<DiscoveredDartElement>()
-        .firstWhereOrNull((c) => c.dartElement == element);
+    final existing = state.definedElements.firstWhereOrNull(
+        (existing) => existing.dartElementName == element.name);
 
-    if (discovered != null) {
-      return resolveReferencedElement(owner, discovered.ownId);
+    if (existing != null) {
+      return resolveReferencedElement(owner, existing.ownId);
     } else {
       return InvalidReferenceResult(
         InvalidReferenceError.noElementWichSuchName,
@@ -170,7 +168,7 @@ class DriftResolver {
     for (final available in driver.cache.crawl(file)) {
       final localElementIds = {
         ...available.analysis.keys,
-        ...available.definedElements.map((e) => e.$1),
+        ...available.definedElements.map((e) => e.ownId),
       };
 
       for (final definedLocally in localElementIds) {

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

@@ -6,6 +6,7 @@ import 'package:package_config/package_config.dart';
 
 import '../../analysis/driver/driver.dart';
 import '../../analysis/options.dart';
+import '../../analysis/resolver/intermediate_state.dart';
 import '../../writer/import_manager.dart';
 import '../../writer/writer.dart';
 import 'backend.dart';
@@ -49,6 +50,8 @@ class DriftDiscover extends Builder {
               {
                 'kind': entry.kind.name,
                 'name': entry.ownId.name,
+                if (entry is DiscoveredDartElement)
+                  'dart_name:': entry.dartElement.name,
               }
           ]
         }),

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

@@ -126,11 +126,14 @@ class BuildCacheReader implements AnalysisResultCacheReader {
       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']]!,
-        },
+            ExistingDriftElement(
+              ownId: DriftElementId(uri, element['name'] as String),
+              kind: DriftElementKind.byName[element['kind']]!,
+              dartElementName: element['dart_name'] as String?,
+            ),
+        ],
       );
     }