add assist to move arm condition to match guard

Author: unrealhoang

crates/ra_assists/src/lib.rs

@@ -120,6 +120,7 @@ fn all_assists<DB: HirDatabase>() -> &'static [fn(AssistCtx<DB>) -> Option<Assis
         add_missing_impl_members::add_missing_default_members,
         inline_local_variable::inline_local_varialbe,
         move_guard_to_arm_body_and_back::move_guard_to_arm_body,
+        move_guard_to_arm_body_and_back::move_arm_cond_to_match_guard,
     ]
 }
 

crates/ra_assists/src/move_guard_to_arm_body_and_back.rs

@@ -2,7 +2,7 @@ use hir::db::HirDatabase;
 use ra_syntax::{
     TextUnit,
     SyntaxElement,
-    ast::{MatchArm, AstNode, AstToken},
+    ast::{MatchArm, AstNode, AstToken, IfExpr},
     ast,
 };
 
@@ -39,6 +39,41 @@ pub(crate) fn move_guard_to_arm_body(mut ctx: AssistCtx<impl HirDatabase>) -> Op
     ctx.build()
 }
 
+pub(crate) fn move_arm_cond_to_match_guard(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
+    let match_arm: &MatchArm = ctx.node_at_offset::<MatchArm>()?;
+    let last_match_pat = match_arm.pats().last()?;
+
+    let arm_body = match_arm.expr()?;
+    let if_expr: &IfExpr = IfExpr::cast(arm_body.syntax())?;
+    let cond = if_expr.condition()?;
+    let then_block = if_expr.then_branch()?;
+
+    // Not support if with else branch
+    if let Some(_) = if_expr.else_branch() {
+        return None;
+    }
+
+    let buf = format!(" if {}", cond.syntax().text());
+
+    ctx.add_action(
+        AssistId("move_arm_cond_to_match_guard"),
+        "move condition to match guard",
+        |edit| {
+            edit.target(last_match_pat.syntax().range());
+
+            if then_block.statements().next().is_none() && then_block.expr().is_some() {
+                edit.replace(if_expr.syntax().range(), then_block.expr().unwrap().syntax().text());
+            } else {
+                edit.replace(if_expr.syntax().range(), then_block.syntax().text());
+            }
+
+            edit.insert(last_match_pat.syntax().range().end(), buf);
+            edit.set_cursor(last_match_pat.syntax().range().end() + TextUnit::from(1));
+        },
+    );
+    ctx.build()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -112,4 +147,118 @@ mod tests {
             "#,
         );
     }
+
+    #[test]
+    fn move_arm_cond_to_match_guard_works() {
+        check_assist(
+            move_arm_cond_to_match_guard,
+            r#"
+            fn f() {
+                let t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' => if chars.clone().next() == Some('\n') { <|>false },
+                    _ => true
+                }
+            }
+            "#,
+            r#"
+            fn f() {
+                let t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' <|>if chars.clone().next() == Some('\n') => false,
+                    _ => true
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn move_arm_cond_to_match_guard_if_let_works() {
+        check_assist(
+            move_arm_cond_to_match_guard,
+            r#"
+            fn f() {
+                let t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' => if let Some(_) = chars.clone().next() { <|>false },
+                    _ => true
+                }
+            }
+            "#,
+            r#"
+            fn f() {
+                let t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' <|>if let Some(_) = chars.clone().next() => false,
+                    _ => true
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn move_arm_cond_to_match_guard_if_empty_body_works() {
+        check_assist(
+            move_arm_cond_to_match_guard,
+            r#"
+            fn f() {
+                let t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' => if let Some(_) = chars.clone().next() { <|> },
+                    _ => true
+                }
+            }
+            "#,
+            r#"
+            fn f() {
+                let t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' <|>if let Some(_) = chars.clone().next() => {  },
+                    _ => true
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn move_arm_cond_to_match_guard_if_multiline_body_works() {
+        check_assist(
+            move_arm_cond_to_match_guard,
+            r#"
+            fn f() {
+                let mut t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' => if let Some(_) = chars.clone().next() {
+                        t = 'e';<|>
+                        false
+                    },
+                    _ => true
+                }
+            }
+            "#,
+            r#"
+            fn f() {
+                let mut t = 'a';
+                let chars = "abcd";
+                match t {
+                    '\r' <|>if let Some(_) = chars.clone().next() => {
+                        t = 'e';
+                        false
+                    },
+                    _ => true
+                }
+            }
+            "#,
+        );
+    }
 }