feat: support user defined macro snippet for completion

crates/hir-def/src/attrs.rs

@@ -98,7 +98,7 @@ fn extract_ra_completions(attr_flags: &mut AttrFlags, tt: ast::TokenTree) {
 
 fn extract_ra_macro_style(attr_flags: &mut AttrFlags, tt: ast::TokenTree) {
     let tt = TokenTreeChildren::new(&tt);
-    if let Ok(NodeOrToken::Token(option)) = Itertools::exactly_one(tt)
+    if let Ok(NodeOrToken::Token(option)) = Itertools::exactly_one(tt.clone())
         && option.kind().is_any_identifier()
     {
         match option.text() {
@@ -107,6 +107,11 @@ fn extract_ra_macro_style(attr_flags: &mut AttrFlags, tt: ast::TokenTree) {
             "parentheses" => attr_flags.insert(AttrFlags::MACRO_STYLE_PARENTHESES),
             _ => {}
         }
+    } else if let Some([NodeOrToken::Token(kind), NodeOrToken::Token(eq), _]) = tt.collect_array()
+        && kind.text() == "snippet"
+        && eq.kind() == T![=]
+    {
+        attr_flags.insert(AttrFlags::HAS_MACRO_STYLE_SNIPPET);
     }
 }
 
@@ -353,6 +358,7 @@ bitflags::bitflags! {
         const DIAGNOSTIC_DO_NOT_RECOMMEND = 1 << 51;
 
         const HAS_RUSTC_MUST_IMPLEMENT_ONE_OF = 1 << 52;
+        const HAS_MACRO_STYLE_SNIPPET = 1 << 53;
     }
 }
 
@@ -1094,6 +1100,33 @@ impl AttrFlags {
         )
     }
 
+    pub fn macro_style_snippet(db: &dyn SourceDatabase, owner: MacroId) -> Option<&str> {
+        if !AttrFlags::query(db, owner.into()).contains(AttrFlags::HAS_MACRO_STYLE_SNIPPET) {
+            return None;
+        }
+
+        return macro_style_snippet(db, owner).as_ref().map(SmolStr::as_str);
+
+        #[salsa::tracked(returns(ref))]
+        fn macro_style_snippet(db: &dyn SourceDatabase, owner: MacroId) -> Option<SmolStr> {
+            collect_attrs(db, owner.into(), |attr| {
+                if let ast::Meta::TokenTreeMeta(attr) = attr
+                    && attr.path().is2("rust_analyzer", "macro_style")
+                    && let Some(tt) = attr.token_tree()
+                {
+                    for atom in DocAtom::parse(tt) {
+                        if let DocAtom::KeyValue { key, value } = atom
+                            && key == "snippet"
+                        {
+                            return ControlFlow::Break(value);
+                        }
+                    }
+                }
+                ControlFlow::Continue(())
+            })
+        }
+    }
+
     #[inline]
     pub fn field_docs(db: &dyn SourceDatabase, field: FieldId) -> Option<&Docs> {
         return fields_docs(db, field.parent).get(field.local_id).and_then(|it| it.as_deref());

crates/hir-expand/src/attrs.rs

@@ -36,18 +36,20 @@ use crate::{
 };
 
 pub trait AstPathExt {
-    fn is1(&self, segment: &str) -> bool;
+    fn is1(&self, segment: &str) -> bool {
+        self.as_one_segment().is_some_and(|it| it == segment)
+    }
+
+    fn is2(&self, qualifier: &str, segment: &str) -> bool {
+        matches!(self.as_up_to_two_segment(), Some((a, Some(b))) if a == qualifier && b == segment)
+    }
 
     fn as_one_segment(&self) -> Option<SmolStr>;
 
     fn as_up_to_two_segment(&self) -> Option<(SmolStr, Option<SmolStr>)>;
 }
 
 impl AstPathExt for ast::Path {
-    fn is1(&self, segment: &str) -> bool {
-        self.as_one_segment().is_some_and(|it| it == segment)
-    }
-
     fn as_one_segment(&self) -> Option<SmolStr> {
         Some(self.as_single_name_ref()?.text().into())
     }
@@ -60,10 +62,6 @@ impl AstPathExt for ast::Path {
 }
 
 impl AstPathExt for Option<ast::Path> {
-    fn is1(&self, segment: &str) -> bool {
-        self.as_ref().is_some_and(|it| it.is1(segment))
-    }
-
     fn as_one_segment(&self) -> Option<SmolStr> {
         self.as_ref().and_then(|it| it.as_one_segment())
     }

crates/hir/src/lib.rs

@@ -3291,9 +3291,10 @@ impl Macro {
         matches!(self.kind(db), MacroKind::Derive | MacroKind::DeriveBuiltIn)
     }
 
-    pub fn preferred_brace_style(&self, db: &dyn HirDatabase) -> Option<MacroBraces> {
+    pub fn preferred_brace_style<'db>(&self, db: &'db dyn HirDatabase) -> Option<MacroBraces<'db>> {
         let attrs = self.attrs(db);
         MacroBraces::extract(attrs.attrs)
+            .or_else(|| AttrFlags::macro_style_snippet(db, self.id).map(MacroBraces::Snippet))
     }
 }
 
@@ -3309,20 +3310,22 @@ impl Macro {
 //  - `braces` for `{...}` style.
 //  - `brackets` for `[...]` style.
 //  - `parentheses` for `(...)` style.
+//  - `snippet = "..."` for custom snippet, e.g `"${path}!($1, $2)"`.
 //
 // Malformed attributes will be ignored without warnings.
 //
 // Note that users have no way to override this attribute, so be careful and only include things
 // users definitely do not want to be completed!
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum MacroBraces {
+pub enum MacroBraces<'a> {
     Braces,
     Brackets,
     Parentheses,
+    Snippet(&'a str),
 }
 
-impl MacroBraces {
+impl MacroBraces<'static> {
     fn extract(attrs: AttrFlags) -> Option<Self> {
         if attrs.contains(AttrFlags::MACRO_STYLE_BRACES) {
             Some(Self::Braces)

crates/ide-completion/src/render/macro_.rs

@@ -1,6 +1,6 @@
 //! Renderer for macro invocations.
 
-use hir::{HirDisplay, db::HirDatabase};
+use hir::{HirDisplay, MacroBraces, db::HirDatabase};
 use ide_db::{SymbolKind, documentation::Documentation};
 use syntax::{SmolStr, ToSmolStr, format_smolstr};
 
@@ -31,6 +31,8 @@ pub(crate) fn render_macro_pat(
     render(ctx, false, false, false, name, macro_)
 }
 
+const PATH_PREFIX: &str = "${path}";
+
 fn render(
     ctx @ RenderContext { completion, .. }: RenderContext<'_, '_>,
     is_use_path: bool,
@@ -50,12 +52,18 @@ fn render(
         (name.as_str(), name.display(ctx.db(), completion.edition).to_smolstr());
     let docs = ctx.docs(macro_);
     let is_fn_like = macro_.is_fn_like(completion.db);
-    let (bra, ket) = if is_fn_like {
-        guess_macro_braces(ctx.db(), macro_, name, docs.as_ref())
+    let (bra, ket, custom) = if is_fn_like {
+        match guess_macro_braces(ctx.db(), macro_, name, docs.as_ref()) {
+            MacroBraces::Braces => (" {", "}", ""),
+            MacroBraces::Brackets => ("[", "]", ""),
+            MacroBraces::Parentheses => ("(", ")", ""),
+            MacroBraces::Snippet(snippet) => ("", "", snippet),
+        }
     } else {
-        ("", "")
+        ("", "", "")
     };
 
+    let force_custom = !custom.is_empty() && !custom.starts_with(PATH_PREFIX);
     let needs_bang = is_fn_like && !is_use_path && !has_macro_bang;
 
     let mut item = CompletionItem::new(
@@ -70,8 +78,15 @@ fn render(
         .set_relevance(ctx.completion_relevance());
 
     match ctx.snippet_cap() {
-        Some(cap) if needs_bang && !has_call_parens => {
-            let snippet = format!("{escaped_name}!{bra}$0{ket}");
+        Some(cap) if needs_bang && (!has_call_parens || force_custom) => {
+            let snippet = if !custom.is_empty() {
+                match custom.strip_prefix(PATH_PREFIX) {
+                    Some(custom) => format!("{escaped_name}{custom}"),
+                    None => custom.into(),
+                }
+            } else {
+                format!("{escaped_name}!{bra}$0{ket}")
+            };
             let lookup = banged_name(name);
             item.insert_snippet(cap, snippet).lookup_by(lookup);
         }
@@ -112,18 +127,14 @@ fn banged_name(name: &str) -> SmolStr {
     SmolStr::from_iter([name, "!"])
 }
 
-fn guess_macro_braces(
-    db: &dyn HirDatabase,
+fn guess_macro_braces<'db>(
+    db: &'db dyn HirDatabase,
     macro_: hir::Macro,
     macro_name: &str,
     docs: Option<&Documentation<'_>>,
-) -> (&'static str, &'static str) {
+) -> MacroBraces<'db> {
     if let Some(style) = macro_.preferred_brace_style(db) {
-        return match style {
-            hir::MacroBraces::Braces => (" {", "}"),
-            hir::MacroBraces::Brackets => ("[", "]"),
-            hir::MacroBraces::Parentheses => ("(", ")"),
-        };
+        return style;
     }
 
     let orig_name = macro_.name(db);
@@ -148,12 +159,12 @@ fn guess_macro_braces(
 
     // Insert a space before `{}`.
     // We prefer the last one when some votes equal.
-    let (_vote, (bra, ket)) = votes
+    let (_vote, style) = votes
         .iter()
-        .zip(&[(" {", "}"), ("[", "]"), ("(", ")")])
+        .zip(&[MacroBraces::Braces, MacroBraces::Brackets, MacroBraces::Parentheses])
         .max_by_key(|&(&vote, _)| vote)
         .unwrap();
-    (*bra, *ket)
+    *style
 }
 
 #[cfg(test)]
@@ -260,6 +271,77 @@ fn main() { bar![$0] }
         );
     }
 
+    #[test]
+    fn custom_macro_snippets() {
+        check_edit(
+            "foo!",
+            r#"
+#[rust_analyzer::macro_style(snippet = "${path}!{\n    $1\n}")]
+macro_rules! foo { () => {} }
+
+fn main() { f$0 }
+"#,
+            r#"
+#[rust_analyzer::macro_style(snippet = "${path}!{\n    $1\n}")]
+macro_rules! foo { () => {} }
+
+fn main() { foo!{
+    $1
+} }
+"#,
+        );
+
+        check_edit(
+            "assert_eq!",
+            r#"
+#[rust_analyzer::macro_style(snippet = "${path}!($1, $2)")]
+macro_rules! assert_eq { ($($t:tt)*) => {} }
+
+fn main() { a$0 }
+"#,
+            r#"
+#[rust_analyzer::macro_style(snippet = "${path}!($1, $2)")]
+macro_rules! assert_eq { ($($t:tt)*) => {} }
+
+fn main() { assert_eq!($1, $2) }
+"#,
+        );
+
+        check_edit(
+            "div!",
+            r#"
+macro_rules! html {
+    (<div>$($t:tt)*) => {
+        /* ... */
+    };
+    ($($t:tt)*) => {{
+        #[rust_analyzer::macro_style(snippet = "<div>\n    $1\n</div>")]
+        macro_rules! div { () => {} }
+        $($t)*
+    }}
+}
+
+fn main() { html!(d$0) }
+"#,
+            r#"
+macro_rules! html {
+    (<div>$($t:tt)*) => {
+        /* ... */
+    };
+    ($($t:tt)*) => {{
+        #[rust_analyzer::macro_style(snippet = "<div>\n    $1\n</div>")]
+        macro_rules! div { () => {} }
+        $($t)*
+    }}
+}
+
+fn main() { html!(<div>
+    $1
+</div>) }
+"#,
+        );
+    }
+
     #[test]
     fn guesses_macro_braces() {
         check_edit(