crates/ide_assists/src/utils/gen_trait_fn_body.rs

   1 //! This module contains functions to generate default trait impl function bodies where possible.
   2
   3 use syntax::{
   4     ast::{self, edit::AstNodeEdit, make, AstNode, NameOwner},
   5     ted,
   6 };
   7
   8 /// Generate custom trait bodies where possible.
   9 ///
  10 /// Returns `Option` so that we can use `?` rather than `if let Some`. Returning
  11 /// `None` means that generating a custom trait body failed, and the body will remain
  12 /// as `todo!` instead.
  13 pub(crate) fn gen_trait_fn_body(
  14     func: &ast::Fn,
  15     trait_path: &ast::Path,
  16     adt: &ast::Adt,
  17 ) -> Option<()> {
  18     match trait_path.segment()?.name_ref()?.text().as_str() {
  19         "Clone" => gen_clone_impl(adt, func),
  20         "Debug" => gen_debug_impl(adt, func),
  21         "Default" => gen_default_impl(adt, func),
  22         "Hash" => gen_hash_impl(adt, func),
  23         "PartialEq" => gen_partial_eq(adt, func),
  24         _ => None,
  25     }
  26 }
  27
  28 /// Generate a `Clone` impl based on the fields and members of the target type.
  29 fn gen_clone_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
  30     fn gen_clone_call(target: ast::Expr) -> ast::Expr {
  31         let method = make::name_ref("clone");
  32         make::expr_method_call(target, method, make::arg_list(None))
  33     }
  34     let expr = match adt {
  35         // `Clone` cannot be derived for unions, so no default impl can be provided.
  36         ast::Adt::Union(_) => return None,
  37         ast::Adt::Enum(enum_) => {
  38             let list = enum_.variant_list()?;
  39             let mut arms = vec![];
  40             for variant in list.variants() {
  41                 let name = variant.name()?;
  42                 let variant_name = make::ext::path_from_idents(["Self", &format!("{}", name)])?;
  43
  44                 match variant.field_list() {
  45                     // => match self { Self::Name { x } => Self::Name { x: x.clone() } }
  46                     Some(ast::FieldList::RecordFieldList(list)) => {
  47                         let mut pats = vec![];
  48                         let mut fields = vec![];
  49                         for field in list.fields() {
  50                             let field_name = field.name()?;
  51                             let pat = make::ident_pat(false, false, field_name.clone());
  52                             pats.push(pat.into());
  53
  54                             let path = make::ext::ident_path(&field_name.to_string());
  55                             let method_call = gen_clone_call(make::expr_path(path));
  56                             let name_ref = make::name_ref(&field_name.to_string());
  57                             let field = make::record_expr_field(name_ref, Some(method_call));
  58                             fields.push(field);
  59                         }
  60                         let pat = make::record_pat(variant_name.clone(), pats.into_iter());
  61                         let fields = make::record_expr_field_list(fields);
  62                         let record_expr = make::record_expr(variant_name, fields).into();
  63                         arms.push(make::match_arm(Some(pat.into()), None, record_expr));
  64                     }
  65
  66                     // => match self { Self::Name(arg1) => Self::Name(arg1.clone()) }
  67                     Some(ast::FieldList::TupleFieldList(list)) => {
  68                         let mut pats = vec![];
  69                         let mut fields = vec![];
  70                         for (i, _) in list.fields().enumerate() {
  71                             let field_name = format!("arg{}", i);
  72                             let pat = make::ident_pat(false, false, make::name(&field_name));
  73                             pats.push(pat.into());
  74
  75                             let f_path = make::expr_path(make::ext::ident_path(&field_name));
  76                             fields.push(gen_clone_call(f_path));
  77                         }
  78                         let pat = make::tuple_struct_pat(variant_name.clone(), pats.into_iter());
  79                         let struct_name = make::expr_path(variant_name);
  80                         let tuple_expr = make::expr_call(struct_name, make::arg_list(fields));
  81                         arms.push(make::match_arm(Some(pat.into()), None, tuple_expr));
  82                     }
  83
  84                     // => match self { Self::Name => Self::Name }
  85                     None => {
  86                         let pattern = make::path_pat(variant_name.clone());
  87                         let variant_expr = make::expr_path(variant_name);
  88                         arms.push(make::match_arm(Some(pattern.into()), None, variant_expr));
  89                     }
  90                 }
  91             }
  92
  93             let match_target = make::expr_path(make::ext::ident_path("self"));
  94             let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
  95             make::expr_match(match_target, list)
  96         }
  97         ast::Adt::Struct(strukt) => {
  98             match strukt.field_list() {
  99                 // => Self { name: self.name.clone() }
 100                 Some(ast::FieldList::RecordFieldList(field_list)) => {
 101                     let mut fields = vec![];
 102                     for field in field_list.fields() {
 103                         let base = make::expr_path(make::ext::ident_path("self"));
 104                         let target = make::expr_field(base, &field.name()?.to_string());
 105                         let method_call = gen_clone_call(target);
 106                         let name_ref = make::name_ref(&field.name()?.to_string());
 107                         let field = make::record_expr_field(name_ref, Some(method_call));
 108                         fields.push(field);
 109                     }
 110                     let struct_name = make::ext::ident_path("Self");
 111                     let fields = make::record_expr_field_list(fields);
 112                     make::record_expr(struct_name, fields).into()
 113                 }
 114                 // => Self(self.0.clone(), self.1.clone())
 115                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 116                     let mut fields = vec![];
 117                     for (i, _) in field_list.fields().enumerate() {
 118                         let f_path = make::expr_path(make::ext::ident_path("self"));
 119                         let target = make::expr_field(f_path, &format!("{}", i)).into();
 120                         fields.push(gen_clone_call(target));
 121                     }
 122                     let struct_name = make::expr_path(make::ext::ident_path("Self"));
 123                     make::expr_call(struct_name, make::arg_list(fields))
 124                 }
 125                 // => Self { }
 126                 None => {
 127                     let struct_name = make::ext::ident_path("Self");
 128                     let fields = make::record_expr_field_list(None);
 129                     make::record_expr(struct_name, fields).into()
 130                 }
 131             }
 132         }
 133     };
 134     let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 135     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 136     Some(())
 137 }
 138
 139 /// Generate a `Debug` impl based on the fields and members of the target type.
 140 fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 141     let annotated_name = adt.name()?;
 142     match adt {
 143         // `Debug` cannot be derived for unions, so no default impl can be provided.
 144         ast::Adt::Union(_) => None,
 145
 146         // => match self { Self::Variant => write!(f, "Variant") }
 147         ast::Adt::Enum(enum_) => {
 148             let list = enum_.variant_list()?;
 149             let mut arms = vec![];
 150             for variant in list.variants() {
 151                 let name = variant.name()?;
 152                 let variant_name =
 153                     make::path_pat(make::ext::path_from_idents(["Self", &format!("{}", name)])?);
 154
 155                 let target = make::expr_path(make::ext::ident_path("f").into());
 156                 let fmt_string = make::expr_literal(&(format!("\"{}\"", name))).into();
 157                 let args = make::arg_list(vec![target, fmt_string]);
 158                 let macro_name = make::expr_path(make::ext::ident_path("write"));
 159                 let macro_call = make::expr_macro_call(macro_name, args);
 160
 161                 arms.push(make::match_arm(Some(variant_name.into()), None, macro_call.into()));
 162             }
 163
 164             let match_target = make::expr_path(make::ext::ident_path("self"));
 165             let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
 166             let match_expr = make::expr_match(match_target, list);
 167
 168             let body = make::block_expr(None, Some(match_expr));
 169             let body = body.indent(ast::edit::IndentLevel(1));
 170             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 171             Some(())
 172         }
 173
 174         ast::Adt::Struct(strukt) => {
 175             let name = format!("\"{}\"", annotated_name);
 176             let args = make::arg_list(Some(make::expr_literal(&name).into()));
 177             let target = make::expr_path(make::ext::ident_path("f"));
 178
 179             let expr = match strukt.field_list() {
 180                 // => f.debug_struct("Name").finish()
 181                 None => make::expr_method_call(target, make::name_ref("debug_struct"), args),
 182
 183                 // => f.debug_struct("Name").field("foo", &self.foo).finish()
 184                 Some(ast::FieldList::RecordFieldList(field_list)) => {
 185                     let method = make::name_ref("debug_struct");
 186                     let mut expr = make::expr_method_call(target, method, args);
 187                     for field in field_list.fields() {
 188                         let name = field.name()?;
 189                         let f_name = make::expr_literal(&(format!("\"{}\"", name))).into();
 190                         let f_path = make::expr_path(make::ext::ident_path("self"));
 191                         let f_path = make::expr_ref(f_path, false);
 192                         let f_path = make::expr_field(f_path, &format!("{}", name)).into();
 193                         let args = make::arg_list(vec![f_name, f_path]);
 194                         expr = make::expr_method_call(expr, make::name_ref("field"), args);
 195                     }
 196                     expr
 197                 }
 198
 199                 // => f.debug_tuple("Name").field(self.0).finish()
 200                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 201                     let method = make::name_ref("debug_tuple");
 202                     let mut expr = make::expr_method_call(target, method, args);
 203                     for (i, _) in field_list.fields().enumerate() {
 204                         let f_path = make::expr_path(make::ext::ident_path("self"));
 205                         let f_path = make::expr_ref(f_path, false);
 206                         let f_path = make::expr_field(f_path, &format!("{}", i)).into();
 207                         let method = make::name_ref("field");
 208                         expr = make::expr_method_call(expr, method, make::arg_list(Some(f_path)));
 209                     }
 210                     expr
 211                 }
 212             };
 213
 214             let method = make::name_ref("finish");
 215             let expr = make::expr_method_call(expr, method, make::arg_list(None));
 216             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 217             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 218             Some(())
 219         }
 220     }
 221 }
 222
 223 /// Generate a `Debug` impl based on the fields and members of the target type.
 224 fn gen_default_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 225     fn gen_default_call() -> Option<ast::Expr> {
 226         let fn_name = make::ext::path_from_idents(["Default", "default"])?;
 227         Some(make::expr_call(make::expr_path(fn_name), make::arg_list(None)))
 228     }
 229     match adt {
 230         // `Debug` cannot be derived for unions, so no default impl can be provided.
 231         ast::Adt::Union(_) => None,
 232         // Deriving `Debug` for enums is not stable yet.
 233         ast::Adt::Enum(_) => None,
 234         ast::Adt::Struct(strukt) => {
 235             let expr = match strukt.field_list() {
 236                 Some(ast::FieldList::RecordFieldList(field_list)) => {
 237                     let mut fields = vec![];
 238                     for field in field_list.fields() {
 239                         let method_call = gen_default_call()?;
 240                         let name_ref = make::name_ref(&field.name()?.to_string());
 241                         let field = make::record_expr_field(name_ref, Some(method_call));
 242                         fields.push(field);
 243                     }
 244                     let struct_name = make::ext::ident_path("Self");
 245                     let fields = make::record_expr_field_list(fields);
 246                     make::record_expr(struct_name, fields).into()
 247                 }
 248                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 249                     let struct_name = make::expr_path(make::ext::ident_path("Self"));
 250                     let fields = field_list
 251                         .fields()
 252                         .map(|_| gen_default_call())
 253                         .collect::<Option<Vec<ast::Expr>>>()?;
 254                     make::expr_call(struct_name, make::arg_list(fields))
 255                 }
 256                 None => {
 257                     let struct_name = make::ext::ident_path("Self");
 258                     let fields = make::record_expr_field_list(None);
 259                     make::record_expr(struct_name, fields).into()
 260                 }
 261             };
 262             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 263             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 264             Some(())
 265         }
 266     }
 267 }
 268
 269 /// Generate a `Hash` impl based on the fields and members of the target type.
 270 fn gen_hash_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 271     fn gen_hash_call(target: ast::Expr) -> ast::Stmt {
 272         let method = make::name_ref("hash");
 273         let arg = make::expr_path(make::ext::ident_path("state"));
 274         let expr = make::expr_method_call(target, method, make::arg_list(Some(arg)));
 275         make::expr_stmt(expr).into()
 276     }
 277
 278     let body = match adt {
 279         // `Hash` cannot be derived for unions, so no default impl can be provided.
 280         ast::Adt::Union(_) => return None,
 281
 282         // => std::mem::discriminant(self).hash(state);
 283         ast::Adt::Enum(_) => {
 284             let fn_name = make_discriminant()?;
 285
 286             let arg = make::expr_path(make::ext::ident_path("self"));
 287             let fn_call = make::expr_call(fn_name, make::arg_list(Some(arg)));
 288             let stmt = gen_hash_call(fn_call);
 289
 290             make::block_expr(Some(stmt), None).indent(ast::edit::IndentLevel(1))
 291         }
 292         ast::Adt::Struct(strukt) => match strukt.field_list() {
 293             // => self.<field>.hash(state);
 294             Some(ast::FieldList::RecordFieldList(field_list)) => {
 295                 let mut stmts = vec![];
 296                 for field in field_list.fields() {
 297                     let base = make::expr_path(make::ext::ident_path("self"));
 298                     let target = make::expr_field(base, &field.name()?.to_string());
 299                     stmts.push(gen_hash_call(target));
 300                 }
 301                 make::block_expr(stmts, None).indent(ast::edit::IndentLevel(1))
 302             }
 303
 304             // => self.<field_index>.hash(state);
 305             Some(ast::FieldList::TupleFieldList(field_list)) => {
 306                 let mut stmts = vec![];
 307                 for (i, _) in field_list.fields().enumerate() {
 308                     let base = make::expr_path(make::ext::ident_path("self"));
 309                     let target = make::expr_field(base, &format!("{}", i));
 310                     stmts.push(gen_hash_call(target));
 311                 }
 312                 make::block_expr(stmts, None).indent(ast::edit::IndentLevel(1))
 313             }
 314
 315             // No fields in the body means there's nothing to hash.
 316             None => return None,
 317         },
 318     };
 319
 320     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 321     Some(())
 322 }
 323
 324 /// Generate a `PartialEq` impl based on the fields and members of the target type.
 325 fn gen_partial_eq(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 326     fn gen_eq_chain(expr: Option<ast::Expr>, cmp: ast::Expr) -> Option<ast::Expr> {
 327         match expr {
 328             Some(expr) => Some(make::expr_op(ast::BinOp::BooleanAnd, expr, cmp)),
 329             None => Some(cmp),
 330         }
 331     }
 332
 333     fn gen_record_pat_field(field_name: &str, pat_name: &str) -> ast::RecordPatField {
 334         let pat = make::ext::simple_ident_pat(make::name(&pat_name));
 335         let name_ref = make::name_ref(field_name);
 336         make::record_pat_field(name_ref, pat.into())
 337     }
 338
 339     fn gen_record_pat(record_name: ast::Path, fields: Vec<ast::RecordPatField>) -> ast::RecordPat {
 340         let list = make::record_pat_field_list(fields);
 341         make::record_pat_with_fields(record_name, list)
 342     }
 343
 344     fn gen_variant_path(variant: &ast::Variant) -> Option<ast::Path> {
 345         make::ext::path_from_idents(["Self", &variant.name()?.to_string()])
 346     }
 347
 348     fn gen_tuple_field(field_name: &String) -> ast::Pat {
 349         ast::Pat::IdentPat(make::ident_pat(false, false, make::name(field_name)))
 350     }
 351
 352     // FIXME: return `None` if the trait carries a generic type; we can only
 353     // generate this code `Self` for the time being.
 354
 355     let body = match adt {
 356         // `Hash` cannot be derived for unions, so no default impl can be provided.
 357         ast::Adt::Union(_) => return None,
 358
 359         ast::Adt::Enum(enum_) => {
 360             // => std::mem::discriminant(self) == std::mem::discriminant(other)
 361             let lhs_name = make::expr_path(make::ext::ident_path("self"));
 362             let lhs = make::expr_call(make_discriminant()?, make::arg_list(Some(lhs_name.clone())));
 363             let rhs_name = make::expr_path(make::ext::ident_path("other"));
 364             let rhs = make::expr_call(make_discriminant()?, make::arg_list(Some(rhs_name.clone())));
 365             let eq_check = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 366
 367             let mut case_count = 0;
 368             let mut arms = vec![];
 369             for variant in enum_.variant_list()?.variants() {
 370                 case_count += 1;
 371                 match variant.field_list() {
 372                     // => (Self::Bar { bin: l_bin }, Self::Bar { bin: r_bin }) => l_bin == r_bin,
 373                     Some(ast::FieldList::RecordFieldList(list)) => {
 374                         let mut expr = None;
 375                         let mut l_fields = vec![];
 376                         let mut r_fields = vec![];
 377
 378                         for field in list.fields() {
 379                             let field_name = field.name()?.to_string();
 380
 381                             let l_name = &format!("l_{}", field_name);
 382                             l_fields.push(gen_record_pat_field(&field_name, &l_name));
 383
 384                             let r_name = &format!("r_{}", field_name);
 385                             r_fields.push(gen_record_pat_field(&field_name, &r_name));
 386
 387                             let lhs = make::expr_path(make::ext::ident_path(l_name));
 388                             let rhs = make::expr_path(make::ext::ident_path(r_name));
 389                             let cmp = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 390                             expr = gen_eq_chain(expr, cmp);
 391                         }
 392
 393                         let left = gen_record_pat(gen_variant_path(&variant)?, l_fields);
 394                         let right = gen_record_pat(gen_variant_path(&variant)?, r_fields);
 395                         let tuple = make::tuple_pat(vec![left.into(), right.into()]);
 396
 397                         if let Some(expr) = expr {
 398                             arms.push(make::match_arm(Some(tuple.into()), None, expr));
 399                         }
 400                     }
 401
 402                     Some(ast::FieldList::TupleFieldList(list)) => {
 403                         let mut expr = None;
 404                         let mut l_fields = vec![];
 405                         let mut r_fields = vec![];
 406
 407                         for (i, _) in list.fields().enumerate() {
 408                             let field_name = format!("{}", i);
 409
 410                             let l_name = format!("l{}", field_name);
 411                             l_fields.push(gen_tuple_field(&l_name));
 412
 413                             let r_name = format!("r{}", field_name);
 414                             r_fields.push(gen_tuple_field(&r_name));
 415
 416                             let lhs = make::expr_path(make::ext::ident_path(&l_name));
 417                             let rhs = make::expr_path(make::ext::ident_path(&r_name));
 418                             let cmp = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 419                             expr = gen_eq_chain(expr, cmp);
 420                         }
 421
 422                         let left = make::tuple_struct_pat(gen_variant_path(&variant)?, l_fields);
 423                         let right = make::tuple_struct_pat(gen_variant_path(&variant)?, r_fields);
 424                         let tuple = make::tuple_pat(vec![left.into(), right.into()]);
 425
 426                         if let Some(expr) = expr {
 427                             arms.push(make::match_arm(Some(tuple.into()), None, expr));
 428                         }
 429                     }
 430                     None => continue,
 431                 }
 432             }
 433
 434             let expr = match arms.len() {
 435                 0 => eq_check,
 436                 _ => {
 437                     if case_count > arms.len() {
 438                         let lhs = make::wildcard_pat().into();
 439                         arms.push(make::match_arm(Some(lhs), None, eq_check));
 440                     }
 441
 442                     let match_target = make::expr_tuple(vec![lhs_name, rhs_name]);
 443                     let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
 444                     make::expr_match(match_target, list)
 445                 }
 446             };
 447
 448             make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1))
 449         }
 450         ast::Adt::Struct(strukt) => match strukt.field_list() {
 451             Some(ast::FieldList::RecordFieldList(field_list)) => {
 452                 let mut expr = None;
 453                 for field in field_list.fields() {
 454                     let lhs = make::expr_path(make::ext::ident_path("self"));
 455                     let lhs = make::expr_field(lhs, &field.name()?.to_string());
 456                     let rhs = make::expr_path(make::ext::ident_path("other"));
 457                     let rhs = make::expr_field(rhs, &field.name()?.to_string());
 458                     let cmp = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 459                     expr = gen_eq_chain(expr, cmp);
 460                 }
 461                 make::block_expr(None, expr).indent(ast::edit::IndentLevel(1))
 462             }
 463
 464             Some(ast::FieldList::TupleFieldList(field_list)) => {
 465                 let mut expr = None;
 466                 for (i, _) in field_list.fields().enumerate() {
 467                     let idx = format!("{}", i);
 468                     let lhs = make::expr_path(make::ext::ident_path("self"));
 469                     let lhs = make::expr_field(lhs, &idx);
 470                     let rhs = make::expr_path(make::ext::ident_path("other"));
 471                     let rhs = make::expr_field(rhs, &idx);
 472                     let cmp = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 473                     expr = gen_eq_chain(expr, cmp);
 474                 }
 475                 make::block_expr(None, expr).indent(ast::edit::IndentLevel(1))
 476             }
 477
 478             // No fields in the body means there's nothing to hash.
 479             None => {
 480                 let expr = make::expr_literal("true").into();
 481                 make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1))
 482             }
 483         },
 484     };
 485
 486     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 487     Some(())
 488 }
 489
 490 fn make_discriminant() -> Option<ast::Expr> {
 491     Some(make::expr_path(make::ext::path_from_idents(["core", "mem", "discriminant"])?))
 492 }