Fix annotating top-level types and extensions

By not using a `SimpleIdentifier` AST node for class and extension
names, the code used to detect imports for Dart code was broken.
Closes #2423

drift_dev/CHANGELOG.md

@@ -1,6 +1,7 @@
-## 2.9.0-dev
+## 2.8.1
 
 - Add documentation comments for comments on columns in drift files.
+- Fix modular generation not generating imports correctly with the latest analyzer.
 
 ## 2.8.0
 

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

@@ -220,7 +220,8 @@ class DartTopLevelSymbol {
     return DartTopLevelSymbol(name, _driftUri);
   }
 
-  factory DartTopLevelSymbol.topLevelElement(Element element) {
+  factory DartTopLevelSymbol.topLevelElement(Element element,
+      [String? elementName]) {
     assert(element.library?.topLevelElements.contains(element) == true);
 
     // We're using this to recover the right import URI when using
@@ -231,7 +232,8 @@ class DartTopLevelSymbol {
       sourceUri = AssetId.resolve(sourceUri).uri;
     }
 
-    return DartTopLevelSymbol(element.name ?? '(???)', sourceUri);
+    return DartTopLevelSymbol(
+        elementName ?? element.name ?? '(???)', sourceUri);
   }
 
   factory DartTopLevelSymbol.fromJson(Map json) =>
@@ -434,6 +436,15 @@ class _AddFromAst extends GeneralizingAstVisitor<void> {
 
   _AddFromAst(this._builder, this._excluding);
 
+  void _addTopLevelReference(Element? element, Token name2) {
+    if (element == null) {
+      _builder.addText(name2.lexeme);
+    } else {
+      _builder.addTopLevel(
+          DartTopLevelSymbol.topLevelElement(element, name2.lexeme));
+    }
+  }
+
   @override
   void visitNode(AstNode node) {
     if (_excluding.contains(node)) return;
@@ -470,6 +481,18 @@ class _AddFromAst extends GeneralizingAstVisitor<void> {
     _builder.addText(')');
   }
 
+  @override
+  void visitExtensionOverride(ExtensionOverride node) {
+    _addTopLevelReference(node.element, node.name); // Transform identifier
+    node.typeArguments?.accept(this);
+    node.argumentList.accept(this);
+  }
+
+  @override
+  void visitNamedType(NamedType node) {
+    _addTopLevelReference(node.element, node.name2);
+  }
+
   @override
   void visitPrefixedIdentifier(PrefixedIdentifier node) {
     final targetOfPrefix = node.prefix.staticElement;

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

@@ -64,7 +64,7 @@ class DriftAnalyzer extends Builder {
         final version = await buildStep.languageVersionForPackage ??
             _languageVersionForGeneralizedTypedefs;
 
-        final imports = LibraryInputManager();
+        final imports = LibraryImportManager();
         final writer = Writer(
           options,
           generationOptions: GenerationOptions(imports: imports),

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

@@ -365,7 +365,7 @@ class _DriftBuildRun {
       );
       writer = Writer(options, generationOptions: generationOptions);
     } else {
-      final imports = LibraryInputManager(buildStep.allowedOutputs.single.uri);
+      final imports = LibraryImportManager(buildStep.allowedOutputs.single.uri);
       final generationOptions = GenerationOptions(
         imports: imports,
         isModular: true,

drift_dev/lib/src/writer/import_manager.dart

@@ -14,7 +14,9 @@ class ImportManagerForPartFiles extends ImportManager {
   }
 }
 
-class LibraryInputManager extends ImportManager {
+/// An [ImportManager] for generation contexts that create standalone Dart
+/// libraries capable of managing their own imports.
+class LibraryImportManager extends ImportManager {
   static final _dartCore = Uri.parse('dart:core');
 
   final Map<Uri, String> _importAliases = {};
@@ -27,7 +29,7 @@ class LibraryInputManager extends ImportManager {
 
   TextEmitter? emitter;
 
-  LibraryInputManager([this._outputUri]);
+  LibraryImportManager([this._outputUri]);
 
   void linkToWriter(Writer writer) {
     emitter = writer.leaf();

drift_dev/pubspec.yaml

@@ -1,6 +1,6 @@
 name: drift_dev
 description: Dev-dependency for users of drift. Contains the generator and development tools.
-version: 2.9.0-dev
+version: 2.8.1-dev
 repository: https://github.com/simolus3/drift
 homepage: https://drift.simonbinder.eu/
 issue_tracker: https://github.com/simolus3/drift/issues

drift_dev/test/analysis/results/dart_test.dart

@@ -0,0 +1,56 @@
+import 'package:drift_dev/src/analysis/options.dart';
+import 'package:drift_dev/src/analysis/results/dart.dart';
+import 'package:drift_dev/src/writer/writer.dart';
+import 'package:test/test.dart';
+
+import '../test_utils.dart';
+
+void main() {
+  late TestBackend tester;
+
+  setUpAll(() => tester = TestBackend({}));
+  tearDownAll(() => tester.dispose());
+
+  group('from AST', () {
+    final testUri = Uri.parse('package:a/test.dart');
+
+    Future<void> checkTransformation(String sourceExpression,
+        String expectedResult, Map<String, String> expectedImports) async {
+      final expression =
+          await tester.resolveExpression(testUri, sourceExpression, const []);
+      final annotated = AnnotatedDartCode.ast(expression);
+
+      final imports = TestImportManager();
+      final writer = Writer(
+        const DriftOptions.defaults(),
+        generationOptions: GenerationOptions(
+          imports: imports,
+        ),
+      );
+
+      final code = writer.dartCode(annotated);
+
+      expect(code, expectedResult);
+      expectedImports.forEach((alias, import) {
+        expect(imports.importAliases[Uri.parse(import)], alias);
+      });
+    }
+
+    test('constructor invocation', () async {
+      await checkTransformation('const Duration(seconds: 12)',
+          'const i0.Duration(seconds: 12)', {'i0': 'dart:core'});
+    });
+
+    test('static invocation', () async {
+      await checkTransformation(
+          'Uri.parse("")', 'i0.Uri.parse("")', {'i0': 'dart:core'});
+    });
+
+    test('explicit extension invocation', () async {
+      await checkTransformation(
+          'IterableExtensions<String>([]).firstOrNull',
+          'i0.IterableExtensions<i1.String>([]).firstOrNull',
+          {'i0': 'dart:collection', 'i1': 'dart:core'});
+    });
+  });
+}

drift_dev/test/analysis/test_utils.dart

@@ -16,6 +16,7 @@ import 'package:drift_dev/src/analysis/driver/error.dart';
 import 'package:drift_dev/src/analysis/driver/state.dart';
 import 'package:drift_dev/src/analysis/results/results.dart';
 import 'package:drift_dev/src/analysis/options.dart';
+import 'package:drift_dev/src/writer/import_manager.dart';
 import 'package:logging/logging.dart';
 import 'package:package_config/package_config.dart';
 import 'package:path/path.dart' as p;
@@ -212,6 +213,16 @@ class TestBackend extends DriftBackend {
   }
 }
 
+class TestImportManager extends ImportManager {
+  final Map<Uri, String> importAliases = {};
+
+  @override
+  String? prefixFor(Uri definitionUri, String elementName) {
+    return importAliases.putIfAbsent(
+        definitionUri, () => 'i${importAliases.length}');
+  }
+}
+
 Matcher get hasNoErrors =>
     isA<FileState>().having((e) => e.allErrors, 'allErrors', isEmpty);
 

drift_dev/test/writer/import_manager_test.dart

@@ -9,11 +9,11 @@ void main() {
   group('LibraryInputManager', () {
     final sourceUri = AssetId('a', 'example/main.dart').uri;
 
-    late LibraryInputManager imports;
+    late LibraryImportManager imports;
     late Writer writer;
 
     setUp(() {
-      imports = LibraryInputManager(sourceUri);
+      imports = LibraryImportManager(sourceUri);
       final generationOptions =
           GenerationOptions(imports: imports, isModular: true);
       writer = Writer(const DriftOptions.defaults(),

drift_dev/test/writer/utils/hash_code_test.dart

@@ -9,7 +9,7 @@ void main() {
   late Writer writer;
 
   setUp(() {
-    final imports = LibraryInputManager(Uri.parse('drift:test'));
+    final imports = LibraryImportManager(Uri.parse('drift:test'));
     final generationOptions =
         GenerationOptions(imports: imports, isModular: true);
     writer = Writer(const DriftOptions.defaults(),

drift_dev/test/writer/utils/override_equals_test.dart

@@ -8,7 +8,7 @@ void main() {
   late Writer writer;
 
   setUp(() {
-    final imports = LibraryInputManager(Uri.parse('drift:test'));
+    final imports = LibraryImportManager(Uri.parse('drift:test'));
     final generationOptions =
         GenerationOptions(imports: imports, isModular: true);
     writer = Writer(const DriftOptions.defaults(),