Merge pull request #18524 from Giga-Bowser/migrate-wrap-unwrap-return

internal: Migrate `(un)wrap_return_type` assists to use `SyntaxEditor`
This commit is contained in:
Lukas Wirth 2025-01-09 07:05:48 +00:00 committed by GitHub
commit 1e9c7dd7ee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 295 additions and 121 deletions

View file

@ -189,7 +189,7 @@ pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext<'_>) -> Opti
/// This will create a turbofish generic arg list corresponding to the number of arguments
fn get_fish_head(make: &SyntaxFactory, number_of_arguments: usize) -> ast::GenericArgList {
let args = (0..number_of_arguments).map(|_| make::type_arg(make::ty_placeholder()).into());
make.turbofish_generic_arg_list(args)
make.generic_arg_list(args, true)
}
#[cfg(test)]

View file

@ -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::<ast::RetType>()?;
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,110 @@ 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();
let mut final_placeholder = None;
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::<ast::ReturnExpr, ast::StmtList>::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;
}
let new_tail_expr = make.expr_unit();
editor.replace(path_expr.syntax(), new_tail_expr.syntax());
if let Some(cap) = ctx.config.snippet_cap {
editor.add_annotation(
new_tail_expr.syntax(),
builder.make_placeholder_snippet(cap),
);
final_placeholder = Some(new_tail_expr);
}
}
_ => (),
}
}
if let Some(cap) = ctx.config.snippet_cap {
if let Some(final_placeholder) = final_placeholder {
editor.add_annotation(final_placeholder.syntax(), builder.make_tabstop_after(cap));
}
}
editor.add_mappings(make.finish_with_mappings());
builder.add_file_edits(ctx.file_id(), editor);
})
}
@ -168,12 +206,12 @@ impl UnwrapperKind {
fn tail_cb_impl(acc: &mut Vec<ast::Expr>, 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 +276,7 @@ fn foo() -> Option<()$0> {
}
"#,
r#"
fn foo() {
}
fn foo() {}
"#,
"Unwrap Option return type",
);
@ -254,8 +291,7 @@ fn foo() -> Option<()$0>{
}
"#,
r#"
fn foo() {
}
fn foo() {}
"#,
"Unwrap Option return type",
);
@ -280,7 +316,42 @@ fn foo() -> i32 {
if true {
42
} else {
()
${1:()}$0
}
}
"#,
"Unwrap Option return type",
);
}
#[test]
fn unwrap_option_return_type_multi_none() {
check_assist_by_label(
unwrap_return_type,
r#"
//- minicore: option
fn foo() -> Option<i3$02> {
if false {
return None;
}
if true {
Some(42)
} else {
None
}
}
"#,
r#"
fn foo() -> i32 {
if false {
return ${1:()};
}
if true {
42
} else {
${2:()}$0
}
}
"#,
@ -1262,8 +1333,7 @@ fn foo() -> Result<(), Box<dyn Error$0>> {
}
"#,
r#"
fn foo() {
}
fn foo() {}
"#,
"Unwrap Result return type",
);
@ -1278,8 +1348,7 @@ fn foo() -> Result<(), Box<dyn Error$0>>{
}
"#,
r#"
fn foo() {
}
fn foo() {}
"#,
"Unwrap Result return type",
);

View file

@ -6,10 +6,9 @@ use ide_db::{
famous_defs::FamousDefs,
syntax_helpers::node_ext::{for_each_tail_expr, walk_expr},
};
use itertools::Itertools;
use syntax::{
ast::{self, make, Expr, HasGenericParams},
match_ast, ted, AstNode, ToSmolStr,
ast::{self, syntax_factory::SyntaxFactory, Expr, HasGenericArgs, HasGenericParams},
match_ast, AstNode,
};
use crate::{AssistContext, AssistId, AssistKind, Assists};
@ -43,11 +42,11 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
pub(crate) fn wrap_return_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
let ret_type = ctx.find_node_at_offset::<ast::RetType>()?;
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,
Expr::BlockExpr(block) => block.into(),
// closures require a block when a return type is specified
_ => return None,
},
@ -75,56 +74,65 @@ pub(crate) fn wrap_return_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
kind.assist_id(),
kind.label(),
type_ref.syntax().text_range(),
|edit| {
let alias = wrapper_alias(ctx, &core_wrapper, type_ref, kind.symbol());
let new_return_ty =
alias.unwrap_or_else(|| kind.wrap_type(type_ref)).clone_for_update();
let body = edit.make_mut(ast::Expr::BlockExpr(body.clone()));
|builder| {
let mut editor = builder.make_editor(&parent);
let make = SyntaxFactory::new();
let alias = wrapper_alias(ctx, &make, &core_wrapper, type_ref, kind.symbol());
let new_return_ty = alias.unwrap_or_else(|| match kind {
WrapperKind::Option => make.ty_option(type_ref.clone()),
WrapperKind::Result => make.ty_result(type_ref.clone(), make.ty_infer().into()),
});
let mut exprs_to_wrap = Vec::new();
let tail_cb = &mut |e: &_| tail_cb_impl(&mut exprs_to_wrap, e);
walk_expr(&body, &mut |expr| {
walk_expr(&body_expr, &mut |expr| {
if let 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);
for ret_expr_arg in exprs_to_wrap {
let happy_wrapped = make::expr_call(
make::expr_path(make::ext::ident_path(kind.happy_ident())),
make::arg_list(iter::once(ret_expr_arg.clone())),
)
.clone_for_update();
ted::replace(ret_expr_arg.syntax(), happy_wrapped.syntax());
let happy_wrapped = make.expr_call(
make.expr_path(make.ident_path(kind.happy_ident())),
make.arg_list(iter::once(ret_expr_arg.clone())),
);
editor.replace(ret_expr_arg.syntax(), happy_wrapped.syntax());
}
let old_return_ty = edit.make_mut(type_ref.clone());
ted::replace(old_return_ty.syntax(), new_return_ty.syntax());
editor.replace(type_ref.syntax(), new_return_ty.syntax());
if let WrapperKind::Result = kind {
// Add a placeholder snippet at the first generic argument that doesn't equal the return type.
// This is normally the error type, but that may not be the case when we inserted a type alias.
let args =
new_return_ty.syntax().descendants().find_map(ast::GenericArgList::cast);
let error_type_arg = args.and_then(|list| {
list.generic_args().find(|arg| match arg {
ast::GenericArg::TypeArg(_) => {
arg.syntax().text() != type_ref.syntax().text()
}
ast::GenericArg::LifetimeArg(_) => false,
_ => true,
})
let args = new_return_ty
.path()
.unwrap()
.segment()
.unwrap()
.generic_arg_list()
.unwrap();
let error_type_arg = args.generic_args().find(|arg| match arg {
ast::GenericArg::TypeArg(_) => {
arg.syntax().text() != type_ref.syntax().text()
}
ast::GenericArg::LifetimeArg(_) => false,
_ => true,
});
if let Some(error_type_arg) = error_type_arg {
if let Some(cap) = ctx.config.snippet_cap {
edit.add_placeholder_snippet(cap, error_type_arg);
editor.add_annotation(
error_type_arg.syntax(),
builder.make_placeholder_snippet(cap),
);
}
}
}
editor.add_mappings(make.finish_with_mappings());
builder.add_file_edits(ctx.file_id(), editor);
},
);
}
@ -176,22 +184,16 @@ impl WrapperKind {
WrapperKind::Result => hir::sym::Result.clone(),
}
}
fn wrap_type(&self, type_ref: &ast::Type) -> ast::Type {
match self {
WrapperKind::Option => make::ext::ty_option(type_ref.clone()),
WrapperKind::Result => make::ext::ty_result(type_ref.clone(), make::ty_placeholder()),
}
}
}
// Try to find an wrapper type alias in the current scope (shadowing the default).
fn wrapper_alias(
ctx: &AssistContext<'_>,
make: &SyntaxFactory,
core_wrapper: &hir::Enum,
ret_type: &ast::Type,
wrapper: hir::Symbol,
) -> Option<ast::Type> {
) -> Option<ast::PathType> {
let wrapper_path = hir::ModPath::from_segments(
hir::PathKind::Plain,
iter::once(hir::Name::new_symbol_root(wrapper)),
@ -207,25 +209,28 @@ fn wrapper_alias(
})
.find_map(|alias| {
let mut inserted_ret_type = false;
let generic_params = alias
.source(ctx.db())?
.value
.generic_param_list()?
.generic_params()
.map(|param| match param {
// Replace the very first type parameter with the functions return type.
ast::GenericParam::TypeParam(_) if !inserted_ret_type => {
inserted_ret_type = true;
ret_type.to_smolstr()
let generic_args =
alias.source(ctx.db())?.value.generic_param_list()?.generic_params().map(|param| {
match param {
// Replace the very first type parameter with the function's return type.
ast::GenericParam::TypeParam(_) if !inserted_ret_type => {
inserted_ret_type = true;
make.type_arg(ret_type.clone()).into()
}
ast::GenericParam::LifetimeParam(_) => {
make.lifetime_arg(make.lifetime("'_")).into()
}
_ => make.type_arg(make.ty_infer().into()).into(),
}
ast::GenericParam::LifetimeParam(_) => make::lifetime("'_").to_smolstr(),
_ => make::ty_placeholder().to_smolstr(),
})
.join(", ");
});
let name = alias.name(ctx.db());
let name = name.as_str();
Some(make::ty(&format!("{name}<{generic_params}>")))
let generic_arg_list = make.generic_arg_list(generic_args, false);
let path = make.path_unqualified(
make.path_segment_generics(make.name_ref(name.as_str()), generic_arg_list),
);
Some(make.ty_path(path))
})
})
}

View file

@ -599,7 +599,7 @@ fn gen_partial_ord(adt: &ast::Adt, func: &ast::Fn, trait_ref: Option<TraitRef>)
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(),

View file

@ -558,8 +558,8 @@ pub fn expr_const_value(text: &str) -> ast::ConstArg {
ast_from_text(&format!("trait Foo<const N: usize = {text}> {{}}"))
}
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())

View file

@ -1,6 +1,9 @@
//! Wrappers over [`make`] constructors
use crate::{
ast::{self, make, HasGenericArgs, HasGenericParams, HasName, HasTypeBounds, HasVisibility},
ast::{
self, make, HasArgList, HasGenericArgs, HasGenericParams, HasName, HasTypeBounds,
HasVisibility,
},
syntax_editor::SyntaxMappingBuilder,
AstNode, NodeOrToken, SyntaxKind, SyntaxNode, SyntaxToken,
};
@ -16,6 +19,10 @@ impl SyntaxFactory {
make::name_ref(name).clone_for_update()
}
pub fn lifetime(&self, text: &str) -> ast::Lifetime {
make::lifetime(text).clone_for_update()
}
pub fn ty(&self, text: &str) -> ast::Type {
make::ty(text).clone_for_update()
}
@ -28,6 +35,20 @@ impl SyntaxFactory {
ast
}
pub fn ty_path(&self, path: ast::Path) -> ast::PathType {
let ast::Type::PathType(ast) = make::ty_path(path.clone()).clone_for_update() else {
unreachable!()
};
if let Some(mut mapping) = self.mappings() {
let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
builder.map_node(path.syntax().clone(), ast.path().unwrap().syntax().clone());
builder.finish(&mut mapping);
}
ast
}
pub fn type_param(
&self,
name: ast::Name,
@ -190,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<Item = ast::Expr>) -> ast::TupleExpr {
@ -253,6 +274,37 @@ impl SyntaxFactory {
ast
}
pub fn expr_call(&self, expr: ast::Expr, arg_list: ast::ArgList) -> ast::CallExpr {
// FIXME: `make::expr_call`` should return a `CallExpr`, not just an `Expr`
let ast::Expr::CallExpr(ast) =
make::expr_call(expr.clone(), arg_list.clone()).clone_for_update()
else {
unreachable!()
};
if let Some(mut mapping) = self.mappings() {
let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
builder.map_node(expr.syntax().clone(), ast.expr().unwrap().syntax().clone());
builder.map_node(arg_list.syntax().clone(), ast.arg_list().unwrap().syntax().clone());
builder.finish(&mut mapping);
}
ast
}
pub fn arg_list(&self, args: impl IntoIterator<Item = ast::Expr>) -> ast::ArgList {
let (args, input) = iterator_input(args);
let ast = make::arg_list(args).clone_for_update();
if let Some(mut mapping) = self.mappings() {
let mut builder = SyntaxMappingBuilder::new(ast.syntax.clone());
builder.map_children(input.into_iter(), ast.args().map(|it| it.syntax().clone()));
builder.finish(&mut mapping);
}
ast
}
pub fn expr_ref(&self, expr: ast::Expr, exclusive: bool) -> ast::Expr {
let ast::Expr::RefExpr(ast) = make::expr_ref(expr.clone(), exclusive).clone_for_update()
else {
@ -428,6 +480,30 @@ impl SyntaxFactory {
ast
}
pub fn type_arg(&self, ty: ast::Type) -> ast::TypeArg {
let ast = make::type_arg(ty.clone()).clone_for_update();
if let Some(mut mapping) = self.mappings() {
let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
builder.map_node(ty.syntax().clone(), ast.ty().unwrap().syntax().clone());
builder.finish(&mut mapping);
}
ast
}
pub fn lifetime_arg(&self, lifetime: ast::Lifetime) -> ast::LifetimeArg {
let ast = make::lifetime_arg(lifetime.clone()).clone_for_update();
if let Some(mut mapping) = self.mappings() {
let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
builder.map_node(lifetime.syntax().clone(), ast.lifetime().unwrap().syntax().clone());
builder.finish(&mut mapping);
}
ast
}
pub fn item_const(
&self,
visibility: Option<ast::Visibility>,
@ -495,12 +571,17 @@ impl SyntaxFactory {
ast
}
pub fn turbofish_generic_arg_list(
pub fn generic_arg_list(
&self,
generic_args: impl IntoIterator<Item = ast::GenericArg>,
is_turbo: bool,
) -> ast::GenericArgList {
let (generic_args, input) = iterator_input(generic_args);
let ast = make::turbofish_generic_arg_list(generic_args.clone()).clone_for_update();
let ast = if is_turbo {
make::turbofish_generic_arg_list(generic_args).clone_for_update()
} else {
make::generic_arg_list(generic_args).clone_for_update()
};
if let Some(mut mapping) = self.mappings() {
let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
@ -753,12 +834,31 @@ impl SyntaxFactory {
// `ext` constructors
impl SyntaxFactory {
pub fn ident_path(&self, ident: &str) -> ast::Path {
self.path_unqualified(self.path_segment(self.name_ref(ident)))
}
pub fn expr_unit(&self) -> ast::Expr {
self.expr_tuple([]).into()
}
pub fn ident_path(&self, ident: &str) -> ast::Path {
self.path_unqualified(self.path_segment(self.name_ref(ident)))
pub fn ty_option(&self, t: ast::Type) -> ast::PathType {
let generic_arg_list = self.generic_arg_list([self.type_arg(t).into()], false);
let path = self.path_unqualified(
self.path_segment_generics(self.name_ref("Option"), generic_arg_list),
);
self.ty_path(path)
}
pub fn ty_result(&self, t: ast::Type, e: ast::Type) -> ast::PathType {
let generic_arg_list =
self.generic_arg_list([self.type_arg(t).into(), self.type_arg(e).into()], false);
let path = self.path_unqualified(
self.path_segment_generics(self.name_ref("Result"), generic_arg_list),
);
self.ty_path(path)
}
}