Add a "Add attribute" assist

crates/ide-assists/src/handlers/add_attribute.rs

@@ -0,0 +1,196 @@
+use ide_db::assists::{AssistId, AssistKind, GroupLabel};
+use syntax::{
+    ast::{self, HasAttrs},
+    match_ast, AstNode, SyntaxKind, TextSize,
+};
+
+use crate::assist_context::{AssistContext, Assists};
+
+// Assist: add_attribute
+//
+// Adds commonly used attributes to items.
+//
+// ```
+// struct Point {
+//     x: u32,
+//     y: u32,$0
+// }
+// ```
+// ->add_derive
+// ```
+// #[derive($0)]
+// struct Point {
+//     x: u32,
+//     y: u32,
+// }
+// ```
+pub(crate) fn add_attribute(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let cap = ctx.config.snippet_cap?;
+
+    let (attr_owner, attrs) = ctx
+        .find_node_at_offset::<ast::AnyHasAttrs>()?
+        .syntax()
+        .ancestors()
+        .filter_map(ast::AnyHasAttrs::cast)
+        .find_map(|attr_owner| {
+            let node = attr_owner.syntax();
+            match_ast! {
+                match node {
+                    ast::Adt(_) => Some((attr_owner, ADT_ATTRS)),
+                    ast::Fn(_) => Some((attr_owner, FN_ATTRS)),
+                    _ => None,
+                }
+            }
+        })?;
+
+    let offset = attr_insertion_offset(&attr_owner)?;
+
+    for tpl in attrs {
+        let existing_offset = attr_owner.attrs().find_map(|attr| {
+            if attr.simple_name()? == tpl.name {
+                match attr.token_tree() {
+                    Some(tt) => {
+                        // Attribute like `#[derive(...)]`, position cursor right before the `)`
+                        return Some(tt.syntax().text_range().end() - TextSize::of(')'));
+                    }
+                    None => {
+                        // `#[key = value]`
+                        let tok = attr.syntax().last_token()?;
+                        if tok.kind() == SyntaxKind::R_BRACK {
+                            return Some(tok.text_range().end() - TextSize::of(']'));
+                        }
+                    }
+                }
+            }
+            None
+        });
+        acc.add_group(
+            &GroupLabel("Add attribute".into()),
+            AssistId(tpl.id, AssistKind::Generate),
+            format!("Add `#[{}]`", tpl.name),
+            attr_owner.syntax().text_range(),
+            |b| match existing_offset {
+                Some(offset) => {
+                    b.insert_snippet(cap, offset, "$0");
+                }
+                None => {
+                    b.insert_snippet(cap, offset, format!("#[{}]\n", tpl.snippet));
+                }
+            },
+        );
+    }
+
+    Some(())
+}
+
+fn attr_insertion_offset(nominal: &ast::AnyHasAttrs) -> Option<TextSize> {
+    let non_ws_child = nominal
+        .syntax()
+        .children_with_tokens()
+        .find(|it| it.kind() != SyntaxKind::COMMENT && it.kind() != SyntaxKind::WHITESPACE)?;
+    Some(non_ws_child.text_range().start())
+}
+
+static ADT_ATTRS: &[AttrTemplate] = &[
+    AttrTemplate { id: "add_derive", name: "derive", snippet: "derive($0)" },
+    AttrTemplate { id: "add_must_use", name: "must_use", snippet: "must_use$0" },
+];
+
+static FN_ATTRS: &[AttrTemplate] = &[
+    AttrTemplate { id: "add_inline", name: "inline", snippet: "inline$0" },
+    AttrTemplate { id: "add_must_use", name: "must_use", snippet: "must_use$0" },
+];
+
+struct AttrTemplate {
+    /// Assist ID.
+    id: &'static str,
+    /// User-facing attribute name.
+    name: &'static str,
+    /// Snippet to insert.
+    snippet: &'static str,
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist_by_label, check_assist_target};
+
+    use super::add_attribute;
+
+    fn check_derive(ra_fixture_before: &str, ra_fixture_after: &str) {
+        check_assist_by_label(
+            add_attribute,
+            ra_fixture_before,
+            ra_fixture_after,
+            "Add `#[derive]`",
+        );
+    }
+
+    fn check_must_use(ra_fixture_before: &str, ra_fixture_after: &str) {
+        check_assist_by_label(
+            add_attribute,
+            ra_fixture_before,
+            ra_fixture_after,
+            "Add `#[must_use]`",
+        );
+    }
+
+    #[test]
+    fn add_derive_new() {
+        check_derive("struct Foo { a: i32, $0}", "#[derive($0)]\nstruct Foo { a: i32, }");
+        check_derive("struct Foo { $0 a: i32, }", "#[derive($0)]\nstruct Foo {  a: i32, }");
+    }
+
+    #[test]
+    fn add_derive_existing() {
+        check_derive(
+            "#[derive(Clone)]\nstruct Foo { a: i32$0, }",
+            "#[derive(Clone$0)]\nstruct Foo { a: i32, }",
+        );
+    }
+
+    #[test]
+    fn add_derive_new_with_doc_comment() {
+        check_derive(
+            "
+/// `Foo` is a pretty important struct.
+/// It does stuff.
+struct Foo { a: i32$0, }
+            ",
+            "
+/// `Foo` is a pretty important struct.
+/// It does stuff.
+#[derive($0)]
+struct Foo { a: i32, }
+            ",
+        );
+    }
+
+    #[test]
+    fn add_derive_target() {
+        check_assist_target(
+            add_attribute,
+            r#"
+struct SomeThingIrrelevant;
+/// `Foo` is a pretty important struct.
+/// It does stuff.
+struct Foo { a: i32$0, }
+struct EvenMoreIrrelevant;
+            "#,
+            "/// `Foo` is a pretty important struct.
+/// It does stuff.
+struct Foo { a: i32, }",
+        );
+    }
+
+    #[test]
+    fn insert_must_use() {
+        check_must_use("struct S$0;", "#[must_use$0]\nstruct S;");
+        check_must_use("$0fn f() {}", "#[must_use$0]\nfn f() {}");
+
+        check_must_use(r#"#[must_use = "bla"] struct S$0;"#, r#"#[must_use = "bla"$0] struct S;"#);
+        check_must_use(r#"#[must_use = ] struct S$0;"#, r#"#[must_use = $0] struct S;"#);
+
+        check_must_use(r#"#[must_use = "bla"] $0fn f() {}"#, r#"#[must_use = "bla"$0] fn f() {}"#);
+        check_must_use(r#"#[must_use = ] $0fn f() {}"#, r#"#[must_use = $0] fn f() {}"#);
+    }
+}

crates/ide-assists/src/handlers/generate_derive.rs

@@ -1,132 +0,0 @@
-use syntax::{
-    ast::{self, AstNode, HasAttrs},
-    SyntaxKind::{COMMENT, WHITESPACE},
-    TextSize,
-};
-
-use crate::{AssistContext, AssistId, AssistKind, Assists};
-
-// Assist: generate_derive
-//
-// Adds a new `#[derive()]` clause to a struct or enum.
-//
-// ```
-// struct Point {
-//     x: u32,
-//     y: u32,$0
-// }
-// ```
-// ->
-// ```
-// #[derive($0)]
-// struct Point {
-//     x: u32,
-//     y: u32,
-// }
-// ```
-pub(crate) fn generate_derive(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let cap = ctx.config.snippet_cap?;
-    let nominal = ctx.find_node_at_offset::<ast::Adt>()?;
-    let node_start = derive_insertion_offset(&nominal)?;
-    let target = nominal.syntax().text_range();
-    acc.add(
-        AssistId("generate_derive", AssistKind::Generate),
-        "Add `#[derive]`",
-        target,
-        |builder| {
-            let derive_attr = nominal
-                .attrs()
-                .filter_map(|x| x.as_simple_call())
-                .filter(|(name, _arg)| name == "derive")
-                .map(|(_name, arg)| arg)
-                .next();
-            match derive_attr {
-                None => {
-                    builder.insert_snippet(cap, node_start, "#[derive($0)]\n");
-                }
-                Some(tt) => {
-                    // Just move the cursor.
-                    builder.insert_snippet(
-                        cap,
-                        tt.syntax().text_range().end() - TextSize::of(')'),
-                        "$0",
-                    )
-                }
-            };
-        },
-    )
-}
-
-// Insert `derive` after doc comments.
-fn derive_insertion_offset(nominal: &ast::Adt) -> Option<TextSize> {
-    let non_ws_child = nominal
-        .syntax()
-        .children_with_tokens()
-        .find(|it| it.kind() != COMMENT && it.kind() != WHITESPACE)?;
-    Some(non_ws_child.text_range().start())
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::tests::{check_assist, check_assist_target};
-
-    use super::*;
-
-    #[test]
-    fn add_derive_new() {
-        check_assist(
-            generate_derive,
-            "struct Foo { a: i32, $0}",
-            "#[derive($0)]\nstruct Foo { a: i32, }",
-        );
-        check_assist(
-            generate_derive,
-            "struct Foo { $0 a: i32, }",
-            "#[derive($0)]\nstruct Foo {  a: i32, }",
-        );
-    }
-
-    #[test]
-    fn add_derive_existing() {
-        check_assist(
-            generate_derive,
-            "#[derive(Clone)]\nstruct Foo { a: i32$0, }",
-            "#[derive(Clone$0)]\nstruct Foo { a: i32, }",
-        );
-    }
-
-    #[test]
-    fn add_derive_new_with_doc_comment() {
-        check_assist(
-            generate_derive,
-            "
-/// `Foo` is a pretty important struct.
-/// It does stuff.
-struct Foo { a: i32$0, }
-            ",
-            "
-/// `Foo` is a pretty important struct.
-/// It does stuff.
-#[derive($0)]
-struct Foo { a: i32, }
-            ",
-        );
-    }
-
-    #[test]
-    fn add_derive_target() {
-        check_assist_target(
-            generate_derive,
-            "
-struct SomeThingIrrelevant;
-/// `Foo` is a pretty important struct.
-/// It does stuff.
-struct Foo { a: i32$0, }
-struct EvenMoreIrrelevant;
-            ",
-            "/// `Foo` is a pretty important struct.
-/// It does stuff.
-struct Foo { a: i32, }",
-        );
-    }
-}

crates/ide-assists/src/handlers/generate_getter.rs

@@ -38,7 +38,6 @@ use crate::{
 // }
 //
 // impl Person {
-//     #[must_use]
 //     fn $0name(&self) -> &str {
 //         self.name.as_ref()
 //     }
@@ -64,7 +63,6 @@ pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<
 // }
 //
 // impl Person {
-//     #[must_use]
 //     fn $0name_mut(&mut self) -> &mut String {
 //         &mut self.name
 //     }
@@ -133,8 +131,7 @@ pub(crate) fn generate_getter_impl(
 
             format_to!(
                 buf,
-                "    #[must_use]
+                "    {}fn {}(&{}self) -> {} {{
-    {}fn {}(&{}self) -> {} {{
         {}
     }}",
                 vis,
@@ -182,7 +179,6 @@ struct Context {
 }
 
 impl Context {
-    #[must_use]
     fn $0data(&self) -> &Data {
         &self.data
     }
@@ -203,7 +199,6 @@ struct Context {
 }
 
 impl Context {
-    #[must_use]
     fn $0data_mut(&mut self) -> &mut Data {
         &mut self.data
     }
@@ -237,7 +232,6 @@ struct Context {
 }
 
 impl Context {
-    #[must_use]
     fn data_mut(&mut self) -> &mut Data {
         &mut self.data
     }
@@ -261,7 +255,6 @@ pub(crate) struct Context {
 }
 
 impl Context {
-    #[must_use]
     pub(crate) fn $0data(&self) -> &Data {
         &self.data
     }
@@ -281,7 +274,6 @@ struct Context {
 }
 
 impl Context {
-    #[must_use]
     fn data(&self) -> &Data {
         &self.data
     }
@@ -294,12 +286,10 @@ struct Context {
 }
 
 impl Context {
-    #[must_use]
     fn data(&self) -> &Data {
         &self.data
     }
 
-    #[must_use]
     fn $0count(&self) -> &usize {
         &self.count
     }
@@ -325,7 +315,6 @@ pub struct String;
 struct S { foo: String }
 
 impl S {
-    #[must_use]
     fn $0foo(&self) -> &String {
         &self.foo
     }
@@ -349,7 +338,6 @@ struct S { foo: $0bool }
 struct S { foo: bool }
 
 impl S {
-    #[must_use]
     fn $0foo(&self) -> bool {
         self.foo
     }
@@ -382,7 +370,6 @@ impl AsRef<str> for String {
 struct S { foo: String }
 
 impl S {
-    #[must_use]
     fn $0foo(&self) -> &str {
         self.foo.as_ref()
     }
@@ -419,7 +406,6 @@ impl<T> AsRef<T> for Box<T> {
 struct S { foo: Box<Sweets> }
 
 impl S {
-    #[must_use]
     fn $0foo(&self) -> &Sweets {
         self.foo.as_ref()
     }
@@ -452,7 +438,6 @@ impl<T> AsRef<[T]> for Vec<T> {
 struct S { foo: Vec<()> }
 
 impl S {
-    #[must_use]
     fn $0foo(&self) -> &[()] {
         self.foo.as_ref()
     }
@@ -475,7 +460,6 @@ struct Failure;
 struct S { foo: Option<Failure> }
 
 impl S {
-    #[must_use]
     fn $0foo(&self) -> Option<&Failure> {
         self.foo.as_ref()
     }
@@ -498,7 +482,6 @@ struct Context {
 }
 
 impl Context {
-    #[must_use]
     fn $0data(&self) -> Result<&bool, &i32> {
         self.data.as_ref()
     }

crates/ide-assists/src/lib.rs

@@ -106,6 +106,8 @@ mod handlers {
     mod add_explicit_type;
     mod add_lifetime_to_type;
     mod add_missing_impl_members;
+    mod add_missing_match_arms;
+    mod add_attribute;
     mod add_turbo_fish;
     mod apply_demorgan;
     mod auto_import;
@@ -125,7 +127,6 @@ mod handlers {
     mod extract_struct_from_enum_variant;
     mod extract_type_alias;
     mod extract_variable;
-    mod add_missing_match_arms;
     mod fix_visibility;
     mod flip_binexpr;
     mod flip_comma;
@@ -134,7 +135,6 @@ mod handlers {
     mod generate_default_from_enum_variant;
     mod generate_default_from_new;
     mod generate_deref;
-    mod generate_derive;
     mod generate_documentation_template;
     mod generate_enum_is_method;
     mod generate_enum_projection_method;
@@ -190,6 +190,7 @@ mod handlers {
     pub(crate) fn all() -> &'static [Handler] {
         &[
             // These are alphabetic for the foolish consistency
+            add_attribute::add_attribute,
             add_explicit_type::add_explicit_type,
             add_missing_match_arms::add_missing_match_arms,
             add_lifetime_to_type::add_lifetime_to_type,
@@ -219,7 +220,6 @@ mod handlers {
             generate_constant::generate_constant,
             generate_default_from_enum_variant::generate_default_from_enum_variant,
             generate_default_from_new::generate_default_from_new,
-            generate_derive::generate_derive,
             generate_documentation_template::generate_documentation_template,
             generate_enum_is_method::generate_enum_is_method,
             generate_enum_projection_method::generate_enum_as_method,

crates/ide-assists/src/tests.rs

@@ -251,6 +251,7 @@ pub fn test_some_range(a: int) -> bool {
         Extract into variable
         Extract into function
         Replace if let with match
+        Add attribute
     "#]]
     .assert_eq(&expected);
 }

crates/ide-assists/src/tests/generated.rs

@@ -2,6 +2,26 @@
 
 use super::check_doc_test;
 
+#[test]
+fn doctest_add_attribute() {
+    check_doc_test(
+        "add_derive",
+        r#####"
+struct Point {
+    x: u32,
+    y: u32,$0
+}
+"#####,
+        r#####"
+#[derive($0)]
+struct Point {
+    x: u32,
+    y: u32,
+}
+"#####,
+    )
+}
+
 #[test]
 fn doctest_add_explicit_type() {
     check_doc_test(
@@ -822,26 +842,6 @@ impl core::ops::Deref for B {
     )
 }
 
-#[test]
-fn doctest_generate_derive() {
-    check_doc_test(
-        "generate_derive",
-        r#####"
-struct Point {
-    x: u32,
-    y: u32,$0
-}
-"#####,
-        r#####"
-#[derive($0)]
-struct Point {
-    x: u32,
-    y: u32,
-}
-"#####,
-    )
-}
-
 #[test]
 fn doctest_generate_documentation_template() {
     check_doc_test(
@@ -1037,7 +1037,6 @@ struct Person {
 }
 
 impl Person {
-    #[must_use]
     fn $0name(&self) -> &str {
         self.name.as_ref()
     }
@@ -1061,7 +1060,6 @@ struct Person {
 }
 
 impl Person {
-    #[must_use]
     fn $0name_mut(&mut self) -> &mut String {
         &mut self.name
     }

crates/ide-assists/src/tests/sourcegen.rs

@@ -31,7 +31,7 @@ r#####"
 }}
 "######,
                     &test_id,
-                    &assist.id,
+                    &section.assist_id,
                     reveal_hash_comments(&section.before),
                     reveal_hash_comments(&section.after)
                 );
@@ -61,6 +61,7 @@ r#####"
 }
 #[derive(Debug)]
 struct Section {
+    assist_id: String,
     doc: String,
     before: String,
     after: String,
@@ -111,11 +112,13 @@ impl Assist {
 
                     let before = take_until(lines.by_ref(), "```");
 
-                    assert_eq!(lines.next().unwrap().as_str(), "->");
+                    let arrow = lines.next().unwrap();
+                    assert!(arrow.starts_with("->"));
+                    let id = if arrow[2..].is_empty() { &assist.id } else { &arrow[2..] };
                     assert_eq!(lines.next().unwrap().as_str(), "```");
                     let after = take_until(lines.by_ref(), "```");
 
-                    assist.sections.push(Section { doc, before, after });
+                    assist.sections.push(Section { assist_id: id.to_string(), doc, before, after });
                 }
 
                 acc.push(assist)