Merge pull request #22335 from A4-Tacks/lazy-method-param

fix: add param on result methods for replace_method_eager_lazy

Author: ChayimFriedman2

crates/ide-assists/src/handlers/replace_method_eager_lazy.rs

@@ -1,8 +1,8 @@
 use hir::Semantics;
-use ide_db::{RootDatabase, assists::AssistId, defs::Definition};
+use ide_db::{RootDatabase, assists::AssistId, defs::Definition, famous_defs::FamousDefs};
 use syntax::{
     AstNode,
-    ast::{self, Expr, HasArgList, make, syntax_factory::SyntaxFactory},
+    ast::{self, Expr, HasArgList, syntax_factory::SyntaxFactory},
 };
 
 use crate::{AssistContext, Assists, utils::wrap_paren_in_call};
@@ -64,9 +64,26 @@ pub(crate) fn replace_with_lazy_method(
         format!("Replace {method_name} with {method_name_lazy}"),
         call.syntax().text_range(),
         |builder| {
-            let closured = into_closure(&last_arg, &method_name_lazy);
-            builder.replace(method_name.syntax().text_range(), method_name_lazy);
-            builder.replace_ast(last_arg, closured);
+            let editor = builder.make_editor(call.syntax());
+            let add_param = match &*method_name_lazy {
+                "and_then" => true,
+                "or_else" | "unwrap_or_else" => {
+                    FamousDefs(&ctx.sema, scope.krate()).core_result_Result().is_some_and(
+                        |result| result.ty(ctx.db()).could_unify_with(ctx.db(), &receiver_ty),
+                    )
+                }
+                _ => false,
+            };
+            let closured = into_closure(&last_arg, add_param, editor.make());
+            editor.replace(method_name.syntax(), editor.make().name(&method_name_lazy).syntax());
+            editor.replace(last_arg.syntax(), closured.syntax());
+            if let Some(cap) = ctx.config.snippet_cap
+                && let ast::Expr::ClosureExpr(closured) = closured
+                && let Some(param) = closured.param_list().and_then(|it| it.params().next())
+            {
+                editor.add_annotation(param.syntax(), builder.make_placeholder_snippet(cap));
+            }
+            builder.add_file_edits(ctx.vfs_file_id(), editor);
         },
     )
 }
@@ -83,7 +100,7 @@ fn lazy_method_name(name: &str) -> String {
     }
 }
 
-fn into_closure(param: &Expr, name_lazy: &str) -> Expr {
+fn into_closure(param: &Expr, add_param: bool, make: &SyntaxFactory) -> Expr {
     (|| {
         if let ast::Expr::CallExpr(call) = param {
             if call.arg_list()?.args().count() == 0 { Some(call.expr()?) } else { None }
@@ -92,9 +109,8 @@ fn into_closure(param: &Expr, name_lazy: &str) -> Expr {
         }
     })()
     .unwrap_or_else(|| {
-        let pats = (name_lazy == "and_then")
-            .then(|| make::untyped_param(make::ext::simple_ident_pat(make::name("it")).into()));
-        make::expr_closure(pats, param.clone()).into()
+        let pats = add_param.then(|| make.untyped_param(make.wildcard_pat().into()));
+        make.expr_closure(pats, param.clone()).into()
     })
 }
 
@@ -156,14 +172,16 @@ pub(crate) fn replace_with_eager_method(
         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, &ctx.sema);
-            builder.replace_ast(last_arg, called);
+            let editor = builder.make_editor(call.syntax());
+            let called = into_call(&last_arg, &ctx.sema, editor.make());
+            editor.replace(method_name.syntax(), editor.make().name(method_name_eager).syntax());
+            editor.replace(last_arg.syntax(), called.syntax());
+            builder.add_file_edits(ctx.vfs_file_id(), editor);
         },
     )
 }
 
-fn into_call(param: &Expr, sema: &Semantics<'_, RootDatabase>) -> Expr {
+fn into_call(param: &Expr, sema: &Semantics<'_, RootDatabase>, make: &SyntaxFactory) -> Expr {
     (|| {
         if let ast::Expr::ClosureExpr(closure) = param {
             let mut params = closure.param_list()?.params();
@@ -183,8 +201,8 @@ fn into_call(param: &Expr, sema: &Semantics<'_, RootDatabase>) -> Expr {
         }
     })()
     .unwrap_or_else(|| {
-        let callable = wrap_paren_in_call(param.clone(), &SyntaxFactory::without_mappings());
-        make::expr_call(callable, make::arg_list(Vec::new())).into()
+        let callable = wrap_paren_in_call(param.clone(), make);
+        make.expr_call(callable, make.arg_list(Vec::new())).into()
     })
 }
 
@@ -213,7 +231,7 @@ mod tests {
         check_assist(
             replace_with_lazy_method,
             r#"
-//- minicore: option, fn
+//- minicore: option, result, fn
 fn foo() {
     let foo = Some(1);
     return foo.unwrap_$0or(2);
@@ -228,6 +246,26 @@ fn foo() {
         )
     }
 
+    #[test]
+    fn replace_or_with_or_else_with_parameter() {
+        check_assist(
+            replace_with_lazy_method,
+            r#"
+//- minicore: option, result, fn
+fn foo() {
+    let foo = Ok(1);
+    return foo.unwrap_$0or(2);
+}
+"#,
+            r#"
+fn foo() {
+    let foo = Ok(1);
+    return foo.unwrap_or_else(|${0:_}| 2);
+}
+"#,
+        )
+    }
+
     #[test]
     fn replace_or_with_or_else_call() {
         check_assist(
@@ -358,7 +396,7 @@ fn foo() {
             r#"
 fn foo() {
     let foo = Some("foo");
-    return foo.and_then(|it| Some("bar"));
+    return foo.and_then(|${0:_}| Some("bar"));
 }
 "#,
         )

crates/test-utils/src/minicore.rs

@@ -1713,6 +1713,43 @@ pub mod result {
         #[lang = "Err"]
         Err(E),
     }
+    impl<T, E> Result<T, E> {
+        pub const fn or<F>(self, res: Result<T, F>) -> Result<T, F> {
+            match self {
+                Ok(v) => Ok(v),
+                Err(_) => res,
+            }
+        }
+
+        pub const fn unwrap_or(self, default: T) -> T {
+            match self {
+                Ok(t) => t,
+                Err(_) => default,
+            }
+        }
+
+        // region:fn
+        pub const fn or_else<F, O>(self, op: O) -> Result<T, F>
+        where
+            O: FnOnce(E) -> Result<T, F>,
+        {
+            match self {
+                Ok(t) => Ok(t),
+                Err(e) => op(e),
+            }
+        }
+
+        pub const fn unwrap_or_else<F>(self, op: F) -> T
+        where
+            F: FnOnce(E) -> T,
+        {
+            match self {
+                Ok(t) => t,
+                Err(e) => op(e),
+            }
+        }
+        // endregion:fn
+    }
 }
 // endregion:result