Move ws insert rendering for macro expansion into ide_db

10 files changed, 139 insertions, 95 deletions

diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index ed1b2f64fd1..9302eaf913b 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -1163,9 +1163,12 @@ impl<'a> SemanticsScope<'a> {
         Some(Crate { id: self.resolver.krate()? })
     }
 
+    pub fn in_macro_file(&self) -> bool {
+        self.file_id.is_macro()
+    }
+
     /// Note: `FxHashSet<TraitId>` should be treated as an opaque type, passed into `Type
-    // FIXME: rename to visible_traits to not repeat scope?
-    pub fn traits_in_scope(&self) -> FxHashSet<TraitId> {
+    pub fn visible_traits(&self) -> FxHashSet<TraitId> {
         let resolver = &self.resolver;
         resolver.traits_in_scope(self.db.upcast())
     }
diff --git a/crates/ide/src/expand_macro.rs b/crates/ide/src/expand_macro.rs
index ee49732240e..7f57d1d6389 100644
--- a/crates/ide/src/expand_macro.rs
+++ b/crates/ide/src/expand_macro.rs
@@ -1,9 +1,10 @@
-use std::iter;
-
 use hir::Semantics;
-use ide_db::{helpers::pick_best_token, RootDatabase};
+use ide_db::{
+    helpers::{pick_best_token, render_macro_node::render_with_ws_inserted},
+    RootDatabase,
+};
 use itertools::Itertools;
-use syntax::{ast, ted, AstNode, NodeOrToken, SyntaxKind, SyntaxNode, WalkEvent, T};
+use syntax::{ast, ted, AstNode, SyntaxKind, SyntaxNode};
 
 use crate::FilePosition;
 
@@ -49,7 +50,7 @@ pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<
             let expansions = sema.expand_derive_macro(&attr)?;
             Some(ExpandedMacro {
                 name: tt,
-                expansion: expansions.into_iter().map(insert_whitespaces).join(""),
+                expansion: expansions.into_iter().map(render_with_ws_inserted).join(""),
             })
         } else {
             None
@@ -82,7 +83,7 @@ pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<
     // FIXME:
     // macro expansion may lose all white space information
     // But we hope someday we can use ra_fmt for that
-    let expansion = insert_whitespaces(expanded?);
+    let expansion = render_with_ws_inserted(expanded?).to_string();
     Some(ExpandedMacro { name: name.unwrap_or_else(|| "???".to_owned()), expansion })
 }
 
@@ -122,84 +123,6 @@ fn expand<T: AstNode>(
     Some(expanded)
 }
 
-// FIXME: It would also be cool to share logic here and in the mbe tests,
-// which are pretty unreadable at the moment.
-fn insert_whitespaces(syn: SyntaxNode) -> String {
-    use SyntaxKind::*;
-    let mut res = String::new();
-
-    let mut indent = 0;
-    let mut last: Option<SyntaxKind> = None;
-
-    for event in syn.preorder_with_tokens() {
-        let token = match event {
-            WalkEvent::Enter(NodeOrToken::Token(token)) => token,
-            WalkEvent::Leave(NodeOrToken::Node(node))
-                if matches!(node.kind(), ATTR | MATCH_ARM | STRUCT | ENUM | UNION | FN | IMPL) =>
-            {
-                res.push('\n');
-                res.extend(iter::repeat(" ").take(2 * indent));
-                continue;
-            }
-            _ => continue,
-        };
-        let is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
-            token.next_token().map(|it| f(it.kind())).unwrap_or(default)
-        };
-        let is_last =
-            |f: fn(SyntaxKind) -> bool, default| -> bool { last.map(f).unwrap_or(default) };
-
-        match token.kind() {
-            k if is_text(k) && is_next(|it| !it.is_punct(), true) => {
-                res.push_str(token.text());
-                res.push(' ');
-            }
-            L_CURLY if is_next(|it| it != R_CURLY, true) => {
-                indent += 1;
-                if is_last(is_text, false) {
-                    res.push(' ');
-                }
-                res.push_str("{\n");
-                res.extend(iter::repeat(" ").take(2 * indent));
-            }
-            R_CURLY if is_last(|it| it != L_CURLY, true) => {
-                indent = indent.saturating_sub(1);
-                res.push('\n');
-                res.extend(iter::repeat(" ").take(2 * indent));
-                res.push_str("}");
-            }
-            R_CURLY => {
-                res.push_str("}\n");
-                res.extend(iter::repeat(" ").take(2 * indent));
-            }
-            LIFETIME_IDENT if is_next(|it| it == IDENT || it == MUT_KW, true) => {
-                res.push_str(token.text());
-                res.push(' ');
-            }
-            AS_KW => {
-                res.push_str(token.text());
-                res.push(' ');
-            }
-            T![;] => {
-                res.push_str(";\n");
-                res.extend(iter::repeat(" ").take(2 * indent));
-            }
-            T![->] => res.push_str(" -> "),
-            T![=] => res.push_str(" = "),
-            T![=>] => res.push_str(" => "),
-            _ => res.push_str(token.text()),
-        }
-
-        last = Some(token.kind());
-    }
-
-    return res;
-
-    fn is_text(k: SyntaxKind) -> bool {
-        k.is_keyword() || k.is_literal() || k == IDENT
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use expect_test::{expect, Expect};
diff --git a/crates/ide_assists/src/handlers/add_missing_impl_members.rs b/crates/ide_assists/src/handlers/add_missing_impl_members.rs
index a145598c791..6b4640ce7d7 100644
--- a/crates/ide_assists/src/handlers/add_missing_impl_members.rs
+++ b/crates/ide_assists/src/handlers/add_missing_impl_members.rs
@@ -124,6 +124,9 @@ fn add_missing_impl_members_inner(
             impl_def.clone(),
             target_scope,
         );
+        // if target_scope.in_macro_file() {
+
+        // }
         match ctx.config.snippet_cap {
             None => builder.replace(target, new_impl_def.to_string()),
             Some(cap) => {
diff --git a/crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs b/crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs
index eca6d047a1a..7fbbdb4f5eb 100644
--- a/crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs
+++ b/crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs
@@ -148,7 +148,7 @@ fn is_ref_and_impls_iter_method(
     let ty = sema.type_of_expr(&expr_behind_ref)?.adjusted();
     let scope = sema.scope(iterable.syntax());
     let krate = scope.module()?.krate();
-    let traits_in_scope = scope.traits_in_scope();
+    let traits_in_scope = scope.visible_traits();
     let iter_trait = FamousDefs(sema, Some(krate)).core_iter_Iterator()?;
 
     let has_wanted_method = ty
diff --git a/crates/ide_assists/src/handlers/generate_is_empty_from_len.rs b/crates/ide_assists/src/handlers/generate_is_empty_from_len.rs
index d831289775a..15025cf0d0e 100644
--- a/crates/ide_assists/src/handlers/generate_is_empty_from_len.rs
+++ b/crates/ide_assists/src/handlers/generate_is_empty_from_len.rs
@@ -92,7 +92,7 @@ fn get_impl_method(
     let scope = ctx.sema.scope(impl_.syntax());
     let krate = impl_def.module(db).krate();
     let ty = impl_def.self_ty(db);
-    let traits_in_scope = scope.traits_in_scope();
+    let traits_in_scope = scope.visible_traits();
     ty.iterate_method_candidates(db, krate, &traits_in_scope, Some(fn_name), |_, func| Some(func))
 }
 
diff --git a/crates/ide_completion/src/completions/dot.rs b/crates/ide_completion/src/completions/dot.rs
index e01e9c9fa7e..e7371270fbf 100644
--- a/crates/ide_completion/src/completions/dot.rs
+++ b/crates/ide_completion/src/completions/dot.rs
@@ -79,7 +79,7 @@ fn complete_methods(
 ) {
     if let Some(krate) = ctx.krate {
         let mut seen_methods = FxHashSet::default();
-        let traits_in_scope = ctx.scope.traits_in_scope();
+        let traits_in_scope = ctx.scope.visible_traits();
         receiver.iterate_method_candidates(ctx.db, krate, &traits_in_scope, None, |_ty, func| {
             if func.self_param(ctx.db).is_some() && seen_methods.insert(func.name(ctx.db)) {
                 f(func);
diff --git a/crates/ide_completion/src/completions/qualified_path.rs b/crates/ide_completion/src/completions/qualified_path.rs
index b5c3d83c168..782615a8848 100644
--- a/crates/ide_completion/src/completions/qualified_path.rs
+++ b/crates/ide_completion/src/completions/qualified_path.rs
@@ -152,9 +152,7 @@ pub(crate) fn complete_qualified_path(acc: &mut Completions, ctx: &CompletionCon
             }
         }
         hir::PathResolution::Def(
-            def
-            @
-            (hir::ModuleDef::Adt(_)
+            def @ (hir::ModuleDef::Adt(_)
             | hir::ModuleDef::TypeAlias(_)
             | hir::ModuleDef::BuiltinType(_)),
         ) => {
@@ -187,7 +185,7 @@ pub(crate) fn complete_qualified_path(acc: &mut Completions, ctx: &CompletionCon
 
             let krate = ctx.krate;
             if let Some(krate) = krate {
-                let traits_in_scope = ctx.scope.traits_in_scope();
+                let traits_in_scope = ctx.scope.visible_traits();
                 ty.iterate_path_candidates(ctx.db, krate, &traits_in_scope, None, |_ty, item| {
                     add_assoc_item(acc, ctx, item);
                     None::<()>
@@ -220,7 +218,7 @@ pub(crate) fn complete_qualified_path(acc: &mut Completions, ctx: &CompletionCon
                     add_enum_variants(acc, ctx, e);
                 }
 
-                let traits_in_scope = ctx.scope.traits_in_scope();
+                let traits_in_scope = ctx.scope.visible_traits();
                 let mut seen = FxHashSet::default();
                 ty.iterate_path_candidates(ctx.db, krate, &traits_in_scope, None, |_ty, item| {
                     // We might iterate candidates of a trait multiple times here, so deduplicate
diff --git a/crates/ide_db/src/helpers.rs b/crates/ide_db/src/helpers.rs
index 1b9cb7ff51c..b8a8723a2e5 100644
--- a/crates/ide_db/src/helpers.rs
+++ b/crates/ide_db/src/helpers.rs
@@ -4,6 +4,7 @@ pub mod generated_lints;
 pub mod import_assets;
 pub mod insert_use;
 pub mod merge_imports;
+pub mod render_macro_node;
 pub mod node_ext;
 pub mod rust_doc;
 
diff --git a/crates/ide_db/src/helpers/render_macro_node.rs b/crates/ide_db/src/helpers/render_macro_node.rs
new file mode 100644
index 00000000000..7c45d281535
--- /dev/null
+++ b/crates/ide_db/src/helpers/render_macro_node.rs
@@ -0,0 +1,116 @@
+use syntax::{
+    ast::make,
+    ted::{self, Position},
+    NodeOrToken,
+    SyntaxKind::{self, *},
+    SyntaxNode, SyntaxToken, WalkEvent, T,
+};
+
+// FIXME: It would also be cool to share logic here and in the mbe tests,
+// which are pretty unreadable at the moment.
+/// Renders a [`SyntaxNode`] with whitespace inserted between tokens that require them.
+pub fn render_with_ws_inserted(syn: SyntaxNode) -> SyntaxNode {
+    let mut indent = 0;
+    let mut last: Option<SyntaxKind> = None;
+    let mut mods = Vec::new();
+    let syn = syn.clone_subtree().clone_for_update();
+
+    let before = Position::before;
+    let after = Position::after;
+
+    let do_indent = |pos: fn(_) -> Position, token: &SyntaxToken, indent| {
+        (pos(token.clone()), make::tokens::whitespace(&" ".repeat(2 * indent)))
+    };
+    let do_ws = |pos: fn(_) -> Position, token: &SyntaxToken| {
+        (pos(token.clone()), make::tokens::single_space())
+    };
+    let do_nl = |pos: fn(_) -> Position, token: &SyntaxToken| {
+        (pos(token.clone()), make::tokens::single_newline())
+    };
+
+    for event in syn.preorder_with_tokens() {
+        let token = match event {
+            WalkEvent::Enter(NodeOrToken::Token(token)) => token,
+            WalkEvent::Leave(NodeOrToken::Node(node))
+                if matches!(node.kind(), ATTR | MATCH_ARM | STRUCT | ENUM | UNION | FN | IMPL) =>
+            {
+                if indent > 0 {
+                    mods.push((
+                        Position::after(node.clone()),
+                        make::tokens::whitespace(&" ".repeat(2 * indent)),
+                    ));
+                }
+                mods.push((Position::after(node), make::tokens::single_newline()));
+                continue;
+            }
+            _ => continue,
+        };
+        let tok = &token;
+
+        let is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
+            tok.next_token().map(|it| f(it.kind())).unwrap_or(default)
+        };
+        let is_last =
+            |f: fn(SyntaxKind) -> bool, default| -> bool { last.map(f).unwrap_or(default) };
+
+        match tok.kind() {
+            k if is_text(k) && is_next(|it| !it.is_punct(), true) => {
+                mods.push(do_ws(after, tok));
+            }
+            L_CURLY if is_next(|it| it != R_CURLY, true) => {
+                indent += 1;
+                if is_last(is_text, false) {
+                    mods.push(do_ws(before, tok));
+                }
+
+                if indent > 0 {
+                    mods.push(do_indent(after, tok, indent));
+                }
+                mods.push(do_nl(after, &tok));
+            }
+            R_CURLY if is_last(|it| it != L_CURLY, true) => {
+                indent = indent.saturating_sub(1);
+
+                if indent > 0 {
+                    mods.push(do_indent(before, tok, indent));
+                }
+                mods.push(do_nl(before, tok));
+            }
+            R_CURLY => {
+                if indent > 0 {
+                    mods.push(do_indent(after, tok, indent));
+                }
+                mods.push(do_nl(after, tok));
+            }
+            LIFETIME_IDENT if is_next(|it| it == IDENT || it == MUT_KW, true) => {
+                mods.push(do_ws(after, tok));
+            }
+            AS_KW => {
+                mods.push(do_ws(after, tok));
+            }
+            T![;] => {
+                if indent > 0 {
+                    mods.push(do_indent(after, tok, indent));
+                }
+                mods.push(do_nl(after, tok));
+            }
+            T![->] | T![=] | T![=>] => {
+                mods.push(do_ws(before, tok));
+                mods.push(do_ws(after, tok));
+            }
+            _ => (),
+        }
+
+        last = Some(tok.kind());
+    }
+
+    for (pos, insert) in mods {
+        ted::insert(pos, insert);
+    }
+
+    syn
+}
+
+fn is_text(k: SyntaxKind) -> bool {
+    k.is_keyword() || k.is_literal() || k == IDENT
+}
diff --git a/crates/ide_ssr/src/resolving.rs b/crates/ide_ssr/src/resolving.rs
index 84e5f82604b..7902295d290 100644
--- a/crates/ide_ssr/src/resolving.rs
+++ b/crates/ide_ssr/src/resolving.rs
@@ -222,7 +222,7 @@ impl<'db> ResolutionScope<'db> {
             adt.ty(self.scope.db).iterate_path_candidates(
                 self.scope.db,
                 self.scope.module()?.krate(),
-                &self.scope.traits_in_scope(),
+                &self.scope.visible_traits(),
                 None,
                 |_ty, assoc_item| {
                     let item_name = assoc_item.name(self.scope.db)?;