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 the current match arm with the following if their bodies are identical.
  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| match arm.expr() {
  48             Some(expr) if arm.guard().is_none() => {
  49                 expr.syntax().text() == current_expr.syntax().text()
  50             }
  51             _ => false,
  52         })
  53         .collect::<Vec<_>>();
  54
  55     if arms_to_merge.len() <= 1 {
  56         return None;
  57     }
  58
  59     acc.add(
  60         AssistId("merge_match_arms", AssistKind::RefactorRewrite),
  61         "Merge match arms",
  62         current_text_range,
  63         |edit| {
  64             let pats = if arms_to_merge.iter().any(contains_placeholder) {
  65                 "_".into()
  66             } else {
  67                 arms_to_merge
  68                     .iter()
  69                     .filter_map(ast::MatchArm::pat)
  70                     .map(|x| x.syntax().to_string())
  71                     .collect::<Vec<String>>()
  72                     .join(" | ")
  73             };
  74
  75             let arm = format!("{} => {}", pats, current_expr.syntax().text());
  76
  77             if let [first, .., last] = &*arms_to_merge {
  78                 let start = first.syntax().text_range().start();
  79                 let end = last.syntax().text_range().end();
  80
  81                 edit.replace(TextRange::new(start, end), arm);
  82             }
  83         },
  84     )
  85 }
  86
  87 fn contains_placeholder(a: &ast::MatchArm) -> bool {
  88     matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
  89 }
  90
  91 #[cfg(test)]
  92 mod tests {
  93     use crate::tests::{check_assist, check_assist_not_applicable};
  94
  95     use super::*;
  96
  97     #[test]
  98     fn merge_match_arms_single_patterns() {
  99         check_assist(
 100             merge_match_arms,
 101             r#"
 102 #[derive(Debug)]
 103 enum X { A, B, C }
 104
 105 fn main() {
 106     let x = X::A;
 107     let y = match x {
 108         X::A => { 1i32$0 }
 109         X::B => { 1i32 }
 110         X::C => { 2i32 }
 111     }
 112 }
 113 "#,
 114             r#"
 115 #[derive(Debug)]
 116 enum X { A, B, C }
 117
 118 fn main() {
 119     let x = X::A;
 120     let y = match x {
 121         X::A | X::B => { 1i32 }
 122         X::C => { 2i32 }
 123     }
 124 }
 125 "#,
 126         );
 127     }
 128
 129     #[test]
 130     fn merge_match_arms_multiple_patterns() {
 131         check_assist(
 132             merge_match_arms,
 133             r#"
 134 #[derive(Debug)]
 135 enum X { A, B, C, D, E }
 136
 137 fn main() {
 138     let x = X::A;
 139     let y = match x {
 140         X::A | X::B => {$0 1i32 },
 141         X::C | X::D => { 1i32 },
 142         X::E => { 2i32 },
 143     }
 144 }
 145 "#,
 146             r#"
 147 #[derive(Debug)]
 148 enum X { A, B, C, D, E }
 149
 150 fn main() {
 151     let x = X::A;
 152     let y = match x {
 153         X::A | X::B | X::C | X::D => { 1i32 },
 154         X::E => { 2i32 },
 155     }
 156 }
 157 "#,
 158         );
 159     }
 160
 161     #[test]
 162     fn merge_match_arms_placeholder_pattern() {
 163         check_assist(
 164             merge_match_arms,
 165             r#"
 166 #[derive(Debug)]
 167 enum X { A, B, C, D, E }
 168
 169 fn main() {
 170     let x = X::A;
 171     let y = match x {
 172         X::A => { 1i32 },
 173         X::B => { 2i$032 },
 174         _ => { 2i32 }
 175     }
 176 }
 177 "#,
 178             r#"
 179 #[derive(Debug)]
 180 enum X { A, B, C, D, E }
 181
 182 fn main() {
 183     let x = X::A;
 184     let y = match x {
 185         X::A => { 1i32 },
 186         _ => { 2i32 }
 187     }
 188 }
 189 "#,
 190         );
 191     }
 192
 193     #[test]
 194     fn merges_all_subsequent_arms() {
 195         check_assist(
 196             merge_match_arms,
 197             r#"
 198 enum X { A, B, C, D, E }
 199
 200 fn main() {
 201     match X::A {
 202         X::A$0 => 92,
 203         X::B => 92,
 204         X::C => 92,
 205         X::D => 62,
 206         _ => panic!(),
 207     }
 208 }
 209 "#,
 210             r#"
 211 enum X { A, B, C, D, E }
 212
 213 fn main() {
 214     match X::A {
 215         X::A | X::B | X::C => 92,
 216         X::D => 62,
 217         _ => panic!(),
 218     }
 219 }
 220 "#,
 221         )
 222     }
 223
 224     #[test]
 225     fn merge_match_arms_rejects_guards() {
 226         check_assist_not_applicable(
 227             merge_match_arms,
 228             r#"
 229 #[derive(Debug)]
 230 enum X {
 231     A(i32),
 232     B,
 233     C
 234 }
 235
 236 fn main() {
 237     let x = X::A;
 238     let y = match x {
 239         X::A(a) if a > 5 => { $01i32 },
 240         X::B => { 1i32 },
 241         X::C => { 2i32 }
 242     }
 243 }
 244 "#,
 245         );
 246     }
 247 }