From 31204e3590a59a6c0cbae53d111b699d92fb229f Mon Sep 17 00:00:00 2001 From: Jesse Bakker Date: Sat, 2 Jan 2021 01:55:56 +0100 Subject: [PATCH 1/2] Add extract-assignment assist --- .../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 diff --git a/crates/assists/src/handlers/extract_assignment.rs b/crates/assists/src/handlers/extract_assignment.rs new file mode 100644 index 0000000000..6042977057 --- /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 fdec886e9d..212464f859 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 d3dfe24e7c..b91a816e89 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( From bfe6a8e71afc0d3bee47261f83647b28eca0aae6 Mon Sep 17 00:00:00 2001 From: Jesse Bakker Date: Sat, 2 Jan 2021 15:33:23 +0100 Subject: [PATCH 2/2] Add support for MatchExpr to extract_assigment assist --- .../src/handlers/extract_assignment.rs | 101 ++++++++++++++++-- 1 file changed, 94 insertions(+), 7 deletions(-) diff --git a/crates/assists/src/handlers/extract_assignment.rs b/crates/assists/src/handlers/extract_assignment.rs index 6042977057..281cf5d24e 100644 --- a/crates/assists/src/handlers/extract_assignment.rs +++ b/crates/assists/src/handlers/extract_assignment.rs @@ -40,25 +40,52 @@ use crate::{ 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 (old_stmt, new_stmt) = if let Some(if_expr) = ctx.find_node_at_offset::() { + ( + ast::Expr::cast(if_expr.syntax().to_owned())?, + exprify_if(&if_expr, &name)?.indent(if_expr.indent_level()), + ) + } else if let Some(match_expr) = ctx.find_node_at_offset::() { + (ast::Expr::cast(match_expr.syntax().to_owned())?, exprify_match(&match_expr, &name)?) + } else { + return None; + }; - 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(), + old_stmt.syntax().text_range(), move |edit| { - edit.replace(if_statement.syntax().text_range(), format!("{} = {};", name, expr_stmt)); + edit.replace(old_stmt.syntax().text_range(), format!("{} = {};", name, expr_stmt)); }, ) } +fn exprify_match(match_expr: &ast::MatchExpr, name: &hir::Name) -> Option { + 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, name)?.indent(block.indent_level()); + Some(arm.replace_descendant(block, new_block)) + } else { + None + } + }) + .collect::>>()?; + 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, 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::Block(ref 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( @@ -97,7 +124,7 @@ mod tests { use crate::tests::{check_assist, check_assist_not_applicable}; #[test] - fn test_extract_assignment() { + fn test_extract_assignment_if() { check_assist( extract_assigment, r#" @@ -123,6 +150,45 @@ fn foo() { ); } + #[test] + fn test_extract_assignment_match() { + check_assist( + extract_assigment, + 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_extract_assignment_not_last_not_applicable() { check_assist_not_applicable( @@ -222,7 +288,7 @@ fn foo() { } #[test] - fn extract_assignment_missing_assigment_not_applicable() { + fn extract_assignment_if_missing_assigment_not_applicable() { check_assist_not_applicable( extract_assigment, r#" @@ -232,6 +298,27 @@ fn foo() { if true { <|>a = 2; } else {} +}"#, + ) + } + + #[test] + fn extract_assignment_match_missing_assigment_not_applicable() { + check_assist_not_applicable( + extract_assigment, + r#" +fn foo() { + let mut a = 1; + + match 1 { + 1 => { + <|>a = 2; + }, + 2 => { + a = 3; + }, + 3 => {}, + } }"#, ) }