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         "Debug" => gen_debug_impl(adt, func),
  20         "Default" => gen_default_impl(adt, func),
  21         _ => None,
  22     }
  23 }
  24
  25 /// Generate a `Debug` impl based on the fields and members of the target type.
  26 fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
  27     let annotated_name = adt.name()?;
  28     match adt {
  29         // `Debug` cannot be derived for unions, so no default impl can be provided.
  30         ast::Adt::Union(_) => None,
  31
  32         // => match self { Self::Variant => write!(f, "Variant") }
  33         ast::Adt::Enum(enum_) => {
  34             let list = enum_.variant_list()?;
  35             let mut arms = vec![];
  36             for variant in list.variants() {
  37                 let name = variant.name()?;
  38                 let left = make::ext::ident_path("Self");
  39                 let right = make::ext::ident_path(&format!("{}", name));
  40                 let variant_name = make::path_pat(make::path_concat(left, right));
  41
  42                 let target = make::expr_path(make::ext::ident_path("f").into());
  43                 let fmt_string = make::expr_literal(&(format!("\"{}\"", name))).into();
  44                 let args = make::arg_list(vec![target, fmt_string]);
  45                 let macro_name = make::expr_path(make::ext::ident_path("write"));
  46                 let macro_call = make::expr_macro_call(macro_name, args);
  47
  48                 arms.push(make::match_arm(Some(variant_name.into()), None, macro_call.into()));
  49             }
  50
  51             let match_target = make::expr_path(make::ext::ident_path("self"));
  52             let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
  53             let match_expr = make::expr_match(match_target, list);
  54
  55             let body = make::block_expr(None, Some(match_expr));
  56             let body = body.indent(ast::edit::IndentLevel(1));
  57             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
  58             Some(())
  59         }
  60
  61         ast::Adt::Struct(strukt) => {
  62             let name = format!("\"{}\"", annotated_name);
  63             let args = make::arg_list(Some(make::expr_literal(&name).into()));
  64             let target = make::expr_path(make::ext::ident_path("f"));
  65
  66             let expr = match strukt.field_list() {
  67                 // => f.debug_struct("Name").finish()
  68                 None => make::expr_method_call(target, make::name_ref("debug_struct"), args),
  69
  70                 // => f.debug_struct("Name").field("foo", &self.foo).finish()
  71                 Some(ast::FieldList::RecordFieldList(field_list)) => {
  72                     let method = make::name_ref("debug_struct");
  73                     let mut expr = make::expr_method_call(target, method, args);
  74                     for field in field_list.fields() {
  75                         let name = field.name()?;
  76                         let f_name = make::expr_literal(&(format!("\"{}\"", name))).into();
  77                         let f_path = make::expr_path(make::ext::ident_path("self"));
  78                         let f_path = make::expr_ref(f_path, false);
  79                         let f_path = make::expr_field(f_path, &format!("{}", name)).into();
  80                         let args = make::arg_list(vec![f_name, f_path]);
  81                         expr = make::expr_method_call(expr, make::name_ref("field"), args);
  82                     }
  83                     expr
  84                 }
  85
  86                 // => f.debug_tuple("Name").field(self.0).finish()
  87                 Some(ast::FieldList::TupleFieldList(field_list)) => {
  88                     let method = make::name_ref("debug_tuple");
  89                     let mut expr = make::expr_method_call(target, method, args);
  90                     for (idx, _) in field_list.fields().enumerate() {
  91                         let f_path = make::expr_path(make::ext::ident_path("self"));
  92                         let f_path = make::expr_ref(f_path, false);
  93                         let f_path = make::expr_field(f_path, &format!("{}", idx)).into();
  94                         let method = make::name_ref("field");
  95                         expr = make::expr_method_call(expr, method, make::arg_list(Some(f_path)));
  96                     }
  97                     expr
  98                 }
  99             };
 100
 101             let method = make::name_ref("finish");
 102             let expr = make::expr_method_call(expr, method, make::arg_list(None));
 103             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 104             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 105             Some(())
 106         }
 107     }
 108 }
 109
 110 /// Generate a `Debug` impl based on the fields and members of the target type.
 111 fn gen_default_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 112     fn gen_default_call() -> ast::Expr {
 113         let trait_name = make::ext::ident_path("Default");
 114         let method_name = make::ext::ident_path("default");
 115         let fn_name = make::expr_path(make::path_concat(trait_name, method_name));
 116         make::expr_call(fn_name, make::arg_list(None))
 117     }
 118     match adt {
 119         // `Debug` cannot be derived for unions, so no default impl can be provided.
 120         ast::Adt::Union(_) => None,
 121         // Deriving `Debug` for enums is not stable yet.
 122         ast::Adt::Enum(_) => None,
 123         ast::Adt::Struct(strukt) => {
 124             let expr = match strukt.field_list() {
 125                 Some(ast::FieldList::RecordFieldList(field_list)) => {
 126                     let mut fields = vec![];
 127                     for field in field_list.fields() {
 128                         let method_call = gen_default_call();
 129                         let name_ref = make::name_ref(&field.name()?.to_string());
 130                         let field = make::record_expr_field(name_ref, Some(method_call));
 131                         fields.push(field);
 132                     }
 133                     let struct_name = make::ext::ident_path("Self");
 134                     let fields = make::record_expr_field_list(fields);
 135                     make::record_expr(struct_name, fields).into()
 136                 }
 137                 Some(ast::FieldList::TupleFieldList(field_list)) => {
 138                     let struct_name = make::expr_path(make::ext::ident_path("Self"));
 139                     let fields = field_list.fields().map(|_| gen_default_call());
 140                     make::expr_call(struct_name, make::arg_list(fields))
 141                 }
 142                 None => {
 143                     let struct_name = make::ext::ident_path("Self");
 144                     let fields = make::record_expr_field_list(None);
 145                     make::record_expr(struct_name, fields).into()
 146                 }
 147             };
 148             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
 149             ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
 150             Some(())
 151         }
 152     }
 153 }