change as requested

Author: Austaras

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

@@ -30,33 +30,33 @@ use crate::{AssistContext, Assists};
 pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
     let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
 
-    is_option_or_result(call.receiver()?, ctx)?;
+    let kind = is_option_or_result(call.receiver()?, ctx)?;
 
     let (name, arg_list) = (call.name_ref()?, call.arg_list()?);
 
+    let mut map_or = false;
+
     let replace = match &*name.text() {
         "unwrap_or" => "unwrap_or_else".to_string(),
-        "ok_or" => "ok_or_else".to_string(),
+        "or" => "or_else".to_string(),
+        "ok_or" if kind == Kind::Option => "ok_or_else".to_string(),
+        "map_or" => {
+            map_or = true;
+            "map_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()));
+            let param = into_closure(first);
             make::arg_list(vec![param])
         }
+        [first, second] if map_or => {
+            let param = into_closure(first);
+            make::arg_list(vec![param, second.clone()])
+        }
         _ => return None,
     };
 
@@ -71,6 +71,21 @@ pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>
     )
 }
 
+fn into_closure(param: &Expr) -> Expr {
+    (|| {
+        if let ast::Expr::CallExpr(call) = param {
+            if call.arg_list()?.args().count() == 0 {
+                Some(call.expr()?.clone())
+            } else {
+                None
+            }
+        } else {
+            None
+        }
+    })()
+    .unwrap_or_else(|| make::expr_closure(None, param.clone()))
+}
+
 // Assist: replace_or_else_with_or
 //
 // Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`.
@@ -92,33 +107,32 @@ pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>
 pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
     let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
 
-    is_option_or_result(call.receiver()?, ctx)?;
+    let kind = is_option_or_result(call.receiver()?, ctx)?;
 
     let (name, arg_list) = (call.name_ref()?, call.arg_list()?);
 
+    let mut map_or = false;
     let replace = match &*name.text() {
         "unwrap_or_else" => "unwrap_or".to_string(),
-        "ok_or_else" => "ok_or".to_string(),
+        "or_else" => "or".to_string(),
+        "ok_or_else" if kind == Kind::Option => "ok_or".to_string(),
+        "map_or_else" => {
+            map_or = true;
+            "map_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())));
+            let param = into_call(first);
             make::arg_list(vec![param])
         }
+        [first, second] if map_or => {
+            let param = into_call(first);
+            make::arg_list(vec![param, second.clone()])
+        }
         _ => return None,
     };
 
@@ -133,14 +147,35 @@ pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>
     )
 }
 
-fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<()> {
+fn into_call(param: &Expr) -> Expr {
+    (|| {
+        if let ast::Expr::ClosureExpr(closure) = param {
+            if closure.param_list()?.params().count() == 0 {
+                Some(closure.body()?.clone())
+            } else {
+                None
+            }
+        } else {
+            None
+        }
+    })()
+    .unwrap_or_else(|| make::expr_call(param.clone(), make::arg_list(Vec::new())))
+}
+
+#[derive(PartialEq, Eq)]
+enum Kind {
+    Option,
+    Result,
+}
+
+fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<Kind> {
     let ty = ctx.sema.type_of_expr(&receiver)?.adjusted().as_adt()?.as_enum()?;
     let option_enum =
         FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_option_Option();
 
     if let Some(option_enum) = option_enum {
         if ty == option_enum {
-            return Some(());
+            return Some(Kind::Option);
         }
     }
 
@@ -149,7 +184,7 @@ fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<()> {
 
     if let Some(result_enum) = result_enum {
         if ty == result_enum {
-            return Some(());
+            return Some(Kind::Result);
         }
     }
 
@@ -294,6 +329,26 @@ fn foo() {
         )
     }
 
+    #[test]
+    fn replace_or_else_with_or_map() {
+        check_assist(
+            replace_or_else_with_or,
+            r#"
+//- minicore: result
+fn foo() {
+    let foo = Ok("foo");
+    return foo.map$0_or_else(|| 42, |v| v.len());
+}
+"#,
+            r#"
+fn foo() {
+    let foo = Ok("foo");
+    return foo.map_or(42, |v| v.len());
+}
+"#,
+        )
+    }
+
     #[test]
     fn replace_or_else_with_or_not_applicable() {
         check_assist_not_applicable(