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 = find_struct_impl(
  40         ctx,
  41         &ast::Adt::Struct(strukt.clone()),
  42         format!("set_{}", fn_name).as_str(),
  43     )?;
  44
  45     let target = field.syntax().text_range();
  46     acc.add_group(
  47         &GroupLabel("Generate getter/setter".to_owned()),
  48         AssistId("generate_setter", AssistKind::Generate),
  49         "Generate a setter method",
  50         target,
  51         |builder| {
  52             let mut buf = String::with_capacity(512);
  53
  54             if impl_def.is_some() {
  55                 buf.push('\n');
  56             }
  57
  58             let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
  59             format_to!(
  60                 buf,
  61                 "    {}fn set_{}(&mut self, {}: {}) {{
  62         self.{} = {};
  63     }}",
  64                 vis,
  65                 fn_name,
  66                 fn_name,
  67                 field_ty,
  68                 fn_name,
  69                 fn_name,
  70             );
  71
  72             let start_offset = impl_def
  73                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
  74                 .unwrap_or_else(|| {
  75                     buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
  76                     strukt.syntax().text_range().end()
  77                 });
  78
  79             builder.insert(start_offset, buf);
  80         },
  81     )
  82 }
  83
  84 #[cfg(test)]
  85 mod tests {
  86     use crate::tests::{check_assist, check_assist_not_applicable};
  87
  88     use super::*;
  89
  90     fn check_not_applicable(ra_fixture: &str) {
  91         check_assist_not_applicable(generate_setter, ra_fixture)
  92     }
  93
  94     #[test]
  95     fn test_generate_setter_from_field() {
  96         check_assist(
  97             generate_setter,
  98             r#"
  99 struct Person<T: Clone> {
 100     dat$0a: T,
 101 }"#,
 102             r#"
 103 struct Person<T: Clone> {
 104     data: T,
 105 }
 106
 107 impl<T: Clone> Person<T> {
 108     fn set_data(&mut self, data: T) {
 109         self.data = data;
 110     }
 111 }"#,
 112         );
 113     }
 114
 115     #[test]
 116     fn test_generate_setter_already_implemented() {
 117         check_not_applicable(
 118             r#"
 119 struct Person<T: Clone> {
 120     dat$0a: T,
 121 }
 122
 123 impl<T: Clone> Person<T> {
 124     fn set_data(&mut self, data: T) {
 125         self.data = data;
 126     }
 127 }"#,
 128         );
 129     }
 130
 131     #[test]
 132     fn test_generate_setter_from_field_with_visibility_marker() {
 133         check_assist(
 134             generate_setter,
 135             r#"
 136 pub(crate) struct Person<T: Clone> {
 137     dat$0a: T,
 138 }"#,
 139             r#"
 140 pub(crate) struct Person<T: Clone> {
 141     data: T,
 142 }
 143
 144 impl<T: Clone> Person<T> {
 145     pub(crate) fn set_data(&mut self, data: T) {
 146         self.data = data;
 147     }
 148 }"#,
 149         );
 150     }
 151
 152     #[test]
 153     fn test_multiple_generate_setter() {
 154         check_assist(
 155             generate_setter,
 156             r#"
 157 struct Context<T: Clone> {
 158     data: T,
 159     cou$0nt: usize,
 160 }
 161
 162 impl<T: Clone> Context<T> {
 163     fn set_data(&mut self, data: T) {
 164         self.data = data;
 165     }
 166 }"#,
 167             r#"
 168 struct Context<T: Clone> {
 169     data: T,
 170     count: usize,
 171 }
 172
 173 impl<T: Clone> Context<T> {
 174     fn set_data(&mut self, data: T) {
 175         self.data = data;
 176     }
 177
 178     fn set_count(&mut self, count: usize) {
 179         self.count = count;
 180     }
 181 }"#,
 182         );
 183     }
 184 }