1 files changed, 400 insertions, 0 deletions

diff --git a/crates/assists/src/handlers/pull_assignment_up.rs b/crates/assists/src/handlers/pull_assignment_up.rs
new file mode 100644
index 000000000..560d93e10
--- /dev/null
+++ b/crates/assists/src/handlers/pull_assignment_up.rs
@@ -0,0 +1,400 @@
+use syntax::{
+    ast::{self, edit::AstNodeEdit, make},
+    AstNode,
+};
+use test_utils::mark;
+
+use crate::{
+    assist_context::{AssistContext, Assists},
+    AssistId, AssistKind,
+};
+
+// Assist: pull_assignment_up
+//
+// Extracts variable assignment 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 pull_assignment_up(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let assign_expr = ctx.find_node_at_offset::<ast::BinExpr>()?;
+    let name_expr = if assign_expr.op_kind()? == ast::BinOp::Assignment {
+        assign_expr.lhs()?
+    } else {
+        return None;
+    };
+
+    let (old_stmt, new_stmt) = if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
+        (
+            ast::Expr::cast(if_expr.syntax().to_owned())?,
+            exprify_if(&if_expr, &ctx.sema, &name_expr)?.indent(if_expr.indent_level()),
+        )
+    } else if let Some(match_expr) = ctx.find_node_at_offset::<ast::MatchExpr>() {
+        (
+            ast::Expr::cast(match_expr.syntax().to_owned())?,
+            exprify_match(&match_expr, &ctx.sema, &name_expr)?,
+        )
+    } else {
+        return None;
+    };
+
+    let expr_stmt = make::expr_stmt(new_stmt);
+
+    acc.add(
+        AssistId("pull_assignment_up", AssistKind::RefactorExtract),
+        "Pull assignment up",
+        old_stmt.syntax().text_range(),
+        move |edit| {
+            edit.replace(old_stmt.syntax().text_range(), format!("{} = {};", name_expr, expr_stmt));
+        },
+    )
+}
+
+fn exprify_match(
+    match_expr: &ast::MatchExpr,
+    sema: &hir::Semantics<ide_db::RootDatabase>,
+    name: &ast::Expr,
+) -> Option<ast::Expr> {
+    let new_arm_list = match_expr
+        .match_arm_list()?
+        .arms()
+        .map(|arm| {
+            if let ast::Expr::BlockExpr(block) = arm.expr()? {
+                let new_block = exprify_block(&block, sema, name)?.indent(block.indent_level());
+                Some(arm.replace_descendant(block, new_block))
+            } else {
+                None
+            }
+        })
+        .collect::<Option<Vec<_>>>()?;
+    let new_arm_list = match_expr
+        .match_arm_list()?
+        .replace_descendants(match_expr.match_arm_list()?.arms().zip(new_arm_list));
+    Some(make::expr_match(match_expr.expr()?, new_arm_list))
+}
+
+fn exprify_if(
+    statement: &ast::IfExpr,
+    sema: &hir::Semantics<ide_db::RootDatabase>,
+    name: &ast::Expr,
+) -> Option<ast::Expr> {
+    let then_branch = exprify_block(&statement.then_branch()?, sema, name)?;
+    let else_branch = match statement.else_branch()? {
+        ast::ElseBranch::Block(ref block) => {
+            ast::ElseBranch::Block(exprify_block(block, sema, name)?)
+        }
+        ast::ElseBranch::IfExpr(expr) => {
+            mark::hit!(test_pull_assignment_up_chained_if);
+            ast::ElseBranch::IfExpr(ast::IfExpr::cast(
+                exprify_if(&expr, sema, name)?.syntax().to_owned(),
+            )?)
+        }
+    };
+    Some(make::expr_if(statement.condition()?, then_branch, Some(else_branch)))
+}
+
+fn exprify_block(
+    block: &ast::BlockExpr,
+    sema: &hir::Semantics<ide_db::RootDatabase>,
+    name: &ast::Expr,
+) -> Option<ast::BlockExpr> {
+    if block.expr().is_some() {
+        return None;
+    }
+
+    let mut stmts: Vec<_> = block.statements().collect();
+    let stmt = stmts.pop()?;
+
+    if let ast::Stmt::ExprStmt(stmt) = stmt {
+        if let ast::Expr::BinExpr(expr) = stmt.expr()? {
+            if expr.op_kind()? == ast::BinOp::Assignment && is_equivalent(sema, &expr.lhs()?, 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
+}
+
+fn is_equivalent(
+    sema: &hir::Semantics<ide_db::RootDatabase>,
+    expr0: &ast::Expr,
+    expr1: &ast::Expr,
+) -> bool {
+    match (expr0, expr1) {
+        (ast::Expr::FieldExpr(field_expr0), ast::Expr::FieldExpr(field_expr1)) => {
+            mark::hit!(test_pull_assignment_up_field_assignment);
+            sema.resolve_field(field_expr0) == sema.resolve_field(field_expr1)
+        }
+        (ast::Expr::PathExpr(path0), ast::Expr::PathExpr(path1)) => {
+            let path0 = path0.path();
+            let path1 = path1.path();
+            if let (Some(path0), Some(path1)) = (path0, path1) {
+                sema.resolve_path(&path0) == sema.resolve_path(&path1)
+            } else {
+                false
+            }
+        }
+        _ => false,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    #[test]
+    fn test_pull_assignment_up_if() {
+        check_assist(
+            pull_assignment_up,
+            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_pull_assignment_up_match() {
+        check_assist(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    match 1 {
+        1 => {
+            <|>a = 2;
+        },
+        2 => {
+            a = 3;
+        },
+        3 => {
+            a = 4;
+        }
+    }
+}"#,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    a = match 1 {
+        1 => {
+            2
+        },
+        2 => {
+            3
+        },
+        3 => {
+            4
+        }
+    };
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_pull_assignment_up_not_last_not_applicable() {
+        check_assist_not_applicable(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    if true {
+        <|>a = 2;
+        b = a;
+    } else {
+        a = 3;
+    }
+}"#,
+        )
+    }
+
+    #[test]
+    fn test_pull_assignment_up_chained_if() {
+        mark::check!(test_pull_assignment_up_chained_if);
+        check_assist(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    if true {
+        <|>a = 2;
+    } else if false {
+        a = 3;
+    } else {
+        a = 4;
+    }
+}"#,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    a = if true {
+        2
+    } else if false {
+        3
+    } else {
+        4
+    };
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_pull_assignment_up_retains_stmts() {
+        check_assist(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    if true {
+        let b = 2;
+        <|>a = 2;
+    } else {
+        let b = 3;
+        a = 3;
+    }
+}"#,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    a = if true {
+        let b = 2;
+        2
+    } else {
+        let b = 3;
+        3
+    };
+}"#,
+        )
+    }
+
+    #[test]
+    fn pull_assignment_up_let_stmt_not_applicable() {
+        check_assist_not_applicable(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    let b = if true {
+        <|>a = 2
+    } else {
+        a = 3
+    };
+}"#,
+        )
+    }
+
+    #[test]
+    fn pull_assignment_up_if_missing_assigment_not_applicable() {
+        check_assist_not_applicable(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    if true {
+        <|>a = 2;
+    } else {}
+}"#,
+        )
+    }
+
+    #[test]
+    fn pull_assignment_up_match_missing_assigment_not_applicable() {
+        check_assist_not_applicable(
+            pull_assignment_up,
+            r#"
+fn foo() {
+    let mut a = 1;
+
+    match 1 {
+        1 => {
+            <|>a = 2;
+        },
+        2 => {
+            a = 3;
+        },
+        3 => {},
+    }
+}"#,
+        )
+    }
+
+    #[test]
+    fn test_pull_assignment_up_field_assignment() {
+        mark::check!(test_pull_assignment_up_field_assignment);
+        check_assist(
+            pull_assignment_up,
+            r#"
+struct A(usize);
+
+fn foo() {
+    let mut a = A(1);
+
+    if true {
+        <|>a.0 = 2;
+    } else {
+        a.0 = 3;
+    }
+}"#,
+            r#"
+struct A(usize);
+
+fn foo() {
+    let mut a = A(1);
+
+    a.0 = if true {
+        2
+    } else {
+        3
+    };
+}"#,
+        )
+    }
+}