feat: Add a "Unmerge match arm" assist to split or-patterns inside match expressions

3 files changed, 321 insertions, 0 deletions

diff --git a/crates/ide-assists/src/handlers/unmerge_match_arm.rs b/crates/ide-assists/src/handlers/unmerge_match_arm.rs
new file mode 100644
index 00000000000..9565f0ee6f2
--- /dev/null
+++ b/crates/ide-assists/src/handlers/unmerge_match_arm.rs
@@ -0,0 +1,293 @@
+use syntax::{
+    algo::neighbor,
+    ast::{self, edit::IndentLevel, make, AstNode},
+    ted::{self, Position},
+    Direction, SyntaxKind, T,
+};
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: unmerge_match_arm
+//
+// Splits the current match with a `|` pattern into two arms with identical bodies.
+//
+// ```
+// enum Action { Move { distance: u32 }, Stop }
+//
+// fn handle(action: Action) {
+//     match action {
+//         Action::Move(..) $0| Action::Stop => foo(),
+//     }
+// }
+// ```
+// ->
+// ```
+// enum Action { Move { distance: u32 }, Stop }
+//
+// fn handle(action: Action) {
+//     match action {
+//         Action::Move(..) => foo(),
+//         Action::Stop => foo(),
+//     }
+// }
+// ```
+pub(crate) fn unmerge_match_arm(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let pipe_token = ctx.find_token_syntax_at_offset(T![|])?;
+    let or_pat = ast::OrPat::cast(pipe_token.parent()?)?.clone_for_update();
+    let match_arm = ast::MatchArm::cast(or_pat.syntax().parent()?)?;
+    let match_arm_body = match_arm.expr()?;
+
+    // We don't need to check for leading pipe because it is directly under `MatchArm`
+    // without `OrPat`.
+
+    let new_parent = match_arm.syntax().parent()?;
+    let old_parent_range = new_parent.text_range();
+
+    acc.add(
+        AssistId("unmerge_match_arm", AssistKind::RefactorRewrite),
+        "Unmerge match arm",
+        pipe_token.text_range(),
+        |edit| {
+            let pats_after = pipe_token
+                .siblings_with_tokens(Direction::Next)
+                .filter_map(|it| ast::Pat::cast(it.into_node()?));
+            // FIXME: We should add a leading pipe if the original arm has one.
+            let new_match_arm = make::match_arm(
+                pats_after,
+                match_arm.guard().and_then(|guard| guard.condition()),
+                match_arm_body,
+            )
+            .clone_for_update();
+
+            let mut pipe_index = pipe_token.index();
+            if pipe_token
+                .prev_sibling_or_token()
+                .map_or(false, |it| it.kind() == SyntaxKind::WHITESPACE)
+            {
+                pipe_index -= 1;
+            }
+            or_pat.syntax().splice_children(
+                pipe_index..or_pat.syntax().children_with_tokens().count(),
+                Vec::new(),
+            );
+
+            let mut insert_after_old_arm = Vec::new();
+
+            // A comma can be:
+            //  - After the arm. In this case we always want to insert a comma after the newly
+            //    inserted arm.
+            //  - Missing after the arm, with no arms after. In this case we want to insert a
+            //    comma before the newly inserted arm. It can not be necessary if there arm
+            //    body is a block, but we don't bother to check that.
+            //  - Missing after the arm with arms after, if the arm body is a block. In this case
+            //    we don't want to insert a comma at all.
+            let has_comma_after =
+                std::iter::successors(match_arm.syntax().last_child_or_token(), |it| {
+                    it.prev_sibling_or_token()
+                })
+                .map(|it| it.kind())
+                .skip_while(|it| it.is_trivia())
+                .next()
+                    == Some(T![,]);
+            let has_arms_after = neighbor(&match_arm, Direction::Next).is_some();
+            if !has_comma_after && !has_arms_after {
+                insert_after_old_arm.push(make::token(T![,]).into());
+            }
+
+            let indent = IndentLevel::from_node(match_arm.syntax());
+            insert_after_old_arm.push(make::tokens::whitespace(&format!("\n{indent}")).into());
+
+            insert_after_old_arm.push(new_match_arm.syntax().clone().into());
+
+            ted::insert_all_raw(Position::after(match_arm.syntax()), insert_after_old_arm);
+
+            if has_comma_after {
+                ted::insert_raw(
+                    Position::last_child_of(new_match_arm.syntax()),
+                    make::token(T![,]),
+                );
+            }
+
+            edit.replace(old_parent_range, new_parent.to_string());
+        },
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn unmerge_match_arm_single_pipe() {
+        check_assist(
+            unmerge_match_arm,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A $0| X::B => { 1i32 }
+        X::C => { 2i32 }
+    };
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A => { 1i32 }
+        X::B => { 1i32 }
+        X::C => { 2i32 }
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unmerge_match_arm_guard() {
+        check_assist(
+            unmerge_match_arm,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A $0| X::B if true => { 1i32 }
+        _ => { 2i32 }
+    };
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A if true => { 1i32 }
+        X::B if true => { 1i32 }
+        _ => { 2i32 }
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unmerge_match_arm_leading_pipe() {
+        check_assist_not_applicable(
+            unmerge_match_arm,
+            r#"
+
+fn main() {
+    let y = match 0 {
+        |$0 0 => { 1i32 }
+        1 => { 2i32 }
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unmerge_match_arm_multiple_pipes() {
+        check_assist(
+            unmerge_match_arm,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C, D, E }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A | X::B |$0 X::C | X::D => 1i32,
+        X::E => 2i32,
+    };
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C, D, E }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A | X::B => 1i32,
+        X::C | X::D => 1i32,
+        X::E => 2i32,
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unmerge_match_arm_inserts_comma_if_required() {
+        check_assist(
+            unmerge_match_arm,
+            r#"
+#[derive(Debug)]
+enum X { A, B }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A $0| X::B => 1i32
+    };
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B }
+
+fn main() {
+    let x = X::A;
+    let y = match x {
+        X::A => 1i32,
+        X::B => 1i32
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unmerge_match_arm_inserts_comma_if_had_after() {
+        check_assist(
+            unmerge_match_arm,
+            r#"
+#[derive(Debug)]
+enum X { A, B }
+
+fn main() {
+    let x = X::A;
+    match x {
+        X::A $0| X::B => {},
+    }
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B }
+
+fn main() {
+    let x = X::A;
+    match x {
+        X::A => {},
+        X::B => {},
+    }
+}
+"#,
+        );
+    }
+}
diff --git a/crates/ide-assists/src/lib.rs b/crates/ide-assists/src/lib.rs
index 7fb35143fa2..c558553b1cb 100644
--- a/crates/ide-assists/src/lib.rs
+++ b/crates/ide-assists/src/lib.rs
@@ -185,6 +185,7 @@ mod handlers {
     mod replace_string_with_char;
     mod replace_turbofish_with_explicit_type;
     mod split_import;
+    mod unmerge_match_arm;
     mod sort_items;
     mod toggle_ignore;
     mod unmerge_use;
@@ -278,6 +279,7 @@ mod handlers {
             sort_items::sort_items,
             split_import::split_import,
             toggle_ignore::toggle_ignore,
+            unmerge_match_arm::unmerge_match_arm,
             unmerge_use::unmerge_use,
             unnecessary_async::unnecessary_async,
             unwrap_block::unwrap_block,
diff --git a/crates/ide-assists/src/tests/generated.rs b/crates/ide-assists/src/tests/generated.rs
index 22319f36134..7b2c16806b2 100644
--- a/crates/ide-assists/src/tests/generated.rs
+++ b/crates/ide-assists/src/tests/generated.rs
@@ -2208,6 +2208,32 @@ fn arithmetics {
 }
 
 #[test]
+fn doctest_unmerge_match_arm() {
+    check_doc_test(
+        "unmerge_match_arm",
+        r#####"
+enum Action { Move { distance: u32 }, Stop }
+
+fn handle(action: Action) {
+    match action {
+        Action::Move(..) $0| Action::Stop => foo(),
+    }
+}
+"#####,
+        r#####"
+enum Action { Move { distance: u32 }, Stop }
+
+fn handle(action: Action) {
+    match action {
+        Action::Move(..) => foo(),
+        Action::Stop => foo(),
+    }
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_unmerge_use() {
     check_doc_test(
         "unmerge_use",