crates/ide_assists/src/handlers/merge_match_arms.rs

   1 use std::iter::successors;
   2
   3 use syntax::{
   4     algo::neighbor,
   5     ast::{self, AstNode},
   6     Direction,
   7 };
   8
   9 use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};
  10
  11 // Assist: merge_match_arms
  12 //
  13 // Merges identical match arms.
  14 //
  15 // ```
  16 // enum Action { Move { distance: u32 }, Stop }
  17 //
  18 // fn handle(action: Action) {
  19 //     match action {
  20 //         $0Action::Move(..) => foo(),
  21 //         Action::Stop => foo(),
  22 //     }
  23 // }
  24 // ```
  25 // ->
  26 // ```
  27 // enum Action { Move { distance: u32 }, Stop }
  28 //
  29 // fn handle(action: Action) {
  30 //     match action {
  31 //         Action::Move(..) | Action::Stop => foo(),
  32 //     }
  33 // }
  34 // ```
  35 pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  36     let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
  37     // Don't try to handle arms with guards for now - can add support for this later
  38     if current_arm.guard().is_some() {
  39         return None;
  40     }
  41     let current_expr = current_arm.expr()?;
  42     let current_text_range = current_arm.syntax().text_range();
  43
  44     // We check if the following match arms match this one. We could, but don't,
  45     // compare to the previous match arm as well.
  46     let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
  47         .take_while(|arm| {
  48             if arm.guard().is_some() {
  49                 return false;
  50             }
  51             match arm.expr() {
  52                 Some(expr) => expr.syntax().text() == current_expr.syntax().text(),
  53                 None => false,
  54             }
  55         })
  56         .collect::<Vec<_>>();
  57
  58     if arms_to_merge.len() <= 1 {
  59         return None;
  60     }
  61
  62     acc.add(
  63         AssistId("merge_match_arms", AssistKind::RefactorRewrite),
  64         "Merge match arms",
  65         current_text_range,
  66         |edit| {
  67             let pats = if arms_to_merge.iter().any(contains_placeholder) {
  68                 "_".into()
  69             } else {
  70                 arms_to_merge
  71                     .iter()
  72                     .filter_map(ast::MatchArm::pat)
  73                     .map(|x| x.syntax().to_string())
  74                     .collect::<Vec<String>>()
  75                     .join(" | ")
  76             };
  77
  78             let arm = format!("{} => {}", pats, current_expr.syntax().text());
  79
  80             let start = arms_to_merge.first().unwrap().syntax().text_range().start();
  81             let end = arms_to_merge.last().unwrap().syntax().text_range().end();
  82
  83             edit.replace(TextRange::new(start, end), arm);
  84         },
  85     )
  86 }
  87
  88 fn contains_placeholder(a: &ast::MatchArm) -> bool {
  89     matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
  90 }
  91
  92 #[cfg(test)]
  93 mod tests {
  94     use crate::tests::{check_assist, check_assist_not_applicable};
  95
  96     use super::*;
  97
  98     #[test]
  99     fn merge_match_arms_single_patterns() {
 100         check_assist(
 101             merge_match_arms,
 102             r#"
 103             #[derive(Debug)]
 104             enum X { A, B, C }
 105
 106             fn main() {
 107                 let x = X::A;
 108                 let y = match x {
 109                     X::A => { 1i32$0 }
 110                     X::B => { 1i32 }
 111                     X::C => { 2i32 }
 112                 }
 113             }
 114             "#,
 115             r#"
 116             #[derive(Debug)]
 117             enum X { A, B, C }
 118
 119             fn main() {
 120                 let x = X::A;
 121                 let y = match x {
 122                     X::A | X::B => { 1i32 }
 123                     X::C => { 2i32 }
 124                 }
 125             }
 126             "#,
 127         );
 128     }
 129
 130     #[test]
 131     fn merge_match_arms_multiple_patterns() {
 132         check_assist(
 133             merge_match_arms,
 134             r#"
 135             #[derive(Debug)]
 136             enum X { A, B, C, D, E }
 137
 138             fn main() {
 139                 let x = X::A;
 140                 let y = match x {
 141                     X::A | X::B => {$0 1i32 },
 142                     X::C | X::D => { 1i32 },
 143                     X::E => { 2i32 },
 144                 }
 145             }
 146             "#,
 147             r#"
 148             #[derive(Debug)]
 149             enum X { A, B, C, D, E }
 150
 151             fn main() {
 152                 let x = X::A;
 153                 let y = match x {
 154                     X::A | X::B | X::C | X::D => { 1i32 },
 155                     X::E => { 2i32 },
 156                 }
 157             }
 158             "#,
 159         );
 160     }
 161
 162     #[test]
 163     fn merge_match_arms_placeholder_pattern() {
 164         check_assist(
 165             merge_match_arms,
 166             r#"
 167             #[derive(Debug)]
 168             enum X { A, B, C, D, E }
 169
 170             fn main() {
 171                 let x = X::A;
 172                 let y = match x {
 173                     X::A => { 1i32 },
 174                     X::B => { 2i$032 },
 175                     _ => { 2i32 }
 176                 }
 177             }
 178             "#,
 179             r#"
 180             #[derive(Debug)]
 181             enum X { A, B, C, D, E }
 182
 183             fn main() {
 184                 let x = X::A;
 185                 let y = match x {
 186                     X::A => { 1i32 },
 187                     _ => { 2i32 }
 188                 }
 189             }
 190             "#,
 191         );
 192     }
 193
 194     #[test]
 195     fn merges_all_subsequent_arms() {
 196         check_assist(
 197             merge_match_arms,
 198             r#"
 199             enum X { A, B, C, D, E }
 200
 201             fn main() {
 202                 match X::A {
 203                     X::A$0 => 92,
 204                     X::B => 92,
 205                     X::C => 92,
 206                     X::D => 62,
 207                     _ => panic!(),
 208                 }
 209             }
 210             "#,
 211             r#"
 212             enum X { A, B, C, D, E }
 213
 214             fn main() {
 215                 match X::A {
 216                     X::A | X::B | X::C => 92,
 217                     X::D => 62,
 218                     _ => panic!(),
 219                 }
 220             }
 221             "#,
 222         )
 223     }
 224
 225     #[test]
 226     fn merge_match_arms_rejects_guards() {
 227         check_assist_not_applicable(
 228             merge_match_arms,
 229             r#"
 230             #[derive(Debug)]
 231             enum X {
 232                 A(i32),
 233                 B,
 234                 C
 235             }
 236
 237             fn main() {
 238                 let x = X::A;
 239                 let y = match x {
 240                     X::A(a) if a > 5 => { $01i32 },
 241                     X::B => { 1i32 },
 242                     X::C => { 2i32 }
 243                 }
 244             }
 245             "#,
 246         );
 247     }
 248 }