1 files changed, 137 insertions, 52 deletions

diff --git a/crates/ide_assists/src/handlers/generate_getter.rs b/crates/ide_assists/src/handlers/generate_getter.rs
index df7d1bb95..9faaaf284 100644
--- a/crates/ide_assists/src/handlers/generate_getter.rs
+++ b/crates/ide_assists/src/handlers/generate_getter.rs
@@ -23,12 +23,46 @@ use crate::{
 //
 // impl Person {
 //     /// Get a reference to the person's name.
-//     fn name(&self) -> &String {
+//     fn $0name(&self) -> &String {
 //         &self.name
 //     }
 // }
 // ```
 pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    generate_getter_impl(acc, ctx, false)
+}
+
+// Assist: generate_getter_mut
+//
+// Generate a mut getter method.
+//
+// ```
+// struct Person {
+//     nam$0e: String,
+// }
+// ```
+// ->
+// ```
+// struct Person {
+//     name: String,
+// }
+//
+// impl Person {
+//     /// Get a mutable reference to the person's name.
+//     fn $0name_mut(&mut self) -> &mut String {
+//         &mut self.name
+//     }
+// }
+// ```
+pub(crate) fn generate_getter_mut(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    generate_getter_impl(acc, ctx, true)
+}
+
+pub(crate) fn generate_getter_impl(
+    acc: &mut Assists,
+    ctx: &AssistContext,
+    mutable: bool,
+) -> Option<()> {
     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
     let field = ctx.find_node_at_offset::<ast::RecordField>()?;
 
@@ -37,22 +71,26 @@ pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<
     let field_ty = field.ty()?;
 
     // Return early if we've found an existing fn
-    let fn_name = to_lower_snake_case(&field_name.to_string());
+    let mut fn_name = to_lower_snake_case(&field_name.to_string());
+    if mutable {
+        format_to!(fn_name, "_mut");
+    }
     let impl_def = find_struct_impl(&ctx, &ast::Adt::Struct(strukt.clone()), fn_name.as_str())?;
 
+    let (id, label) = if mutable {
+        ("generate_getter_mut", "Generate a mut getter method")
+    } else {
+        ("generate_getter", "Generate a getter method")
+    };
     let target = field.syntax().text_range();
     acc.add_group(
         &GroupLabel("Generate getter/setter".to_owned()),
-        AssistId("generate_getter", AssistKind::Generate),
-        "Generate a getter method",
+        AssistId(id, AssistKind::Generate),
+        label,
         target,
         |builder| {
             let mut buf = String::with_capacity(512);
 
-            let fn_name_spaced = fn_name.replace('_', " ");
-            let strukt_name_spaced =
-                to_lower_snake_case(&strukt_name.to_string()).replace('_', " ");
-
             if impl_def.is_some() {
                 buf.push('\n');
             }
@@ -60,16 +98,18 @@ pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<
             let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
             format_to!(
                 buf,
-                "    /// Get a reference to the {}'s {}.
-    {}fn {}(&self) -> &{} {{
-        &self.{}
+                "    /// Get a {}reference to the {}'s {}.
+    {}fn {}(&{mut_}self) -> &{mut_}{} {{
+        &{mut_}self.{}
     }}",
-                strukt_name_spaced,
-                fn_name_spaced,
+                mutable.then(|| "mutable ").unwrap_or_default(),
+                to_lower_snake_case(&strukt_name.to_string()).replace('_', " "),
+                fn_name.trim_end_matches("_mut").replace('_', " "),
                 vis,
                 fn_name,
                 field_ty,
-                fn_name,
+                field_name,
+                mut_ = mutable.then(|| "mut ").unwrap_or_default(),
             );
 
             let start_offset = impl_def
@@ -79,7 +119,12 @@ pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<
                     strukt.syntax().text_range().end()
                 });
 
-            builder.insert(start_offset, buf);
+            match ctx.config.snippet_cap {
+                Some(cap) => {
+                    builder.insert_snippet(cap, start_offset, buf.replacen("fn ", "fn $0", 1))
+                }
+                None => builder.insert(start_offset, buf),
+            }
         },
     )
 }
@@ -90,45 +135,81 @@ mod tests {
 
     use super::*;
 
-    fn check_not_applicable(ra_fixture: &str) {
-        check_assist_not_applicable(generate_getter, ra_fixture)
-    }
-
     #[test]
     fn test_generate_getter_from_field() {
         check_assist(
             generate_getter,
             r#"
-struct Context<T: Clone> {
-    dat$0a: T,
-}"#,
+struct Context {
+    dat$0a: Data,
+}
+"#,
             r#"
-struct Context<T: Clone> {
-    data: T,
+struct Context {
+    data: Data,
 }
 
-impl<T: Clone> Context<T> {
+impl Context {
     /// Get a reference to the context's data.
-    fn data(&self) -> &T {
+    fn $0data(&self) -> &Data {
         &self.data
     }
-}"#,
+}
+"#,
+        );
+
+        check_assist(
+            generate_getter_mut,
+            r#"
+struct Context {
+    dat$0a: Data,
+}
+"#,
+            r#"
+struct Context {
+    data: Data,
+}
+
+impl Context {
+    /// Get a mutable reference to the context's data.
+    fn $0data_mut(&mut self) -> &mut Data {
+        &mut self.data
+    }
+}
+"#,
         );
     }
 
     #[test]
     fn test_generate_getter_already_implemented() {
-        check_not_applicable(
+        check_assist_not_applicable(
+            generate_getter,
             r#"
-struct Context<T: Clone> {
-    dat$0a: T,
+struct Context {
+    dat$0a: Data,
 }
 
-impl<T: Clone> Context<T> {
-    fn data(&self) -> &T {
+impl Context {
+    fn data(&self) -> &Data {
         &self.data
     }
-}"#,
+}
+"#,
+        );
+
+        check_assist_not_applicable(
+            generate_getter_mut,
+            r#"
+struct Context {
+    dat$0a: Data,
+}
+
+impl Context {
+    fn data_mut(&mut self) -> &mut Data {
+        &mut self.data
+    }
+}
+"#,
         );
     }
 
@@ -137,20 +218,22 @@ impl<T: Clone> Context<T> {
         check_assist(
             generate_getter,
             r#"
-pub(crate) struct Context<T: Clone> {
-    dat$0a: T,
-}"#,
+pub(crate) struct Context {
+    dat$0a: Data,
+}
+"#,
             r#"
-pub(crate) struct Context<T: Clone> {
-    data: T,
+pub(crate) struct Context {
+    data: Data,
 }
 
-impl<T: Clone> Context<T> {
+impl Context {
     /// Get a reference to the context's data.
-    pub(crate) fn data(&self) -> &T {
+    pub(crate) fn $0data(&self) -> &Data {
         &self.data
     }
-}"#,
+}
+"#,
         );
     }
 
@@ -159,34 +242,36 @@ impl<T: Clone> Context<T> {
         check_assist(
             generate_getter,
             r#"
-struct Context<T: Clone> {
-    data: T,
+struct Context {
+    data: Data,
     cou$0nt: usize,
 }
 
-impl<T: Clone> Context<T> {
+impl Context {
     /// Get a reference to the context's data.
-    fn data(&self) -> &T {
+    fn data(&self) -> &Data {
         &self.data
     }
-}"#,
+}
+"#,
             r#"
-struct Context<T: Clone> {
-    data: T,
+struct Context {
+    data: Data,
     count: usize,
 }
 
-impl<T: Clone> Context<T> {
+impl Context {
     /// Get a reference to the context's data.
-    fn data(&self) -> &T {
+    fn data(&self) -> &Data {
         &self.data
     }
 
     /// Get a reference to the context's count.
-    fn count(&self) -> &usize {
+    fn $0count(&self) -> &usize {
         &self.count
     }
-}"#,
+}
+"#,
         );
     }
 }