rust/crates/assists/src/handlers/extract_assignment.rs

use hir::AsName;
use syntax::{
    ast::{self, edit::AstNodeEdit, make},
    AstNode,
};
use test_utils::mark;

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::<ast::NameRef>()?.as_name();

    let if_statement = ctx.find_node_at_offset::<ast::IfExpr>()?;

    let new_stmt = exprify_if(&if_statement, &name)?.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: &ast::IfExpr, name: &hir::Name) -> Option<ast::Expr> {
    let then_branch = exprify_block(&statement.then_branch()?, name)?;
    let else_branch = match statement.else_branch()? {
        ast::ElseBranch::Block(block) => ast::ElseBranch::Block(exprify_block(&block, name)?),
        ast::ElseBranch::IfExpr(expr) => {
            mark::hit!(test_extract_assigment_chained_if);
            ast::ElseBranch::IfExpr(ast::IfExpr::cast(
                exprify_if(&expr, name)?.syntax().to_owned(),
            )?)
        }
    };
    Some(make::expr_if(statement.condition()?, then_branch, Some(else_branch)))
}

fn exprify_block(block: &ast::BlockExpr, name: &hir::Name) -> 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
                && &expr.lhs()?.name_ref()?.as_name() == 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;
    }
}"#,
        )
    }

    #[test]
    fn test_extract_assignment_chained_if() {
        mark::check!(test_extract_assigment_chained_if);
        check_assist(
            extract_assigment,
            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_extract_assigment_retains_stmts() {
        check_assist(
            extract_assigment,
            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 extract_assignment_let_stmt_not_applicable() {
        check_assist_not_applicable(
            extract_assigment,
            r#"
fn foo() {
    let mut a = 1;

    let b = if true {
        <|>a = 2
    } else {
        a = 3
    };
}"#,
        )
    }

    #[test]
    fn extract_assignment_missing_assigment_not_applicable() {
        check_assist_not_applicable(
            extract_assigment,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        <|>a = 2;
    } else {}
}"#,
        )
    }
}