diff --git a/crates/ra_assists/src/assists/early_return.rs b/crates/ra_assists/src/assists/early_return.rs index ad6c5695a3a..570a07a20c5 100644 --- a/crates/ra_assists/src/assists/early_return.rs +++ b/crates/ra_assists/src/assists/early_return.rs @@ -3,9 +3,10 @@ use hir::db::HirDatabase; use ra_syntax::{ algo::replace_children, - ast::{self, edit::IndentLevel, make}, + ast::{self, edit::IndentLevel, make, Block, Pat::TupleStructPat}, AstNode, SyntaxKind::{FN_DEF, LOOP_EXPR, L_CURLY, R_CURLY, WHILE_EXPR, WHITESPACE}, + SyntaxNode, }; use crate::{ @@ -37,7 +38,23 @@ // ``` pub(crate) fn convert_to_guarded_return(ctx: AssistCtx) -> Option { let if_expr: ast::IfExpr = ctx.find_node_at_offset()?; - let expr = if_expr.condition()?.expr()?; + let cond = if_expr.condition()?; + let mut if_let_ident: Option = None; + + // Check if there is an IfLet that we can handle. + match cond.pat() { + None => {} // No IfLet, supported. + Some(TupleStructPat(ref pat)) if pat.args().count() == 1usize => match &pat.path() { + Some(p) => match p.qualifier() { + None => if_let_ident = Some(p.syntax().text().to_string()), + _ => return None, + }, + _ => return None, + }, + _ => return None, // Unsupported IfLet. + }; + + let expr = cond.expr()?; let then_block = if_expr.then_branch()?.block()?; if if_expr.else_branch().is_some() { return None; @@ -63,8 +80,8 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx) -> Opt 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;"), + WHILE_EXPR | LOOP_EXPR => Some("continue"), + FN_DEF => Some("return"), _ => None, }?; @@ -77,34 +94,58 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx) -> Opt ctx.add_assist(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, - ); + let new_block = match if_let_ident { + None => { + // If. + let early_expression = &(early_expression.to_owned() + ";"); + let new_expr = + if_indent_level.increase_indent(make::if_expression(&expr, early_expression)); + replace(new_expr, &then_block, &parent_block, &if_expr) + } + Some(if_let_ident) => { + // If-let. + let new_expr = if_indent_level.increase_indent(make::let_match_early( + expr, + &if_let_ident, + early_expression, + )); + replace(new_expr, &then_block, &parent_block, &if_expr) + } + }; edit.target(if_expr.syntax().text_range()); edit.replace_ast(parent_block, ast::Block::cast(new_block).unwrap()); edit.set_cursor(cursor_position); + + fn replace( + new_expr: impl AstNode, + then_block: &Block, + parent_block: &Block, + if_expr: &ast::IfExpr, + ) -> SyntaxNode { + 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 then_statements = new_expr.syntax().children_with_tokens().chain( + then_block_items + .syntax() + .children_with_tokens() + .skip(1) + .take_while(|i| *i != end_of_then), + ); + replace_children( + &parent_block.syntax(), + RangeInclusive::new( + if_expr.clone().syntax().clone().into(), + if_expr.syntax().clone().into(), + ), + &mut then_statements, + ) + } }) } @@ -143,6 +184,68 @@ fn main() { ); } + #[test] + fn convert_let_inside_fn() { + check_assist( + convert_to_guarded_return, + r#" + fn main(n: Option) { + bar(); + if<|> let Some(n) = n { + foo(n); + + //comment + bar(); + } + } + "#, + r#" + fn main(n: Option) { + bar(); + le<|>t n = match n { + Some(it) => it, + None => return, + }; + foo(n); + + //comment + bar(); + } + "#, + ); + } + + #[test] + fn convert_let_ok_inside_fn() { + check_assist( + convert_to_guarded_return, + r#" + fn main(n: Option) { + bar(); + if<|> let Ok(n) = n { + foo(n); + + //comment + bar(); + } + } + "#, + r#" + fn main(n: Option) { + bar(); + le<|>t n = match n { + Ok(it) => it, + None => return, + }; + foo(n); + + //comment + bar(); + } + "#, + ); + } + #[test] fn convert_inside_while() { check_assist( @@ -171,6 +274,35 @@ fn main() { ); } + #[test] + fn convert_let_inside_while() { + check_assist( + convert_to_guarded_return, + r#" + fn main() { + while true { + if<|> let Some(n) = n { + foo(n); + bar(); + } + } + } + "#, + r#" + fn main() { + while true { + le<|>t n = match n { + Some(it) => it, + None => continue, + }; + foo(n); + bar(); + } + } + "#, + ); + } + #[test] fn convert_inside_loop() { check_assist( @@ -199,6 +331,35 @@ fn main() { ); } + #[test] + fn convert_let_inside_loop() { + check_assist( + convert_to_guarded_return, + r#" + fn main() { + loop { + if<|> let Some(n) = n { + foo(n); + bar(); + } + } + } + "#, + r#" + fn main() { + loop { + le<|>t n = match n { + Some(it) => it, + None => continue, + }; + foo(n); + bar(); + } + } + "#, + ); + } + #[test] fn ignore_already_converted_if() { check_assist_not_applicable( diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs index 3d5f18bfae6..95062ef6c46 100644 --- a/crates/ra_syntax/src/ast/make.rs +++ b/crates/ra_syntax/src/ast/make.rs @@ -110,6 +110,23 @@ fn from_text(text: &str) -> ast::MatchArmList { } } +pub fn let_match_early(expr: ast::Expr, path: &str, early_expression: &str) -> ast::LetStmt { + return from_text(&format!( + r#"let {} = match {} {{ + {}(it) => it, + None => {}, +}};"#, + expr.syntax().text(), + expr.syntax().text(), + path, + early_expression + )); + + fn from_text(text: &str) -> ast::LetStmt { + ast_from_text(&format!("fn f() {{ {} }}", text)) + } +} + pub fn where_pred(path: ast::Path, bounds: impl Iterator) -> ast::WherePred { let bounds = bounds.map(|b| b.syntax().to_string()).join(" + "); return from_text(&format!("{}: {}", path.syntax(), bounds));