Auto merge of rust-lang#14266 - Veykril:generalize-eager-lazy, r=Veykril

feature: Make replace_or_with_or_else assists more generally applicable

Author: bors

crates/hir/src/lib.rs

@@ -1677,6 +1677,10 @@ impl Function {
             .collect()
     }
 
+    pub fn num_params(self, db: &dyn HirDatabase) -> usize {
+        db.function_data(self.id).params.len()
+    }
+
     pub fn method_params(self, db: &dyn HirDatabase) -> Option<Vec<Param>> {
         if self.self_param(db).is_none() {
             return None;
@@ -3857,11 +3861,13 @@ impl Type {
     }
 }
 
+// FIXME: Document this
 #[derive(Debug)]
 pub struct Callable {
     ty: Type,
     sig: CallableSig,
     callee: Callee,
+    /// Whether this is a method that was called with method call syntax.
     pub(crate) is_bound_method: bool,
 }
 
@@ -3895,14 +3901,14 @@ impl Callable {
             Other => CallableKind::Other,
         }
     }
-    pub fn receiver_param(&self, db: &dyn HirDatabase) -> Option<ast::SelfParam> {
+    pub fn receiver_param(&self, db: &dyn HirDatabase) -> Option<(ast::SelfParam, Type)> {
         let func = match self.callee {
             Callee::Def(CallableDefId::FunctionId(it)) if self.is_bound_method => it,
             _ => return None,
         };
         let src = func.lookup(db.upcast()).source(db.upcast());
         let param_list = src.value.param_list()?;
-        param_list.self_param()
+        Some((param_list.self_param()?, self.ty.derived(self.sig.params()[0].clone())))
     }
     pub fn n_params(&self) -> usize {
         self.sig.params().len() - if self.is_bound_method { 1 } else { 0 }

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