crates/ide_assists/src/handlers/generate_getter.rs

   1 use stdx::{format_to, to_lower_snake_case};
   2 use syntax::ast::{self, AstNode, NameOwner, VisibilityOwner};
   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_getter
  10 //
  11 // Generate a getter 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 //     /// Get a reference to the person's name.
  26 //     fn name(&self) -> &String {
  27 //         &self.name
  28 //     }
  29 // }
  30 // ```
  31 pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  32     generate_getter_impl(acc, ctx, false)
  33 }
  34
  35 // Assist: generate_getter_mut
  36 //
  37 // Generate a mut getter method.
  38 //
  39 // ```
  40 // struct Person {
  41 //     nam$0e: String,
  42 // }
  43 // ```
  44 // ->
  45 // ```
  46 // struct Person {
  47 //     name: String,
  48 // }
  49 //
  50 // impl Person {
  51 //     /// Get a mutable reference to the person's name.
  52 //     fn name_mut(&mut self) -> &mut String {
  53 //         &mut self.name
  54 //     }
  55 // }
  56 // ```
  57 pub(crate) fn generate_getter_mut(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  58     generate_getter_impl(acc, ctx, true)
  59 }
  60
  61 pub(crate) fn generate_getter_impl(
  62     acc: &mut Assists,
  63     ctx: &AssistContext,
  64     mutable: bool,
  65 ) -> Option<()> {
  66     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  67     let field = ctx.find_node_at_offset::<ast::RecordField>()?;
  68
  69     let strukt_name = strukt.name()?;
  70     let field_name = field.name()?;
  71     let field_ty = field.ty()?;
  72
  73     // Return early if we've found an existing fn
  74     let mut fn_name = to_lower_snake_case(&field_name.to_string());
  75     if mutable {
  76         format_to!(fn_name, "_mut");
  77     }
  78     let impl_def = find_struct_impl(&ctx, &ast::Adt::Struct(strukt.clone()), fn_name.as_str())?;
  79
  80     let (id, label) = if mutable {
  81         ("generate_getter_mut", "Generate a mut getter method")
  82     } else {
  83         ("generate_getter", "Generate a getter method")
  84     };
  85     let target = field.syntax().text_range();
  86     acc.add_group(
  87         &GroupLabel("Generate getter/setter".to_owned()),
  88         AssistId(id, AssistKind::Generate),
  89         label,
  90         target,
  91         |builder| {
  92             let mut buf = String::with_capacity(512);
  93
  94             if impl_def.is_some() {
  95                 buf.push('\n');
  96             }
  97
  98             let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
  99             format_to!(
 100                 buf,
 101                 "    /// Get a {}reference to the {}'s {}.
 102     {}fn {}(&{mut_}self) -> &{mut_}{} {{
 103         &{mut_}self.{}
 104     }}",
 105                 mutable.then(|| "mutable ").unwrap_or_default(),
 106                 to_lower_snake_case(&strukt_name.to_string()).replace('_', " "),
 107                 fn_name.trim_end_matches("_mut").replace('_', " "),
 108                 vis,
 109                 fn_name,
 110                 field_ty,
 111                 field_name,
 112                 mut_ = mutable.then(|| "mut ").unwrap_or_default(),
 113             );
 114
 115             let start_offset = impl_def
 116                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
 117                 .unwrap_or_else(|| {
 118                     buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
 119                     strukt.syntax().text_range().end()
 120                 });
 121
 122             builder.insert(start_offset, buf);
 123         },
 124     )
 125 }
 126
 127 #[cfg(test)]
 128 mod tests {
 129     use crate::tests::{check_assist, check_assist_not_applicable};
 130
 131     use super::*;
 132
 133     #[test]
 134     fn test_generate_getter_from_field() {
 135         check_assist(
 136             generate_getter,
 137             r#"
 138 struct Context<T: Clone> {
 139     dat$0a: T,
 140 }"#,
 141             r#"
 142 struct Context<T: Clone> {
 143     data: T,
 144 }
 145
 146 impl<T: Clone> Context<T> {
 147     /// Get a reference to the context's data.
 148     fn data(&self) -> &T {
 149         &self.data
 150     }
 151 }"#,
 152         );
 153
 154         check_assist(
 155             generate_getter_mut,
 156             r#"
 157 struct Context<T: Clone> {
 158     dat$0a: T,
 159 }"#,
 160             r#"
 161 struct Context<T: Clone> {
 162     data: T,
 163 }
 164
 165 impl<T: Clone> Context<T> {
 166     /// Get a mutable reference to the context's data.
 167     fn data_mut(&mut self) -> &mut T {
 168         &mut self.data
 169     }
 170 }"#,
 171         );
 172     }
 173
 174     #[test]
 175     fn test_generate_getter_already_implemented() {
 176         check_assist_not_applicable(
 177             generate_getter,
 178             r#"
 179 struct Context<T: Clone> {
 180     dat$0a: T,
 181 }
 182
 183 impl<T: Clone> Context<T> {
 184     fn data(&self) -> &T {
 185         &self.data
 186     }
 187 }"#,
 188         );
 189
 190         check_assist_not_applicable(
 191             generate_getter_mut,
 192             r#"
 193 struct Context<T: Clone> {
 194     dat$0a: T,
 195 }
 196
 197 impl<T: Clone> Context<T> {
 198     fn data_mut(&mut self) -> &mut T {
 199         &mut self.data
 200     }
 201 }"#,
 202         );
 203     }
 204
 205     #[test]
 206     fn test_generate_getter_from_field_with_visibility_marker() {
 207         check_assist(
 208             generate_getter,
 209             r#"
 210 pub(crate) struct Context<T: Clone> {
 211     dat$0a: T,
 212 }"#,
 213             r#"
 214 pub(crate) struct Context<T: Clone> {
 215     data: T,
 216 }
 217
 218 impl<T: Clone> Context<T> {
 219     /// Get a reference to the context's data.
 220     pub(crate) fn data(&self) -> &T {
 221         &self.data
 222     }
 223 }"#,
 224         );
 225     }
 226
 227     #[test]
 228     fn test_multiple_generate_getter() {
 229         check_assist(
 230             generate_getter,
 231             r#"
 232 struct Context<T: Clone> {
 233     data: T,
 234     cou$0nt: usize,
 235 }
 236
 237 impl<T: Clone> Context<T> {
 238     /// Get a reference to the context's data.
 239     fn data(&self) -> &T {
 240         &self.data
 241     }
 242 }"#,
 243             r#"
 244 struct Context<T: Clone> {
 245     data: T,
 246     count: usize,
 247 }
 248
 249 impl<T: Clone> Context<T> {
 250     /// Get a reference to the context's data.
 251     fn data(&self) -> &T {
 252         &self.data
 253     }
 254
 255     /// Get a reference to the context's count.
 256     fn count(&self) -> &usize {
 257         &self.count
 258     }
 259 }"#,
 260         );
 261     }
 262 }