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

   1 use syntax::{
   2     ast::{
   3         self,
   4         edit::{AstNodeEdit, IndentLevel},
   5     },
   6     AstNode, SyntaxKind, TextRange, T,
   7 };
   8
   9 use crate::{AssistContext, AssistId, AssistKind, Assists};
  10
  11 // Assist: unwrap_block
  12 //
  13 // This assist removes if...else, for, while and loop control statements to just keep the body.
  14 //
  15 // ```
  16 // fn foo() {
  17 //     if true {$0
  18 //         println!("foo");
  19 //     }
  20 // }
  21 // ```
  22 // ->
  23 // ```
  24 // fn foo() {
  25 //     println!("foo");
  26 // }
  27 // ```
  28 pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  29     let assist_id = AssistId("unwrap_block", AssistKind::RefactorRewrite);
  30     let assist_label = "Unwrap block";
  31
  32     let l_curly_token = ctx.find_token_syntax_at_offset(T!['{'])?;
  33     let mut block = ast::BlockExpr::cast(l_curly_token.parent_ancestors().nth(1)?)?;
  34     let target = block.syntax().text_range();
  35     let mut parent = block.syntax().parent()?;
  36     if ast::MatchArm::can_cast(parent.kind()) {
  37         parent = parent.ancestors().find(|it| ast::MatchExpr::can_cast(it.kind()))?
  38     }
  39
  40     if matches!(parent.kind(), SyntaxKind::STMT_LIST | SyntaxKind::EXPR_STMT) {
  41         return acc.add(assist_id, assist_label, target, |builder| {
  42             builder.replace(block.syntax().text_range(), update_expr_string(block.to_string()));
  43         });
  44     }
  45
  46     let parent = ast::Expr::cast(parent)?;
  47
  48     match parent.clone() {
  49         ast::Expr::ForExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::LoopExpr(_) => (),
  50         ast::Expr::MatchExpr(_) => block = block.dedent(IndentLevel(1)),
  51         ast::Expr::IfExpr(if_expr) => {
  52             let then_branch = if_expr.then_branch()?;
  53             if then_branch == block {
  54                 if let Some(ancestor) = if_expr.syntax().parent().and_then(ast::IfExpr::cast) {
  55                     // For `else if` blocks
  56                     let ancestor_then_branch = ancestor.then_branch()?;
  57
  58                     return acc.add(assist_id, assist_label, target, |edit| {
  59                         let range_to_del_else_if = TextRange::new(
  60                             ancestor_then_branch.syntax().text_range().end(),
  61                             l_curly_token.text_range().start(),
  62                         );
  63                         let range_to_del_rest = TextRange::new(
  64                             then_branch.syntax().text_range().end(),
  65                             if_expr.syntax().text_range().end(),
  66                         );
  67
  68                         edit.delete(range_to_del_rest);
  69                         edit.delete(range_to_del_else_if);
  70                         edit.replace(
  71                             target,
  72                             update_expr_string_without_newline(then_branch.to_string()),
  73                         );
  74                     });
  75                 }
  76             } else {
  77                 return acc.add(assist_id, assist_label, target, |edit| {
  78                     let range_to_del = TextRange::new(
  79                         then_branch.syntax().text_range().end(),
  80                         l_curly_token.text_range().start(),
  81                     );
  82
  83                     edit.delete(range_to_del);
  84                     edit.replace(target, update_expr_string_without_newline(block.to_string()));
  85                 });
  86             }
  87         }
  88         _ => return None,
  89     };
  90
  91     acc.add(assist_id, assist_label, target, |builder| {
  92         builder.replace(parent.syntax().text_range(), update_expr_string(block.to_string()));
  93     })
  94 }
  95
  96 fn update_expr_string(expr_string: String) -> String {
  97     update_expr_string_with_pat(expr_string, &[' ', '\n'])
  98 }
  99
 100 fn update_expr_string_without_newline(expr_string: String) -> String {
 101     update_expr_string_with_pat(expr_string, &[' '])
 102 }
 103
 104 fn update_expr_string_with_pat(expr_str: String, whitespace_pat: &[char]) -> String {
 105     // Remove leading whitespace, index [1..] to remove the leading '{',
 106     // then continue to remove leading whitespace.
 107     let expr_str =
 108         expr_str.trim_start_matches(whitespace_pat)[1..].trim_start_matches(whitespace_pat);
 109
 110     // Remove trailing whitespace, index [..expr_str.len() - 1] to remove the trailing '}',
 111     // then continue to remove trailing whitespace.
 112     let expr_str = expr_str.trim_end_matches(whitespace_pat);
 113     let expr_str = expr_str[..expr_str.len() - 1].trim_end_matches(whitespace_pat);
 114
 115     expr_str
 116         .lines()
 117         .map(|line| line.replacen("    ", "", 1)) // Delete indentation
 118         .collect::<Vec<String>>()
 119         .join("\n")
 120 }
 121
 122 #[cfg(test)]
 123 mod tests {
 124     use crate::tests::{check_assist, check_assist_not_applicable};
 125
 126     use super::*;
 127
 128     #[test]
 129     fn unwrap_tail_expr_block() {
 130         check_assist(
 131             unwrap_block,
 132             r#"
 133 fn main() {
 134     $0{
 135         92
 136     }
 137 }
 138 "#,
 139             r#"
 140 fn main() {
 141     92
 142 }
 143 "#,
 144         )
 145     }
 146
 147     #[test]
 148     fn unwrap_stmt_expr_block() {
 149         check_assist(
 150             unwrap_block,
 151             r#"
 152 fn main() {
 153     $0{
 154         92;
 155     }
 156     ()
 157 }
 158 "#,
 159             r#"
 160 fn main() {
 161     92;
 162     ()
 163 }
 164 "#,
 165         );
 166         // Pedantically, we should add an `;` here...
 167         check_assist(
 168             unwrap_block,
 169             r#"
 170 fn main() {
 171     $0{
 172         92
 173     }
 174     ()
 175 }
 176 "#,
 177             r#"
 178 fn main() {
 179     92
 180     ()
 181 }
 182 "#,
 183         );
 184     }
 185
 186     #[test]
 187     fn simple_if() {
 188         check_assist(
 189             unwrap_block,
 190             r#"
 191 fn main() {
 192     bar();
 193     if true {$0
 194         foo();
 195
 196         // comment
 197         bar();
 198     } else {
 199         println!("bar");
 200     }
 201 }
 202 "#,
 203             r#"
 204 fn main() {
 205     bar();
 206     foo();
 207
 208     // comment
 209     bar();
 210 }
 211 "#,
 212         );
 213     }
 214
 215     #[test]
 216     fn simple_if_else() {
 217         check_assist(
 218             unwrap_block,
 219             r#"
 220 fn main() {
 221     bar();
 222     if true {
 223         foo();
 224
 225         // comment
 226         bar();
 227     } else {$0
 228         println!("bar");
 229     }
 230 }
 231 "#,
 232             r#"
 233 fn main() {
 234     bar();
 235     if true {
 236         foo();
 237
 238         // comment
 239         bar();
 240     }
 241     println!("bar");
 242 }
 243 "#,
 244         );
 245     }
 246
 247     #[test]
 248     fn simple_if_else_if() {
 249         check_assist(
 250             unwrap_block,
 251             r#"
 252 fn main() {
 253     // bar();
 254     if true {
 255         println!("true");
 256
 257         // comment
 258         // bar();
 259     } else if false {$0
 260         println!("bar");
 261     } else {
 262         println!("foo");
 263     }
 264 }
 265 "#,
 266             r#"
 267 fn main() {
 268     // bar();
 269     if true {
 270         println!("true");
 271
 272         // comment
 273         // bar();
 274     }
 275     println!("bar");
 276 }
 277 "#,
 278         );
 279     }
 280
 281     #[test]
 282     fn simple_if_else_if_nested() {
 283         check_assist(
 284             unwrap_block,
 285             r#"
 286 fn main() {
 287     // bar();
 288     if true {
 289         println!("true");
 290
 291         // comment
 292         // bar();
 293     } else if false {
 294         println!("bar");
 295     } else if true {$0
 296         println!("foo");
 297     }
 298 }
 299 "#,
 300             r#"
 301 fn main() {
 302     // bar();
 303     if true {
 304         println!("true");
 305
 306         // comment
 307         // bar();
 308     } else if false {
 309         println!("bar");
 310     }
 311     println!("foo");
 312 }
 313 "#,
 314         );
 315     }
 316
 317     #[test]
 318     fn simple_if_else_if_nested_else() {
 319         check_assist(
 320             unwrap_block,
 321             r#"
 322 fn main() {
 323     // bar();
 324     if true {
 325         println!("true");
 326
 327         // comment
 328         // bar();
 329     } else if false {
 330         println!("bar");
 331     } else if true {
 332         println!("foo");
 333     } else {$0
 334         println!("else");
 335     }
 336 }
 337 "#,
 338             r#"
 339 fn main() {
 340     // bar();
 341     if true {
 342         println!("true");
 343
 344         // comment
 345         // bar();
 346     } else if false {
 347         println!("bar");
 348     } else if true {
 349         println!("foo");
 350     }
 351     println!("else");
 352 }
 353 "#,
 354         );
 355     }
 356
 357     #[test]
 358     fn simple_if_else_if_nested_middle() {
 359         check_assist(
 360             unwrap_block,
 361             r#"
 362 fn main() {
 363     // bar();
 364     if true {
 365         println!("true");
 366
 367         // comment
 368         // bar();
 369     } else if false {
 370         println!("bar");
 371     } else if true {$0
 372         println!("foo");
 373     } else {
 374         println!("else");
 375     }
 376 }
 377 "#,
 378             r#"
 379 fn main() {
 380     // bar();
 381     if true {
 382         println!("true");
 383
 384         // comment
 385         // bar();
 386     } else if false {
 387         println!("bar");
 388     }
 389     println!("foo");
 390 }
 391 "#,
 392         );
 393     }
 394
 395     #[test]
 396     fn simple_if_bad_cursor_position() {
 397         check_assist_not_applicable(
 398             unwrap_block,
 399             r#"
 400 fn main() {
 401     bar();$0
 402     if true {
 403         foo();
 404
 405         // comment
 406         bar();
 407     } else {
 408         println!("bar");
 409     }
 410 }
 411 "#,
 412         );
 413     }
 414
 415     #[test]
 416     fn simple_for() {
 417         check_assist(
 418             unwrap_block,
 419             r#"
 420 fn main() {
 421     for i in 0..5 {$0
 422         if true {
 423             foo();
 424
 425             // comment
 426             bar();
 427         } else {
 428             println!("bar");
 429         }
 430     }
 431 }
 432 "#,
 433             r#"
 434 fn main() {
 435     if true {
 436         foo();
 437
 438         // comment
 439         bar();
 440     } else {
 441         println!("bar");
 442     }
 443 }
 444 "#,
 445         );
 446     }
 447
 448     #[test]
 449     fn simple_if_in_for() {
 450         check_assist(
 451             unwrap_block,
 452             r#"
 453 fn main() {
 454     for i in 0..5 {
 455         if true {$0
 456             foo();
 457
 458             // comment
 459             bar();
 460         } else {
 461             println!("bar");
 462         }
 463     }
 464 }
 465 "#,
 466             r#"
 467 fn main() {
 468     for i in 0..5 {
 469         foo();
 470
 471         // comment
 472         bar();
 473     }
 474 }
 475 "#,
 476         );
 477     }
 478
 479     #[test]
 480     fn simple_loop() {
 481         check_assist(
 482             unwrap_block,
 483             r#"
 484 fn main() {
 485     loop {$0
 486         if true {
 487             foo();
 488
 489             // comment
 490             bar();
 491         } else {
 492             println!("bar");
 493         }
 494     }
 495 }
 496 "#,
 497             r#"
 498 fn main() {
 499     if true {
 500         foo();
 501
 502         // comment
 503         bar();
 504     } else {
 505         println!("bar");
 506     }
 507 }
 508 "#,
 509         );
 510     }
 511
 512     #[test]
 513     fn simple_while() {
 514         check_assist(
 515             unwrap_block,
 516             r#"
 517 fn main() {
 518     while true {$0
 519         if true {
 520             foo();
 521
 522             // comment
 523             bar();
 524         } else {
 525             println!("bar");
 526         }
 527     }
 528 }
 529 "#,
 530             r#"
 531 fn main() {
 532     if true {
 533         foo();
 534
 535         // comment
 536         bar();
 537     } else {
 538         println!("bar");
 539     }
 540 }
 541 "#,
 542         );
 543     }
 544
 545     #[test]
 546     fn unwrap_match_arm() {
 547         check_assist(
 548             unwrap_block,
 549             r#"
 550 fn main() {
 551     match rel_path {
 552         Ok(rel_path) => {$0
 553             let rel_path = RelativePathBuf::from_path(rel_path).ok()?;
 554             Some((*id, rel_path))
 555         }
 556         Err(_) => None,
 557     }
 558 }
 559 "#,
 560             r#"
 561 fn main() {
 562     let rel_path = RelativePathBuf::from_path(rel_path).ok()?;
 563     Some((*id, rel_path))
 564 }
 565 "#,
 566         );
 567     }
 568
 569     #[test]
 570     fn simple_if_in_while_bad_cursor_position() {
 571         check_assist_not_applicable(
 572             unwrap_block,
 573             r#"
 574 fn main() {
 575     while true {
 576         if true {
 577             foo();$0
 578
 579             // comment
 580             bar();
 581         } else {
 582             println!("bar");
 583         }
 584     }
 585 }
 586 "#,
 587         );
 588     }
 589
 590     #[test]
 591     fn simple_single_line() {
 592         check_assist(
 593             unwrap_block,
 594             r#"
 595 fn main() {
 596     {$0 0 }
 597 }
 598 "#,
 599             r#"
 600 fn main() {
 601     0
 602 }
 603 "#,
 604         );
 605     }
 606
 607     #[test]
 608     fn simple_nested_block() {
 609         check_assist(
 610             unwrap_block,
 611             r#"
 612 fn main() {
 613     $0{
 614         {
 615             3
 616         }
 617     }
 618 }
 619 "#,
 620             r#"
 621 fn main() {
 622     {
 623         3
 624     }
 625 }
 626 "#,
 627         );
 628     }
 629
 630     #[test]
 631     fn nested_single_line() {
 632         check_assist(
 633             unwrap_block,
 634             r#"
 635 fn main() {
 636     {$0 { println!("foo"); } }
 637 }
 638 "#,
 639             r#"
 640 fn main() {
 641     { println!("foo"); }
 642 }
 643 "#,
 644         );
 645
 646         check_assist(
 647             unwrap_block,
 648             r#"
 649 fn main() {
 650     {$0 { 0 } }
 651 }
 652 "#,
 653             r#"
 654 fn main() {
 655     { 0 }
 656 }
 657 "#,
 658         );
 659     }
 660
 661     #[test]
 662     fn simple_if_single_line() {
 663         check_assist(
 664             unwrap_block,
 665             r#"
 666 fn main() {
 667     if true {$0 /* foo */ foo() } else { bar() /* bar */}
 668 }
 669 "#,
 670             r#"
 671 fn main() {
 672     /* foo */ foo()
 673 }
 674 "#,
 675         );
 676     }
 677
 678     #[test]
 679     fn if_single_statement() {
 680         check_assist(
 681             unwrap_block,
 682             r#"
 683 fn main() {
 684     if true {$0
 685         return 3;
 686     }
 687 }
 688 "#,
 689             r#"
 690 fn main() {
 691     return 3;
 692 }
 693 "#,
 694         );
 695     }
 696
 697     #[test]
 698     fn multiple_statements() {
 699         check_assist(
 700             unwrap_block,
 701             r#"
 702 fn main() -> i32 {
 703     if 2 > 1 {$0
 704         let a = 5;
 705         return 3;
 706     }
 707     5
 708 }
 709 "#,
 710             r#"
 711 fn main() -> i32 {
 712     let a = 5;
 713     return 3;
 714     5
 715 }
 716 "#,
 717         );
 718     }
 719 }