[DAS] Fixes import prefix go to definition (multiple)

[email protected]

Fixes #56931

Change-Id: Ida0cb58cdb4267af31a87a29862a3f4f2e28907f
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/405201
Auto-Submit: Felipe Morschel <[email protected]>
Commit-Queue: Samuel Rawlins <[email protected]>
Reviewed-by: Keerti Parthasarathy <[email protected]>
Reviewed-by: Samuel Rawlins <[email protected]>

Author: FMorschel

pkg/analysis_server/test/lsp/definition_test.dart

@@ -119,6 +119,14 @@ void [!^f!]() {}
     await testContents(contents);
   }
 
+  Future<void> test_atDeclaration_importPrefix() async {
+    var contents = '''
+import 'dart:math' as [!^math!];
+''';
+
+    await testContents(contents);
+  }
+
   Future<void> test_atDeclaration_method() async {
     var contents = '''
 class A {
@@ -185,6 +193,15 @@ extension StringExtension on String {
     await testContents(contents);
   }
 
+  Future<void> test_comment_importPrefix() async {
+    var contents = '''
+/// This is a comment for [^math]
+import 'dart:math' as [!math!];
+''';
+
+    await testContents(contents);
+  }
+
   Future<void> test_comment_instanceMember_qualified() async {
     var contents = '''
 /// [A.myFi^eld].
@@ -450,6 +467,91 @@ foo(Object pair) {
     await testContents(contents);
   }
 
+  Future<void> test_importPrefix() async {
+    var contents = '''
+import 'dart:math' as [!math!];
+
+^math.Random? r;
+''';
+
+    await testContents(contents);
+  }
+
+  Future<void> test_importPrefix_multiple() async {
+    setLocationLinkSupport();
+
+    var code = TestCode.parse('''
+import 'dart:math' as /*[0*/math/*0]*/;
+import 'dart:async' as /*[1*/math/*1]*/;
+
+/*[2*/^math/*2]*/.Random? r;
+''');
+
+    await initialize();
+    await openFile(mainFileUri, code.code);
+    var res = await getDefinitionAsLocationLinks(
+      mainFileUri,
+      code.position.position,
+    );
+
+    expect(res, hasLength(2));
+    for (var (index, loc) in res.indexed) {
+      expect(loc.originSelectionRange, equals(code.ranges.last.range));
+      expect(loc.targetRange, equals(code.ranges[index].range));
+      expect(loc.targetSelectionRange, equals(code.ranges[index].range));
+    }
+  }
+
+  Future<void> test_importPrefix_multiple_alone() async {
+    var code = TestCode.parse('''
+import 'dart:math' as /*[0*/math/*0]*/;
+import 'dart:async' as /*[1*/math/*1]*/;
+
+void foo() {
+  // ignore: prefix_identifier_not_followed_by_dot
+  /*[2*/^math/*2]*/;
+}
+''');
+
+    await initialize();
+    await openFile(mainFileUri, code.code);
+    var res = await getDefinitionAsLocation(
+      mainFileUri,
+      code.position.position,
+    );
+
+    expect(res, hasLength(2));
+    for (var (index, loc) in res.indexed) {
+      expect(loc.range, equals(code.ranges[index].range));
+    }
+  }
+
+  Future<void> test_importPrefix_multiple_comment() async {
+    setLocationLinkSupport();
+
+    var code = TestCode.parse('''
+import 'dart:math' as /*[0*/math/*0]*/;
+import 'dart:async' as /*[1*/math/*1]*/;
+
+/// This is a comment that talks about [/*[2*/^math/*2]*/].
+math.Random? r;
+''');
+
+    await initialize();
+    await openFile(mainFileUri, code.code);
+    var res = await getDefinitionAsLocationLinks(
+      mainFileUri,
+      code.position.position,
+    );
+
+    expect(res, hasLength(2));
+    for (var (index, loc) in res.indexed) {
+      expect(loc.originSelectionRange, equals(code.ranges.last.range));
+      expect(loc.targetRange, equals(code.ranges[index].range));
+      expect(loc.targetSelectionRange, equals(code.ranges[index].range));
+    }
+  }
+
   Future<void> test_locationLink_class() async {
     setLocationLinkSupport();
 

pkg/analyzer_plugin/lib/src/utilities/navigation/navigation_dart.dart

@@ -136,6 +136,13 @@ class _DartNavigationCollector {
     var offset = nodeOrToken.offset;
     var length = nodeOrToken.length;
 
+    _addRegionForFragmentRange(offset, length, fragment);
+  }
+
+  void _addRegionForFragmentRange(
+      int? offset, int? length, Fragment? fragment) {
+    if (offset == null || length == null || fragment == null) return;
+
     // If this fragment is for a synthetic element, use the first fragment for
     // the non-synthetic element.
     if (fragment.element.isSynthetic) {
@@ -424,7 +431,11 @@ class _DartNavigationComputerVisitor extends RecursiveAstVisitor<void> {
 
   @override
   void visitImportPrefixReference(ImportPrefixReference node) {
-    computer._addRegionForElement(node.name, node.element2);
+    var element = node.element2;
+    if (element == null) return;
+    for (var fragment in element.fragments) {
+      computer._addRegionForFragment(node.name, fragment);
+    }
   }
 
   @override
@@ -559,7 +570,17 @@ class _DartNavigationComputerVisitor extends RecursiveAstVisitor<void> {
   @override
   void visitSimpleIdentifier(SimpleIdentifier node) {
     var element = node.writeOrReadElement2;
-    computer._addRegionForElement(node, element);
+    if (element case PrefixElement2(:var fragments, :var name3)) {
+      for (var fragment in fragments) {
+        computer._addRegionForFragmentRange(
+          node.offset,
+          name3?.length,
+          fragment,
+        );
+      }
+    } else {
+      computer._addRegionForElement(node, element);
+    }
   }
 
   @override