Merge pull request #21508 from A4-Tacks/move-guard-else-if

Support else-branch for move_guard

Author: ChayimFriedman2

crates/ide-assists/src/handlers/move_guard.rs

@@ -1,6 +1,7 @@
-use itertools::Itertools;
+use itertools::{Itertools, chain};
 use syntax::{
     SyntaxKind::WHITESPACE,
+    TextRange,
     ast::{
         AstNode, BlockExpr, ElseBranch, Expr, IfExpr, MatchArm, Pat, edit::AstNodeEdit, make,
         prec::ExprPrecedence, syntax_factory::SyntaxFactory,
@@ -44,13 +45,26 @@ pub(crate) fn move_guard_to_arm_body(acc: &mut Assists, ctx: &AssistContext<'_>)
         cov_mark::hit!(move_guard_inapplicable_in_arm_body);
         return None;
     }
-    let space_before_guard = guard.syntax().prev_sibling_or_token();
+    let rest_arms = rest_arms(&match_arm, ctx.selection_trimmed())?;
+    let space_before_delete = chain(
+        guard.syntax().prev_sibling_or_token(),
+        rest_arms.iter().filter_map(|it| it.syntax().prev_sibling_or_token()),
+    );
     let space_after_arrow = match_arm.fat_arrow_token()?.next_sibling_or_token();
 
-    let guard_condition = guard.condition()?.reset_indent();
     let arm_expr = match_arm.expr()?;
-    let then_branch = crate::utils::wrap_block(&arm_expr);
-    let if_expr = make::expr_if(guard_condition, then_branch, None).indent(arm_expr.indent_level());
+    let if_branch = chain([&match_arm], &rest_arms)
+        .rfold(None, |else_branch, arm| {
+            if let Some(guard) = arm.guard() {
+                let then_branch = crate::utils::wrap_block(&arm.expr()?);
+                let guard_condition = guard.condition()?.reset_indent();
+                Some(make::expr_if(guard_condition, then_branch, else_branch).into())
+            } else {
+                arm.expr().map(|it| crate::utils::wrap_block(&it).into())
+            }
+        })?
+        .indent(arm_expr.indent_level());
+    let ElseBranch::IfExpr(if_expr) = if_branch else { return None };
 
     let target = guard.syntax().text_range();
     acc.add(
@@ -59,10 +73,13 @@ pub(crate) fn move_guard_to_arm_body(acc: &mut Assists, ctx: &AssistContext<'_>)
         target,
         |builder| {
             let mut edit = builder.make_editor(match_arm.syntax());
-            if let Some(element) = space_before_guard
-                && element.kind() == WHITESPACE
-            {
-                edit.delete(element);
+            for element in space_before_delete {
+                if element.kind() == WHITESPACE {
+                    edit.delete(element);
+                }
+            }
+            for rest_arm in &rest_arms {
+                edit.delete(rest_arm.syntax());
             }
             if let Some(element) = space_after_arrow
                 && element.kind() == WHITESPACE
@@ -221,6 +238,25 @@ pub(crate) fn move_arm_cond_to_match_guard(
     )
 }
 
+fn rest_arms(match_arm: &MatchArm, selection: TextRange) -> Option<Vec<MatchArm>> {
+    match_arm
+        .parent_match()
+        .match_arm_list()?
+        .arms()
+        .skip_while(|it| it != match_arm)
+        .skip(1)
+        .take_while(move |it| {
+            selection.is_empty() || crate::utils::is_selected(it, selection, false)
+        })
+        .take_while(move |it| {
+            it.pat()
+                .zip(match_arm.pat())
+                .is_some_and(|(a, b)| a.syntax().text() == b.syntax().text())
+        })
+        .collect::<Vec<_>>()
+        .into()
+}
+
 // Parses an if-else-if chain to get the conditions and the then branches until we encounter an else
 // branch or the end.
 fn parse_if_chain(if_expr: IfExpr) -> Option<(Vec<(Expr, BlockExpr)>, Option<BlockExpr>)> {
@@ -344,6 +380,115 @@ fn main() {
         );
     }
 
+    #[test]
+    fn move_multiple_guard_to_arm_body_works() {
+        check_assist(
+            move_guard_to_arm_body,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 $0if x % 3 == 0 => false,
+        x @ 0..30 if x % 2 == 0 => true,
+        _ => false
+    }
+}
+"#,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 => if x % 3 == 0 {
+            false
+        } else if x % 2 == 0 {
+            true
+        },
+        _ => false
+    }
+}
+"#,
+        );
+
+        check_assist(
+            move_guard_to_arm_body,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 $0if x % 3 == 0 => false,
+        x @ 0..30 if x % 2 == 0 => true,
+        x @ 0..30 => false,
+        _ => true
+    }
+}
+"#,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 => if x % 3 == 0 {
+            false
+        } else if x % 2 == 0 {
+            true
+        } else {
+            false
+        },
+        _ => true
+    }
+}
+"#,
+        );
+
+        check_assist(
+            move_guard_to_arm_body,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 if x % 3 == 0 => false,
+        x @ 0..30 $0if x % 2 == 0$0 => true,
+        x @ 0..30 => false,
+        _ => true
+    }
+}
+"#,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 if x % 3 == 0 => false,
+        x @ 0..30 => if x % 2 == 0 {
+            true
+        },
+        x @ 0..30 => false,
+        _ => true
+    }
+}
+"#,
+        );
+
+        check_assist(
+            move_guard_to_arm_body,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 $0if x % 3 == 0 => false,
+        x @ 0..30 $0if x % 2 == 0 => true,
+        x @ 0..30 => false,
+        _ => true
+    }
+}
+"#,
+            r#"
+fn main() {
+    match 92 {
+        x @ 0..30 => if x % 3 == 0 {
+            false
+        } else if x % 2 == 0 {
+            true
+        },
+        x @ 0..30 => false,
+        _ => true
+    }
+}
+"#,
+        );
+    }
+
     #[test]
     fn move_guard_to_block_arm_body_works() {
         check_assist(