Add extract-assignment assist

Author: Jesse-Bakker

crates/assists/src/handlers/extract_assignment.rs

@@ -0,0 +1,137 @@
+use syntax::{
+    ast::{edit::AstNodeEdit, make, BinOp, BlockExpr, ElseBranch, Expr, IfExpr, NameRef, Stmt},
+    AstNode,
+};
+
+use crate::{
+    assist_context::{AssistContext, Assists},
+    AssistId, AssistKind,
+};
+
+// Assist: extract_assignment
+//
+// Extracts variable assigment to outside an if or match statement.
+//
+// ```
+// fn main() {
+//     let mut foo = 6;
+//
+//     if true {
+//         <|>foo = 5;
+//     } else {
+//         foo = 4;
+//     }
+// }
+// ```
+// ->
+// ```
+// fn main() {
+//     let mut foo = 6;
+//
+//     foo = if true {
+//         5
+//     } else {
+//         4
+//     }
+// }
+// ```
+pub(crate) fn extract_assigment(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let name = ctx.find_node_at_offset::<NameRef>()?.to_string();
+
+    let if_statement = ctx.find_node_at_offset::<IfExpr>()?;
+
+    let new_stmt =
+        exprify_if(&if_statement, &name.to_string())?.indent(if_statement.indent_level());
+    let expr_stmt = make::expr_stmt(new_stmt);
+
+    acc.add(
+        AssistId("extract_assignment", AssistKind::RefactorExtract),
+        "Extract assignment",
+        if_statement.syntax().text_range(),
+        move |edit| {
+            edit.replace(if_statement.syntax().text_range(), format!("{} = {}", name, expr_stmt));
+        },
+    )
+}
+
+fn exprify_if(statement: &IfExpr, name: &str) -> Option<Expr> {
+    let then_branch = exprify_block(statement.then_branch()?, name)?;
+    let else_branch = match statement.else_branch()? {
+        ElseBranch::Block(block) => ElseBranch::Block(exprify_block(block, name)?),
+        ElseBranch::IfExpr(expr) => {
+            ElseBranch::IfExpr(IfExpr::cast(exprify_if(&expr, name)?.syntax().to_owned())?)
+        }
+    };
+    Some(make::expr_if(statement.condition()?, then_branch, Some(else_branch)))
+}
+
+fn exprify_block(block: BlockExpr, name: &str) -> Option<BlockExpr> {
+    if block.expr().is_some() {
+        return None;
+    }
+
+    let mut stmts: Vec<_> = block.statements().collect();
+    let stmt = stmts.pop()?;
+
+    if let Stmt::ExprStmt(stmt) = stmt {
+        if let Expr::BinExpr(expr) = stmt.expr()? {
+            if expr.op_kind()? == BinOp::Assignment && expr.lhs()?.name_ref()?.to_string() == name {
+                // The last statement in the block is an assignment to the name we want
+                return Some(make::block_expr(stmts, Some(expr.rhs()?)));
+            }
+        }
+    }
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    #[test]
+    fn test_extract_assignment() {
+        check_assist(
+            extract_assigment,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    if true {
+        <|>a = 2;
+    } else {
+        a = 3;
+    }
+}"#,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    a = if true {
+        2
+    } else {
+        3
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_extract_assignment_not_last_not_applicable() {
+        check_assist_not_applicable(
+            extract_assigment,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    if true {
+        <|>a = 2;
+        b = a;
+    } else {
+        a = 3;
+    }
+}"#,
+        )
+    }
+}

crates/assists/src/lib.rs

@@ -116,6 +116,7 @@ mod handlers {
     mod convert_integer_literal;
     mod early_return;
     mod expand_glob_import;
+    mod extract_assignment;
     mod extract_module_to_file;
     mod extract_struct_from_enum_variant;
     mod extract_variable;
@@ -167,6 +168,7 @@ mod handlers {
             convert_integer_literal::convert_integer_literal,
             early_return::convert_to_guarded_return,
             expand_glob_import::expand_glob_import,
+            extract_assignment::extract_assigment,
             extract_module_to_file::extract_module_to_file,
             extract_struct_from_enum_variant::extract_struct_from_enum_variant,
             extract_variable::extract_variable,

crates/assists/src/tests/generated.rs

@@ -237,6 +237,35 @@ fn qux(bar: Bar, baz: Baz) {}
     )
 }
 
+#[test]
+fn doctest_extract_assignment() {
+    check_doc_test(
+        "extract_assignment",
+        r#####"
+fn main() {
+    let mut foo = 6;
+
+    if true {
+        <|>foo = 5;
+    } else {
+        foo = 4;
+    }
+}
+"#####,
+        r#####"
+fn main() {
+    let mut foo = 6;
+
+    foo = if true {
+        5
+    } else {
+        4
+    }
+}
+"#####,
+    )
+}
+
 #[test]
 fn doctest_extract_module_to_file() {
     check_doc_test(