Merge #10686

10686: internal: Add `Semantics::original_ast_node` for upmapping nodes out of macro files r=Veykril a=Veykril

Fixes trying to insert imports into macro expanded files which then do text edits on very wrong text ranges.

Co-authored-by: Lukas Wirth <lukastw97@gmail.com>

Author: bors[bot]

crates/hir/src/semantics.rs

@@ -208,6 +208,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.original_range_opt(node)
     }
 
+    pub fn original_ast_node<N: AstNode>(&self, node: N) -> Option<N> {
+        self.imp.original_ast_node(node)
+    }
+
     pub fn diagnostics_display_range(&self, diagnostics: InFile<SyntaxNodePtr>) -> FileRange {
         self.imp.diagnostics_display_range(diagnostics)
     }
@@ -660,6 +664,11 @@ impl<'db> SemanticsImpl<'db> {
         node.as_ref().original_file_range_opt(self.db.upcast())
     }
 
+    fn original_ast_node<N: AstNode>(&self, node: N) -> Option<N> {
+        let file = self.find_file(node.syntax().clone());
+        file.with_value(node).original_ast_node(self.db.upcast()).map(|it| it.value)
+    }
+
     fn diagnostics_display_range(&self, src: InFile<SyntaxNodePtr>) -> FileRange {
         let root = self.db.parse_or_expand(src.file_id).unwrap();
         let node = src.value.to_node(&root);

crates/hir_expand/src/lib.rs

@@ -24,9 +24,9 @@ use std::{hash::Hash, iter, sync::Arc};
 
 use base_db::{impl_intern_key, salsa, CrateId, FileId, FileRange};
 use syntax::{
-    algo::skip_trivia_token,
+    algo::{self, skip_trivia_token},
     ast::{self, AstNode, HasAttrs},
-    Direction, SyntaxNode, SyntaxToken, TextRange,
+    Direction, SyntaxNode, SyntaxToken,
 };
 
 use crate::{
@@ -600,13 +600,15 @@ impl<'a> InFile<&'a SyntaxNode> {
 
     /// Attempts to map the syntax node back up its macro calls.
     pub fn original_file_range_opt(self, db: &dyn db::AstDatabase) -> Option<FileRange> {
-        match original_range_opt(db, self) {
-            Some(range) => {
-                let original_file = range.file_id.original_file(db);
-                if range.file_id != original_file.into() {
+        match ascend_node_border_tokens(db, self) {
+            Some(InFile { file_id, value: (first, last) }) => {
+                let original_file = file_id.original_file(db);
+                let range = first.text_range().cover(last.text_range());
+                if file_id != original_file.into() {
                     tracing::error!("Failed mapping up more for {:?}", range);
+                    return None;
                 }
-                Some(FileRange { file_id: original_file, range: range.value })
+                Some(FileRange { file_id: original_file, range })
             }
             _ if !self.file_id.is_macro() => Some(FileRange {
                 file_id: self.file_id.original_file(db),
@@ -617,28 +619,29 @@ impl<'a> InFile<&'a SyntaxNode> {
     }
 }
 
-fn original_range_opt(
+fn ascend_node_border_tokens(
     db: &dyn db::AstDatabase,
-    node: InFile<&SyntaxNode>,
-) -> Option<InFile<TextRange>> {
-    let expansion = node.file_id.expansion_info(db)?;
+    InFile { file_id, value: node }: InFile<&SyntaxNode>,
+) -> Option<InFile<(SyntaxToken, SyntaxToken)>> {
+    let expansion = file_id.expansion_info(db)?;
 
     // the input node has only one token ?
-    let single = skip_trivia_token(node.value.first_token()?, Direction::Next)?
-        == skip_trivia_token(node.value.last_token()?, Direction::Prev)?;
+    let first = skip_trivia_token(node.first_token()?, Direction::Next)?;
+    let last = skip_trivia_token(node.last_token()?, Direction::Prev)?;
+    let is_single_token = first == last;
 
-    node.value.descendants().find_map(|it| {
+    node.descendants().find_map(|it| {
         let first = skip_trivia_token(it.first_token()?, Direction::Next)?;
-        let first = ascend_call_token(db, &expansion, node.with_value(first))?;
+        let first = ascend_call_token(db, &expansion, InFile::new(file_id, first))?;
 
         let last = skip_trivia_token(it.last_token()?, Direction::Prev)?;
-        let last = ascend_call_token(db, &expansion, node.with_value(last))?;
+        let last = ascend_call_token(db, &expansion, InFile::new(file_id, last))?;
 
-        if (!single && first == last) || (first.file_id != last.file_id) {
+        if (!is_single_token && first == last) || (first.file_id != last.file_id) {
             return None;
         }
 
-        Some(first.with_value(first.value.text_range().cover(last.value.text_range())))
+        Some(InFile::new(first.file_id, (first.value, last.value)))
     })
 }
 
@@ -674,6 +677,23 @@ impl<N: AstNode> InFile<N> {
         self.value.syntax().descendants().filter_map(T::cast).map(move |n| self.with_value(n))
     }
 
+    pub fn original_ast_node(self, db: &dyn db::AstDatabase) -> Option<InFile<N>> {
+        match ascend_node_border_tokens(db, self.syntax()) {
+            Some(InFile { file_id, value: (first, last) }) => {
+                let original_file = file_id.original_file(db);
+                if file_id != original_file.into() {
+                    let range = first.text_range().cover(last.text_range());
+                    tracing::error!("Failed mapping up more for {:?}", range);
+                    return None;
+                }
+                let anc = algo::least_common_ancestor(&first.parent()?, &last.parent()?)?;
+                Some(InFile::new(file_id, anc.ancestors().find_map(N::cast)?))
+            }
+            _ if !self.file_id.is_macro() => Some(self),
+            _ => None,
+        }
+    }
+
     pub fn syntax(&self) -> InFile<&SyntaxNode> {
         self.with_value(self.value.syntax())
     }

crates/ide_assists/src/handlers/auto_import.rs

@@ -95,7 +95,7 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
         NodeOrToken::Token(token) => token.text_range(),
     };
     let group_label = group_label(import_assets.import_candidate());
-    let scope = ImportScope::find_insert_use_container_with_macros(
+    let scope = ImportScope::find_insert_use_container(
         &match syntax_under_caret {
             NodeOrToken::Node(it) => it,
             NodeOrToken::Token(it) => it.parent()?,
@@ -164,6 +164,60 @@ mod tests {
 
     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
 
+    #[test]
+    fn not_applicable_if_scope_inside_macro() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+mod bar {
+    pub struct Baz;
+}
+macro_rules! foo {
+    ($it:ident) => {
+        mod __ {
+            fn __(x: $it) {}
+        }
+    };
+}
+foo! {
+    Baz$0
+}
+",
+        );
+    }
+
+    #[test]
+    fn applicable_in_attributes() {
+        check_assist(
+            auto_import,
+            r"
+//- proc_macros: identity
+#[proc_macros::identity]
+mod foo {
+    mod bar {
+        const _: Baz$0 = ();
+    }
+}
+mod baz {
+    pub struct Baz;
+}
+",
+            r"
+#[proc_macros::identity]
+mod foo {
+    mod bar {
+        use crate::baz::Baz;
+
+        const _: Baz = ();
+    }
+}
+mod baz {
+    pub struct Baz;
+}
+",
+        );
+    }
+
     #[test]
     fn applicable_when_found_an_import_partial() {
         check_assist(

crates/ide_assists/src/handlers/extract_function.rs

@@ -91,7 +91,7 @@ pub(crate) fn extract_function(acc: &mut Assists, ctx: &AssistContext) -> Option
 
     let target_range = body.text_range();
 
-    let scope = ImportScope::find_insert_use_container_with_macros(&node, &ctx.sema)?;
+    let scope = ImportScope::find_insert_use_container(&node, &ctx.sema)?;
 
     acc.add(
         AssistId("extract_function", crate::AssistKind::RefactorExtract),

crates/ide_assists/src/handlers/extract_struct_from_enum_variant.rs

@@ -314,7 +314,7 @@ fn process_references(
                 if let Some(mut mod_path) = mod_path {
                     mod_path.pop_segment();
                     mod_path.push_segment(variant_hir_name.clone());
-                    let scope = ImportScope::find_insert_use_container(&scope_node)?;
+                    let scope = ImportScope::find_insert_use_container(&scope_node, &ctx.sema)?;
                     visited_modules.insert(module);
                     return Some((segment, scope_node, Some((scope, mod_path))));
                 }

crates/ide_assists/src/handlers/replace_qualified_name_with_use.rs

@@ -70,7 +70,7 @@ pub(crate) fn replace_qualified_name_with_use(
         })
         .flatten();
 
-    let scope = ImportScope::find_insert_use_container_with_macros(path.syntax(), &ctx.sema)?;
+    let scope = ImportScope::find_insert_use_container(path.syntax(), &ctx.sema)?;
     let target = path.syntax().text_range();
     acc.add(
         AssistId("replace_qualified_name_with_use", AssistKind::RefactorRewrite),

crates/ide_assists/src/tests.rs

@@ -2,7 +2,7 @@ mod sourcegen;
 mod generated;
 
 use expect_test::expect;
-use hir::Semantics;
+use hir::{db::DefDatabase, Semantics};
 use ide_db::{
     base_db::{fixture::WithFixture, FileId, FileRange, SourceDatabaseExt},
     helpers::{
@@ -117,7 +117,8 @@ enum ExpectedResult<'a> {
 
 #[track_caller]
 fn check(handler: Handler, before: &str, expected: ExpectedResult, assist_label: Option<&str>) {
-    let (db, file_with_caret_id, range_or_offset) = RootDatabase::with_range_or_offset(before);
+    let (mut db, file_with_caret_id, range_or_offset) = RootDatabase::with_range_or_offset(before);
+    db.set_enable_proc_attr_macros(true);
     let text_without_caret = db.file_text(file_with_caret_id).to_string();
 
     let frange = FileRange { file_id: file_with_caret_id, range: range_or_offset.into() };

crates/ide_completion/src/completions/attribute/derive.rs

@@ -98,7 +98,7 @@ fn flyimport_attribute(ctx: &CompletionContext, acc: &mut Completions) -> Option
         &ctx.sema,
         parent.clone(),
     )?;
-    let import_scope = ImportScope::find_insert_use_container_with_macros(&parent, &ctx.sema)?;
+    let import_scope = ImportScope::find_insert_use_container(&parent, &ctx.sema)?;
     acc.add_all(
         import_assets
             .search_for_imports(&ctx.sema, ctx.config.insert_use.prefix_kind)

crates/ide_completion/src/completions/flyimport.rs

@@ -129,7 +129,7 @@ pub(crate) fn import_on_the_fly(acc: &mut Completions, ctx: &CompletionContext)
 
     let user_input_lowercased = potential_import_name.to_lowercase();
     let import_assets = import_assets(ctx, potential_import_name)?;
-    let import_scope = ImportScope::find_insert_use_container_with_macros(
+    let import_scope = ImportScope::find_insert_use_container(
         &position_for_import(ctx, Some(import_assets.import_candidate()))?,
         &ctx.sema,
     )?;

crates/ide_completion/src/completions/postfix.rs

@@ -244,8 +244,7 @@ fn add_custom_postfix_completions(
     postfix_snippet: impl Fn(&str, &str, &str) -> Builder,
     receiver_text: &str,
 ) -> Option<()> {
-    let import_scope =
-        ImportScope::find_insert_use_container_with_macros(&ctx.token.parent()?, &ctx.sema)?;
+    let import_scope = ImportScope::find_insert_use_container(&ctx.token.parent()?, &ctx.sema)?;
     ctx.config.postfix_snippets().filter(|(_, snip)| snip.scope == SnippetScope::Expr).for_each(
         |(trigger, snippet)| {
             let imports = match snippet.imports(ctx, &import_scope) {