From f12d5e3f201c3f238f69f66b73d3846fd351c403 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <cmrs@cmrs.io>
Date: Mon, 18 Nov 2024 23:07:19 +0100
Subject: [PATCH] improve selection handling for `merge_match_arms`

---
 crates/ide-assists/src/assist_context.rs      |   3 +
 .../src/handlers/merge_match_arms.rs          | 101 +++++++++++++++++-
 2 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/crates/ide-assists/src/assist_context.rs b/crates/ide-assists/src/assist_context.rs
index 1d2d3350f7..0146369f29 100644
--- a/crates/ide-assists/src/assist_context.rs
+++ b/crates/ide-assists/src/assist_context.rs
@@ -116,6 +116,9 @@ impl<'a> AssistContext<'a> {
     pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
         find_node_at_offset(self.source_file.syntax(), self.offset())
     }
+    pub(crate) fn find_node_at_trimmed_offset<N: AstNode>(&self) -> Option<N> {
+        find_node_at_offset(self.source_file.syntax(), self.trimmed_range.start())
+    }
     pub(crate) fn find_node_at_range<N: AstNode>(&self) -> Option<N> {
         find_node_at_range(self.source_file.syntax(), self.trimmed_range)
     }
diff --git a/crates/ide-assists/src/handlers/merge_match_arms.rs b/crates/ide-assists/src/handlers/merge_match_arms.rs
index 4608e9494b..f83de931ea 100644
--- a/crates/ide-assists/src/handlers/merge_match_arms.rs
+++ b/crates/ide-assists/src/handlers/merge_match_arms.rs
@@ -34,7 +34,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};
 // }
 // ```
 pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
+    let current_arm = ctx.find_node_at_trimmed_offset::<ast::MatchArm>()?;
     // Don't try to handle arms with guards for now - can add support for this later
     if current_arm.guard().is_some() {
         return None;
@@ -42,12 +42,21 @@ pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
     let current_expr = current_arm.expr()?;
     let current_text_range = current_arm.syntax().text_range();
     let current_arm_types = get_arm_types(ctx, &current_arm);
+    let multi_arm_selection = !ctx.has_empty_selection()
+        && ctx.selection_trimmed().end() > current_arm.syntax().text_range().end();
 
     // We check if the following match arms match this one. We could, but don't,
     // compare to the previous match arm as well.
     let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
         .take_while(|arm| match arm.expr() {
             Some(expr) if arm.guard().is_none() => {
+                // don't include match arms that start after our selection
+                if multi_arm_selection
+                    && arm.syntax().text_range().start() >= ctx.selection_trimmed().end()
+                {
+                    return false;
+                }
+
                 let same_text = expr.syntax().text() == current_expr.syntax().text();
                 if !same_text {
                     return false;
@@ -298,6 +307,96 @@ fn main() {
         )
     }
 
+    #[test]
+    fn merge_match_arms_selection_has_leading_whitespace() {
+        check_assist(
+            merge_match_arms,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    match X::A {
+    $0    X::A => 0,
+        X::B => 0,$0
+        X::C => 1,
+    }
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    match X::A {
+        X::A | X::B => 0,
+        X::C => 1,
+    }
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn merge_match_arms_stops_at_end_of_selection() {
+        check_assist(
+            merge_match_arms,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    match X::A {
+    $0    X::A => 0,
+        X::B => 0,
+        $0X::C => 0,
+    }
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    match X::A {
+        X::A | X::B => 0,
+        X::C => 0,
+    }
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn merge_match_arms_works_despite_accidental_selection() {
+        check_assist(
+            merge_match_arms,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    match X::A {
+        X::$0A$0 => 0,
+        X::B => 0,
+        X::C => 1,
+    }
+}
+"#,
+            r#"
+#[derive(Debug)]
+enum X { A, B, C }
+
+fn main() {
+    match X::A {
+        X::A | X::B => 0,
+        X::C => 1,
+    }
+}
+"#,
+        );
+    }
+
     #[test]
     fn merge_match_arms_rejects_guards() {
         check_assist_not_applicable(