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::useless_type_special_case,
   6     utils::{find_impl_block_end, find_struct_impl, generate_impl_text},
   7     AssistContext, AssistId, AssistKind, Assists, GroupLabel,
   8 };
   9
  10 // Assist: generate_getter
  11 //
  12 // Generate a 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 reference to the person's name.
  27 //     fn $0name(&self) -> &str {
  28 //         self.name.as_str()
  29 //     }
  30 // }
  31 // ```
  32 pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  33     generate_getter_impl(acc, ctx, false)
  34 }
  35
  36 // Assist: generate_getter_mut
  37 //
  38 // Generate a mut getter method.
  39 //
  40 // ```
  41 // struct Person {
  42 //     nam$0e: String,
  43 // }
  44 // ```
  45 // ->
  46 // ```
  47 // struct Person {
  48 //     name: String,
  49 // }
  50 //
  51 // impl Person {
  52 //     /// Get a mutable reference to the person's name.
  53 //     fn $0name_mut(&mut self) -> &mut String {
  54 //         &mut self.name
  55 //     }
  56 // }
  57 // ```
  58 pub(crate) fn generate_getter_mut(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  59     generate_getter_impl(acc, ctx, true)
  60 }
  61
  62 pub(crate) fn generate_getter_impl(
  63     acc: &mut Assists,
  64     ctx: &AssistContext,
  65     mutable: bool,
  66 ) -> Option<()> {
  67     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  68     let field = ctx.find_node_at_offset::<ast::RecordField>()?;
  69
  70     let strukt_name = strukt.name()?;
  71     let field_name = field.name()?;
  72     let field_ty = field.ty()?;
  73
  74     // Return early if we've found an existing fn
  75     let mut fn_name = to_lower_snake_case(&field_name.to_string());
  76     if mutable {
  77         format_to!(fn_name, "_mut");
  78     }
  79     let impl_def = find_struct_impl(ctx, &ast::Adt::Struct(strukt.clone()), fn_name.as_str())?;
  80
  81     let (id, label) = if mutable {
  82         ("generate_getter_mut", "Generate a mut getter method")
  83     } else {
  84         ("generate_getter", "Generate a getter method")
  85     };
  86     let target = field.syntax().text_range();
  87     acc.add_group(
  88         &GroupLabel("Generate getter/setter".to_owned()),
  89         AssistId(id, AssistKind::Generate),
  90         label,
  91         target,
  92         |builder| {
  93             let mut buf = String::with_capacity(512);
  94
  95             if impl_def.is_some() {
  96                 buf.push('\n');
  97             }
  98
  99             let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
 100             let (ty, body) = if mutable {
 101                 (format!("&mut {}", field_ty), format!("&mut self.{}", field_name))
 102             } else {
 103                 useless_type_special_case(&field_name.to_string(), &field_ty.to_string())
 104                     .unwrap_or_else(|| (format!("&{}", field_ty), format!("&self.{}", field_name)))
 105             };
 106
 107             format_to!(
 108                 buf,
 109                 "    /// Get a {}reference to the {}'s {}.
 110     {}fn {}(&{}self) -> {} {{
 111         {}
 112     }}",
 113                 mutable.then(|| "mutable ").unwrap_or_default(),
 114                 to_lower_snake_case(&strukt_name.to_string()).replace('_', " "),
 115                 fn_name.trim_end_matches("_mut").replace('_', " "),
 116                 vis,
 117                 fn_name,
 118                 mutable.then(|| "mut ").unwrap_or_default(),
 119                 ty,
 120                 body,
 121             );
 122
 123             let start_offset = impl_def
 124                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
 125                 .unwrap_or_else(|| {
 126                     buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
 127                     strukt.syntax().text_range().end()
 128                 });
 129
 130             match ctx.config.snippet_cap {
 131                 Some(cap) => {
 132                     builder.insert_snippet(cap, start_offset, buf.replacen("fn ", "fn $0", 1))
 133                 }
 134                 None => builder.insert(start_offset, buf),
 135             }
 136         },
 137     )
 138 }
 139
 140 #[cfg(test)]
 141 mod tests {
 142     use crate::tests::{check_assist, check_assist_not_applicable};
 143
 144     use super::*;
 145
 146     #[test]
 147     fn test_generate_getter_from_field() {
 148         check_assist(
 149             generate_getter,
 150             r#"
 151 struct Context {
 152     dat$0a: Data,
 153 }
 154 "#,
 155             r#"
 156 struct Context {
 157     data: Data,
 158 }
 159
 160 impl Context {
 161     /// Get a reference to the context's data.
 162     fn $0data(&self) -> &Data {
 163         &self.data
 164     }
 165 }
 166 "#,
 167         );
 168
 169         check_assist(
 170             generate_getter_mut,
 171             r#"
 172 struct Context {
 173     dat$0a: Data,
 174 }
 175 "#,
 176             r#"
 177 struct Context {
 178     data: Data,
 179 }
 180
 181 impl Context {
 182     /// Get a mutable reference to the context's data.
 183     fn $0data_mut(&mut self) -> &mut Data {
 184         &mut self.data
 185     }
 186 }
 187 "#,
 188         );
 189     }
 190
 191     #[test]
 192     fn test_generate_getter_already_implemented() {
 193         check_assist_not_applicable(
 194             generate_getter,
 195             r#"
 196 struct Context {
 197     dat$0a: Data,
 198 }
 199
 200 impl Context {
 201     fn data(&self) -> &Data {
 202         &self.data
 203     }
 204 }
 205 "#,
 206         );
 207
 208         check_assist_not_applicable(
 209             generate_getter_mut,
 210             r#"
 211 struct Context {
 212     dat$0a: Data,
 213 }
 214
 215 impl Context {
 216     fn data_mut(&mut self) -> &mut Data {
 217         &mut self.data
 218     }
 219 }
 220 "#,
 221         );
 222     }
 223
 224     #[test]
 225     fn test_generate_getter_from_field_with_visibility_marker() {
 226         check_assist(
 227             generate_getter,
 228             r#"
 229 pub(crate) struct Context {
 230     dat$0a: Data,
 231 }
 232 "#,
 233             r#"
 234 pub(crate) struct Context {
 235     data: Data,
 236 }
 237
 238 impl Context {
 239     /// Get a reference to the context's data.
 240     pub(crate) fn $0data(&self) -> &Data {
 241         &self.data
 242     }
 243 }
 244 "#,
 245         );
 246     }
 247
 248     #[test]
 249     fn test_multiple_generate_getter() {
 250         check_assist(
 251             generate_getter,
 252             r#"
 253 struct Context {
 254     data: Data,
 255     cou$0nt: usize,
 256 }
 257
 258 impl Context {
 259     /// Get a reference to the context's data.
 260     fn data(&self) -> &Data {
 261         &self.data
 262     }
 263 }
 264 "#,
 265             r#"
 266 struct Context {
 267     data: Data,
 268     count: usize,
 269 }
 270
 271 impl Context {
 272     /// Get a reference to the context's data.
 273     fn data(&self) -> &Data {
 274         &self.data
 275     }
 276
 277     /// Get a reference to the context's count.
 278     fn $0count(&self) -> &usize {
 279         &self.count
 280     }
 281 }
 282 "#,
 283         );
 284     }
 285
 286     #[test]
 287     fn test_special_cases() {
 288         cov_mark::check!(useless_type_special_case);
 289         check_assist(
 290             generate_getter,
 291             r#"
 292 struct S { foo: $0String }
 293 "#,
 294             r#"
 295 struct S { foo: String }
 296
 297 impl S {
 298     /// Get a reference to the s's foo.
 299     fn $0foo(&self) -> &str {
 300         self.foo.as_str()
 301     }
 302 }
 303 "#,
 304         );
 305         check_assist(
 306             generate_getter,
 307             r#"
 308 struct S { foo: $0Box<Sweets> }
 309 "#,
 310             r#"
 311 struct S { foo: Box<Sweets> }
 312
 313 impl S {
 314     /// Get a reference to the s's foo.
 315     fn $0foo(&self) -> &Sweets {
 316         self.foo.as_ref()
 317     }
 318 }
 319 "#,
 320         );
 321         check_assist(
 322             generate_getter,
 323             r#"
 324 struct S { foo: $0Vec<()> }
 325 "#,
 326             r#"
 327 struct S { foo: Vec<()> }
 328
 329 impl S {
 330     /// Get a reference to the s's foo.
 331     fn $0foo(&self) -> &[()] {
 332         self.foo.as_slice()
 333     }
 334 }
 335 "#,
 336         );
 337         check_assist(
 338             generate_getter,
 339             r#"
 340 struct S { foo: $0Option<Failure> }
 341 "#,
 342             r#"
 343 struct S { foo: Option<Failure> }
 344
 345 impl S {
 346     /// Get a reference to the s's foo.
 347     fn $0foo(&self) -> Option<&Failure> {
 348         self.foo.as_ref()
 349     }
 350 }
 351 "#,
 352         );
 353     }
 354 }