crates/ra_assists/src/fill_match_arms.rs

   1 use std::fmt::Write;
   2
   3 use hir::{
   4     AdtDef, Ty, FieldSource, source_binder,
   5     db::HirDatabase,
   6 };
   7 use ra_syntax::ast::{self, AstNode};
   8
   9 use crate::{AssistCtx, Assist};
  10
  11 pub(crate) fn fill_match_arms(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
  12     let match_expr = ctx.node_at_offset::<ast::MatchExpr>()?;
  13
  14     // We already have some match arms, so we don't provide any assists.
  15     match match_expr.match_arm_list() {
  16         Some(arm_list) if arm_list.arms().count() > 0 => {
  17             return None;
  18         }
  19         _ => {}
  20     }
  21
  22     let expr = match_expr.expr()?;
  23     let function =
  24         source_binder::function_from_child_node(ctx.db, ctx.frange.file_id, expr.syntax())?;
  25     let infer_result = function.infer(ctx.db);
  26     let syntax_mapping = function.body_syntax_mapping(ctx.db);
  27     let node_expr = syntax_mapping.node_expr(expr)?;
  28     let match_expr_ty = infer_result[node_expr].clone();
  29     let enum_def = match match_expr_ty {
  30         Ty::Adt { def_id: AdtDef::Enum(e), .. } => e,
  31         _ => return None,
  32     };
  33     let enum_name = enum_def.name(ctx.db)?;
  34     let db = ctx.db;
  35
  36     ctx.add_action("fill match arms", |edit| {
  37         let mut buf = format!("match {} {{\n", expr.syntax().text().to_string());
  38         let variants = enum_def.variants(db);
  39         for variant in variants {
  40             let name = match variant.name(db) {
  41                 Some(it) => it,
  42                 None => continue,
  43             };
  44             write!(&mut buf, "    {}::{}", enum_name, name.to_string()).unwrap();
  45
  46             let pat = variant
  47                 .fields(db)
  48                 .into_iter()
  49                 .map(|field| {
  50                     let name = field.name(db).to_string();
  51                     let (_, source) = field.source(db);
  52                     match source {
  53                         FieldSource::Named(_) => name,
  54                         FieldSource::Pos(_) => "_".to_string(),
  55                     }
  56                 })
  57                 .collect::<Vec<_>>();
  58
  59             match pat.first().map(|s| s.as_str()) {
  60                 Some("_") => write!(&mut buf, "({})", pat.join(", ")).unwrap(),
  61                 Some(_) => write!(&mut buf, "{{{}}}", pat.join(", ")).unwrap(),
  62                 None => (),
  63             };
  64
  65             buf.push_str(" => (),\n");
  66         }
  67         buf.push_str("}");
  68         edit.target(match_expr.syntax().range());
  69         edit.set_cursor(expr.syntax().range().start());
  70         edit.replace_node_and_indent(match_expr.syntax(), buf);
  71     });
  72
  73     ctx.build()
  74 }
  75
  76 #[cfg(test)]
  77 mod tests {
  78     use crate::helpers::{check_assist, check_assist_target};
  79
  80     use super::fill_match_arms;
  81
  82     #[test]
  83     fn fill_match_arms_empty_body() {
  84         check_assist(
  85             fill_match_arms,
  86             r#"
  87             enum A {
  88                 As,
  89                 Bs,
  90                 Cs(String),
  91                 Ds(String, String),
  92                 Es{x: usize, y: usize}
  93             }
  94
  95             fn main() {
  96                 let a = A::As;
  97                 match a<|> {}
  98             }
  99             "#,
 100             r#"
 101             enum A {
 102                 As,
 103                 Bs,
 104                 Cs(String),
 105                 Ds(String, String),
 106                 Es{x: usize, y: usize}
 107             }
 108
 109             fn main() {
 110                 let a = A::As;
 111                 match <|>a {
 112                     A::As => (),
 113                     A::Bs => (),
 114                     A::Cs(_) => (),
 115                     A::Ds(_, _) => (),
 116                     A::Es{x, y} => (),
 117                 }
 118             }
 119             "#,
 120         );
 121     }
 122     #[test]
 123     fn fill_match_arms_no_body() {
 124         check_assist(
 125             fill_match_arms,
 126             r#"
 127             enum E { X, Y}
 128
 129             fn main() {
 130                 match E::X<|>
 131             }
 132             "#,
 133             r#"
 134             enum E { X, Y}
 135
 136             fn main() {
 137                 match <|>E::X {
 138                     E::X => (),
 139                     E::Y => (),
 140                 }
 141             }
 142             "#,
 143         );
 144     }
 145
 146     #[test]
 147     fn fill_match_arms_target() {
 148         check_assist_target(
 149             fill_match_arms,
 150             r#"
 151             enum E { X, Y}
 152
 153             fn main() {
 154                 match E::X<|> {}
 155             }
 156             "#,
 157             "match E::X {}",
 158         );
 159     }
 160 }