fix: Fix import insertion not being fully cfg aware

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() {
+                let scope = builder.make_import_scope_mut(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)),
-                };
                 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() {
+                        let scope = builder.make_import_scope_mut(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)),
-                        };
                         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) {}
+"#,
+        );
+    }
 }

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 {
+                if let Some((scope, path)) = import_data {
-                    let scope = match import_scope {
+                    let scope = edit.make_import_scope_mut(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)),
-                    };
                     delayed_mutations.push((scope, path));
                 }
             },

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);
 "#,
         );
     }

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 }
 "#,
         );
     }

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 {
+                let scope = builder.make_import_scope_mut(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 control_flow_enum =
                     FamousDefs(&ctx.sema, module.krate()).core_ops_ControlFlow();
 

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 {
+            let scope = builder.make_import_scope_mut(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)),
-            };
             shorten_paths(scope.as_syntax_node(), &original_path);
             let path = drop_generic_args(&original_path);
             let edition = ctx

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 {
+            let scope = edit.make_import_scope_mut(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)),
-            };
             ide_db::imports::insert_use::insert_use(&scope, import, &ctx.config.insert_use);
         }
     }

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 {
+        // The closest block expression ancestor
-            attrs.attrs().any(|attr| attr.as_simple_call().is_some_and(|(ident, _)| ident == "cfg"))
+        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));
+                return Some(ImportScope { kind: ImportScopeKind::File(file), required_cfgs });
-            } else if let Some(item) = ast::Item::cast(syntax) {
+            } else if let Some(module) = ast::Module::cast(syntax.clone()) {
-                return match item {
+                // early return is important here, if we can't find the original module
-                    ast::Item::Const(konst) if contains_cfg_attr(&konst) => {
+                // in the input there is no way for us to insert an import anywhere.
-                        // FIXME: Instead of bailing out with None, we should note down that
+                return sema
-                        // this import needs an attribute added
+                    .original_ast_node(module)?
-                        match sema.original_ast_node(konst)?.body()? {
+                    .item_list()
-                            ast::Expr::BlockExpr(block) => block,
+                    .map(ImportScopeKind::Module)
-                            _ => return None,
+                    .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
+                if has_attrs
-                        // this import needs an attribute added
+                    .attrs()
-                        sema.original_ast_node(func)?.body()?.stmt_list().map(ImportScope::Block)
+                    .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,
+                        });
                     }
-                    ast::Item::Static(statik) if contains_cfg_attr(&statik) => {
+                    required_cfgs.extend(has_attrs.attrs().filter(|attr| {
-                        // FIXME: Instead of bailing out with None, we should note down that
+                        attr.as_simple_call().is_some_and(|(ident, _)| ident == "cfg")
-                        // 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)
-                    }
-                    _ => continue,
-                };
             }
         }
         None
     }
 
     pub fn as_syntax_node(&self) -> &SyntaxNode {
-        match self {
+        match &self.kind {
-            ImportScope::File(file) => file.syntax(),
+            ImportScopeKind::File(file) => file.syntax(),
-            ImportScope::Module(item_list) => item_list.syntax(),
+            ImportScopeKind::Module(item_list) => item_list.syntax(),
-            ImportScope::Block(block) => block.syntax(),
+            ImportScopeKind::Block(block) => block.syntax(),
         }
     }
 
     pub fn clone_for_update(&self) -> Self {
-        match self {
+        Self {
-            ImportScope::File(file) => ImportScope::File(file.clone_for_update()),
+            kind: match &self.kind {
-            ImportScope::Module(item_list) => ImportScope::Module(item_list.clone_for_update()),
+                ImportScopeKind::File(file) => ImportScopeKind::File(file.clone_for_update()),
-            ImportScope::Block(block) => ImportScope::Block(block.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 {
+    let mut use_stmts = match &scope.kind {
-        ImportScope::File(f) => f.items(),
+        ImportScopeKind::File(f) => f.items(),
-        ImportScope::Module(m) => m.items(),
+        ImportScopeKind::Module(m) => m.items(),
-        ImportScope::Block(b) => b.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 {
+    let l_curly = match &scope.kind {
-        ImportScope::File(_) => None,
+        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

crates/ide-db/src/imports/insert_use/tests.rs

@@ -23,7 +23,7 @@ struct Struct;
 }
 
 #[test]
-fn respects_cfg_attr_fn() {
+fn respects_cfg_attr_fn_body() {
     check(
         r"bar::Bar",
         r#"
@@ -40,6 +40,25 @@ fn foo() {
     );
 }
 
+#[test]
+fn respects_cfg_attr_fn_sig() {
+    check(
+        r"bar::Bar",
+        r#"
+#[cfg(test)]
+fn foo($0) {}
+"#,
+        r#"
+#[cfg(test)]
+use bar::Bar;
+
+#[cfg(test)]
+fn foo() {}
+"#,
+        ImportGranularity::Crate,
+    );
+}
+
 #[test]
 fn respects_cfg_attr_const() {
     check(
@@ -58,6 +77,51 @@ const FOO: Bar = {
     );
 }
 
+#[test]
+fn respects_cfg_attr_impl() {
+    check(
+        r"bar::Bar",
+        r#"
+#[cfg(test)]
+impl () {$0}
+"#,
+        r#"
+#[cfg(test)]
+use bar::Bar;
+
+#[cfg(test)]
+impl () {}
+"#,
+        ImportGranularity::Crate,
+    );
+}
+
+#[test]
+fn respects_cfg_attr_multiple_layers() {
+    check(
+        r"bar::Bar",
+        r#"
+#[cfg(test)]
+impl () {
+    #[cfg(test2)]
+    fn f($0) {}
+}
+"#,
+        r#"
+#[cfg(test)]
+#[cfg(test2)]
+use bar::Bar;
+
+#[cfg(test)]
+impl () {
+    #[cfg(test2)]
+    fn f() {}
+}
+"#,
+        ImportGranularity::Crate,
+    );
+}
+
 #[test]
 fn insert_skips_lone_glob_imports() {
     check(
@@ -813,7 +877,7 @@ use {std::io};",
 }
 
 #[test]
-fn merge_groups_skip_attributed() {
+fn merge_groups_cfg_vs_no_cfg() {
     check_crate(
         "std::io",
         r#"
@@ -836,6 +900,25 @@ use {std::io};
     );
 }
 
+#[test]
+fn merge_groups_cfg_matching() {
+    check_crate(
+        "std::io",
+        r#"
+#[cfg(feature = "gated")] use std::fmt::{Result, Display};
+
+#[cfg(feature = "gated")]
+fn f($0) {}
+"#,
+        r#"
+#[cfg(feature = "gated")] use std::{fmt::{Display, Result}, io};
+
+#[cfg(feature = "gated")]
+fn f() {}
+"#,
+    );
+}
+
 #[test]
 fn split_out_merge() {
     // FIXME: This is suboptimal, we want to get `use std::fmt::{self, Result}`
@@ -1259,12 +1342,14 @@ fn check_with_config(
     };
     let sema = &Semantics::new(&db);
     let source_file = sema.parse(file_id);
-    let syntax = source_file.syntax().clone_for_update();
     let file = pos
-        .and_then(|pos| syntax.token_at_offset(pos.expect_offset()).next()?.parent())
+        .and_then(|pos| source_file.syntax().token_at_offset(pos.expect_offset()).next()?.parent())
         .and_then(|it| ImportScope::find_insert_use_container(&it, sema))
-        .or_else(|| ImportScope::from(syntax))
+        .unwrap_or_else(|| ImportScope {
-        .unwrap();
+            kind: ImportScopeKind::File(source_file),
+            required_cfgs: vec![],
+        })
+        .clone_for_update();
     let path = ast::SourceFile::parse(&format!("use {path};"), span::Edition::CURRENT)
         .tree()
         .syntax()
@@ -1349,7 +1434,7 @@ fn check_merge_only_fail(ra_fixture0: &str, ra_fixture1: &str, mb: MergeBehavior
 }
 
 fn check_guess(#[rust_analyzer::rust_fixture] ra_fixture: &str, expected: ImportGranularityGuess) {
-    let syntax = ast::SourceFile::parse(ra_fixture, span::Edition::CURRENT).tree().syntax().clone();
+    let syntax = ast::SourceFile::parse(ra_fixture, span::Edition::CURRENT).tree();
-    let file = ImportScope::from(syntax).unwrap();
+    let file = ImportScope { kind: ImportScopeKind::File(syntax), required_cfgs: vec![] };
     assert_eq!(super::guess_granularity_from_scope(&file), expected);
 }

crates/ide-db/src/source_change.rs

@@ -5,6 +5,7 @@
 
 use std::{collections::hash_map::Entry, fmt, iter, mem};
 
+use crate::imports::insert_use::{ImportScope, ImportScopeKind};
 use crate::text_edit::{TextEdit, TextEditBuilder};
 use crate::{SnippetCap, assists::Command, syntax_helpers::tree_diff::diff};
 use base_db::AnchoredPathBuf;
@@ -367,6 +368,17 @@ impl SourceChangeBuilder {
     pub fn make_mut<N: AstNode>(&mut self, node: N) -> N {
         self.mutated_tree.get_or_insert_with(|| TreeMutator::new(node.syntax())).make_mut(&node)
     }
+
+    pub fn make_import_scope_mut(&mut self, scope: ImportScope) -> ImportScope {
+        ImportScope {
+            kind: match scope.kind.clone() {
+                ImportScopeKind::File(it) => ImportScopeKind::File(self.make_mut(it)),
+                ImportScopeKind::Module(it) => ImportScopeKind::Module(self.make_mut(it)),
+                ImportScopeKind::Block(it) => ImportScopeKind::Block(self.make_mut(it)),
+            },
+            required_cfgs: scope.required_cfgs.iter().map(|it| self.make_mut(it.clone())).collect(),
+        }
+    }
     /// Returns a copy of the `node`, suitable for mutation.
     ///
     /// Syntax trees in rust-analyzer are typically immutable, and mutating

crates/ide-diagnostics/src/handlers/json_is_not_rust.rs

@@ -137,11 +137,7 @@ pub(crate) fn json_in_items(
                     )
                     .with_fixes(Some(vec![{
                         let mut scb = SourceChangeBuilder::new(vfs_file_id);
-                        let scope = match import_scope {
+                        let scope = scb.make_import_scope_mut(import_scope);
-                            ImportScope::File(it) => ImportScope::File(scb.make_mut(it)),
-                            ImportScope::Module(it) => ImportScope::Module(scb.make_mut(it)),
-                            ImportScope::Block(it) => ImportScope::Block(scb.make_mut(it)),
-                        };
                         let current_module = semantics_scope.module();
 
                         let cfg = ImportPathConfig {

crates/syntax/src/ted.rs

@@ -207,5 +207,12 @@ fn ws_between(left: &SyntaxElement, right: &SyntaxElement) -> Option<SyntaxToken
         }
         return Some(make::tokens::whitespace(&format!("\n{indent}")));
     }
+    if left.kind() == SyntaxKind::ATTR {
+        let mut indent = IndentLevel::from_element(right);
+        if right.kind() == SyntaxKind::ATTR {
+            indent.0 = IndentLevel::from_element(left).0.max(indent.0);
+        }
+        return Some(make::tokens::whitespace(&format!("\n{indent}")));
+    }
     Some(make::tokens::single_space())
 }