|
| 1 | +use ide_db::assists::{AssistId, AssistKind, GroupLabel}; |
| 2 | +use syntax::{ |
| 3 | + ast::{self, HasAttrs}, |
| 4 | + match_ast, AstNode, SyntaxKind, TextSize, |
| 5 | +}; |
| 6 | + |
| 7 | +use crate::assist_context::{AssistContext, Assists}; |
| 8 | + |
| 9 | +// Assist: add_attribute |
| 10 | +// |
| 11 | +// Adds commonly used attributes to items. |
| 12 | +// |
| 13 | +// ``` |
| 14 | +// struct Point { |
| 15 | +// x: u32, |
| 16 | +// y: u32,$0 |
| 17 | +// } |
| 18 | +// ``` |
| 19 | +// ->add_derive |
| 20 | +// ``` |
| 21 | +// #[derive($0)] |
| 22 | +// struct Point { |
| 23 | +// x: u32, |
| 24 | +// y: u32, |
| 25 | +// } |
| 26 | +// ``` |
| 27 | +pub(crate) fn add_attribute(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { |
| 28 | + let cap = ctx.config.snippet_cap?; |
| 29 | + |
| 30 | + let (attr_owner, attrs) = ctx |
| 31 | + .find_node_at_offset::<ast::AnyHasAttrs>()? |
| 32 | + .syntax() |
| 33 | + .ancestors() |
| 34 | + .filter_map(ast::AnyHasAttrs::cast) |
| 35 | + .find_map(|attr_owner| { |
| 36 | + let node = attr_owner.syntax(); |
| 37 | + match_ast! { |
| 38 | + match node { |
| 39 | + ast::Adt(_) => Some((attr_owner, ADT_ATTRS)), |
| 40 | + ast::Fn(_) => Some((attr_owner, FN_ATTRS)), |
| 41 | + _ => None, |
| 42 | + } |
| 43 | + } |
| 44 | + })?; |
| 45 | + |
| 46 | + let offset = attr_insertion_offset(&attr_owner)?; |
| 47 | + |
| 48 | + for tpl in attrs { |
| 49 | + let existing_offset = attr_owner.attrs().find_map(|attr| { |
| 50 | + if attr.simple_name()? == tpl.name { |
| 51 | + match attr.token_tree() { |
| 52 | + Some(tt) => { |
| 53 | + // Attribute like `#[derive(...)]`, position cursor right before the `)` |
| 54 | + return Some(tt.syntax().text_range().end() - TextSize::of(')')); |
| 55 | + } |
| 56 | + None => { |
| 57 | + // `#[key = value]` |
| 58 | + let tok = attr.syntax().last_token()?; |
| 59 | + if tok.kind() == SyntaxKind::R_BRACK { |
| 60 | + return Some(tok.text_range().end() - TextSize::of(']')); |
| 61 | + } |
| 62 | + } |
| 63 | + } |
| 64 | + } |
| 65 | + None |
| 66 | + }); |
| 67 | + acc.add_group( |
| 68 | + &GroupLabel("Add attribute".into()), |
| 69 | + AssistId(tpl.id, AssistKind::Generate), |
| 70 | + format!("Add `#[{}]`", tpl.name), |
| 71 | + attr_owner.syntax().text_range(), |
| 72 | + |b| match existing_offset { |
| 73 | + Some(offset) => { |
| 74 | + b.insert_snippet(cap, offset, "$0"); |
| 75 | + } |
| 76 | + None => { |
| 77 | + b.insert_snippet(cap, offset, format!("#[{}]\n", tpl.snippet)); |
| 78 | + } |
| 79 | + }, |
| 80 | + ); |
| 81 | + } |
| 82 | + |
| 83 | + Some(()) |
| 84 | +} |
| 85 | + |
| 86 | +fn attr_insertion_offset(nominal: &ast::AnyHasAttrs) -> Option<TextSize> { |
| 87 | + let non_ws_child = nominal |
| 88 | + .syntax() |
| 89 | + .children_with_tokens() |
| 90 | + .find(|it| it.kind() != SyntaxKind::COMMENT && it.kind() != SyntaxKind::WHITESPACE)?; |
| 91 | + Some(non_ws_child.text_range().start()) |
| 92 | +} |
| 93 | + |
| 94 | +static ADT_ATTRS: &[AttrTemplate] = &[ |
| 95 | + AttrTemplate { id: "add_derive", name: "derive", snippet: "derive($0)" }, |
| 96 | + AttrTemplate { id: "add_must_use", name: "must_use", snippet: "must_use$0" }, |
| 97 | +]; |
| 98 | + |
| 99 | +static FN_ATTRS: &[AttrTemplate] = &[ |
| 100 | + AttrTemplate { id: "add_inline", name: "inline", snippet: "inline$0" }, |
| 101 | + AttrTemplate { id: "add_must_use", name: "must_use", snippet: "must_use$0" }, |
| 102 | +]; |
| 103 | + |
| 104 | +struct AttrTemplate { |
| 105 | + /// Assist ID. |
| 106 | + id: &'static str, |
| 107 | + /// User-facing attribute name. |
| 108 | + name: &'static str, |
| 109 | + /// Snippet to insert. |
| 110 | + snippet: &'static str, |
| 111 | +} |
| 112 | + |
| 113 | +#[cfg(test)] |
| 114 | +mod tests { |
| 115 | + use crate::tests::{check_assist_by_label, check_assist_target}; |
| 116 | + |
| 117 | + use super::add_attribute; |
| 118 | + |
| 119 | + fn check_derive(ra_fixture_before: &str, ra_fixture_after: &str) { |
| 120 | + check_assist_by_label( |
| 121 | + add_attribute, |
| 122 | + ra_fixture_before, |
| 123 | + ra_fixture_after, |
| 124 | + "Add `#[derive]`", |
| 125 | + ); |
| 126 | + } |
| 127 | + |
| 128 | + fn check_must_use(ra_fixture_before: &str, ra_fixture_after: &str) { |
| 129 | + check_assist_by_label( |
| 130 | + add_attribute, |
| 131 | + ra_fixture_before, |
| 132 | + ra_fixture_after, |
| 133 | + "Add `#[must_use]`", |
| 134 | + ); |
| 135 | + } |
| 136 | + |
| 137 | + #[test] |
| 138 | + fn add_derive_new() { |
| 139 | + check_derive("struct Foo { a: i32, $0}", "#[derive($0)]\nstruct Foo { a: i32, }"); |
| 140 | + check_derive("struct Foo { $0 a: i32, }", "#[derive($0)]\nstruct Foo { a: i32, }"); |
| 141 | + } |
| 142 | + |
| 143 | + #[test] |
| 144 | + fn add_derive_existing() { |
| 145 | + check_derive( |
| 146 | + "#[derive(Clone)]\nstruct Foo { a: i32$0, }", |
| 147 | + "#[derive(Clone$0)]\nstruct Foo { a: i32, }", |
| 148 | + ); |
| 149 | + } |
| 150 | + |
| 151 | + #[test] |
| 152 | + fn add_derive_new_with_doc_comment() { |
| 153 | + check_derive( |
| 154 | + " |
| 155 | +/// `Foo` is a pretty important struct. |
| 156 | +/// It does stuff. |
| 157 | +struct Foo { a: i32$0, } |
| 158 | + ", |
| 159 | + " |
| 160 | +/// `Foo` is a pretty important struct. |
| 161 | +/// It does stuff. |
| 162 | +#[derive($0)] |
| 163 | +struct Foo { a: i32, } |
| 164 | + ", |
| 165 | + ); |
| 166 | + } |
| 167 | + |
| 168 | + #[test] |
| 169 | + fn add_derive_target() { |
| 170 | + check_assist_target( |
| 171 | + add_attribute, |
| 172 | + r#" |
| 173 | +struct SomeThingIrrelevant; |
| 174 | +/// `Foo` is a pretty important struct. |
| 175 | +/// It does stuff. |
| 176 | +struct Foo { a: i32$0, } |
| 177 | +struct EvenMoreIrrelevant; |
| 178 | + "#, |
| 179 | + "/// `Foo` is a pretty important struct. |
| 180 | +/// It does stuff. |
| 181 | +struct Foo { a: i32, }", |
| 182 | + ); |
| 183 | + } |
| 184 | + |
| 185 | + #[test] |
| 186 | + fn insert_must_use() { |
| 187 | + check_must_use("struct S$0;", "#[must_use$0]\nstruct S;"); |
| 188 | + check_must_use("$0fn f() {}", "#[must_use$0]\nfn f() {}"); |
| 189 | + |
| 190 | + check_must_use(r#"#[must_use = "bla"] struct S$0;"#, r#"#[must_use = "bla"$0] struct S;"#); |
| 191 | + check_must_use(r#"#[must_use = ] struct S$0;"#, r#"#[must_use = $0] struct S;"#); |
| 192 | + |
| 193 | + check_must_use(r#"#[must_use = "bla"] $0fn f() {}"#, r#"#[must_use = "bla"$0] fn f() {}"#); |
| 194 | + check_must_use(r#"#[must_use = ] $0fn f() {}"#, r#"#[must_use = $0] fn f() {}"#); |
| 195 | + } |
| 196 | +} |
0 commit comments