src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_setter.rs

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