turn `unwrap_or` into `unwrap_or_else` and vice versa

2 files changed, 233 insertions, 0 deletions

diff --git a/crates/ide-assists/src/handlers/replace_or_with_or_else.rs b/crates/ide-assists/src/handlers/replace_or_with_or_else.rs
new file mode 100644
index 00000000000..b5b5798b8bc
--- /dev/null
+++ b/crates/ide-assists/src/handlers/replace_or_with_or_else.rs
@@ -0,0 +1,230 @@
+use ide_db::assists::{AssistId, AssistKind};
+use syntax::{
+    ast::{self, make, HasArgList},
+    AstNode,
+};
+
+use crate::{AssistContext, Assists};
+
+// Assist: replace_or_with_or_else
+//
+// Replace `unwrap_or` with `unwrap_or_else` and `ok_or` with `ok_or_else`.
+//
+// ```
+// let a = Some(1);
+// a.unwra$0p_or(2);
+// ```
+// ->
+// ```
+// let a = Some(1);
+// a.unwrap_or_else(|| 2);
+// ```
+pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
+    let (name, arg_list) = (call.name_ref()?, call.arg_list()?);
+
+    let replace = match &*name.text() {
+        "unwrap_or" => "unwrap_or_else".to_string(),
+        "ok_or" => "ok_or_else".to_string(),
+        _ => return None,
+    };
+
+    let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
+        [] => make::arg_list(Vec::new()),
+        [first] => {
+            let param = (|| {
+                if let ast::Expr::CallExpr(call) = first {
+                    if call.arg_list()?.args().count() == 0 {
+                        Some(call.expr()?.clone())
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })()
+            .unwrap_or_else(|| make::expr_closure(None, first.clone()));
+            make::arg_list(vec![param])
+        }
+        _ => return None,
+    };
+
+    acc.add(
+        AssistId("replace_or_with_or_else", AssistKind::RefactorRewrite),
+        "Replace unwrap_or or ok_or with lazy version",
+        call.syntax().text_range(),
+        |builder| {
+            builder.replace(name.syntax().text_range(), replace);
+            builder.replace_ast(arg_list, arg)
+        },
+    )
+}
+
+// Assist: replace_or_else_with_or
+//
+// Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`.
+//
+// ```
+// let a = Some(1);
+// a.unwra$0p_or_else(|| 2);
+// ```
+// ->
+// ```
+// let a = Some(1);
+// a.unwrap_or(2);
+// ```
+pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
+
+    let (name, arg_list) = (call.name_ref()?, call.arg_list()?);
+
+    let replace = match &*name.text() {
+        "unwrap_or_else" => "unwrap_or".to_string(),
+        "ok_or_else" => "ok_or".to_string(),
+        _ => return None,
+    };
+
+    let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
+        [] => make::arg_list(Vec::new()),
+        [first] => {
+            let param = (|| {
+                if let ast::Expr::ClosureExpr(closure) = first {
+                    if closure.param_list()?.params().count() == 0 {
+                        Some(closure.body()?.clone())
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })()
+            .unwrap_or_else(|| make::expr_call(first.clone(), make::arg_list(Vec::new())));
+            make::arg_list(vec![param])
+        }
+        _ => return None,
+    };
+
+    acc.add(
+        AssistId("replace_or_else_with_or", AssistKind::RefactorRewrite),
+        "Replace unwrap_or_else or ok_or_else with eager version",
+        call.syntax().text_range(),
+        |builder| {
+            builder.replace(name.syntax().text_range(), replace);
+            builder.replace_ast(arg_list, arg)
+        },
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::check_assist;
+
+    use super::*;
+
+    #[test]
+    fn replace_or_with_or_else_simple() {
+        check_assist(
+            replace_or_with_or_else,
+            r#"
+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_or_with_or_else,
+            r#"
+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_or_with_or_else,
+            r#"
+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_or_else_with_or,
+            r#"
+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_or_else_with_or,
+            r#"
+fn foo() {
+    let foo = Some(1);
+    return foo.unwrap_$0or_else(x);
+}
+"#,
+            r#"
+fn foo() {
+    let foo = Some(1);
+    return foo.unwrap_or(x());
+}
+"#,
+        )
+    }
+}
diff --git a/crates/ide-assists/src/lib.rs b/crates/ide-assists/src/lib.rs
index 7fb35143fa2..94fe387efe9 100644
--- a/crates/ide-assists/src/lib.rs
+++ b/crates/ide-assists/src/lib.rs
@@ -179,6 +179,7 @@ mod handlers {
     mod replace_try_expr_with_match;
     mod replace_derive_with_manual_impl;
     mod replace_if_let_with_match;
+    mod replace_or_with_or_else;
     mod introduce_named_generic;
     mod replace_let_with_if_let;
     mod replace_qualified_name_with_use;
@@ -273,6 +274,8 @@ mod handlers {
             replace_if_let_with_match::replace_if_let_with_match,
             replace_if_let_with_match::replace_match_with_if_let,
             replace_let_with_if_let::replace_let_with_if_let,
+            replace_or_with_or_else::replace_or_else_with_or,
+            replace_or_with_or_else::replace_or_with_or_else,
             replace_turbofish_with_explicit_type::replace_turbofish_with_explicit_type,
             replace_qualified_name_with_use::replace_qualified_name_with_use,
             sort_items::sort_items,