[analysis_server] Copy the original imports from source when moving top-levels to new files

Previously we would just import the library containing the declaration (which for other packages could be in `src/`).

See #30310 (comment).

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

Author: DanTup

pkg/analysis_server/lib/src/services/refactoring/move_top_level_to_file.dart

@@ -182,12 +182,17 @@ class MoveTopLevelToFile extends RefactoringProducer {
   void _addImportsForMovingDeclarations(
       DartFileEditBuilder builder, ImportAnalyzer analyzer) {
     for (var entry in analyzer.movingReferences.entries) {
-      var library = entry.key.library;
-      if (library != null && !library.isDartCore) {
-        var uri = library.source.uri;
-        for (var prefix in entry.value) {
-          builder.importLibrary(uri, prefix: prefix.isEmpty ? null : prefix);
+      var imports = entry.value;
+      for (var import in imports) {
+        var library = import.importedLibrary;
+        if (library == null || library.isDartCore) {
+          continue;
         }
+        // TODO(dantup): This does not support show/hide. We should be able to
+        //  pass them in (and have them merge with any existing imports or
+        //  pending imports).
+        builder.importLibrary(library.source.uri,
+            prefix: import.prefix?.element.name);
       }
     }
   }

pkg/analysis_server/lib/src/utilities/import_analyzer.dart

@@ -25,20 +25,20 @@ class ImportAnalyzer {
   final Set<Element> stayingDeclarations = {};
 
   /// A map from the elements referenced by the declarations to be moved to the
-  /// set of prefixes used to reference those declarations.
-  final Map<Element, Set<String>> movingReferences = {};
+  /// set of imports used to reference those declarations.
+  final Map<Element, Set<LibraryImportElement>> movingReferences = {};
 
   /// A map from the elements referenced by the declarations that are staying to
-  /// the set of prefixes used to reference those declarations.
-  final Map<Element, Set<String>> stayingReferences = {};
+  /// the set of imports used to reference those declarations.
+  final Map<Element, Set<LibraryImportElement>> stayingReferences = {};
 
   /// Analyze the given library [result] to find the declarations and references
   /// being moved and that are staying. The declarations being moved are in the
   /// file at the given [path] in the given [range].
   ImportAnalyzer(this.result, String path, List<SourceRange> ranges) {
     for (var unit in result.units) {
       var finder = _ReferenceFinder(
-          _ElementRecorder(this, path == unit.path ? ranges : []));
+          unit, _ElementRecorder(this, path == unit.path ? ranges : []));
       unit.unit.accept(finder);
     }
     // Remove references that will be within the same file.
@@ -102,23 +102,27 @@ class _ElementRecorder {
     }
   }
 
-  /// Record that the [element] is referenced in the library at the
-  /// [referenceOffset]. [prefix] is the prefixused to reference the element, or
-  /// `null` if no prefix was used.
-  void recordReference(
-      Element referencedElement, int referenceOffset, PrefixElement? prefix) {
+  /// Record that [referencedElement] is referenced in the library at the
+  /// [referenceOffset]. [import] is the specific import used to reference the
+  /// including any prefix, show, hide.
+  void recordReference(Element referencedElement, int referenceOffset,
+      LibraryImportElement? import) {
     if (referencedElement is PropertyAccessorElement &&
         referencedElement.isSynthetic) {
       referencedElement = referencedElement.variable;
     }
     if (_isBeingMoved(referenceOffset)) {
-      var prefixes =
+      var imports =
           analyzer.movingReferences.putIfAbsent(referencedElement, () => {});
-      prefixes.add(prefix?.name ?? '');
+      if (import != null) {
+        imports.add(import);
+      }
     } else {
-      var prefixes =
+      var imports =
           analyzer.stayingReferences.putIfAbsent(referencedElement, () => {});
-      prefixes.add(prefix?.name ?? '');
+      if (import != null) {
+        imports.add(import);
+      }
     }
   }
 
@@ -139,8 +143,24 @@ class _ReferenceFinder extends RecursiveAstVisitor<void> {
   /// sent.
   final _ElementRecorder recorder;
 
+  /// The unit being searched for references.
+  final ResolvedUnitResult unit;
+
+  /// A mapping of prefixes to the imports with those prefixes. An
+  /// empty string is used for unprefixed imports.
+  ///
+  /// Library imports are ordered the same as they appear in the source file
+  /// (since this is a [LinkedHashSet]).
+  final _importsByPrefix = <String, Set<LibraryImportElement>>{};
+
   /// Initialize a newly created finder to send information to the [recorder].
-  _ReferenceFinder(this.recorder);
+  _ReferenceFinder(this.unit, this.recorder) {
+    for (var import in unit.libraryElement.libraryImports) {
+      _importsByPrefix
+          .putIfAbsent(import.prefix?.element.name ?? '', () => {})
+          .add(import);
+    }
+  }
 
   @override
   void visitAssignmentExpression(AssignmentExpression node) {
@@ -249,6 +269,43 @@ class _ReferenceFinder extends RecursiveAstVisitor<void> {
     super.visitTopLevelVariableDeclaration(node);
   }
 
+  /// Finds the [LibraryImportElement] that is used to import [element] for use
+  /// in [node].
+  LibraryImportElement? _getImportForElement(AstNode? node, Element element) {
+    var prefix = _getPrefixFromExpression(node)?.name;
+    var elementName = element.name;
+    // We cannot locate imports for unnamed elements.
+    if (elementName == null) {
+      return null;
+    }
+
+    var import = _importsByPrefix[prefix ?? '']?.where((import) {
+      // Check if this import is providing our element with the correct
+      // prefix/name.
+      var exportedElement = prefix != null
+          ? import.namespace.getPrefixed(prefix, elementName)
+          : import.namespace.get(elementName);
+      return exportedElement == element;
+    }).firstOrNull;
+
+    // Extensions can be used without a prefix, so we can use any import that
+    // brings in the extension.
+    if (import == null && prefix == null && element is ExtensionElement) {
+      import = _importsByPrefix.values
+          .expand((imports) => imports)
+          .where((import) =>
+              // Because we don't know what prefix we're looking for (any is
+              // allowed), use the imports own prefix when checking for the
+              // element.
+              import.namespace.getPrefixed(
+                  import.prefix?.element.name ?? '', elementName) ==
+              element)
+          .firstOrNull;
+    }
+
+    return import;
+  }
+
   /// Return the prefix used in [node].
   PrefixElement? _getPrefixFromExpression(AstNode? node) {
     if (node is PrefixedIdentifier) {
@@ -291,8 +348,9 @@ class _ReferenceFinder extends RecursiveAstVisitor<void> {
     if (!element.isInterestingReference) {
       return;
     }
-    recorder.recordReference(
-        element, node.offset, _getPrefixFromExpression(prefixNode));
+
+    var import = _getImportForElement(prefixNode, element);
+    recorder.recordReference(element, node.offset, import);
   }
 }
 

pkg/analysis_server/test/src/services/refactoring/move_top_level_to_file_test.dart

@@ -186,6 +186,37 @@ class A {}
     expect(content[newFilePath], expectedNewFileContent);
   }
 
+  Future<void> test_imports_declarationInSrc() async {
+    var libFilePath = join(projectFolderPath, 'lib', 'a.dart');
+    var srcFilePath = join(projectFolderPath, 'lib', 'src', 'a.dart');
+    addSource(libFilePath, 'export "src/a.dart";');
+    addSource(srcFilePath, 'class A {}');
+    var originalSource = '''
+import 'package:test/a.dart';
+
+A? staying;
+A? mov^ing;
+''';
+    var modifiedSource = '''
+import 'package:test/a.dart';
+
+A? staying;
+''';
+    var declarationName = 'moving';
+    var newFileName = 'moving.dart';
+    var newFileContent = '''
+import 'package:test/a.dart';
+
+A? moving;
+''';
+    await _singleDeclaration(
+        originalSource: originalSource,
+        modifiedSource: modifiedSource,
+        declarationName: declarationName,
+        newFileName: newFileName,
+        newFileContent: newFileContent);
+  }
+
   Future<void> test_imports_extensionMethod() async {
     var otherFilePath = '$projectFolderPath/lib/extensions.dart';
     var otherFileContent = '''
@@ -433,10 +464,8 @@ class A {}
 ''';
     var movingDeclarationName = 'moving';
     var newFileName = 'moving.dart';
-    // The prefix is not included because it's not used (the only use of
-    // extensions.dart is the extension method).
     var newFileContent = '''
-import 'package:test/extensions.dart';
+import 'package:test/extensions.dart' as other;
 import 'package:test/main.dart';
 
 void moving() {
@@ -479,10 +508,8 @@ class A {}
 ''';
     var movingDeclarationName = 'moving';
     var newFileName = 'moving.dart';
-    // The prefix is not included because it's not used (the only use of
-    // extensions.dart is the extension method).
     var newFileContent = '''
-import 'package:test/extensions.dart';
+import 'package:test/extensions.dart' as other;
 import 'package:test/main.dart';
 
 void moving() {