reduce duplication

1 files changed, 160 insertions, 13 deletions

diff --git a/crates/ide_assists/src/handlers/generate_getter.rs b/crates/ide_assists/src/handlers/generate_getter.rs
index df7d1bb95..e01985112 100644
--- a/crates/ide_assists/src/handlers/generate_getter.rs
+++ b/crates/ide_assists/src/handlers/generate_getter.rs
@@ -29,6 +29,40 @@ use crate::{
 // }
 // ```
 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 name_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
@@ -190,3 +230,110 @@ impl<T: Clone> Context<T> {
         );
     }
 }
+
+#[cfg(test)]
+mod tests_mut {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    fn check_not_applicable(ra_fixture: &str) {
+        check_assist_not_applicable(generate_getter_mut, ra_fixture)
+    }
+
+    #[test]
+    fn test_generate_getter_mut_from_field() {
+        check_assist(
+            generate_getter_mut,
+            r#"
+struct Context<T: Clone> {
+    dat$0a: T,
+}"#,
+            r#"
+struct Context<T: Clone> {
+    data: T,
+}
+
+impl<T: Clone> Context<T> {
+    /// Get a mutable reference to the context's data.
+    fn data_mut(&mut self) -> &mut T {
+        &mut self.data
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_getter_mut_already_implemented() {
+        check_not_applicable(
+            r#"
+struct Context<T: Clone> {
+    dat$0a: T,
+}
+
+impl<T: Clone> Context<T> {
+    fn data_mut(&mut self) -> &mut T {
+        &mut self.data
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_getter_mut_from_field_with_visibility_marker() {
+        check_assist(
+            generate_getter_mut,
+            r#"
+pub(crate) struct Context<T: Clone> {
+    dat$0a: T,
+}"#,
+            r#"
+pub(crate) struct Context<T: Clone> {
+    data: T,
+}
+
+impl<T: Clone> Context<T> {
+    /// Get a mutable reference to the context's data.
+    pub(crate) fn data_mut(&mut self) -> &mut T {
+        &mut self.data
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_multiple_generate_getter_mut() {
+        check_assist(
+            generate_getter_mut,
+            r#"
+struct Context<T: Clone> {
+    data: T,
+    cou$0nt: usize,
+}
+
+impl<T: Clone> Context<T> {
+    /// Get a mutable reference to the context's data.
+    fn data_mut(&mut self) -> &mut T {
+        &mut self.data
+    }
+}"#,
+            r#"
+struct Context<T: Clone> {
+    data: T,
+    count: usize,
+}
+
+impl<T: Clone> Context<T> {
+    /// Get a mutable reference to the context's data.
+    fn data_mut(&mut self) -> &mut T {
+        &mut self.data
+    }
+
+    /// Get a mutable reference to the context's count.
+    fn count_mut(&mut self) -> &mut usize {
+        &mut self.count
+    }
+}"#,
+        );
+    }
+}