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