Don't insert imports outside of cfg attributed items

4 files changed, 118 insertions, 32 deletions

diff --git a/crates/ide_assists/src/handlers/auto_import.rs b/crates/ide_assists/src/handlers/auto_import.rs
index 6c73481781b..8df8b060d97 100644
--- a/crates/ide_assists/src/handlers/auto_import.rs
+++ b/crates/ide_assists/src/handlers/auto_import.rs
@@ -103,6 +103,7 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
                 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)),
                 };
                 insert_use(&scope, mod_path_to_ast(&import.import_path), &ctx.config.insert_use);
             },
@@ -994,4 +995,62 @@ const _: () = {
 "#,
         );
     }
+
+    #[test]
+    fn respects_cfg_attr() {
+        check_assist(
+            auto_import,
+            r#"
+mod bar {
+    pub struct Bar;
+}
+
+#[cfg(test)]
+fn foo() {
+    Bar$0
+}
+"#,
+            r#"
+mod bar {
+    pub struct Bar;
+}
+
+#[cfg(test)]
+fn foo() {
+use bar::Bar;
+
+    Bar
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn respects_cfg_attr2() {
+        check_assist(
+            auto_import,
+            r#"
+mod bar {
+    pub struct Bar;
+}
+
+#[cfg(test)]
+const FOO: Bar = {
+    Bar$0
+}
+"#,
+            r#"
+mod bar {
+    pub struct Bar;
+}
+
+#[cfg(test)]
+const FOO: Bar = {
+use bar::Bar;
+
+    Bar
+}
+"#,
+        );
+    }
 }
diff --git a/crates/ide_assists/src/handlers/replace_qualified_name_with_use.rs b/crates/ide_assists/src/handlers/replace_qualified_name_with_use.rs
index 26019c793aa..26778ee743e 100644
--- a/crates/ide_assists/src/handlers/replace_qualified_name_with_use.rs
+++ b/crates/ide_assists/src/handlers/replace_qualified_name_with_use.rs
@@ -32,7 +32,6 @@ pub(crate) fn replace_qualified_name_with_use(
 
     let target = path.syntax().text_range();
     let scope = ImportScope::find_insert_use_container_with_macros(path.syntax(), &ctx.sema)?;
-    let syntax = scope.as_syntax_node();
     acc.add(
         AssistId("replace_qualified_name_with_use", AssistKind::RefactorRewrite),
         "Replace qualified path with use",
@@ -40,11 +39,13 @@ 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 syntax = builder.make_syntax_mut(syntax.clone());
-            if let Some(ref import_scope) = ImportScope::from(syntax.clone()) {
-                shorten_paths(&syntax, &path.clone_for_update());
-                insert_use(import_scope, path, &ctx.config.insert_use);
-            }
+            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)),
+            };
+            shorten_paths(scope.as_syntax_node(), &path.clone_for_update());
+            insert_use(&scope, path, &ctx.config.insert_use);
         },
     )
 }
diff --git a/crates/ide_db/src/helpers/insert_use.rs b/crates/ide_db/src/helpers/insert_use.rs
index 4da058cb2d2..e6b4832e7db 100644
--- a/crates/ide_db/src/helpers/insert_use.rs
+++ b/crates/ide_db/src/helpers/insert_use.rs
@@ -5,7 +5,7 @@ use hir::Semantics;
 use syntax::{
     algo,
     ast::{self, make, AstNode, AttrsOwner, ModuleItemOwner, PathSegmentKind, VisibilityOwner},
-    ted, AstToken, Direction, NodeOrToken, SyntaxNode, SyntaxToken,
+    match_ast, ted, AstToken, Direction, NodeOrToken, SyntaxNode, SyntaxToken,
 };
 
 use crate::{
@@ -43,16 +43,32 @@ pub struct InsertUseConfig {
 pub enum ImportScope {
     File(ast::SourceFile),
     Module(ast::ItemList),
+    Block(ast::BlockExpr),
 }
 
 impl ImportScope {
-    pub fn from(syntax: SyntaxNode) -> Option<Self> {
-        if let Some(module) = ast::Module::cast(syntax.clone()) {
-            module.item_list().map(ImportScope::Module)
-        } else if let this @ Some(_) = ast::SourceFile::cast(syntax.clone()) {
-            this.map(ImportScope::File)
-        } else {
-            ast::ItemList::cast(syntax).map(ImportScope::Module)
+    fn from(syntax: SyntaxNode) -> Option<Self> {
+        fn contains_cfg_attr(attrs: &dyn AttrsOwner) -> bool {
+            attrs
+                .attrs()
+                .any(|attr| attr.as_simple_call().map_or(false, |(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().map(ImportScope::Block)).flatten(),
+                ast::Const(konst) => contains_cfg_attr(&konst).then(|| match konst.body()? {
+                    ast::Expr::BlockExpr(block) => Some(block),
+                    _ => None,
+                }).flatten().map(ImportScope::Block),
+                ast::Static(statik) => contains_cfg_attr(&statik).then(|| match statik.body()? {
+                    ast::Expr::BlockExpr(block) => Some(block),
+                    _ => None,
+                }).flatten().map(ImportScope::Block),
+                _ => None,
+
+            }
         }
     }
 
@@ -73,6 +89,7 @@ impl ImportScope {
         match self {
             ImportScope::File(file) => file.syntax(),
             ImportScope::Module(item_list) => item_list.syntax(),
+            ImportScope::Block(block) => block.syntax(),
         }
     }
 
@@ -80,6 +97,7 @@ impl ImportScope {
         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()),
         }
     }
 
@@ -96,6 +114,7 @@ impl ImportScope {
         let mut use_stmts = match self {
             ImportScope::File(f) => f.items(),
             ImportScope::Module(m) => m.items(),
+            ImportScope::Block(b) => b.items(),
         }
         .filter_map(use_stmt);
         let mut res = ImportGranularityGuess::Unknown;
@@ -319,28 +338,29 @@ fn insert_use_(
         ted::insert(ted::Position::after(last_inner_element), make::tokens::single_newline());
         return;
     }
-    match scope {
+    let l_curly = match scope {
         ImportScope::File(_) => {
             cov_mark::hit!(insert_group_empty_file);
             ted::insert(ted::Position::first_child_of(scope_syntax), make::tokens::blank_line());
-            ted::insert(ted::Position::first_child_of(scope_syntax), use_item.syntax())
+            ted::insert(ted::Position::first_child_of(scope_syntax), use_item.syntax());
+            return;
         }
+        // don't insert the imports before the item list/block expr's opening curly brace
+        ImportScope::Module(item_list) => item_list.l_curly_token(),
         // don't insert the imports before the item list's opening curly brace
-        ImportScope::Module(item_list) => match item_list.l_curly_token() {
-            Some(b) => {
-                cov_mark::hit!(insert_group_empty_module);
-                ted::insert(ted::Position::after(&b), make::tokens::single_newline());
-                ted::insert(ted::Position::after(&b), use_item.syntax());
-            }
-            None => {
-                // This should never happens, broken module syntax node
-                ted::insert(
-                    ted::Position::first_child_of(scope_syntax),
-                    make::tokens::blank_line(),
-                );
-                ted::insert(ted::Position::first_child_of(scope_syntax), use_item.syntax());
-            }
-        },
+        ImportScope::Block(block) => block.l_curly_token(),
+    };
+    match l_curly {
+        Some(b) => {
+            cov_mark::hit!(insert_group_empty_module);
+            ted::insert(ted::Position::after(&b), make::tokens::single_newline());
+            ted::insert(ted::Position::after(&b), use_item.syntax());
+        }
+        None => {
+            // This should never happens, broken module syntax node
+            ted::insert(ted::Position::first_child_of(scope_syntax), make::tokens::blank_line());
+            ted::insert(ted::Position::first_child_of(scope_syntax), use_item.syntax());
+        }
     }
 }
 
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
index 3c92a486f9b..e33e5bb037b 100644
--- a/crates/syntax/src/ast/node_ext.rs
+++ b/crates/syntax/src/ast/node_ext.rs
@@ -8,7 +8,7 @@ use parser::SyntaxKind;
 use rowan::{GreenNodeData, GreenTokenData};
 
 use crate::{
-    ast::{self, support, AstNode, AstToken, AttrsOwner, NameOwner, SyntaxNode},
+    ast::{self, support, AstChildren, AstNode, AstToken, AttrsOwner, NameOwner, SyntaxNode},
     NodeOrToken, SmolStr, SyntaxElement, SyntaxToken, TokenText, T,
 };
 
@@ -45,6 +45,12 @@ fn text_of_first_token(node: &SyntaxNode) -> TokenText<'_> {
     }
 }
 
+impl ast::BlockExpr {
+    pub fn items(&self) -> AstChildren<ast::Item> {
+        support::children(self.syntax())
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Macro {
     MacroRules(ast::MacroRules),