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(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.build("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
  74 #[cfg(test)]
  75 mod tests {
  76     use crate::helpers::{check_assist, check_assist_target};
  77
  78     use super::fill_match_arms;
  79
  80     #[test]
  81     fn fill_match_arms_empty_body() {
  82         check_assist(
  83             fill_match_arms,
  84             r#"
  85             enum A {
  86                 As,
  87                 Bs,
  88                 Cs(String),
  89                 Ds(String, String),
  90                 Es{x: usize, y: usize}
  91             }
  92
  93             fn main() {
  94                 let a = A::As;
  95                 match a<|> {}
  96             }
  97             "#,
  98             r#"
  99             enum A {
 100                 As,
 101                 Bs,
 102                 Cs(String),
 103                 Ds(String, String),
 104                 Es{x: usize, y: usize}
 105             }
 106
 107             fn main() {
 108                 let a = A::As;
 109                 match <|>a {
 110                     A::As => (),
 111                     A::Bs => (),
 112                     A::Cs(_) => (),
 113                     A::Ds(_, _) => (),
 114                     A::Es{x, y} => (),
 115                 }
 116             }
 117             "#,
 118         );
 119     }
 120     #[test]
 121     fn fill_match_arms_no_body() {
 122         check_assist(
 123             fill_match_arms,
 124             r#"
 125             enum E { X, Y}
 126
 127             fn main() {
 128                 match E::X<|>
 129             }
 130             "#,
 131             r#"
 132             enum E { X, Y}
 133
 134             fn main() {
 135                 match <|>E::X {
 136                     E::X => (),
 137                     E::Y => (),
 138                 }
 139             }
 140             "#,
 141         );
 142     }
 143
 144     #[test]
 145     fn fill_match_arms_target() {
 146         check_assist_target(
 147             fill_match_arms,
 148             r#"
 149             enum E { X, Y}
 150
 151             fn main() {
 152                 match E::X<|> {}
 153             }
 154             "#,
 155             "match E::X {}",
 156         );
 157     }
 158 }