src/tools/rust-analyzer/crates/ide-assists/src/handlers/unmerge_match_arm.rs

   1 use syntax::{
   2     algo::neighbor,
   3     ast::{self, edit::IndentLevel, make, AstNode},
   4     ted::{self, Position},
   5     Direction, SyntaxKind, T,
   6 };
   7
   8 use crate::{AssistContext, AssistId, AssistKind, Assists};
   9
  10 // Assist: unmerge_match_arm
  11 //
  12 // Splits the current match with a `|` pattern into two arms with identical bodies.
  13 //
  14 // ```
  15 // enum Action { Move { distance: u32 }, Stop }
  16 //
  17 // fn handle(action: Action) {
  18 //     match action {
  19 //         Action::Move(..) $0| Action::Stop => foo(),
  20 //     }
  21 // }
  22 // ```
  23 // ->
  24 // ```
  25 // enum Action { Move { distance: u32 }, Stop }
  26 //
  27 // fn handle(action: Action) {
  28 //     match action {
  29 //         Action::Move(..) => foo(),
  30 //         Action::Stop => foo(),
  31 //     }
  32 // }
  33 // ```
  34 pub(crate) fn unmerge_match_arm(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  35     let pipe_token = ctx.find_token_syntax_at_offset(T![|])?;
  36     let or_pat = ast::OrPat::cast(pipe_token.parent()?)?.clone_for_update();
  37     let match_arm = ast::MatchArm::cast(or_pat.syntax().parent()?)?;
  38     let match_arm_body = match_arm.expr()?;
  39
  40     // We don't need to check for leading pipe because it is directly under `MatchArm`
  41     // without `OrPat`.
  42
  43     let new_parent = match_arm.syntax().parent()?;
  44     let old_parent_range = new_parent.text_range();
  45
  46     acc.add(
  47         AssistId("unmerge_match_arm", AssistKind::RefactorRewrite),
  48         "Unmerge match arm",
  49         pipe_token.text_range(),
  50         |edit| {
  51             let pats_after = pipe_token
  52                 .siblings_with_tokens(Direction::Next)
  53                 .filter_map(|it| ast::Pat::cast(it.into_node()?));
  54             // FIXME: We should add a leading pipe if the original arm has one.
  55             let new_match_arm = make::match_arm(
  56                 pats_after,
  57                 match_arm.guard().and_then(|guard| guard.condition()),
  58                 match_arm_body,
  59             )
  60             .clone_for_update();
  61
  62             let mut pipe_index = pipe_token.index();
  63             if pipe_token
  64                 .prev_sibling_or_token()
  65                 .map_or(false, |it| it.kind() == SyntaxKind::WHITESPACE)
  66             {
  67                 pipe_index -= 1;
  68             }
  69             or_pat.syntax().splice_children(
  70                 pipe_index..or_pat.syntax().children_with_tokens().count(),
  71                 Vec::new(),
  72             );
  73
  74             let mut insert_after_old_arm = Vec::new();
  75
  76             // A comma can be:
  77             //  - After the arm. In this case we always want to insert a comma after the newly
  78             //    inserted arm.
  79             //  - Missing after the arm, with no arms after. In this case we want to insert a
  80             //    comma before the newly inserted arm. It can not be necessary if there arm
  81             //    body is a block, but we don't bother to check that.
  82             //  - Missing after the arm with arms after, if the arm body is a block. In this case
  83             //    we don't want to insert a comma at all.
  84             let has_comma_after =
  85                 std::iter::successors(match_arm.syntax().last_child_or_token(), |it| {
  86                     it.prev_sibling_or_token()
  87                 })
  88                 .map(|it| it.kind())
  89                 .skip_while(|it| it.is_trivia())
  90                 .next()
  91                     == Some(T![,]);
  92             let has_arms_after = neighbor(&match_arm, Direction::Next).is_some();
  93             if !has_comma_after && !has_arms_after {
  94                 insert_after_old_arm.push(make::token(T![,]).into());
  95             }
  96
  97             let indent = IndentLevel::from_node(match_arm.syntax());
  98             insert_after_old_arm.push(make::tokens::whitespace(&format!("\n{indent}")).into());
  99
 100             insert_after_old_arm.push(new_match_arm.syntax().clone().into());
 101
 102             ted::insert_all_raw(Position::after(match_arm.syntax()), insert_after_old_arm);
 103
 104             if has_comma_after {
 105                 ted::insert_raw(
 106                     Position::last_child_of(new_match_arm.syntax()),
 107                     make::token(T![,]),
 108                 );
 109             }
 110
 111             edit.replace(old_parent_range, new_parent.to_string());
 112         },
 113     )
 114 }
 115
 116 #[cfg(test)]
 117 mod tests {
 118     use crate::tests::{check_assist, check_assist_not_applicable};
 119
 120     use super::*;
 121
 122     #[test]
 123     fn unmerge_match_arm_single_pipe() {
 124         check_assist(
 125             unmerge_match_arm,
 126             r#"
 127 #[derive(Debug)]
 128 enum X { A, B, C }
 129
 130 fn main() {
 131     let x = X::A;
 132     let y = match x {
 133         X::A $0| X::B => { 1i32 }
 134         X::C => { 2i32 }
 135     };
 136 }
 137 "#,
 138             r#"
 139 #[derive(Debug)]
 140 enum X { A, B, C }
 141
 142 fn main() {
 143     let x = X::A;
 144     let y = match x {
 145         X::A => { 1i32 }
 146         X::B => { 1i32 }
 147         X::C => { 2i32 }
 148     };
 149 }
 150 "#,
 151         );
 152     }
 153
 154     #[test]
 155     fn unmerge_match_arm_guard() {
 156         check_assist(
 157             unmerge_match_arm,
 158             r#"
 159 #[derive(Debug)]
 160 enum X { A, B, C }
 161
 162 fn main() {
 163     let x = X::A;
 164     let y = match x {
 165         X::A $0| X::B if true => { 1i32 }
 166         _ => { 2i32 }
 167     };
 168 }
 169 "#,
 170             r#"
 171 #[derive(Debug)]
 172 enum X { A, B, C }
 173
 174 fn main() {
 175     let x = X::A;
 176     let y = match x {
 177         X::A if true => { 1i32 }
 178         X::B if true => { 1i32 }
 179         _ => { 2i32 }
 180     };
 181 }
 182 "#,
 183         );
 184     }
 185
 186     #[test]
 187     fn unmerge_match_arm_leading_pipe() {
 188         check_assist_not_applicable(
 189             unmerge_match_arm,
 190             r#"
 191
 192 fn main() {
 193     let y = match 0 {
 194         |$0 0 => { 1i32 }
 195         1 => { 2i32 }
 196     };
 197 }
 198 "#,
 199         );
 200     }
 201
 202     #[test]
 203     fn unmerge_match_arm_multiple_pipes() {
 204         check_assist(
 205             unmerge_match_arm,
 206             r#"
 207 #[derive(Debug)]
 208 enum X { A, B, C, D, E }
 209
 210 fn main() {
 211     let x = X::A;
 212     let y = match x {
 213         X::A | X::B |$0 X::C | X::D => 1i32,
 214         X::E => 2i32,
 215     };
 216 }
 217 "#,
 218             r#"
 219 #[derive(Debug)]
 220 enum X { A, B, C, D, E }
 221
 222 fn main() {
 223     let x = X::A;
 224     let y = match x {
 225         X::A | X::B => 1i32,
 226         X::C | X::D => 1i32,
 227         X::E => 2i32,
 228     };
 229 }
 230 "#,
 231         );
 232     }
 233
 234     #[test]
 235     fn unmerge_match_arm_inserts_comma_if_required() {
 236         check_assist(
 237             unmerge_match_arm,
 238             r#"
 239 #[derive(Debug)]
 240 enum X { A, B }
 241
 242 fn main() {
 243     let x = X::A;
 244     let y = match x {
 245         X::A $0| X::B => 1i32
 246     };
 247 }
 248 "#,
 249             r#"
 250 #[derive(Debug)]
 251 enum X { A, B }
 252
 253 fn main() {
 254     let x = X::A;
 255     let y = match x {
 256         X::A => 1i32,
 257         X::B => 1i32
 258     };
 259 }
 260 "#,
 261         );
 262     }
 263
 264     #[test]
 265     fn unmerge_match_arm_inserts_comma_if_had_after() {
 266         check_assist(
 267             unmerge_match_arm,
 268             r#"
 269 #[derive(Debug)]
 270 enum X { A, B }
 271
 272 fn main() {
 273     let x = X::A;
 274     match x {
 275         X::A $0| X::B => {},
 276     }
 277 }
 278 "#,
 279             r#"
 280 #[derive(Debug)]
 281 enum X { A, B }
 282
 283 fn main() {
 284     let x = X::A;
 285     match x {
 286         X::A => {},
 287         X::B => {},
 288     }
 289 }
 290 "#,
 291         );
 292     }
 293 }