Merge pull request #19890 from Veykril/push-kzzntrpllsqx

Veykril · web-flow · commit 23738b5d1784 · 2025-05-29T13:06:12.000Z
fix: Fix import insertion not being fully cfg aware
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/auto_import.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/auto_import.rs
@@ -128,11 +128,7 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
             format!("Import `{import_name}`"),
             range,
             |builder| {
-                let scope = match scope.clone() {
-                    ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
-                    ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
-                    ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
-                };
+                let scope = builder.make_import_scope_mut(scope.clone());
                 insert_use(&scope, mod_path_to_ast(&import_path, edition), &ctx.config.insert_use);
             },
         );
@@ -153,11 +149,7 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
                     format!("Import `{import_name} as _`"),
                     range,
                     |builder| {
-                        let scope = match scope.clone() {
-                            ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
-                            ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
-                            ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
-                        };
+                        let scope = builder.make_import_scope_mut(scope.clone());
                         insert_use_as_alias(
                             &scope,
                             mod_path_to_ast(&import_path, edition),
@@ -1877,4 +1869,30 @@ fn main() {
 ",
         );
     }
+
+    #[test]
+    fn carries_cfg_attr() {
+        check_assist(
+            auto_import,
+            r#"
+mod m {
+    pub struct S;
+}
+
+#[cfg(test)]
+fn foo(_: S$0) {}
+"#,
+            r#"
+#[cfg(test)]
+use m::S;
+
+mod m {
+    pub struct S;
+}
+
+#[cfg(test)]
+fn foo(_: S) {}
+"#,
+        );
+    }
 }
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_bool_to_enum.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_bool_to_enum.rs
@@ -312,12 +312,8 @@ fn replace_usages(
                 }
 
                 // add imports across modules where needed
-                if let Some((import_scope, path)) = import_data {
-                    let scope = match import_scope {
-                        ImportScope::File(it) => ImportScope::File(edit.make_mut(it)),
-                        ImportScope::Module(it) => ImportScope::Module(edit.make_mut(it)),
-                        ImportScope::Block(it) => ImportScope::Block(edit.make_mut(it)),
-                    };
+                if let Some((scope, path)) = import_data {
+                    let scope = edit.make_import_scope_mut(scope);
                     delayed_mutations.push((scope, path));
                 }
             },
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_named_struct_to_tuple_struct.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_named_struct_to_tuple_struct.rs
@@ -996,7 +996,8 @@ pub struct $0Foo {
 }
 "#,
             r#"
-pub struct Foo(#[my_custom_attr] u32);
+pub struct Foo(#[my_custom_attr]
+u32);
 "#,
         );
     }
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_tuple_struct_to_named_struct.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_tuple_struct_to_named_struct.rs
@@ -923,7 +923,8 @@ where
 pub struct $0Foo(#[my_custom_attr] u32);
 "#,
             r#"
-pub struct Foo { #[my_custom_attr] field1: u32 }
+pub struct Foo { #[my_custom_attr]
+field1: u32 }
 "#,
         );
     }
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/extract_function.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/extract_function.rs
@@ -204,12 +204,7 @@ pub(crate) fn extract_function(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
                 .kind
                 .is_some_and(|kind| matches!(kind, FlowKind::Break(_, _) | FlowKind::Continue(_)))
             {
-                let scope = match scope {
-                    ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
-                    ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
-                    ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
-                };
-
+                let scope = builder.make_import_scope_mut(scope);
                 let control_flow_enum =
                     FamousDefs(&ctx.sema, module.krate()).core_ops_ControlFlow();
 
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/replace_qualified_name_with_use.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/replace_qualified_name_with_use.rs
@@ -81,11 +81,7 @@ pub(crate) fn replace_qualified_name_with_use(
         |builder| {
             // Now that we've brought the name into scope, re-qualify all paths that could be
             // affected (that is, all paths inside the node we added the `use` to).
-            let scope = match scope {
-                ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
-                ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
-                ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
-            };
+            let scope = builder.make_import_scope_mut(scope);
             shorten_paths(scope.as_syntax_node(), &original_path);
             let path = drop_generic_args(&original_path);
             let edition = ctx
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unqualify_method_call.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unqualify_method_call.rs
@@ -1,4 +1,3 @@
-use ide_db::imports::insert_use::ImportScope;
 use syntax::{
     TextRange,
     ast::{self, AstNode, HasArgList, prec::ExprPrecedence},
@@ -114,11 +113,7 @@ fn add_import(
         );
 
         if let Some(scope) = scope {
-            let scope = match scope {
-                ImportScope::File(it) => ImportScope::File(edit.make_mut(it)),
-                ImportScope::Module(it) => ImportScope::Module(edit.make_mut(it)),
-                ImportScope::Block(it) => ImportScope::Block(edit.make_mut(it)),
-            };
+            let scope = edit.make_import_scope_mut(scope);
             ide_db::imports::insert_use::insert_use(&scope, import, &ctx.config.insert_use);
         }
     }
diff --git a/src/tools/rust-analyzer/crates/ide-db/src/imports/insert_use.rs b/src/tools/rust-analyzer/crates/ide-db/src/imports/insert_use.rs
@@ -60,107 +60,87 @@ pub struct InsertUseConfig {
 }
 
 #[derive(Debug, Clone)]
-pub enum ImportScope {
+pub struct ImportScope {
+    pub kind: ImportScopeKind,
+    pub required_cfgs: Vec<ast::Attr>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ImportScopeKind {
     File(ast::SourceFile),
     Module(ast::ItemList),
     Block(ast::StmtList),
 }
 
 impl ImportScope {
-    // FIXME: Remove this?
-    #[cfg(test)]
-    fn from(syntax: SyntaxNode) -> Option<Self> {
-        use syntax::match_ast;
-        fn contains_cfg_attr(attrs: &dyn HasAttrs) -> bool {
-            attrs.attrs().any(|attr| attr.as_simple_call().is_some_and(|(ident, _)| ident == "cfg"))
-        }
-        match_ast! {
-            match syntax {
-                ast::Module(module) => module.item_list().map(ImportScope::Module),
-                ast::SourceFile(file) => Some(ImportScope::File(file)),
-                ast::Fn(func) => contains_cfg_attr(&func).then(|| func.body().and_then(|it| it.stmt_list().map(ImportScope::Block))).flatten(),
-                ast::Const(konst) => contains_cfg_attr(&konst).then(|| match konst.body()? {
-                    ast::Expr::BlockExpr(block) => Some(block),
-                    _ => None,
-                }).flatten().and_then(|it| it.stmt_list().map(ImportScope::Block)),
-                ast::Static(statik) => contains_cfg_attr(&statik).then(|| match statik.body()? {
-                    ast::Expr::BlockExpr(block) => Some(block),
-                    _ => None,
-                }).flatten().and_then(|it| it.stmt_list().map(ImportScope::Block)),
-                _ => None,
-
-            }
-        }
-    }
-
     /// Determines the containing syntax node in which to insert a `use` statement affecting `position`.
     /// Returns the original source node inside attributes.
     pub fn find_insert_use_container(
         position: &SyntaxNode,
         sema: &Semantics<'_, RootDatabase>,
     ) -> Option<Self> {
-        fn contains_cfg_attr(attrs: &dyn HasAttrs) -> bool {
-            attrs.attrs().any(|attr| attr.as_simple_call().is_some_and(|(ident, _)| ident == "cfg"))
-        }
-
+        // The closest block expression ancestor
+        let mut block = None;
+        let mut required_cfgs = Vec::new();
         // Walk up the ancestor tree searching for a suitable node to do insertions on
         // with special handling on cfg-gated items, in which case we want to insert imports locally
         // or FIXME: annotate inserted imports with the same cfg
         for syntax in sema.ancestors_with_macros(position.clone()) {
             if let Some(file) = ast::SourceFile::cast(syntax.clone()) {
-                return Some(ImportScope::File(file));
-            } else if let Some(item) = ast::Item::cast(syntax) {
-                return match item {
-                    ast::Item::Const(konst) if contains_cfg_attr(&konst) => {
-                        // FIXME: Instead of bailing out with None, we should note down that
-                        // this import needs an attribute added
-                        match sema.original_ast_node(konst)?.body()? {
-                            ast::Expr::BlockExpr(block) => block,
-                            _ => return None,
+                return Some(ImportScope { kind: ImportScopeKind::File(file), required_cfgs });
+            } else if let Some(module) = ast::Module::cast(syntax.clone()) {
+                // early return is important here, if we can't find the original module
+                // in the input there is no way for us to insert an import anywhere.
+                return sema
+                    .original_ast_node(module)?
+                    .item_list()
+                    .map(ImportScopeKind::Module)
+                    .map(|kind| ImportScope { kind, required_cfgs });
+            } else if let Some(has_attrs) = ast::AnyHasAttrs::cast(syntax) {
+                if block.is_none() {
+                    if let Some(b) = ast::BlockExpr::cast(has_attrs.syntax().clone()) {
+                        if let Some(b) = sema.original_ast_node(b) {
+                            block = b.stmt_list();
                         }
-                        .stmt_list()
-                        .map(ImportScope::Block)
                     }
-                    ast::Item::Fn(func) if contains_cfg_attr(&func) => {
-                        // FIXME: Instead of bailing out with None, we should note down that
-                        // this import needs an attribute added
-                        sema.original_ast_node(func)?.body()?.stmt_list().map(ImportScope::Block)
-                    }
-                    ast::Item::Static(statik) if contains_cfg_attr(&statik) => {
-                        // FIXME: Instead of bailing out with None, we should note down that
-                        // this import needs an attribute added
-                        match sema.original_ast_node(statik)?.body()? {
-                            ast::Expr::BlockExpr(block) => block,
-                            _ => return None,
-                        }
-                        .stmt_list()
-                        .map(ImportScope::Block)
-                    }
-                    ast::Item::Module(module) => {
-                        // early return is important here, if we can't find the original module
-                        // in the input there is no way for us to insert an import anywhere.
-                        sema.original_ast_node(module)?.item_list().map(ImportScope::Module)
+                }
+                if has_attrs
+                    .attrs()
+                    .any(|attr| attr.as_simple_call().is_some_and(|(ident, _)| ident == "cfg"))
+                {
+                    if let Some(b) = block {
+                        return Some(ImportScope {
+                            kind: ImportScopeKind::Block(b),
+                            required_cfgs,
+                        });
                     }
-                    _ => continue,
-                };
+                    required_cfgs.extend(has_attrs.attrs().filter(|attr| {
+                        attr.as_simple_call().is_some_and(|(ident, _)| ident == "cfg")
+                    }));
+                }
             }
         }
         None
     }
 
     pub fn as_syntax_node(&self) -> &SyntaxNode {
-        match self {
-            ImportScope::File(file) => file.syntax(),
-            ImportScope::Module(item_list) => item_list.syntax(),
-            ImportScope::Block(block) => block.syntax(),
+        match &self.kind {
+            ImportScopeKind::File(file) => file.syntax(),
+            ImportScopeKind::Module(item_list) => item_list.syntax(),
+            ImportScopeKind::Block(block) => block.syntax(),
         }
     }
 
     pub fn clone_for_update(&self) -> Self {
-        match self {
-            ImportScope::File(file) => ImportScope::File(file.clone_for_update()),
-            ImportScope::Module(item_list) => ImportScope::Module(item_list.clone_for_update()),
-            ImportScope::Block(block) => ImportScope::Block(block.clone_for_update()),
+        Self {
+            kind: match &self.kind {
+                ImportScopeKind::File(file) => ImportScopeKind::File(file.clone_for_update()),
+                ImportScopeKind::Module(item_list) => {
+                    ImportScopeKind::Module(item_list.clone_for_update())
+                }
+                ImportScopeKind::Block(block) => ImportScopeKind::Block(block.clone_for_update()),
+            },
+            required_cfgs: self.required_cfgs.iter().map(|attr| attr.clone_for_update()).collect(),
         }
     }
 }
@@ -216,6 +196,11 @@ fn insert_use_with_alias_option(
         use_tree.wrap_in_tree_list();
     }
     let use_item = make::use_(None, use_tree).clone_for_update();
+    for attr in
+        scope.required_cfgs.iter().map(|attr| attr.syntax().clone_subtree().clone_for_update())
+    {
+        ted::insert(ted::Position::first_child_of(use_item.syntax()), attr);
+    }
 
     // merge into existing imports if possible
     if let Some(mb) = mb {
@@ -229,7 +214,6 @@ fn insert_use_with_alias_option(
             }
         }
     }
-
     // either we weren't allowed to merge or there is no import that fits the merge conditions
     // so look for the place we have to insert to
     insert_use_(scope, use_item, cfg.group);
@@ -316,10 +300,10 @@ fn guess_granularity_from_scope(scope: &ImportScope) -> ImportGranularityGuess {
         }
         _ => None,
     };
-    let mut use_stmts = match scope {
-        ImportScope::File(f) => f.items(),
-        ImportScope::Module(m) => m.items(),
-        ImportScope::Block(b) => b.items(),
+    let mut use_stmts = match &scope.kind {
+        ImportScopeKind::File(f) => f.items(),
+        ImportScopeKind::Module(m) => m.items(),
+        ImportScopeKind::Block(b) => b.items(),
     }
     .filter_map(use_stmt);
     let mut res = ImportGranularityGuess::Unknown;
@@ -463,12 +447,12 @@ fn insert_use_(scope: &ImportScope, use_item: ast::Use, group_imports: bool) {
         }
     }
 
-    let l_curly = match scope {
-        ImportScope::File(_) => None,
+    let l_curly = match &scope.kind {
+        ImportScopeKind::File(_) => None,
         // don't insert the imports before the item list/block expr's opening curly brace
-        ImportScope::Module(item_list) => item_list.l_curly_token(),
+        ImportScopeKind::Module(item_list) => item_list.l_curly_token(),
         // don't insert the imports before the item list's opening curly brace
-        ImportScope::Block(block) => block.l_curly_token(),
+        ImportScopeKind::Block(block) => block.l_curly_token(),
     };
     // there are no imports in this file at all
     // so put the import after all inner module attributes and possible license header comments
diff --git a/src/tools/rust-analyzer/crates/ide-db/src/imports/insert_use/tests.rs b/src/tools/rust-analyzer/crates/ide-db/src/imports/insert_use/tests.rs
diff --git a/src/tools/rust-analyzer/crates/ide-db/src/source_change.rs b/src/tools/rust-analyzer/crates/ide-db/src/source_change.rs
diff --git a/src/tools/rust-analyzer/crates/ide-diagnostics/src/handlers/json_is_not_rust.rs b/src/tools/rust-analyzer/crates/ide-diagnostics/src/handlers/json_is_not_rust.rs
diff --git a/src/tools/rust-analyzer/crates/syntax/src/ted.rs b/src/tools/rust-analyzer/crates/syntax/src/ted.rs

Original file line number	Diff line number	Diff line change
`@@ -312,12 +312,8 @@ fn replace_usages(`
`312`	`312`	`}`
`313`	`313`
`314`	`314`	`// add imports across modules where needed`
`315`		`- if let Some((import_scope, path)) = import_data {`
`316`		`- let scope = match import_scope {`
`317`		`- ImportScope::File(it) => ImportScope::File(edit.make_mut(it)),`
`318`		`- ImportScope::Module(it) => ImportScope::Module(edit.make_mut(it)),`
`319`		`- ImportScope::Block(it) => ImportScope::Block(edit.make_mut(it)),`
`320`		`- };`
	`315`	`+ if let Some((scope, path)) = import_data {`
	`316`	`+ let scope = edit.make_import_scope_mut(scope);`
`321`	`317`	`delayed_mutations.push((scope, path));`
`322`	`318`	`}`
`323`	`319`	`},`
Original file line number	Diff line number	Diff line change
`@@ -996,7 +996,8 @@ pub struct $0Foo {`
`996`	`996`	`}`
`997`	`997`	`"#,`
`998`	`998`	`r#"`
`999`		`-pub struct Foo(#[my_custom_attr] u32);`
	`999`	`+pub struct Foo(#[my_custom_attr]`
	`1000`	`+u32);`
`1000`	`1001`	`"#,`
`1001`	`1002`	`);`
`1002`	`1003`	`}`