diff --git a/crates/ra_assists/src/assists/early_return.rs b/crates/ra_assists/src/assists/early_return.rs new file mode 100644 index 00000000000..8c975714c81 --- /dev/null +++ b/crates/ra_assists/src/assists/early_return.rs @@ -0,0 +1,276 @@ +//! FIXME: write short doc here + +use crate::{ + assist_ctx::{Assist, AssistCtx}, + AssistId, +}; +use hir::db::HirDatabase; +use ra_syntax::{ + algo::replace_children, + ast::edit::IndentLevel, + ast::make, + ast::Block, + ast::ContinueExpr, + ast::IfExpr, + ast::ReturnExpr, + AstNode, + SyntaxKind::{FN_DEF, LOOP_EXPR, L_CURLY, R_CURLY, WHILE_EXPR, WHITESPACE}, +}; +use std::ops::RangeInclusive; + +pub(crate) fn convert_to_guarded_return(mut ctx: AssistCtx) -> Option { + let if_expr: IfExpr = ctx.node_at_offset()?; + let expr = if_expr.condition()?.expr()?; + let then_block = if_expr.then_branch()?.block()?; + if if_expr.else_branch().is_some() { + return None; + } + + let parent_block = if_expr.syntax().parent()?.ancestors().find_map(Block::cast)?; + + if parent_block.expr()? != if_expr.clone().into() { + return None; + } + + // check for early return and continue + let first_in_then_block = then_block.syntax().first_child()?.clone(); + if ReturnExpr::can_cast(first_in_then_block.kind()) + || ContinueExpr::can_cast(first_in_then_block.kind()) + || first_in_then_block + .children() + .any(|x| ReturnExpr::can_cast(x.kind()) || ContinueExpr::can_cast(x.kind())) + { + return None; + } + + let parent_container = parent_block.syntax().parent()?.parent()?; + + let early_expression = match parent_container.kind() { + WHILE_EXPR | LOOP_EXPR => Some("continue;"), + FN_DEF => Some("return;"), + _ => None, + }?; + + if then_block.syntax().first_child_or_token().map(|t| t.kind() == L_CURLY).is_none() { + return None; + } + + then_block.syntax().last_child_or_token().filter(|t| t.kind() == R_CURLY)?; + let cursor_position = ctx.frange.range.start(); + + ctx.add_action(AssistId("convert_to_guarded_return"), "convert to guarded return", |edit| { + let if_indent_level = IndentLevel::from_node(&if_expr.syntax()); + let new_if_expr = + if_indent_level.increase_indent(make::if_expression(&expr, early_expression)); + let then_block_items = IndentLevel::from(1).decrease_indent(then_block.clone()); + let end_of_then = then_block_items.syntax().last_child_or_token().unwrap(); + let end_of_then = + if end_of_then.prev_sibling_or_token().map(|n| n.kind()) == Some(WHITESPACE) { + end_of_then.prev_sibling_or_token().unwrap() + } else { + end_of_then + }; + let mut new_if_and_then_statements = new_if_expr.syntax().children_with_tokens().chain( + then_block_items + .syntax() + .children_with_tokens() + .skip(1) + .take_while(|i| *i != end_of_then), + ); + let new_block = replace_children( + &parent_block.syntax(), + RangeInclusive::new( + if_expr.clone().syntax().clone().into(), + if_expr.syntax().clone().into(), + ), + &mut new_if_and_then_statements, + ); + edit.target(if_expr.syntax().text_range()); + edit.replace_ast(parent_block, Block::cast(new_block).unwrap()); + edit.set_cursor(cursor_position); + }); + ctx.build() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::helpers::{check_assist, check_assist_not_applicable}; + + #[test] + fn convert_inside_fn() { + check_assist( + convert_to_guarded_return, + r#" + fn main() { + bar(); + if<|> true { + foo(); + + //comment + bar(); + } + } + "#, + r#" + fn main() { + bar(); + if<|> !true { + return; + } + foo(); + + //comment + bar(); + } + "#, + ); + } + + #[test] + fn convert_inside_while() { + check_assist( + convert_to_guarded_return, + r#" + fn main() { + while true { + if<|> true { + foo(); + bar(); + } + } + } + "#, + r#" + fn main() { + while true { + if<|> !true { + continue; + } + foo(); + bar(); + } + } + "#, + ); + } + + #[test] + fn convert_inside_loop() { + check_assist( + convert_to_guarded_return, + r#" + fn main() { + loop { + if<|> true { + foo(); + bar(); + } + } + } + "#, + r#" + fn main() { + loop { + if<|> !true { + continue; + } + foo(); + bar(); + } + } + "#, + ); + } + + #[test] + fn ignore_already_converted_if() { + check_assist_not_applicable( + convert_to_guarded_return, + r#" + fn main() { + if<|> true { + return; + } + } + "#, + ); + } + + #[test] + fn ignore_already_converted_loop() { + check_assist_not_applicable( + convert_to_guarded_return, + r#" + fn main() { + loop { + if<|> true { + continue; + } + } + } + "#, + ); + } + + #[test] + fn ignore_return() { + check_assist_not_applicable( + convert_to_guarded_return, + r#" + fn main() { + if<|> true { + return + } + } + "#, + ); + } + + #[test] + fn ignore_else_branch() { + check_assist_not_applicable( + convert_to_guarded_return, + r#" + fn main() { + if<|> true { + foo(); + } else { + bar() + } + } + "#, + ); + } + + #[test] + fn ignore_statements_aftert_if() { + check_assist_not_applicable( + convert_to_guarded_return, + r#" + fn main() { + if<|> true { + foo(); + } + bar(); + } + "#, + ); + } + + #[test] + fn ignore_statements_inside_if() { + check_assist_not_applicable( + convert_to_guarded_return, + r#" + fn main() { + if false { + if<|> true { + foo(); + } + } + } + "#, + ); + } +} diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs index d2376c475cb..ab77b46a997 100644 --- a/crates/ra_assists/src/lib.rs +++ b/crates/ra_assists/src/lib.rs @@ -108,6 +108,7 @@ mod assists { mod add_missing_impl_members; mod move_guard; mod move_bounds; + mod early_return; pub(crate) fn all() -> &'static [fn(AssistCtx) -> Option] { &[ @@ -135,6 +136,7 @@ mod assists { raw_string::make_raw_string, raw_string::make_usual_string, raw_string::remove_hash, + early_return::convert_to_guarded_return, ] } } diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs index ea92284b833..47bdbb81a1a 100644 --- a/crates/ra_syntax/src/ast/edit.rs +++ b/crates/ra_syntax/src/ast/edit.rs @@ -284,6 +284,34 @@ impl IndentLevel { .collect(); algo::replace_descendants(&node, &replacements) } + + pub fn decrease_indent(self, node: N) -> N { + N::cast(self._decrease_indent(node.syntax().clone())).unwrap() + } + + fn _decrease_indent(self, node: SyntaxNode) -> SyntaxNode { + let replacements: FxHashMap = node + .descendants_with_tokens() + .filter_map(|el| el.into_token()) + .filter_map(ast::Whitespace::cast) + .filter(|ws| { + let text = ws.syntax().text(); + text.contains('\n') + }) + .map(|ws| { + ( + ws.syntax().clone().into(), + make::tokens::whitespace( + &ws.syntax() + .text() + .replace(&format!("\n{:1$}", "", self.0 as usize * 4), "\n"), + ) + .into(), + ) + }) + .collect(); + algo::replace_descendants(&node, &replacements) + } } // FIXME: replace usages with IndentLevel above diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs index 143835172af..00422ea913a 100644 --- a/crates/ra_syntax/src/ast/make.rs +++ b/crates/ra_syntax/src/ast/make.rs @@ -128,6 +128,14 @@ pub fn where_clause(preds: impl Iterator) -> ast::WhereCl } } +pub fn if_expression(condition: &ast::Expr, statement: &str) -> ast::IfExpr { + return ast_from_text(&format!( + "fn f() {{ if !{} {{\n {}\n}}\n}}", + condition.syntax().text(), + statement + )); +} + fn ast_from_text(text: &str) -> N { let parse = SourceFile::parse(text); let res = parse.tree().syntax().descendants().find_map(N::cast).unwrap();