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         let stmt = make::expr_stmt(expr);
 278         stmt.into()
 279     }
 280
 281     let body = match adt {
 282         // `Hash` cannot be derived for unions, so no default impl can be provided.
 283         ast::Adt::Union(_) => return None,
 284
 285         // => std::mem::discriminant(self).hash(state);
 286         ast::Adt::Enum(_) => {
 287             let root = make::ext::ident_path("core");
 288             let submodule = make::ext::ident_path("mem");
 289             let fn_name = make::ext::ident_path("discriminant");
 290             let fn_name = make::path_concat(submodule, fn_name);
 291             let fn_name = make::expr_path(make::path_concat(root, fn_name));
 292
 293             let arg = make::expr_path(make::ext::ident_path("self"));
 294             let fn_call = make::expr_call(fn_name, make::arg_list(Some(arg)));
 295             let stmt = gen_hash_call(fn_call);
 296
 297             make::block_expr(Some(stmt), None).indent(ast::edit::IndentLevel(1))
 298         }
 299         ast::Adt::Struct(strukt) => match strukt.field_list() {
 300             // => self.<field>.hash(state);
 301             Some(ast::FieldList::RecordFieldList(field_list)) => {
 302                 let mut stmts = vec![];
 303                 for field in field_list.fields() {
 304                     let base = make::expr_path(make::ext::ident_path("self"));
 305                     let target = make::expr_field(base, &field.name()?.to_string());
 306                     stmts.push(gen_hash_call(target));
 307                 }
 308                 make::block_expr(stmts, None).indent(ast::edit::IndentLevel(1))
 309             }
 310
 311             // => self.<field_index>.hash(state);
 312             Some(ast::FieldList::TupleFieldList(field_list)) => {
 313                 let mut stmts = vec![];
 314                 for (i, _) in field_list.fields().enumerate() {
 315                     let base = make::expr_path(make::ext::ident_path("self"));
 316                     let target = make::expr_field(base, &format!("{}", i));
 317                     stmts.push(gen_hash_call(target));
 318                 }
 319                 make::block_expr(stmts, None).indent(ast::edit::IndentLevel(1))
 320             }
 321
 322             // No fields in the body means there's nothing to hash.
 323             None => return None,
 324         },
 325     };
 326
 327     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 328     Some(())
 329 }
 330
 331 /// Generate a `PartialEq` impl based on the fields and members of the target type.
 332 fn gen_partial_eq(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 333     // FIXME: return `None` if the trait carries a generic type; we can only
 334     // generate this code `Self` for the time being.
 335
 336     let body = match adt {
 337         // `Hash` cannot be derived for unions, so no default impl can be provided.
 338         ast::Adt::Union(_) => return None,
 339
 340         // FIXME: generate trait variants
 341         ast::Adt::Enum(_) => todo!(),
 342         ast::Adt::Struct(strukt) => match strukt.field_list() {
 343             // => self.<field>.hash(state);
 344             Some(ast::FieldList::RecordFieldList(field_list)) => {
 345                 let mut expr = None;
 346                 for field in field_list.fields() {
 347                     let lhs = make::expr_path(make::ext::ident_path("self"));
 348                     let lhs = make::expr_field(lhs, &field.name()?.to_string());
 349                     let rhs = make::expr_path(make::ext::ident_path("other"));
 350                     let rhs = make::expr_field(rhs, &field.name()?.to_string());
 351                     let cmp = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 352                     expr = match expr {
 353                         Some(expr) => Some(make::expr_op(ast::BinOp::BooleanAnd, expr, cmp)),
 354                         None => Some(cmp),
 355                     };
 356                 }
 357                 make::block_expr(None, expr).indent(ast::edit::IndentLevel(1))
 358             }
 359
 360             // => self.<field_index>.hash(state);
 361             Some(ast::FieldList::TupleFieldList(field_list)) => {
 362                 let mut expr = None;
 363                 for (i, _) in field_list.fields().enumerate() {
 364                     let idx = format!("{}", i);
 365                     let lhs = make::expr_path(make::ext::ident_path("self"));
 366                     let lhs = make::expr_field(lhs, &idx);
 367                     let rhs = make::expr_path(make::ext::ident_path("other"));
 368                     let rhs = make::expr_field(rhs, &idx);
 369                     let cmp = make::expr_op(ast::BinOp::EqualityTest, lhs, rhs);
 370                     expr = match expr {
 371                         Some(expr) => Some(make::expr_op(ast::BinOp::BooleanAnd, expr, cmp)),
 372                         None => Some(cmp),
 373                     };
 374                 }
 375                 make::block_expr(None, expr).indent(ast::edit::IndentLevel(1))
 376             }
 377
 378             // No fields in the body means there's nothing to hash.
 379             None => {
 380                 let expr = make::expr_literal("true").into();
 381                 make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1))
 382             }
 383         },
 384     };
 385
 386     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 387     Some(())
 388 }