Unnamed repository; edit this file 'description' to name the repository.
Diffstat (limited to 'crates/ide-assists/src/handlers/replace_method_eager_lazy.rs')
| -rw-r--r-- | crates/ide-assists/src/handlers/replace_method_eager_lazy.rs | 310 |
1 files changed, 310 insertions, 0 deletions
diff --git a/crates/ide-assists/src/handlers/replace_method_eager_lazy.rs b/crates/ide-assists/src/handlers/replace_method_eager_lazy.rs new file mode 100644 index 0000000000..a7e3ed793f --- /dev/null +++ b/crates/ide-assists/src/handlers/replace_method_eager_lazy.rs @@ -0,0 +1,310 @@ +use ide_db::assists::{AssistId, AssistKind}; +use syntax::{ + ast::{self, make, Expr, HasArgList}, + AstNode, +}; + +use crate::{AssistContext, Assists}; + +// Assist: replace_with_lazy_method +// +// Replace `unwrap_or` with `unwrap_or_else` and `ok_or` with `ok_or_else`. +// +// ``` +// # //- minicore:option, fn +// fn foo() { +// let a = Some(1); +// a.unwra$0p_or(2); +// } +// ``` +// -> +// ``` +// fn foo() { +// let a = Some(1); +// a.unwrap_or_else(|| 2); +// } +// ``` +pub(crate) fn replace_with_lazy_method(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { + let call: ast::MethodCallExpr = ctx.find_node_at_offset()?; + let scope = ctx.sema.scope(call.syntax())?; + + let last_arg = call.arg_list()?.args().next()?; + let method_name = call.name_ref()?; + + let callable = ctx.sema.resolve_method_call_as_callable(&call)?; + let (_, receiver_ty) = callable.receiver_param(ctx.sema.db)?; + let n_params = callable.n_params() + 1; + + let method_name_lazy = format!( + "{method_name}{}", + if method_name.text().ends_with("or") { "_else" } else { "_with" } + ); + + receiver_ty.iterate_method_candidates_with_traits( + ctx.sema.db, + &scope, + &scope.visible_traits().0, + None, + None, + |func| { + let valid = func.name(ctx.sema.db).as_str() == Some(&*method_name_lazy) + && func.num_params(ctx.sema.db) == n_params + && { + let params = func.params_without_self(ctx.sema.db); + let last_p = params.first()?; + // FIXME: Check that this has the form of `() -> T` where T is the current type of the argument + last_p.ty().impls_fnonce(ctx.sema.db) + }; + valid.then_some(func) + }, + )?; + + acc.add( + AssistId("replace_with_lazy_method", AssistKind::RefactorRewrite), + format!("Replace {method_name} with {method_name_lazy}"), + call.syntax().text_range(), + |builder| { + builder.replace(method_name.syntax().text_range(), method_name_lazy); + let closured = into_closure(&last_arg); + builder.replace_ast(last_arg, closured); + }, + ) +} + +fn into_closure(param: &Expr) -> Expr { + (|| { + if let ast::Expr::CallExpr(call) = param { + if call.arg_list()?.args().count() == 0 { + Some(call.expr()?) + } else { + None + } + } else { + None + } + })() + .unwrap_or_else(|| make::expr_closure(None, param.clone())) +} + +// Assist: replace_with_eager_method +// +// Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`. +// +// ``` +// # //- minicore:option, fn +// fn foo() { +// let a = Some(1); +// a.unwra$0p_or_else(|| 2); +// } +// ``` +// -> +// ``` +// fn foo() { +// let a = Some(1); +// a.unwrap_or(2); +// } +// ``` +pub(crate) fn replace_with_eager_method(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { + let call: ast::MethodCallExpr = ctx.find_node_at_offset()?; + let scope = ctx.sema.scope(call.syntax())?; + + let last_arg = call.arg_list()?.args().next()?; + let method_name = call.name_ref()?; + + let callable = ctx.sema.resolve_method_call_as_callable(&call)?; + let (_, receiver_ty) = callable.receiver_param(ctx.sema.db)?; + let n_params = callable.n_params() + 1; + let params = callable.params(ctx.sema.db); + + // FIXME: Check that the arg is of the form `() -> T` + if !params.first()?.1.impls_fnonce(ctx.sema.db) { + return None; + } + + let method_name_text = method_name.text(); + let method_name_eager = method_name_text + .strip_suffix("_else") + .or_else(|| method_name_text.strip_suffix("_with"))?; + + receiver_ty.iterate_method_candidates_with_traits( + ctx.sema.db, + &scope, + &scope.visible_traits().0, + None, + None, + |func| { + let valid = func.name(ctx.sema.db).as_str() == Some(&*method_name_eager) + && func.num_params(ctx.sema.db) == n_params; + valid.then_some(func) + }, + )?; + + acc.add( + AssistId("replace_with_eager_method", AssistKind::RefactorRewrite), + format!("Replace {method_name} with {method_name_eager}"), + call.syntax().text_range(), + |builder| { + builder.replace(method_name.syntax().text_range(), method_name_eager); + let called = into_call(&last_arg); + builder.replace_ast(last_arg, called); + }, + ) +} + +fn into_call(param: &Expr) -> Expr { + (|| { + if let ast::Expr::ClosureExpr(closure) = param { + if closure.param_list()?.params().count() == 0 { + Some(closure.body()?) + } else { + None + } + } else { + None + } + })() + .unwrap_or_else(|| make::expr_call(param.clone(), make::arg_list(Vec::new()))) +} + +#[cfg(test)] +mod tests { + use crate::tests::check_assist; + + use super::*; + + #[test] + fn replace_or_with_or_else_simple() { + check_assist( + replace_with_lazy_method, + r#" +//- minicore: option, fn +fn foo() { + let foo = Some(1); + return foo.unwrap_$0or(2); +} +"#, + r#" +fn foo() { + let foo = Some(1); + return foo.unwrap_or_else(|| 2); +} +"#, + ) + } + + #[test] + fn replace_or_with_or_else_call() { + check_assist( + replace_with_lazy_method, + r#" +//- minicore: option, fn +fn foo() { + let foo = Some(1); + return foo.unwrap_$0or(x()); +} +"#, + r#" +fn foo() { + let foo = Some(1); + return foo.unwrap_or_else(x); +} +"#, + ) + } + + #[test] + fn replace_or_with_or_else_block() { + check_assist( + replace_with_lazy_method, + r#" +//- minicore: option, fn +fn foo() { + let foo = Some(1); + return foo.unwrap_$0or({ + let mut x = bar(); + for i in 0..10 { + x += i; + } + x + }); +} +"#, + r#" +fn foo() { + let foo = Some(1); + return foo.unwrap_or_else(|| { + let mut x = bar(); + for i in 0..10 { + x += i; + } + x + }); +} +"#, + ) + } + + #[test] + fn replace_or_else_with_or_simple() { + check_assist( + replace_with_eager_method, + r#" +//- minicore: option, fn +fn foo() { + let foo = Some(1); + return foo.unwrap_$0or_else(|| 2); +} +"#, + r#" +fn foo() { + let foo = Some(1); + return foo.unwrap_or(2); +} +"#, + ) + } + + #[test] + fn replace_or_else_with_or_call() { + check_assist( + replace_with_eager_method, + r#" +//- minicore: option, fn +fn foo() { + let foo = Some(1); + return foo.unwrap_$0or_else(x); +} + +fn x() -> i32 { 0 } +"#, + r#" +fn foo() { + let foo = Some(1); + return foo.unwrap_or(x()); +} + +fn x() -> i32 { 0 } +"#, + ) + } + + #[test] + fn replace_or_else_with_or_map() { + check_assist( + replace_with_eager_method, + r#" +//- minicore: option, fn +fn foo() { + let foo = Some("foo"); + return foo.map$0_or_else(|| 42, |v| v.len()); +} +"#, + r#" +fn foo() { + let foo = Some("foo"); + return foo.map_or(42, |v| v.len()); +} +"#, + ) + } +} |