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()); } "#, ) } }