src/tools/rust-analyzer/crates/ide-assists/src/handlers/replace_or_with_or_else.rs

   1 use ide_db::{
   2     assists::{AssistId, AssistKind},
   3     famous_defs::FamousDefs,
   4 };
   5 use syntax::{
   6     ast::{self, make, Expr, HasArgList},
   7     AstNode,
   8 };
   9
  10 use crate::{AssistContext, Assists};
  11
  12 // Assist: replace_or_with_or_else
  13 //
  14 // Replace `unwrap_or` with `unwrap_or_else` and `ok_or` with `ok_or_else`.
  15 //
  16 // ```
  17 // # //- minicore:option
  18 // fn foo() {
  19 //     let a = Some(1);
  20 //     a.unwra$0p_or(2);
  21 // }
  22 // ```
  23 // ->
  24 // ```
  25 // fn foo() {
  26 //     let a = Some(1);
  27 //     a.unwrap_or_else(|| 2);
  28 // }
  29 // ```
  30 pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  31     let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
  32
  33     let kind = is_option_or_result(call.receiver()?, ctx)?;
  34
  35     let (name, arg_list) = (call.name_ref()?, call.arg_list()?);
  36
  37     let mut map_or = false;
  38
  39     let replace = match &*name.text() {
  40         "unwrap_or" => "unwrap_or_else".to_string(),
  41         "or" => "or_else".to_string(),
  42         "ok_or" if kind == Kind::Option => "ok_or_else".to_string(),
  43         "map_or" => {
  44             map_or = true;
  45             "map_or_else".to_string()
  46         }
  47         _ => return None,
  48     };
  49
  50     let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
  51         [] => make::arg_list(Vec::new()),
  52         [first] => {
  53             let param = into_closure(first);
  54             make::arg_list(vec![param])
  55         }
  56         [first, second] if map_or => {
  57             let param = into_closure(first);
  58             make::arg_list(vec![param, second.clone()])
  59         }
  60         _ => return None,
  61     };
  62
  63     acc.add(
  64         AssistId("replace_or_with_or_else", AssistKind::RefactorRewrite),
  65         format!("Replace {} with {}", name.text(), replace),
  66         call.syntax().text_range(),
  67         |builder| {
  68             builder.replace(name.syntax().text_range(), replace);
  69             builder.replace_ast(arg_list, arg)
  70         },
  71     )
  72 }
  73
  74 fn into_closure(param: &Expr) -> Expr {
  75     (|| {
  76         if let ast::Expr::CallExpr(call) = param {
  77             if call.arg_list()?.args().count() == 0 {
  78                 Some(call.expr()?.clone())
  79             } else {
  80                 None
  81             }
  82         } else {
  83             None
  84         }
  85     })()
  86     .unwrap_or_else(|| make::expr_closure(None, param.clone()))
  87 }
  88
  89 // Assist: replace_or_else_with_or
  90 //
  91 // Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`.
  92 //
  93 // ```
  94 // # //- minicore:option
  95 // fn foo() {
  96 //     let a = Some(1);
  97 //     a.unwra$0p_or_else(|| 2);
  98 // }
  99 // ```
 100 // ->
 101 // ```
 102 // fn foo() {
 103 //     let a = Some(1);
 104 //     a.unwrap_or(2);
 105 // }
 106 // ```
 107 pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
 108     let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
 109
 110     let kind = is_option_or_result(call.receiver()?, ctx)?;
 111
 112     let (name, arg_list) = (call.name_ref()?, call.arg_list()?);
 113
 114     let mut map_or = false;
 115     let replace = match &*name.text() {
 116         "unwrap_or_else" => "unwrap_or".to_string(),
 117         "or_else" => "or".to_string(),
 118         "ok_or_else" if kind == Kind::Option => "ok_or".to_string(),
 119         "map_or_else" => {
 120             map_or = true;
 121             "map_or".to_string()
 122         }
 123         _ => return None,
 124     };
 125
 126     let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
 127         [] => make::arg_list(Vec::new()),
 128         [first] => {
 129             let param = into_call(first);
 130             make::arg_list(vec![param])
 131         }
 132         [first, second] if map_or => {
 133             let param = into_call(first);
 134             make::arg_list(vec![param, second.clone()])
 135         }
 136         _ => return None,
 137     };
 138
 139     acc.add(
 140         AssistId("replace_or_else_with_or", AssistKind::RefactorRewrite),
 141         format!("Replace {} with {}", name.text(), replace),
 142         call.syntax().text_range(),
 143         |builder| {
 144             builder.replace(name.syntax().text_range(), replace);
 145             builder.replace_ast(arg_list, arg)
 146         },
 147     )
 148 }
 149
 150 fn into_call(param: &Expr) -> Expr {
 151     (|| {
 152         if let ast::Expr::ClosureExpr(closure) = param {
 153             if closure.param_list()?.params().count() == 0 {
 154                 Some(closure.body()?.clone())
 155             } else {
 156                 None
 157             }
 158         } else {
 159             None
 160         }
 161     })()
 162     .unwrap_or_else(|| make::expr_call(param.clone(), make::arg_list(Vec::new())))
 163 }
 164
 165 #[derive(PartialEq, Eq)]
 166 enum Kind {
 167     Option,
 168     Result,
 169 }
 170
 171 fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<Kind> {
 172     let ty = ctx.sema.type_of_expr(&receiver)?.adjusted().as_adt()?.as_enum()?;
 173     let option_enum =
 174         FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_option_Option();
 175
 176     if let Some(option_enum) = option_enum {
 177         if ty == option_enum {
 178             return Some(Kind::Option);
 179         }
 180     }
 181
 182     let result_enum =
 183         FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_result_Result();
 184
 185     if let Some(result_enum) = result_enum {
 186         if ty == result_enum {
 187             return Some(Kind::Result);
 188         }
 189     }
 190
 191     None
 192 }
 193
 194 #[cfg(test)]
 195 mod tests {
 196     use crate::tests::{check_assist, check_assist_not_applicable};
 197
 198     use super::*;
 199
 200     #[test]
 201     fn replace_or_with_or_else_simple() {
 202         check_assist(
 203             replace_or_with_or_else,
 204             r#"
 205 //- minicore: option
 206 fn foo() {
 207     let foo = Some(1);
 208     return foo.unwrap_$0or(2);
 209 }
 210 "#,
 211             r#"
 212 fn foo() {
 213     let foo = Some(1);
 214     return foo.unwrap_or_else(|| 2);
 215 }
 216 "#,
 217         )
 218     }
 219
 220     #[test]
 221     fn replace_or_with_or_else_call() {
 222         check_assist(
 223             replace_or_with_or_else,
 224             r#"
 225 //- minicore: option
 226 fn foo() {
 227     let foo = Some(1);
 228     return foo.unwrap_$0or(x());
 229 }
 230 "#,
 231             r#"
 232 fn foo() {
 233     let foo = Some(1);
 234     return foo.unwrap_or_else(x);
 235 }
 236 "#,
 237         )
 238     }
 239
 240     #[test]
 241     fn replace_or_with_or_else_block() {
 242         check_assist(
 243             replace_or_with_or_else,
 244             r#"
 245 //- minicore: option
 246 fn foo() {
 247     let foo = Some(1);
 248     return foo.unwrap_$0or({
 249         let mut x = bar();
 250         for i in 0..10 {
 251             x += i;
 252         }
 253         x
 254     });
 255 }
 256 "#,
 257             r#"
 258 fn foo() {
 259     let foo = Some(1);
 260     return foo.unwrap_or_else(|| {
 261         let mut x = bar();
 262         for i in 0..10 {
 263             x += i;
 264         }
 265         x
 266     });
 267 }
 268 "#,
 269         )
 270     }
 271
 272     #[test]
 273     fn replace_or_else_with_or_simple() {
 274         check_assist(
 275             replace_or_else_with_or,
 276             r#"
 277 //- minicore: option
 278 fn foo() {
 279     let foo = Some(1);
 280     return foo.unwrap_$0or_else(|| 2);
 281 }
 282 "#,
 283             r#"
 284 fn foo() {
 285     let foo = Some(1);
 286     return foo.unwrap_or(2);
 287 }
 288 "#,
 289         )
 290     }
 291
 292     #[test]
 293     fn replace_or_else_with_or_call() {
 294         check_assist(
 295             replace_or_else_with_or,
 296             r#"
 297 //- minicore: option
 298 fn foo() {
 299     let foo = Some(1);
 300     return foo.unwrap_$0or_else(x);
 301 }
 302 "#,
 303             r#"
 304 fn foo() {
 305     let foo = Some(1);
 306     return foo.unwrap_or(x());
 307 }
 308 "#,
 309         )
 310     }
 311
 312     #[test]
 313     fn replace_or_else_with_or_result() {
 314         check_assist(
 315             replace_or_else_with_or,
 316             r#"
 317 //- minicore: result
 318 fn foo() {
 319     let foo = Ok(1);
 320     return foo.unwrap_$0or_else(x);
 321 }
 322 "#,
 323             r#"
 324 fn foo() {
 325     let foo = Ok(1);
 326     return foo.unwrap_or(x());
 327 }
 328 "#,
 329         )
 330     }
 331
 332     #[test]
 333     fn replace_or_else_with_or_map() {
 334         check_assist(
 335             replace_or_else_with_or,
 336             r#"
 337 //- minicore: result
 338 fn foo() {
 339     let foo = Ok("foo");
 340     return foo.map$0_or_else(|| 42, |v| v.len());
 341 }
 342 "#,
 343             r#"
 344 fn foo() {
 345     let foo = Ok("foo");
 346     return foo.map_or(42, |v| v.len());
 347 }
 348 "#,
 349         )
 350     }
 351
 352     #[test]
 353     fn replace_or_else_with_or_not_applicable() {
 354         check_assist_not_applicable(
 355             replace_or_else_with_or,
 356             r#"
 357 fn foo() {
 358     let foo = Ok(1);
 359     return foo.unwrap_$0or_else(x);
 360 }
 361 "#,
 362         )
 363     }
 364 }