[analysis_server] Update "Go to Imports" to support enhanced parts

Instead of looking in a single root unit, we now look in each unit up the tree.

Change-Id: I59b0692438dd73fce7be7790ddc24240761d72c3
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/394200
Reviewed-by: Konstantin Shcheglov <[email protected]>
Reviewed-by: Brian Wilkerson <[email protected]>
Commit-Queue: Brian Wilkerson <[email protected]>

Author: DanTup

pkg/analysis_server/lib/src/lsp/handlers/custom/handler_imports.dart

@@ -9,12 +9,11 @@ import 'package:analysis_server/src/lsp/error_or.dart';
 import 'package:analysis_server/src/lsp/handlers/handlers.dart';
 import 'package:analysis_server/src/lsp/mapping.dart';
 import 'package:analyzer/dart/analysis/results.dart';
-import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element2.dart';
+import 'package:analyzer/src/dart/ast/ast.dart';
 import 'package:analyzer/src/dart/ast/element_locator.dart';
-import 'package:analyzer/src/dart/ast/utilities.dart';
-
-typedef _ImportRecord = ({CompilationUnit unit, ImportDirective directive});
+import 'package:analyzer/src/utilities/extensions/ast.dart';
+import 'package:analyzer/src/utilities/extensions/results.dart';
 
 class ImportsHandler
     extends SharedMessageHandler<TextDocumentPositionParams, List<Location>?> {
@@ -60,7 +59,7 @@ class ImportsHandler
     );
 
     return (library, unit, offset).mapResults((library, unit, offset) async {
-      var node = NodeLocator(offset).searchWithin(unit.unit);
+      var node = unit.unit.nodeCovering(offset: offset);
       if (node == null) {
         return success(null);
       }
@@ -89,29 +88,56 @@ class ImportsHandler
         element = enclosingElement;
       }
 
-      var locations = _getImportLocations(element, library, unit, prefix);
+      var locations = _getImportLocations(library, unit, element, prefix);
 
       return success(nullIfEmpty(locations));
     });
   }
 
   /// Returns [Location]s for imports that import the given [element] into
-  /// [unit].
+  /// [unitResult].
   List<Location> _getImportLocations(
-    Element2 element,
     ResolvedLibraryResult libraryResult,
-    ResolvedUnitResult unit,
+    ResolvedUnitResult? unitResult,
+    Element2 element,
     String? prefix,
   ) {
     var elementName = element.name3;
     if (elementName == null) {
       return [];
     }
 
-    var imports = _getImports(libraryResult);
-    var results = <Location>[];
+    // Search in each unit up the chain for related imports.
+    while (unitResult is ResolvedUnitResult) {
+      var results = _getImportsInUnit(
+        unitResult.unit,
+        element,
+        prefix: prefix,
+        elementName: elementName,
+      );
+
+      // Stop searching in the unit where we find any matching imports.
+      if (results.isNotEmpty) {
+        return results;
+      }
+
+      // Otherwise, we continue up the chain.
+      unitResult = libraryResult.parentUnitOf(unitResult);
+    }
+
+    return [];
+  }
 
-    for (var (:unit, :directive) in imports) {
+  /// Gets the locations of all imports that provide [element] with [prefix] in
+  /// [unit].
+  List<Location> _getImportsInUnit(
+    CompilationUnit unit,
+    Element2 element, {
+    required String? prefix,
+    required String elementName,
+  }) {
+    var results = <Location>[];
+    for (var directive in unit.directives.whereType<ImportDirective>()) {
       var import = directive.libraryImport;
       if (import == null) continue;
 
@@ -134,20 +160,6 @@ class ImportsHandler
         results.add(Location(uri: uri, range: range));
       }
     }
-
     return results;
   }
-
-  List<_ImportRecord> _getImports(ResolvedLibraryResult libraryResult) {
-    var imports = <_ImportRecord>[];
-
-    // TODO(dantup): With enhanced parts, we may need to look at more than
-    //  just the first fragment.
-    var unit = libraryResult.units.first.unit;
-    for (var directive in unit.directives.whereType<ImportDirective>()) {
-      imports.add((unit: unit, directive: directive));
-    }
-
-    return imports;
-  }
 }

pkg/analysis_server/test/lsp/import_test.dart

@@ -18,6 +18,18 @@ void main() {
 
 @reflectiveTest
 class ImportTest extends AbstractLspAnalysisServerTest {
+  /// A helper for creating files for multi-level-parts tests.
+  ({Uri uri, String filePath, TestCode code}) createFile(
+    String filename,
+    String content,
+  ) {
+    assert(filename.endsWith('.dart'));
+    var filePath = join(projectFolderPath, 'lib', filename);
+    var code = TestCode.parse(content);
+    newFile(filePath, code.code);
+    return (uri: toUri(filePath), filePath: filePath, code: code);
+  }
+
   Future<void> test_constant() async {
     await _verifyGoToImports(
       TestCode.parse('''
@@ -180,6 +192,131 @@ a.A^? r;
     );
   }
 
+  Future<void> test_import_multiLevelParts() async {
+    // Create a tree of files that all import 'dart:math' and ensure we find
+    // only the import from the parent (not a grandparent, sibling, or child).
+    //
+    //
+    // - root                      has import
+    //     - level1_other          has import
+    //     - level1                has import, is the used reference
+    //         - level2_other      has import
+    //         - level2            has reference
+    //             - level3_other  has import
+
+    createFile('root.dart', '''
+import 'dart:math';
+part 'level1_other.dart';
+part 'level1.dart';
+''');
+
+    createFile('level1_other.dart', '''
+part of 'root.dart';
+import 'dart:math';
+''');
+
+    var level1 = createFile('level1.dart', '''
+part of 'root.dart';
+[!import 'dart:math';!]
+part 'level2_other.dart';
+part 'level2.dart';
+''');
+
+    createFile('level2_other.dart', '''
+part of 'level1.dart';
+import 'dart:math';
+''');
+
+    var level2 = createFile('level2.dart', '''
+part of 'level1.dart';
+part 'level3_other.dart';
+
+Rando^m? r;
+''');
+
+    createFile('level3_other.dart', '''
+part of 'level2.dart';
+import 'dart:math';
+''');
+
+    await _verifyGoToImports(
+      // Test the position in level2.
+      level2.code,
+      fileUri: level2.uri,
+      // Expect only the nearest parent import (level1).
+      expecting: (level1.uri, level1.code.ranges),
+    );
+  }
+
+  Future<void> test_import_multiLevelParts_findsInGrandParent() async {
+    var root = createFile('root.dart', '''
+[!import 'dart:math';!]
+part 'level1.dart';
+''');
+
+    createFile('level1.dart', '''
+part of 'root.dart';
+part 'level2.dart';
+''');
+
+    var level2 = createFile('level2.dart', '''
+part of 'level1.dart';
+Rando^m? r;''');
+
+    await _verifyGoToImports(
+      // Test the position in level2, expect reference in root.
+      level2.code,
+      fileUri: level2.uri,
+      expecting: (root.uri, root.code.ranges),
+    );
+  }
+
+  Future<void> test_import_multiLevelParts_findsInParent() async {
+    createFile('root.dart', '''
+part 'level1.dart';
+''');
+
+    var level1 = createFile('level1.dart', '''
+part of 'root.dart';
+[!import 'dart:math';!]
+part 'level2.dart';
+''');
+
+    var level2 = createFile('level2.dart', '''
+part of 'level1.dart';
+Rando^m? r;''');
+
+    await _verifyGoToImports(
+      // Test the position in level2, expect reference in level1.
+      level2.code,
+      fileUri: level2.uri,
+      expecting: (level1.uri, level1.code.ranges),
+    );
+  }
+
+  Future<void> test_import_multiLevelParts_findsInSelf() async {
+    createFile('root.dart', '''
+part 'level1.dart';
+''');
+
+    createFile('level1.dart', '''
+part of 'root.dart';
+part 'level2.dart';
+''');
+
+    var level2 = createFile('level2.dart', '''
+part of 'level1.dart';
+[!import 'dart:math';!]
+Rando^m? r;''');
+
+    await _verifyGoToImports(
+      // Test the position in level2, expect reference in same.
+      level2.code,
+      fileUri: level2.uri,
+      expecting: (level2.uri, level2.code.ranges),
+    );
+  }
+
   Future<void> test_import_part() async {
     var otherFileUri = Uri.file(join(projectFolderPath, 'lib', 'other.dart'));
     var main = TestCode.parse('''
@@ -195,7 +332,7 @@ part of '$mainFileUri';
 Rando^m? r;
 '''),
       fileUri: otherFileUri,
-      expecting: [_Results(mainFileUri, main.ranges)],
+      expecting: (mainFileUri, main.ranges),
     );
   }
 
@@ -277,38 +414,23 @@ var a = 1.abs().ba^r();
     );
   }
 
-  void _expecting(List<Location>? res, List<_Results>? fileRanges) {
-    List<Location>? expected;
-    if (fileRanges != null && fileRanges.expand((r) => r.ranges).isNotEmpty) {
-      expected = [
-        for (final _Results(:uri, :ranges) in fileRanges)
-          for (final range in ranges) Location(uri: uri, range: range.range),
-      ];
-    }
-
-    expect(res, equals(expected));
-  }
-
   Future<void> _verifyGoToImports(
     TestCode code, {
     Uri? fileUri,
-    List<_Results>? expecting,
+    (Uri, List<TestCodeRange>)? expecting,
   }) async {
+    expecting ??= (fileUri ?? mainFileUri, code.ranges);
+    var (expectedUri, expectedRanges) = expecting;
+    var expectedLocations = [
+      for (var range in expectedRanges)
+        Location(uri: expectedUri, range: range.range),
+    ];
+
     newFile(fromUri(fileUri ?? mainFileUri), code.code);
     await initialize();
     await initialAnalysis;
     var res = await getImports(fileUri ?? mainFileUri, code.position.position);
-    List<_Results>? results;
-    if (expecting == null && code.ranges.isNotEmpty) {
-      results = [_Results(fileUri ?? mainFileUri, code.ranges)];
-    }
-    _expecting(res, expecting ?? results);
-  }
-}
-
-extension type _Results._((Uri uri, List<TestCodeRange> ranges) _r) {
-  _Results(Uri uri, List<TestCodeRange> ranges) : _r = (uri, ranges);
 
-  List<TestCodeRange> get ranges => _r.$2;
-  Uri get uri => _r.$1;
+    expect(res ?? [], expectedLocations);
+  }
 }

pkg/analyzer/lib/src/utilities/extensions/results.dart

@@ -0,0 +1,16 @@
+// Copyright (c) 2024, the Dart project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:analyzer/dart/analysis/results.dart';
+
+extension ResolvedLibraryResultExtension on ResolvedLibraryResult {
+  /// Returns the parent unit (the unit that this is a `part of`) from this
+  /// result.
+  ///
+  /// Returns `null` if this result does not contain the parent unit.
+  ResolvedUnitResult? parentUnitOf(ResolvedUnitResult unit) {
+    var parentPath = unit.libraryFragment.enclosingFragment?.source.fullName;
+    return parentPath != null ? unitWithPath(parentPath) : null;
+  }
+}