crates/ra_assists/src/handlers/fill_match_arms.rs

   1 //! FIXME: write short doc here
   2
   3 use std::iter;
   4
   5 use hir::{db::HirDatabase, Adt, HasSource};
   6 use ra_syntax::ast::{self, edit::IndentLevel, make, AstNode, NameOwner};
   7
   8 use crate::{Assist, AssistCtx, AssistId};
   9
  10 // Assist: fill_match_arms
  11 //
  12 // Adds missing clauses to a `match` expression.
  13 //
  14 // ```
  15 // enum Action { Move { distance: u32 }, Stop }
  16 //
  17 // fn handle(action: Action) {
  18 //     match action {
  19 //         <|>
  20 //     }
  21 // }
  22 // ```
  23 // ->
  24 // ```
  25 // enum Action { Move { distance: u32 }, Stop }
  26 //
  27 // fn handle(action: Action) {
  28 //     match action {
  29 //         Action::Move { distance } => (),
  30 //         Action::Stop => (),
  31 //     }
  32 // }
  33 // ```
  34 pub(crate) fn fill_match_arms(ctx: AssistCtx) -> Option<Assist> {
  35     let match_expr = ctx.find_node_at_offset::<ast::MatchExpr>()?;
  36     let match_arm_list = match_expr.match_arm_list()?;
  37
  38     // We already have some match arms, so we don't provide any assists.
  39     // Unless if there is only one trivial match arm possibly created
  40     // by match postfix complete. Trivial match arm is the catch all arm.
  41     let mut existing_arms = match_arm_list.arms();
  42     if let Some(arm) = existing_arms.next() {
  43         if !is_trivial(&arm) || existing_arms.next().is_some() {
  44             return None;
  45         }
  46     };
  47
  48     let expr = match_expr.expr()?;
  49     let (enum_def, module) = {
  50         let analyzer = ctx.source_analyzer(expr.syntax(), None);
  51         (resolve_enum_def(ctx.db, &analyzer, &expr)?, analyzer.module()?)
  52     };
  53     let variants = enum_def.variants(ctx.db);
  54     if variants.is_empty() {
  55         return None;
  56     }
  57
  58     let db = ctx.db;
  59
  60     ctx.add_assist(AssistId("fill_match_arms"), "Fill match arms", |edit| {
  61         let indent_level = IndentLevel::from_node(match_arm_list.syntax());
  62
  63         let new_arm_list = {
  64             let arms = variants
  65                 .into_iter()
  66                 .filter_map(|variant| build_pat(db, module, variant))
  67                 .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()));
  68             indent_level.increase_indent(make::match_arm_list(arms))
  69         };
  70
  71         edit.target(match_expr.syntax().text_range());
  72         edit.set_cursor(expr.syntax().text_range().start());
  73         edit.replace_ast(match_arm_list, new_arm_list);
  74     })
  75 }
  76
  77 fn is_trivial(arm: &ast::MatchArm) -> bool {
  78     match arm.pat() {
  79         Some(ast::Pat::PlaceholderPat(..)) => true,
  80         _ => false,
  81     }
  82 }
  83
  84 fn resolve_enum_def(
  85     db: &impl HirDatabase,
  86     analyzer: &hir::SourceAnalyzer,
  87     expr: &ast::Expr,
  88 ) -> Option<hir::Enum> {
  89     let expr_ty = analyzer.type_of(db, &expr)?;
  90
  91     let result = expr_ty.autoderef(db).find_map(|ty| match ty.as_adt() {
  92         Some(Adt::Enum(e)) => Some(e),
  93         _ => None,
  94     });
  95     result
  96 }
  97
  98 fn build_pat(
  99     db: &impl HirDatabase,
 100     module: hir::Module,
 101     var: hir::EnumVariant,
 102 ) -> Option<ast::Pat> {
 103     let path = crate::ast_transform::path_to_ast(module.find_use_path(db, var.into())?);
 104
 105     // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
 106     let pat: ast::Pat = match var.source(db).value.kind() {
 107         ast::StructKind::Tuple(field_list) => {
 108             let pats =
 109                 iter::repeat(make::placeholder_pat().into()).take(field_list.fields().count());
 110             make::tuple_struct_pat(path, pats).into()
 111         }
 112         ast::StructKind::Record(field_list) => {
 113             let pats = field_list.fields().map(|f| make::bind_pat(f.name().unwrap()).into());
 114             make::record_pat(path, pats).into()
 115         }
 116         ast::StructKind::Unit => make::path_pat(path),
 117     };
 118
 119     Some(pat)
 120 }
 121
 122 #[cfg(test)]
 123 mod tests {
 124     use crate::helpers::{check_assist, check_assist_target};
 125
 126     use super::fill_match_arms;
 127
 128     #[test]
 129     fn fill_match_arms_empty_body() {
 130         check_assist(
 131             fill_match_arms,
 132             r#"
 133             enum A {
 134                 As,
 135                 Bs,
 136                 Cs(String),
 137                 Ds(String, String),
 138                 Es{ x: usize, y: usize }
 139             }
 140
 141             fn main() {
 142                 let a = A::As;
 143                 match a<|> {}
 144             }
 145             "#,
 146             r#"
 147             enum A {
 148                 As,
 149                 Bs,
 150                 Cs(String),
 151                 Ds(String, String),
 152                 Es{ x: usize, y: usize }
 153             }
 154
 155             fn main() {
 156                 let a = A::As;
 157                 match <|>a {
 158                     A::As => (),
 159                     A::Bs => (),
 160                     A::Cs(_) => (),
 161                     A::Ds(_, _) => (),
 162                     A::Es { x, y } => (),
 163                 }
 164             }
 165             "#,
 166         );
 167     }
 168
 169     #[test]
 170     fn test_fill_match_arm_refs() {
 171         check_assist(
 172             fill_match_arms,
 173             r#"
 174             enum A {
 175                 As,
 176             }
 177
 178             fn foo(a: &A) {
 179                 match a<|> {
 180                 }
 181             }
 182             "#,
 183             r#"
 184             enum A {
 185                 As,
 186             }
 187
 188             fn foo(a: &A) {
 189                 match <|>a {
 190                     A::As => (),
 191                 }
 192             }
 193             "#,
 194         );
 195
 196         check_assist(
 197             fill_match_arms,
 198             r#"
 199             enum A {
 200                 Es{ x: usize, y: usize }
 201             }
 202
 203             fn foo(a: &mut A) {
 204                 match a<|> {
 205                 }
 206             }
 207             "#,
 208             r#"
 209             enum A {
 210                 Es{ x: usize, y: usize }
 211             }
 212
 213             fn foo(a: &mut A) {
 214                 match <|>a {
 215                     A::Es { x, y } => (),
 216                 }
 217             }
 218             "#,
 219         );
 220     }
 221
 222     #[test]
 223     fn fill_match_arms_target() {
 224         check_assist_target(
 225             fill_match_arms,
 226             r#"
 227             enum E { X, Y }
 228
 229             fn main() {
 230                 match E::X<|> {}
 231             }
 232             "#,
 233             "match E::X {}",
 234         );
 235     }
 236
 237     #[test]
 238     fn fill_match_arms_trivial_arm() {
 239         check_assist(
 240             fill_match_arms,
 241             r#"
 242             enum E { X, Y }
 243
 244             fn main() {
 245                 match E::X {
 246                     <|>_ => {},
 247                 }
 248             }
 249             "#,
 250             r#"
 251             enum E { X, Y }
 252
 253             fn main() {
 254                 match <|>E::X {
 255                     E::X => (),
 256                     E::Y => (),
 257                 }
 258             }
 259             "#,
 260         );
 261     }
 262
 263     #[test]
 264     fn fill_match_arms_qualifies_path() {
 265         check_assist(
 266             fill_match_arms,
 267             r#"
 268             mod foo { pub enum E { X, Y } }
 269             use foo::E::X;
 270
 271             fn main() {
 272                 match X {
 273                     <|>
 274                 }
 275             }
 276             "#,
 277             r#"
 278             mod foo { pub enum E { X, Y } }
 279             use foo::E::X;
 280
 281             fn main() {
 282                 match <|>X {
 283                     X => (),
 284                     foo::E::Y => (),
 285                 }
 286             }
 287             "#,
 288         );
 289     }
 290 }