crates/ide_assists/src/handlers/merge_match_arms.rs

   1 use itertools::Itertools;
   2 use std::iter::successors;
   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     let current_arm_types = get_arm_types(&ctx, &current_arm);
  44
  45     // We check if the following match arms match this one. We could, but don't,
  46     // compare to the previous match arm as well.
  47     let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
  48         .take_while(|arm| match arm.expr() {
  49             Some(expr) if arm.guard().is_none() => {
  50                 let same_text = expr.syntax().text() == current_expr.syntax().text();
  51                 if !same_text {
  52                     return false;
  53                 }
  54
  55                 return are_same_types(&current_arm_types, arm, ctx);
  56             }
  57             _ => false,
  58         })
  59         .collect::<Vec<_>>();
  60
  61     if arms_to_merge.len() <= 1 {
  62         return None;
  63     }
  64
  65     acc.add(
  66         AssistId("merge_match_arms", AssistKind::RefactorRewrite),
  67         "Merge match arms",
  68         current_text_range,
  69         |edit| {
  70             let pats = if arms_to_merge.iter().any(contains_placeholder) {
  71                 "_".into()
  72             } else {
  73                 arms_to_merge
  74                     .iter()
  75                     .filter_map(ast::MatchArm::pat)
  76                     .map(|x| x.syntax().to_string())
  77                     .collect::<Vec<String>>()
  78                     .join(" | ")
  79             };
  80
  81             let arm = format!("{} => {},", pats, current_expr.syntax().text());
  82
  83             if let [first, .., last] = &*arms_to_merge {
  84                 let start = first.syntax().text_range().start();
  85                 let end = last.syntax().text_range().end();
  86
  87                 edit.replace(TextRange::new(start, end), arm);
  88             }
  89         },
  90     )
  91 }
  92
  93 fn contains_placeholder(a: &ast::MatchArm) -> bool {
  94     matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
  95 }
  96
  97 fn are_same_types(
  98     current_arm_types: &Vec<Option<hir::TypeInfo>>,
  99     arm: &ast::MatchArm,
 100     ctx: &AssistContext,
 101 ) -> bool {
 102     let arm_types = get_arm_types(&ctx, &arm);
 103     for i in 0..arm_types.len() {
 104         let other_arm_type = &arm_types[i];
 105         let current_arm_type = &current_arm_types[i];
 106         if let (Some(other_arm_type), Some(current_arm_type)) = (other_arm_type, current_arm_type) {
 107             return &other_arm_type.original == &current_arm_type.original;
 108         }
 109     }
 110
 111     return true;
 112 }
 113
 114 fn get_arm_types(ctx: &AssistContext, arm: &ast::MatchArm) -> Vec<Option<hir::TypeInfo>> {
 115     match arm.pat() {
 116         Some(ast::Pat::TupleStructPat(tp)) => {
 117             tp.fields().into_iter().map(|field| ctx.sema.type_of_pat(&field)).collect_vec()
 118         }
 119         _ => Vec::new(),
 120     }
 121 }
 122
 123 #[cfg(test)]
 124 mod tests {
 125     use crate::tests::{check_assist, check_assist_not_applicable};
 126
 127     use super::*;
 128
 129     #[test]
 130     fn merge_match_arms_single_patterns() {
 131         check_assist(
 132             merge_match_arms,
 133             r#"
 134 #[derive(Debug)]
 135 enum X { A, B, C }
 136
 137 fn main() {
 138     let x = X::A;
 139     let y = match x {
 140         X::A => { 1i32$0 }
 141         X::B => { 1i32 }
 142         X::C => { 2i32 }
 143     }
 144 }
 145 "#,
 146             r#"
 147 #[derive(Debug)]
 148 enum X { A, B, C }
 149
 150 fn main() {
 151     let x = X::A;
 152     let y = match x {
 153         X::A | X::B => { 1i32 },
 154         X::C => { 2i32 }
 155     }
 156 }
 157 "#,
 158         );
 159     }
 160
 161     #[test]
 162     fn merge_match_arms_multiple_patterns() {
 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 | X::B => {$0 1i32 },
 173         X::C | X::D => { 1i32 },
 174         X::E => { 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 | X::B | X::C | X::D => { 1i32 },
 186         X::E => { 2i32 },
 187     }
 188 }
 189 "#,
 190         );
 191     }
 192
 193     #[test]
 194     fn merge_match_arms_placeholder_pattern() {
 195         check_assist(
 196             merge_match_arms,
 197             r#"
 198 #[derive(Debug)]
 199 enum X { A, B, C, D, E }
 200
 201 fn main() {
 202     let x = X::A;
 203     let y = match x {
 204         X::A => { 1i32 },
 205         X::B => { 2i$032 },
 206         _ => { 2i32 }
 207     }
 208 }
 209 "#,
 210             r#"
 211 #[derive(Debug)]
 212 enum X { A, B, C, D, E }
 213
 214 fn main() {
 215     let x = X::A;
 216     let y = match x {
 217         X::A => { 1i32 },
 218         _ => { 2i32 },
 219     }
 220 }
 221 "#,
 222         );
 223     }
 224
 225     #[test]
 226     fn merges_all_subsequent_arms() {
 227         check_assist(
 228             merge_match_arms,
 229             r#"
 230 enum X { A, B, C, D, E }
 231
 232 fn main() {
 233     match X::A {
 234         X::A$0 => 92,
 235         X::B => 92,
 236         X::C => 92,
 237         X::D => 62,
 238         _ => panic!(),
 239     }
 240 }
 241 "#,
 242             r#"
 243 enum X { A, B, C, D, E }
 244
 245 fn main() {
 246     match X::A {
 247         X::A | X::B | X::C => 92,
 248         X::D => 62,
 249         _ => panic!(),
 250     }
 251 }
 252 "#,
 253         )
 254     }
 255
 256     #[test]
 257     fn merge_match_arms_rejects_guards() {
 258         check_assist_not_applicable(
 259             merge_match_arms,
 260             r#"
 261 #[derive(Debug)]
 262 enum X {
 263     A(i32),
 264     B,
 265     C
 266 }
 267
 268 fn main() {
 269     let x = X::A;
 270     let y = match x {
 271         X::A(a) if a > 5 => { $01i32 },
 272         X::B => { 1i32 },
 273         X::C => { 2i32 }
 274     }
 275 }
 276 "#,
 277         );
 278     }
 279
 280     #[test]
 281     fn merge_match_arms_different_type() {
 282         check_assist_not_applicable(
 283             merge_match_arms,
 284             r#"//- minicore: result
 285 fn func() {
 286     match Result::<f64, f32>::Ok(0f64) {
 287         Ok(x) => $0x.classify(),
 288         Err(x) => x.classify()
 289     };
 290 }
 291 "#,
 292         );
 293     }
 294
 295     #[test]
 296     fn merge_match_arms_different_type_multiple_fields() {
 297         check_assist_not_applicable(
 298             merge_match_arms,
 299             r#"//- minicore: result
 300 fn func() {
 301     match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
 302         Ok(x) => $0x.1.classify(),
 303         Err(x) => x.1.classify()
 304     };
 305 }
 306 "#,
 307         );
 308     }
 309
 310     #[test]
 311     fn merge_match_arms_same_type_multiple_fields() {
 312         check_assist(
 313             merge_match_arms,
 314             r#"//- minicore: result
 315 fn func() {
 316     match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
 317         Ok(x) => $0x.1.classify(),
 318         Err(x) => x.1.classify()
 319     };
 320 }
 321 "#,
 322             r#"
 323 fn func() {
 324     match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
 325         Ok(x) | Err(x) => x.1.classify(),
 326     };
 327 }
 328 "#,
 329         );
 330     }
 331
 332     #[test]
 333     fn merge_match_arms_same_type_subsequent_arm_with_different_type_in_other() {
 334         check_assist(
 335             merge_match_arms,
 336             r#"
 337 enum MyEnum {
 338     OptionA(f32),
 339     OptionB(f32),
 340     OptionC(f64)
 341 }
 342
 343 fn func(e: MyEnum) {
 344     match e {
 345         MyEnum::OptionA(x) => $0x.classify(),
 346         MyEnum::OptionB(x) => x.classify(),
 347         MyEnum::OptionC(x) => x.classify(),
 348     };
 349 }
 350 "#,
 351             r#"
 352 enum MyEnum {
 353     OptionA(f32),
 354     OptionB(f32),
 355     OptionC(f64)
 356 }
 357
 358 fn func(e: MyEnum) {
 359     match e {
 360         MyEnum::OptionA(x) | MyEnum::OptionB(x) => x.classify(),
 361         MyEnum::OptionC(x) => x.classify(),
 362     };
 363 }
 364 "#,
 365         );
 366     }
 367
 368     #[test]
 369     fn merge_match_arms_same_type_skip_arm_with_different_type_in_between() {
 370         check_assist_not_applicable(
 371             merge_match_arms,
 372             r#"
 373 enum MyEnum {
 374     OptionA(f32),
 375     OptionB(f64),
 376     OptionC(f32)
 377 }
 378
 379 fn func(e: MyEnum) {
 380     match e {
 381         MyEnum::OptionA(x) => $0x.classify(),
 382         MyEnum::OptionB(x) => x.classify(),
 383         MyEnum::OptionC(x) => x.classify(),
 384     };
 385 }
 386 "#,
 387         );
 388     }
 389 }