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         _ => None,
  24     }
  25 }
  26
  27 /// Generate a `Clone` impl based on the fields and members of the target type.
  28 fn gen_clone_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
  29     fn gen_clone_call(target: ast::Expr) -> ast::Expr {
  30         let method = make::name_ref("clone");
  31         make::expr_method_call(target, method, make::arg_list(None))
  32     }
  33     let expr = match adt {
  34         // `Clone` cannot be derived for unions, so no default impl can be provided.
  35         ast::Adt::Union(_) => return None,
  36         ast::Adt::Enum(enum_) => {
  37             let list = enum_.variant_list()?;
  38             let mut arms = vec![];
  39             for variant in list.variants() {
  40                 let name = variant.name()?;
  41                 let left = make::ext::ident_path("Self");
  42                 let right = make::ext::ident_path(&format!("{}", name));
  43                 let variant_name = make::path_concat(left, right);
  44
  45                 match variant.field_list() {
  46                     // => match self { Self::Name { x } => Self::Name { x: x.clone() } }
  47                     Some(ast::FieldList::RecordFieldList(list)) => {
  48                         let mut pats = vec![];
  49                         let mut fields = vec![];
  50                         for field in list.fields() {
  51                             let field_name = field.name()?;
  52                             let pat = make::ident_pat(false, false, field_name.clone());
  53                             pats.push(pat.into());
  54
  55                             let path = make::ext::ident_path(&field_name.to_string());
  56                             let method_call = gen_clone_call(make::expr_path(path));
  57                             let name_ref = make::name_ref(&field_name.to_string());
  58                             let field = make::record_expr_field(name_ref, Some(method_call));
  59                             fields.push(field);
  60                         }
  61                         let pattern = make::record_pat(variant_name.clone(), pats.into_iter());
  62
  63                         let fields = make::record_expr_field_list(fields);
  64                         let record_expr = make::record_expr(variant_name, fields).into();
  65
  66                         arms.push(make::match_arm(Some(pattern.into()), None, record_expr));
  67                     }
  68
  69                     // => match self { Self::Name(arg1) => Self::Name(arg1.clone()) }
  70                     Some(ast::FieldList::TupleFieldList(list)) => todo!(),
  71
  72                     // => match self { Self::Name => Self::Name }
  73                     None => {
  74                         let pattern = make::path_pat(variant_name.clone());
  75                         let variant_expr = make::expr_path(variant_name);
  76                         arms.push(make::match_arm(Some(pattern.into()), None, variant_expr));
  77                     }
  78                 }
  79             }
  80
  81             let match_target = make::expr_path(make::ext::ident_path("self"));
  82             let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
  83             make::expr_match(match_target, list)
  84         }
  85         ast::Adt::Struct(strukt) => {
  86             match strukt.field_list() {
  87                 // => Self { name: self.name.clone() }
  88                 Some(ast::FieldList::RecordFieldList(field_list)) => {
  89                     let mut fields = vec![];
  90                     for field in field_list.fields() {
  91                         let base = make::expr_path(make::ext::ident_path("self"));
  92                         let target = make::expr_field(base, &field.name()?.to_string());
  93                         let method_call = gen_clone_call(target);
  94                         let name_ref = make::name_ref(&field.name()?.to_string());
  95                         let field = make::record_expr_field(name_ref, Some(method_call));
  96                         fields.push(field);
  97                     }
  98                     let struct_name = make::ext::ident_path("Self");
  99                     let fields = make::record_expr_field_list(fields);
 100                     make::record_expr(struct_name, fields).into()
 101                 }
 102                 // => Self(self.0.clone(), self.1.clone())
 103                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 104                     let mut fields = vec![];
 105                     for (i, _) in field_list.fields().enumerate() {
 106                         let f_path = make::expr_path(make::ext::ident_path("self"));
 107                         let target = make::expr_field(f_path, &format!("{}", i)).into();
 108                         fields.push(gen_clone_call(target));
 109                     }
 110                     let struct_name = make::expr_path(make::ext::ident_path("Self"));
 111                     make::expr_call(struct_name, make::arg_list(fields))
 112                 }
 113                 // => Self { }
 114                 None => {
 115                     let struct_name = make::ext::ident_path("Self");
 116                     let fields = make::record_expr_field_list(None);
 117                     make::record_expr(struct_name, fields).into()
 118                 }
 119             }
 120         }
 121     };
 122     let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 123     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 124     Some(())
 125 }
 126
 127 /// Generate a `Debug` impl based on the fields and members of the target type.
 128 fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 129     let annotated_name = adt.name()?;
 130     match adt {
 131         // `Debug` cannot be derived for unions, so no default impl can be provided.
 132         ast::Adt::Union(_) => None,
 133
 134         // => match self { Self::Variant => write!(f, "Variant") }
 135         ast::Adt::Enum(enum_) => {
 136             let list = enum_.variant_list()?;
 137             let mut arms = vec![];
 138             for variant in list.variants() {
 139                 let name = variant.name()?;
 140                 let left = make::ext::ident_path("Self");
 141                 let right = make::ext::ident_path(&format!("{}", name));
 142                 let variant_name = make::path_pat(make::path_concat(left, right));
 143
 144                 let target = make::expr_path(make::ext::ident_path("f").into());
 145                 let fmt_string = make::expr_literal(&(format!("\"{}\"", name))).into();
 146                 let args = make::arg_list(vec![target, fmt_string]);
 147                 let macro_name = make::expr_path(make::ext::ident_path("write"));
 148                 let macro_call = make::expr_macro_call(macro_name, args);
 149
 150                 arms.push(make::match_arm(Some(variant_name.into()), None, macro_call.into()));
 151             }
 152
 153             let match_target = make::expr_path(make::ext::ident_path("self"));
 154             let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
 155             let match_expr = make::expr_match(match_target, list);
 156
 157             let body = make::block_expr(None, Some(match_expr));
 158             let body = body.indent(ast::edit::IndentLevel(1));
 159             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 160             Some(())
 161         }
 162
 163         ast::Adt::Struct(strukt) => {
 164             let name = format!("\"{}\"", annotated_name);
 165             let args = make::arg_list(Some(make::expr_literal(&name).into()));
 166             let target = make::expr_path(make::ext::ident_path("f"));
 167
 168             let expr = match strukt.field_list() {
 169                 // => f.debug_struct("Name").finish()
 170                 None => make::expr_method_call(target, make::name_ref("debug_struct"), args),
 171
 172                 // => f.debug_struct("Name").field("foo", &self.foo).finish()
 173                 Some(ast::FieldList::RecordFieldList(field_list)) => {
 174                     let method = make::name_ref("debug_struct");
 175                     let mut expr = make::expr_method_call(target, method, args);
 176                     for field in field_list.fields() {
 177                         let name = field.name()?;
 178                         let f_name = make::expr_literal(&(format!("\"{}\"", name))).into();
 179                         let f_path = make::expr_path(make::ext::ident_path("self"));
 180                         let f_path = make::expr_ref(f_path, false);
 181                         let f_path = make::expr_field(f_path, &format!("{}", name)).into();
 182                         let args = make::arg_list(vec![f_name, f_path]);
 183                         expr = make::expr_method_call(expr, make::name_ref("field"), args);
 184                     }
 185                     expr
 186                 }
 187
 188                 // => f.debug_tuple("Name").field(self.0).finish()
 189                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 190                     let method = make::name_ref("debug_tuple");
 191                     let mut expr = make::expr_method_call(target, method, args);
 192                     for (i, _) in field_list.fields().enumerate() {
 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!("{}", i)).into();
 196                         let method = make::name_ref("field");
 197                         expr = make::expr_method_call(expr, method, make::arg_list(Some(f_path)));
 198                     }
 199                     expr
 200                 }
 201             };
 202
 203             let method = make::name_ref("finish");
 204             let expr = make::expr_method_call(expr, method, make::arg_list(None));
 205             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 206             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 207             Some(())
 208         }
 209     }
 210 }
 211
 212 /// Generate a `Debug` impl based on the fields and members of the target type.
 213 fn gen_default_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 214     fn gen_default_call() -> ast::Expr {
 215         let trait_name = make::ext::ident_path("Default");
 216         let method_name = make::ext::ident_path("default");
 217         let fn_name = make::expr_path(make::path_concat(trait_name, method_name));
 218         make::expr_call(fn_name, make::arg_list(None))
 219     }
 220     match adt {
 221         // `Debug` cannot be derived for unions, so no default impl can be provided.
 222         ast::Adt::Union(_) => None,
 223         // Deriving `Debug` for enums is not stable yet.
 224         ast::Adt::Enum(_) => None,
 225         ast::Adt::Struct(strukt) => {
 226             let expr = match strukt.field_list() {
 227                 Some(ast::FieldList::RecordFieldList(field_list)) => {
 228                     let mut fields = vec![];
 229                     for field in field_list.fields() {
 230                         let method_call = gen_default_call();
 231                         let name_ref = make::name_ref(&field.name()?.to_string());
 232                         let field = make::record_expr_field(name_ref, Some(method_call));
 233                         fields.push(field);
 234                     }
 235                     let struct_name = make::ext::ident_path("Self");
 236                     let fields = make::record_expr_field_list(fields);
 237                     make::record_expr(struct_name, fields).into()
 238                 }
 239                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 240                     let struct_name = make::expr_path(make::ext::ident_path("Self"));
 241                     let fields = field_list.fields().map(|_| gen_default_call());
 242                     make::expr_call(struct_name, make::arg_list(fields))
 243                 }
 244                 None => {
 245                     let struct_name = make::ext::ident_path("Self");
 246                     let fields = make::record_expr_field_list(None);
 247                     make::record_expr(struct_name, fields).into()
 248                 }
 249             };
 250             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 251             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 252             Some(())
 253         }
 254     }
 255 }
 256
 257 /// Generate a `Hash` impl based on the fields and members of the target type.
 258 fn gen_hash_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 259     fn gen_hash_call(target: ast::Expr) -> ast::Stmt {
 260         let method = make::name_ref("hash");
 261         let arg = make::expr_path(make::ext::ident_path("state"));
 262         let expr = make::expr_method_call(target, method, make::arg_list(Some(arg)));
 263         let stmt = make::expr_stmt(expr);
 264         stmt.into()
 265     }
 266
 267     let body = match adt {
 268         // `Hash` cannot be derived for unions, so no default impl can be provided.
 269         ast::Adt::Union(_) => return None,
 270
 271         // => std::mem::discriminant(self).hash(state);
 272         ast::Adt::Enum(_) => {
 273             let root = make::ext::ident_path("core");
 274             let submodule = make::ext::ident_path("mem");
 275             let fn_name = make::ext::ident_path("discriminant");
 276             let fn_name = make::path_concat(submodule, fn_name);
 277             let fn_name = make::expr_path(make::path_concat(root, fn_name));
 278
 279             let arg = make::expr_path(make::ext::ident_path("self"));
 280             let fn_call = make::expr_call(fn_name, make::arg_list(Some(arg)));
 281             let stmt = gen_hash_call(fn_call);
 282
 283             make::block_expr(Some(stmt), None).indent(ast::edit::IndentLevel(1))
 284         }
 285         ast::Adt::Struct(strukt) => match strukt.field_list() {
 286             // => self.<field>.hash(state);
 287             Some(ast::FieldList::RecordFieldList(field_list)) => {
 288                 let mut stmts = vec![];
 289                 for field in field_list.fields() {
 290                     let base = make::expr_path(make::ext::ident_path("self"));
 291                     let target = make::expr_field(base, &field.name()?.to_string());
 292                     stmts.push(gen_hash_call(target));
 293                 }
 294                 make::block_expr(stmts, None).indent(ast::edit::IndentLevel(1))
 295             }
 296
 297             // => self.<field_index>.hash(state);
 298             Some(ast::FieldList::TupleFieldList(field_list)) => {
 299                 let mut stmts = vec![];
 300                 for (i, _) in field_list.fields().enumerate() {
 301                     let base = make::expr_path(make::ext::ident_path("self"));
 302                     let target = make::expr_field(base, &format!("{}", i));
 303                     stmts.push(gen_hash_call(target));
 304                 }
 305                 make::block_expr(stmts, None).indent(ast::edit::IndentLevel(1))
 306             }
 307
 308             // No fields in the body means there's nothing to hash.
 309             None => return None,
 310         },
 311     };
 312
 313     ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 314     Some(())
 315 }