From 31204e3590a59a6c0cbae53d111b699d92fb229f Mon Sep 17 00:00:00 2001 From: Jesse Bakker Date: Sat, 2 Jan 2021 01:55:56 +0100 Subject: Add extract-assignment assist --- crates/assists/src/handlers/extract_assignment.rs | 238 ++++++++++++++++++++++ crates/assists/src/lib.rs | 2 + crates/assists/src/tests/generated.rs | 29 +++ 3 files changed, 269 insertions(+) create mode 100644 crates/assists/src/handlers/extract_assignment.rs (limited to 'crates') diff --git a/crates/assists/src/handlers/extract_assignment.rs b/crates/assists/src/handlers/extract_assignment.rs new file mode 100644 index 000000000..604297705 --- /dev/null +++ b/crates/assists/src/handlers/extract_assignment.rs @@ -0,0 +1,238 @@ +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 {} +}"#, + ) + } +} diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs index fdec886e9..212464f85 100644 --- a/crates/assists/src/lib.rs +++ b/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, diff --git a/crates/assists/src/tests/generated.rs b/crates/assists/src/tests/generated.rs index d3dfe24e7..b91a816e8 100644 --- a/crates/assists/src/tests/generated.rs +++ b/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( -- cgit v1.2.3