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 $0name(&self) -> &str {
  27 //         self.name.as_str()
  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 $0name_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             let (ty, body) = if mutable {
 100                 (format!("&mut {}", field_ty), format!("&mut self.{}", field_name))
 101             } else {
 102                 useless_type_special_case(&field_name.to_string(), &field_ty)
 103                     .unwrap_or_else(|| (format!("&{}", field_ty), format!("&self.{}", field_name)))
 104             };
 105
 106             format_to!(
 107                 buf,
 108                 "    /// Get a {}reference to the {}'s {}.
 109     {}fn {}(&{}self) -> {} {{
 110         {}
 111     }}",
 112                 mutable.then(|| "mutable ").unwrap_or_default(),
 113                 to_lower_snake_case(&strukt_name.to_string()).replace('_', " "),
 114                 fn_name.trim_end_matches("_mut").replace('_', " "),
 115                 vis,
 116                 fn_name,
 117                 mutable.then(|| "mut ").unwrap_or_default(),
 118                 ty,
 119                 body,
 120             );
 121
 122             let start_offset = impl_def
 123                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
 124                 .unwrap_or_else(|| {
 125                     buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
 126                     strukt.syntax().text_range().end()
 127                 });
 128
 129             match ctx.config.snippet_cap {
 130                 Some(cap) => {
 131                     builder.insert_snippet(cap, start_offset, buf.replacen("fn ", "fn $0", 1))
 132                 }
 133                 None => builder.insert(start_offset, buf),
 134             }
 135         },
 136     )
 137 }
 138
 139 fn useless_type_special_case(field_name: &str, field_ty: &ast::Type) -> Option<(String, String)> {
 140     if field_ty.to_string() == "String" {
 141         cov_mark::hit!(useless_type_special_case);
 142         return Some(("&str".to_string(), format!("self.{}.as_str()", field_name)));
 143     }
 144     if let Some(arg) = ty_ctor(field_ty, "Vec") {
 145         return Some((format!("&[{}]", arg), format!("self.{}.as_slice()", field_name)));
 146     }
 147     if let Some(arg) = ty_ctor(field_ty, "Box") {
 148         return Some((format!("&{}", arg), format!("self.{}.as_ref()", field_name)));
 149     }
 150     if let Some(arg) = ty_ctor(field_ty, "Option") {
 151         return Some((format!("Option<&{}>", arg), format!("self.{}.as_ref()", field_name)));
 152     }
 153     None
 154 }
 155
 156 // FIXME: This should rely on semantic info.
 157 fn ty_ctor(ty: &ast::Type, ctor: &str) -> Option<String> {
 158     let res = ty.to_string().strip_prefix(ctor)?.strip_prefix('<')?.strip_suffix('>')?.to_string();
 159     Some(res)
 160 }
 161
 162 #[cfg(test)]
 163 mod tests {
 164     use crate::tests::{check_assist, check_assist_not_applicable};
 165
 166     use super::*;
 167
 168     #[test]
 169     fn test_generate_getter_from_field() {
 170         check_assist(
 171             generate_getter,
 172             r#"
 173 struct Context {
 174     dat$0a: Data,
 175 }
 176 "#,
 177             r#"
 178 struct Context {
 179     data: Data,
 180 }
 181
 182 impl Context {
 183     /// Get a reference to the context's data.
 184     fn $0data(&self) -> &Data {
 185         &self.data
 186     }
 187 }
 188 "#,
 189         );
 190
 191         check_assist(
 192             generate_getter_mut,
 193             r#"
 194 struct Context {
 195     dat$0a: Data,
 196 }
 197 "#,
 198             r#"
 199 struct Context {
 200     data: Data,
 201 }
 202
 203 impl Context {
 204     /// Get a mutable reference to the context's data.
 205     fn $0data_mut(&mut self) -> &mut Data {
 206         &mut self.data
 207     }
 208 }
 209 "#,
 210         );
 211     }
 212
 213     #[test]
 214     fn test_generate_getter_already_implemented() {
 215         check_assist_not_applicable(
 216             generate_getter,
 217             r#"
 218 struct Context {
 219     dat$0a: Data,
 220 }
 221
 222 impl Context {
 223     fn data(&self) -> &Data {
 224         &self.data
 225     }
 226 }
 227 "#,
 228         );
 229
 230         check_assist_not_applicable(
 231             generate_getter_mut,
 232             r#"
 233 struct Context {
 234     dat$0a: Data,
 235 }
 236
 237 impl Context {
 238     fn data_mut(&mut self) -> &mut Data {
 239         &mut self.data
 240     }
 241 }
 242 "#,
 243         );
 244     }
 245
 246     #[test]
 247     fn test_generate_getter_from_field_with_visibility_marker() {
 248         check_assist(
 249             generate_getter,
 250             r#"
 251 pub(crate) struct Context {
 252     dat$0a: Data,
 253 }
 254 "#,
 255             r#"
 256 pub(crate) struct Context {
 257     data: Data,
 258 }
 259
 260 impl Context {
 261     /// Get a reference to the context's data.
 262     pub(crate) fn $0data(&self) -> &Data {
 263         &self.data
 264     }
 265 }
 266 "#,
 267         );
 268     }
 269
 270     #[test]
 271     fn test_multiple_generate_getter() {
 272         check_assist(
 273             generate_getter,
 274             r#"
 275 struct Context {
 276     data: Data,
 277     cou$0nt: usize,
 278 }
 279
 280 impl Context {
 281     /// Get a reference to the context's data.
 282     fn data(&self) -> &Data {
 283         &self.data
 284     }
 285 }
 286 "#,
 287             r#"
 288 struct Context {
 289     data: Data,
 290     count: usize,
 291 }
 292
 293 impl Context {
 294     /// Get a reference to the context's data.
 295     fn data(&self) -> &Data {
 296         &self.data
 297     }
 298
 299     /// Get a reference to the context's count.
 300     fn $0count(&self) -> &usize {
 301         &self.count
 302     }
 303 }
 304 "#,
 305         );
 306     }
 307
 308     #[test]
 309     fn test_special_cases() {
 310         cov_mark::check!(useless_type_special_case);
 311         check_assist(
 312             generate_getter,
 313             r#"
 314 struct S { foo: $0String }
 315 "#,
 316             r#"
 317 struct S { foo: String }
 318
 319 impl S {
 320     /// Get a reference to the s's foo.
 321     fn $0foo(&self) -> &str {
 322         self.foo.as_str()
 323     }
 324 }
 325 "#,
 326         );
 327         check_assist(
 328             generate_getter,
 329             r#"
 330 struct S { foo: $0Box<Sweets> }
 331 "#,
 332             r#"
 333 struct S { foo: Box<Sweets> }
 334
 335 impl S {
 336     /// Get a reference to the s's foo.
 337     fn $0foo(&self) -> &Sweets {
 338         self.foo.as_ref()
 339     }
 340 }
 341 "#,
 342         );
 343         check_assist(
 344             generate_getter,
 345             r#"
 346 struct S { foo: $0Vec<()> }
 347 "#,
 348             r#"
 349 struct S { foo: Vec<()> }
 350
 351 impl S {
 352     /// Get a reference to the s's foo.
 353     fn $0foo(&self) -> &[()] {
 354         self.foo.as_slice()
 355     }
 356 }
 357 "#,
 358         );
 359         check_assist(
 360             generate_getter,
 361             r#"
 362 struct S { foo: $0Option<Failure> }
 363 "#,
 364             r#"
 365 struct S { foo: Option<Failure> }
 366
 367 impl S {
 368     /// Get a reference to the s's foo.
 369     fn $0foo(&self) -> Option<&Failure> {
 370         self.foo.as_ref()
 371     }
 372 }
 373 "#,
 374         );
 375     }
 376 }