crates/assists/src/handlers/generate_getter_mut.rs

   1 use stdx::{format_to, to_lower_snake_case};
   2 use syntax::ast::VisibilityOwner;
   3 use syntax::ast::{self, AstNode, NameOwner};
   4
   5 use crate::{
   6     utils::{find_impl_block, find_struct_impl, generate_impl_text},
   7     AssistContext, AssistId, AssistKind, Assists,
   8 };
   9
  10 // Assist: generate_getter_mut
  11 //
  12 // Generate a mut getter method.
  13 //
  14 // ```
  15 // struct Person {
  16 //     nam$0e: String,
  17 // }
  18 // ```
  19 // ->
  20 // ```
  21 // struct Person {
  22 //     name: String,
  23 // }
  24 //
  25 // impl Person {
  26 //     /// Get a mutable reference to the person's name.
  27 //     fn name_mut(&mut self) -> &mut String {
  28 //         &mut self.name
  29 //     }
  30 // }
  31 // ```
  32 pub(crate) fn generate_getter_mut(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  33     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  34     let field = ctx.find_node_at_offset::<ast::RecordField>()?;
  35
  36     let strukt_name = strukt.name()?;
  37     let field_name = field.name()?;
  38     let field_ty = field.ty()?;
  39
  40     // Return early if we've found an existing fn
  41     let fn_name = to_lower_snake_case(&field_name.to_string());
  42     let impl_def = find_struct_impl(
  43         &ctx,
  44         &ast::Adt::Struct(strukt.clone()),
  45         format!("{}_mut", fn_name).as_str(),
  46     )?;
  47
  48     let target = field.syntax().text_range();
  49     acc.add(
  50         AssistId("generate_getter_mut", AssistKind::Generate),
  51         "Generate a mut getter method",
  52         target,
  53         |builder| {
  54             let mut buf = String::with_capacity(512);
  55             let fn_name_spaced = fn_name.replace('_', " ");
  56             let strukt_name_spaced =
  57                 to_lower_snake_case(&strukt_name.to_string()).replace('_', " ");
  58
  59             if impl_def.is_some() {
  60                 buf.push('\n');
  61             }
  62
  63             let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
  64             format_to!(
  65                 buf,
  66                 "    /// Get a mutable reference to the {}'s {}.
  67     {}fn {}_mut(&mut self) -> &mut {} {{
  68         &mut self.{}
  69     }}",
  70                 strukt_name_spaced,
  71                 fn_name_spaced,
  72                 vis,
  73                 fn_name,
  74                 field_ty,
  75                 fn_name,
  76             );
  77
  78             let start_offset = impl_def
  79                 .and_then(|impl_def| find_impl_block(impl_def, &mut buf))
  80                 .unwrap_or_else(|| {
  81                     buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
  82                     strukt.syntax().text_range().end()
  83                 });
  84
  85             builder.insert(start_offset, buf);
  86         },
  87     )
  88 }
  89
  90 #[cfg(test)]
  91 mod tests {
  92     use crate::tests::{check_assist, check_assist_not_applicable};
  93
  94     use super::*;
  95
  96     fn check_not_applicable(ra_fixture: &str) {
  97         check_assist_not_applicable(generate_getter_mut, ra_fixture)
  98     }
  99
 100     #[test]
 101     fn test_generate_getter_mut_from_field() {
 102         check_assist(
 103             generate_getter_mut,
 104             r#"
 105 struct Context<T: Clone> {
 106     dat$0a: T,
 107 }"#,
 108             r#"
 109 struct Context<T: Clone> {
 110     data: T,
 111 }
 112
 113 impl<T: Clone> Context<T> {
 114     /// Get a mutable reference to the context's data.
 115     fn data_mut(&mut self) -> &mut T {
 116         &mut self.data
 117     }
 118 }"#,
 119         );
 120     }
 121
 122     #[test]
 123     fn test_generate_getter_mut_already_implemented() {
 124         check_not_applicable(
 125             r#"
 126 struct Context<T: Clone> {
 127     dat$0a: T,
 128 }
 129
 130 impl<T: Clone> Context<T> {
 131     fn data_mut(&mut self) -> &mut T {
 132         &mut self.data
 133     }
 134 }"#,
 135         );
 136     }
 137
 138     #[test]
 139     fn test_generate_getter_mut_from_field_with_visibility_marker() {
 140         check_assist(
 141             generate_getter_mut,
 142             r#"
 143 pub(crate) struct Context<T: Clone> {
 144     dat$0a: T,
 145 }"#,
 146             r#"
 147 pub(crate) struct Context<T: Clone> {
 148     data: T,
 149 }
 150
 151 impl<T: Clone> Context<T> {
 152     /// Get a mutable reference to the context's data.
 153     pub(crate) fn data_mut(&mut self) -> &mut T {
 154         &mut self.data
 155     }
 156 }"#,
 157         );
 158     }
 159 }