crates/ra_assists/src/handlers/unwrap_block.rs

   1 use crate::{Assist, AssistCtx, AssistId};
   2
   3 use ast::LoopBodyOwner;
   4 use ra_fmt::unwrap_trivial_block;
   5 use ra_syntax::{ast, match_ast, AstNode, TextRange, T};
   6
   7 // Assist: unwrap_block
   8 //
   9 // This assist removes if...else, for, while and loop control statements to just keep the body.
  10 //
  11 // ```
  12 // fn foo() {
  13 //     if true {<|>
  14 //         println!("foo");
  15 //     }
  16 // }
  17 // ```
  18 // ->
  19 // ```
  20 // fn foo() {
  21 //     println!("foo");
  22 // }
  23 // ```
  24 pub(crate) fn unwrap_block(ctx: AssistCtx) -> Option<Assist> {
  25     let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
  26     let block = ast::BlockExpr::cast(l_curly_token.parent())?;
  27     let parent = block.syntax().parent()?;
  28     let (expr, expr_to_unwrap) = match_ast! {
  29         match parent {
  30             ast::IfExpr(if_expr) => {
  31                 let expr_to_unwrap = if_expr.blocks().find_map(|expr| extract_expr(ctx.frange.range, expr));
  32                 let expr_to_unwrap = expr_to_unwrap?;
  33                 // Find if we are in a else if block
  34                 let ancestor = if_expr.syntax().parent().and_then(ast::IfExpr::cast);
  35
  36                 match ancestor {
  37                     None => (ast::Expr::IfExpr(if_expr), expr_to_unwrap),
  38                     Some(ancestor) => (ast::Expr::IfExpr(ancestor), expr_to_unwrap),
  39                 }
  40             },
  41             ast::ForExpr(for_expr) => {
  42                 let block_expr = for_expr.loop_body()?;
  43                 let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
  44                 (ast::Expr::ForExpr(for_expr), expr_to_unwrap)
  45             },
  46             ast::WhileExpr(while_expr) => {
  47                 let block_expr = while_expr.loop_body()?;
  48                 let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
  49                 (ast::Expr::WhileExpr(while_expr), expr_to_unwrap)
  50             },
  51             ast::LoopExpr(loop_expr) => {
  52                 let block_expr = loop_expr.loop_body()?;
  53                 let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
  54                 (ast::Expr::LoopExpr(loop_expr), expr_to_unwrap)
  55             },
  56             _ => return None,
  57         }
  58     };
  59
  60     ctx.add_assist(AssistId("unwrap_block"), "Unwrap block", |edit| {
  61         edit.set_cursor(expr.syntax().text_range().start());
  62         edit.target(expr_to_unwrap.syntax().text_range());
  63
  64         let pat_start: &[_] = &[' ', '{', '\n'];
  65         let expr_to_unwrap = expr_to_unwrap.to_string();
  66         let expr_string = expr_to_unwrap.trim_start_matches(pat_start);
  67         let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
  68         expr_string_lines.pop(); // Delete last line
  69
  70         let expr_string = expr_string_lines
  71             .into_iter()
  72             .map(|line| line.replacen("    ", "", 1)) // Delete indentation
  73             .collect::<Vec<String>>()
  74             .join("\n");
  75
  76         edit.replace(expr.syntax().text_range(), expr_string);
  77     })
  78 }
  79
  80 fn extract_expr(cursor_range: TextRange, block: ast::BlockExpr) -> Option<ast::Expr> {
  81     let cursor_in_range = block.l_curly_token()?.text_range().contains_range(cursor_range);
  82
  83     if cursor_in_range {
  84         Some(unwrap_trivial_block(block))
  85     } else {
  86         None
  87     }
  88 }
  89
  90 #[cfg(test)]
  91 mod tests {
  92     use crate::tests::{check_assist, check_assist_not_applicable};
  93
  94     use super::*;
  95
  96     #[test]
  97     fn simple_if() {
  98         check_assist(
  99             unwrap_block,
 100             r#"
 101             fn main() {
 102                 bar();
 103                 if true {<|>
 104                     foo();
 105
 106                     //comment
 107                     bar();
 108                 } else {
 109                     println!("bar");
 110                 }
 111             }
 112             "#,
 113             r#"
 114             fn main() {
 115                 bar();
 116                 <|>foo();
 117
 118                 //comment
 119                 bar();
 120             }
 121             "#,
 122         );
 123     }
 124
 125     #[test]
 126     fn simple_if_else() {
 127         check_assist(
 128             unwrap_block,
 129             r#"
 130             fn main() {
 131                 bar();
 132                 if true {
 133                     foo();
 134
 135                     //comment
 136                     bar();
 137                 } else {<|>
 138                     println!("bar");
 139                 }
 140             }
 141             "#,
 142             r#"
 143             fn main() {
 144                 bar();
 145                 <|>println!("bar");
 146             }
 147             "#,
 148         );
 149     }
 150
 151     #[test]
 152     fn simple_if_else_if() {
 153         check_assist(
 154             unwrap_block,
 155             r#"
 156             fn main() {
 157                 //bar();
 158                 if true {
 159                     println!("true");
 160
 161                     //comment
 162                     //bar();
 163                 } else if false {<|>
 164                     println!("bar");
 165                 } else {
 166                     println!("foo");
 167                 }
 168             }
 169             "#,
 170             r#"
 171             fn main() {
 172                 //bar();
 173                 <|>println!("bar");
 174             }
 175             "#,
 176         );
 177     }
 178
 179     #[test]
 180     fn simple_if_bad_cursor_position() {
 181         check_assist_not_applicable(
 182             unwrap_block,
 183             r#"
 184             fn main() {
 185                 bar();<|>
 186                 if true {
 187                     foo();
 188
 189                     //comment
 190                     bar();
 191                 } else {
 192                     println!("bar");
 193                 }
 194             }
 195             "#,
 196         );
 197     }
 198
 199     #[test]
 200     fn simple_for() {
 201         check_assist(
 202             unwrap_block,
 203             r#"
 204             fn main() {
 205                 for i in 0..5 {<|>
 206                     if true {
 207                         foo();
 208
 209                         //comment
 210                         bar();
 211                     } else {
 212                         println!("bar");
 213                     }
 214                 }
 215             }
 216             "#,
 217             r#"
 218             fn main() {
 219                 <|>if true {
 220                     foo();
 221
 222                     //comment
 223                     bar();
 224                 } else {
 225                     println!("bar");
 226                 }
 227             }
 228             "#,
 229         );
 230     }
 231
 232     #[test]
 233     fn simple_if_in_for() {
 234         check_assist(
 235             unwrap_block,
 236             r#"
 237             fn main() {
 238                 for i in 0..5 {
 239                     if true {<|>
 240                         foo();
 241
 242                         //comment
 243                         bar();
 244                     } else {
 245                         println!("bar");
 246                     }
 247                 }
 248             }
 249             "#,
 250             r#"
 251             fn main() {
 252                 for i in 0..5 {
 253                     <|>foo();
 254
 255                     //comment
 256                     bar();
 257                 }
 258             }
 259             "#,
 260         );
 261     }
 262
 263     #[test]
 264     fn simple_loop() {
 265         check_assist(
 266             unwrap_block,
 267             r#"
 268             fn main() {
 269                 loop {<|>
 270                     if true {
 271                         foo();
 272
 273                         //comment
 274                         bar();
 275                     } else {
 276                         println!("bar");
 277                     }
 278                 }
 279             }
 280             "#,
 281             r#"
 282             fn main() {
 283                 <|>if true {
 284                     foo();
 285
 286                     //comment
 287                     bar();
 288                 } else {
 289                     println!("bar");
 290                 }
 291             }
 292             "#,
 293         );
 294     }
 295
 296     #[test]
 297     fn simple_while() {
 298         check_assist(
 299             unwrap_block,
 300             r#"
 301             fn main() {
 302                 while true {<|>
 303                     if true {
 304                         foo();
 305
 306                         //comment
 307                         bar();
 308                     } else {
 309                         println!("bar");
 310                     }
 311                 }
 312             }
 313             "#,
 314             r#"
 315             fn main() {
 316                 <|>if true {
 317                     foo();
 318
 319                     //comment
 320                     bar();
 321                 } else {
 322                     println!("bar");
 323                 }
 324             }
 325             "#,
 326         );
 327     }
 328
 329     #[test]
 330     fn simple_if_in_while_bad_cursor_position() {
 331         check_assist_not_applicable(
 332             unwrap_block,
 333             r#"
 334             fn main() {
 335                 while true {
 336                     if true {
 337                         foo();<|>
 338
 339                         //comment
 340                         bar();
 341                     } else {
 342                         println!("bar");
 343                     }
 344                 }
 345             }
 346             "#,
 347         );
 348     }
 349 }