crates/ide_assists/src/handlers/merge_match_arms.rs

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