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::()?.as_name(); let if_statement = ctx.find_node_at_offset::()?; 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 { 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 { 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 {} }"#, ) } }