crates/ra_assists/src/handlers/merge_match_arms.rs

   1 use std::iter::successors;
   2
   3 use ra_syntax::{
   4     algo::neighbor,
   5     ast::{self, AstNode},
   6     Direction, TextSize,
   7 };
   8
   9 use crate::{Assist, AssistCtx, AssistId, 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 //         <|>Action::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(ctx: AssistCtx) -> Option<Assist> {
  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     enum CursorPos {
  45         InExpr(TextSize),
  46         InPat(TextSize),
  47     }
  48     let cursor_pos = ctx.frange.range.start();
  49     let cursor_pos = if current_expr.syntax().text_range().contains(cursor_pos) {
  50         CursorPos::InExpr(current_text_range.end() - cursor_pos)
  51     } else {
  52         CursorPos::InPat(cursor_pos)
  53     };
  54
  55     // We check if the following match arms match this one. We could, but don't,
  56     // compare to the previous match arm as well.
  57     let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
  58         .take_while(|arm| {
  59             if arm.guard().is_some() {
  60                 return false;
  61             }
  62             match arm.expr() {
  63                 Some(expr) => expr.syntax().text() == current_expr.syntax().text(),
  64                 None => false,
  65             }
  66         })
  67         .collect::<Vec<_>>();
  68
  69     if arms_to_merge.len() <= 1 {
  70         return None;
  71     }
  72
  73     ctx.add_assist(AssistId("merge_match_arms"), "Merge match arms", current_text_range, |edit| {
  74         let pats = if arms_to_merge.iter().any(contains_placeholder) {
  75             "_".into()
  76         } else {
  77             arms_to_merge
  78                 .iter()
  79                 .filter_map(ast::MatchArm::pat)
  80                 .map(|x| x.syntax().to_string())
  81                 .collect::<Vec<String>>()
  82                 .join(" | ")
  83         };
  84
  85         let arm = format!("{} => {}", pats, current_expr.syntax().text());
  86
  87         let start = arms_to_merge.first().unwrap().syntax().text_range().start();
  88         let end = arms_to_merge.last().unwrap().syntax().text_range().end();
  89
  90         edit.set_cursor(match cursor_pos {
  91             CursorPos::InExpr(back_offset) => start + TextSize::of(&arm) - back_offset,
  92             CursorPos::InPat(offset) => offset,
  93         });
  94         edit.replace(TextRange::new(start, end), arm);
  95     })
  96 }
  97
  98 fn contains_placeholder(a: &ast::MatchArm) -> bool {
  99     match a.pat() {
 100         Some(ra_syntax::ast::Pat::PlaceholderPat(..)) => true,
 101         _ => false,
 102     }
 103 }
 104
 105 #[cfg(test)]
 106 mod tests {
 107     use crate::tests::{check_assist, check_assist_not_applicable};
 108
 109     use super::*;
 110
 111     #[test]
 112     fn merge_match_arms_single_patterns() {
 113         check_assist(
 114             merge_match_arms,
 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 => { 1i32<|> }
 123                     X::B => { 1i32 }
 124                     X::C => { 2i32 }
 125                 }
 126             }
 127             "#,
 128             r#"
 129             #[derive(Debug)]
 130             enum X { A, B, C }
 131
 132             fn main() {
 133                 let x = X::A;
 134                 let y = match x {
 135                     X::A | X::B => { 1i32<|> }
 136                     X::C => { 2i32 }
 137                 }
 138             }
 139             "#,
 140         );
 141     }
 142
 143     #[test]
 144     fn merge_match_arms_multiple_patterns() {
 145         check_assist(
 146             merge_match_arms,
 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 => {<|> 1i32 },
 155                     X::C | X::D => { 1i32 },
 156                     X::E => { 2i32 },
 157                 }
 158             }
 159             "#,
 160             r#"
 161             #[derive(Debug)]
 162             enum X { A, B, C, D, E }
 163
 164             fn main() {
 165                 let x = X::A;
 166                 let y = match x {
 167                     X::A | X::B | X::C | X::D => {<|> 1i32 },
 168                     X::E => { 2i32 },
 169                 }
 170             }
 171             "#,
 172         );
 173     }
 174
 175     #[test]
 176     fn merge_match_arms_placeholder_pattern() {
 177         check_assist(
 178             merge_match_arms,
 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                     X::B => { 2i<|>32 },
 188                     _ => { 2i32 }
 189                 }
 190             }
 191             "#,
 192             r#"
 193             #[derive(Debug)]
 194             enum X { A, B, C, D, E }
 195
 196             fn main() {
 197                 let x = X::A;
 198                 let y = match x {
 199                     X::A => { 1i32 },
 200                     _ => { 2i<|>32 }
 201                 }
 202             }
 203             "#,
 204         );
 205     }
 206
 207     #[test]
 208     fn merges_all_subsequent_arms() {
 209         check_assist(
 210             merge_match_arms,
 211             r#"
 212             enum X { A, B, C, D, E }
 213
 214             fn main() {
 215                 match X::A {
 216                     X::A<|> => 92,
 217                     X::B => 92,
 218                     X::C => 92,
 219                     X::D => 62,
 220                     _ => panic!(),
 221                 }
 222             }
 223             "#,
 224             r#"
 225             enum X { A, B, C, D, E }
 226
 227             fn main() {
 228                 match X::A {
 229                     X::A<|> | X::B | X::C => 92,
 230                     X::D => 62,
 231                     _ => panic!(),
 232                 }
 233             }
 234             "#,
 235         )
 236     }
 237
 238     #[test]
 239     fn merge_match_arms_rejects_guards() {
 240         check_assist_not_applicable(
 241             merge_match_arms,
 242             r#"
 243             #[derive(Debug)]
 244             enum X {
 245                 A(i32),
 246                 B,
 247                 C
 248             }
 249
 250             fn main() {
 251                 let x = X::A;
 252                 let y = match x {
 253                     X::A(a) if a > 5 => { <|>1i32 },
 254                     X::B => { 1i32 },
 255                     X::C => { 2i32 }
 256                 }
 257             }
 258             "#,
 259         );
 260     }
 261 }