From a5a79f5957f2a821ea8562c2b9d04e9304378d7e Mon Sep 17 00:00:00 2001 From: Giga Bowser <45986823+Giga-Bowser@users.noreply.github.com> Date: Sun, 17 Nov 2024 14:03:28 -0500 Subject: [PATCH] internal: Migrate `unwrap_return_type` assist to use `SyntaxEditor` Also changes `make::expr_empty_block()` to return `ast::BlockExpr` instead of `ast::Expr` --- .../src/handlers/unwrap_return_type.rs | 124 ++++++++++-------- .../src/utils/gen_trait_fn_body.rs | 2 +- crates/syntax/src/ast/make.rs | 4 +- .../src/ast/syntax_factory/constructors.rs | 2 +- 4 files changed, 75 insertions(+), 57 deletions(-) diff --git a/crates/ide-assists/src/handlers/unwrap_return_type.rs b/crates/ide-assists/src/handlers/unwrap_return_type.rs index 64d5e2c9b8..6fd46260c0 100644 --- a/crates/ide-assists/src/handlers/unwrap_return_type.rs +++ b/crates/ide-assists/src/handlers/unwrap_return_type.rs @@ -1,11 +1,11 @@ +use either::Either; use ide_db::{ famous_defs::FamousDefs, syntax_helpers::node_ext::{for_each_tail_expr, walk_expr}, }; -use itertools::Itertools; use syntax::{ - ast::{self, Expr, HasGenericArgs}, - match_ast, AstNode, NodeOrToken, SyntaxKind, TextRange, + ast::{self, syntax_factory::SyntaxFactory, HasArgList, HasGenericArgs}, + match_ast, AstNode, NodeOrToken, SyntaxKind, }; use crate::{AssistContext, AssistId, AssistKind, Assists}; @@ -39,11 +39,11 @@ use crate::{AssistContext, AssistId, AssistKind, Assists}; pub(crate) fn unwrap_return_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { let ret_type = ctx.find_node_at_offset::()?; let parent = ret_type.syntax().parent()?; - let body = match_ast! { + let body_expr = match_ast! { match parent { - ast::Fn(func) => func.body()?, + ast::Fn(func) => func.body()?.into(), ast::ClosureExpr(closure) => match closure.body()? { - Expr::BlockExpr(block) => block, + ast::Expr::BlockExpr(block) => block.into(), // closures require a block when a return type is specified _ => return None, }, @@ -65,72 +65,94 @@ pub(crate) fn unwrap_return_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> let happy_type = extract_wrapped_type(type_ref)?; acc.add(kind.assist_id(), kind.label(), type_ref.syntax().text_range(), |builder| { - let body = ast::Expr::BlockExpr(body); + let mut editor = builder.make_editor(&parent); + let make = SyntaxFactory::new(); let mut exprs_to_unwrap = Vec::new(); let tail_cb = &mut |e: &_| tail_cb_impl(&mut exprs_to_unwrap, e); - walk_expr(&body, &mut |expr| { - if let Expr::ReturnExpr(ret_expr) = expr { + walk_expr(&body_expr, &mut |expr| { + if let ast::Expr::ReturnExpr(ret_expr) = expr { if let Some(ret_expr_arg) = &ret_expr.expr() { for_each_tail_expr(ret_expr_arg, tail_cb); } } }); - for_each_tail_expr(&body, tail_cb); + for_each_tail_expr(&body_expr, tail_cb); let is_unit_type = is_unit_type(&happy_type); if is_unit_type { - let mut text_range = ret_type.syntax().text_range(); - if let Some(NodeOrToken::Token(token)) = ret_type.syntax().next_sibling_or_token() { if token.kind() == SyntaxKind::WHITESPACE { - text_range = TextRange::new(text_range.start(), token.text_range().end()); + editor.delete(token); } } - builder.delete(text_range); + editor.delete(ret_type.syntax()); } else { - builder.replace(type_ref.syntax().text_range(), happy_type.syntax().text()); + editor.replace(type_ref.syntax(), happy_type.syntax()); } - for ret_expr_arg in exprs_to_unwrap { - let ret_expr_str = ret_expr_arg.to_string(); + for tail_expr in exprs_to_unwrap { + match &tail_expr { + ast::Expr::CallExpr(call_expr) => { + let ast::Expr::PathExpr(path_expr) = call_expr.expr().unwrap() else { + continue; + }; - let needs_replacing = match kind { - UnwrapperKind::Option => ret_expr_str.starts_with("Some("), - UnwrapperKind::Result => { - ret_expr_str.starts_with("Ok(") || ret_expr_str.starts_with("Err(") - } - }; + let path_str = path_expr.path().unwrap().to_string(); + let needs_replacing = match kind { + UnwrapperKind::Option => path_str == "Some", + UnwrapperKind::Result => path_str == "Ok" || path_str == "Err", + }; - if needs_replacing { - let arg_list = ret_expr_arg.syntax().children().find_map(ast::ArgList::cast); - if let Some(arg_list) = arg_list { + if !needs_replacing { + continue; + } + + let arg_list = call_expr.arg_list().unwrap(); if is_unit_type { - match ret_expr_arg.syntax().prev_sibling_or_token() { - // Useful to delete the entire line without leaving trailing whitespaces - Some(whitespace) => { - let new_range = TextRange::new( - whitespace.text_range().start(), - ret_expr_arg.syntax().text_range().end(), - ); - builder.delete(new_range); + let tail_parent = tail_expr + .syntax() + .parent() + .and_then(Either::::cast) + .unwrap(); + match tail_parent { + Either::Left(ret_expr) => { + editor.replace(ret_expr.syntax(), make.expr_return(None).syntax()) } - None => { - builder.delete(ret_expr_arg.syntax().text_range()); + Either::Right(stmt_list) => { + let new_block = if stmt_list.statements().next().is_none() { + make.expr_empty_block() + } else { + make.block_expr(stmt_list.statements(), None) + }; + editor.replace( + stmt_list.syntax(), + new_block.stmt_list().unwrap().syntax(), + ); } } - } else { - builder.replace( - ret_expr_arg.syntax().text_range(), - arg_list.args().join(", "), - ); + } else if let Some(first_arg) = arg_list.args().next() { + editor.replace(tail_expr.syntax(), first_arg.syntax()); } } - } else if matches!(kind, UnwrapperKind::Option if ret_expr_str == "None") { - builder.replace(ret_expr_arg.syntax().text_range(), "()"); + ast::Expr::PathExpr(path_expr) => { + let UnwrapperKind::Option = kind else { + continue; + }; + + if path_expr.path().unwrap().to_string() != "None" { + continue; + } + + editor.replace(path_expr.syntax(), make.expr_unit().syntax()); + } + _ => (), } } + + editor.add_mappings(make.finish_with_mappings()); + builder.add_file_edits(ctx.file_id(), editor); }) } @@ -168,12 +190,12 @@ impl UnwrapperKind { fn tail_cb_impl(acc: &mut Vec, e: &ast::Expr) { match e { - Expr::BreakExpr(break_expr) => { + ast::Expr::BreakExpr(break_expr) => { if let Some(break_expr_arg) = break_expr.expr() { for_each_tail_expr(&break_expr_arg, &mut |e| tail_cb_impl(acc, e)) } } - Expr::ReturnExpr(_) => { + ast::Expr::ReturnExpr(_) => { // all return expressions have already been handled by the walk loop } e => acc.push(e.clone()), @@ -238,8 +260,7 @@ fn foo() -> Option<()$0> { } "#, r#" -fn foo() { -} +fn foo() {} "#, "Unwrap Option return type", ); @@ -254,8 +275,7 @@ fn foo() -> Option<()$0>{ } "#, r#" -fn foo() { -} +fn foo() {} "#, "Unwrap Option return type", ); @@ -1262,8 +1282,7 @@ fn foo() -> Result<(), Box> { } "#, r#" -fn foo() { -} +fn foo() {} "#, "Unwrap Result return type", ); @@ -1278,8 +1297,7 @@ fn foo() -> Result<(), Box>{ } "#, r#" -fn foo() { -} +fn foo() {} "#, "Unwrap Result return type", ); diff --git a/crates/ide-assists/src/utils/gen_trait_fn_body.rs b/crates/ide-assists/src/utils/gen_trait_fn_body.rs index a9a889acc2..7a9bdfe1ec 100644 --- a/crates/ide-assists/src/utils/gen_trait_fn_body.rs +++ b/crates/ide-assists/src/utils/gen_trait_fn_body.rs @@ -599,7 +599,7 @@ fn gen_partial_ord(adt: &ast::Adt, func: &ast::Fn, trait_ref: Option) let variant_name = make::path_pat(make::ext::path_from_idents(["core", "cmp", "Ordering", "Equal"])?); let lhs = make::tuple_struct_pat(make::ext::path_from_idents(["Some"])?, [variant_name]); - arms.push(make::match_arm(lhs.into(), None, make::expr_empty_block())); + arms.push(make::match_arm(lhs.into(), None, make::expr_empty_block().into())); arms.push(make::match_arm( make::ident_pat(false, false, make::name("ord")).into(), diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs index a920fe4f35..dca231604f 100644 --- a/crates/syntax/src/ast/make.rs +++ b/crates/syntax/src/ast/make.rs @@ -558,8 +558,8 @@ pub fn expr_const_value(text: &str) -> ast::ConstArg { ast_from_text(&format!("trait Foo {{}}")) } -pub fn expr_empty_block() -> ast::Expr { - expr_from_text("{}") +pub fn expr_empty_block() -> ast::BlockExpr { + ast_from_text("const C: () = {};") } pub fn expr_path(path: ast::Path) -> ast::Expr { expr_from_text(&path.to_string()) diff --git a/crates/syntax/src/ast/syntax_factory/constructors.rs b/crates/syntax/src/ast/syntax_factory/constructors.rs index af7b3c8158..572622db54 100644 --- a/crates/syntax/src/ast/syntax_factory/constructors.rs +++ b/crates/syntax/src/ast/syntax_factory/constructors.rs @@ -211,7 +211,7 @@ impl SyntaxFactory { } pub fn expr_empty_block(&self) -> ast::BlockExpr { - ast::BlockExpr { syntax: make::expr_empty_block().syntax().clone_for_update() } + make::expr_empty_block().clone_for_update() } pub fn expr_tuple(&self, fields: impl IntoIterator) -> ast::TupleExpr {