crates/ra_assists/src/assists/early_return.rs

   1 use std::ops::RangeInclusive;
   2
   3 use hir::db::HirDatabase;
   4 use ra_syntax::{
   5     algo::replace_children,
   6     ast::{self, edit::IndentLevel, make},
   7     AstNode,
   8     SyntaxKind::{FN_DEF, LOOP_EXPR, L_CURLY, R_CURLY, WHILE_EXPR, WHITESPACE},
   9 };
  10
  11 use crate::{
  12     assist_ctx::{Assist, AssistCtx},
  13     AssistId,
  14 };
  15
  16 // Assist: convert_to_guarded_return
  17 //
  18 // Replace a large conditional with a guarded return.
  19 //
  20 // ```
  21 // fn main() {
  22 //     <|>if cond {
  23 //         foo();
  24 //         bar();
  25 //     }
  26 // }
  27 // ```
  28 // ->
  29 // ```
  30 // fn main() {
  31 //     if !cond {
  32 //         return;
  33 //     }
  34 //     foo();
  35 //     bar();
  36 // }
  37 // ```
  38 pub(crate) fn convert_to_guarded_return(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
  39     let if_expr: ast::IfExpr = ctx.node_at_offset()?;
  40     let expr = if_expr.condition()?.expr()?;
  41     let then_block = if_expr.then_branch()?.block()?;
  42     if if_expr.else_branch().is_some() {
  43         return None;
  44     }
  45
  46     let parent_block = if_expr.syntax().parent()?.ancestors().find_map(ast::Block::cast)?;
  47
  48     if parent_block.expr()? != if_expr.clone().into() {
  49         return None;
  50     }
  51
  52     // check for early return and continue
  53     let first_in_then_block = then_block.syntax().first_child()?.clone();
  54     if ast::ReturnExpr::can_cast(first_in_then_block.kind())
  55         || ast::ContinueExpr::can_cast(first_in_then_block.kind())
  56         || first_in_then_block
  57             .children()
  58             .any(|x| ast::ReturnExpr::can_cast(x.kind()) || ast::ContinueExpr::can_cast(x.kind()))
  59     {
  60         return None;
  61     }
  62
  63     let parent_container = parent_block.syntax().parent()?.parent()?;
  64
  65     let early_expression = match parent_container.kind() {
  66         WHILE_EXPR | LOOP_EXPR => Some("continue;"),
  67         FN_DEF => Some("return;"),
  68         _ => None,
  69     }?;
  70
  71     if then_block.syntax().first_child_or_token().map(|t| t.kind() == L_CURLY).is_none() {
  72         return None;
  73     }
  74
  75     then_block.syntax().last_child_or_token().filter(|t| t.kind() == R_CURLY)?;
  76     let cursor_position = ctx.frange.range.start();
  77
  78     ctx.add_action(AssistId("convert_to_guarded_return"), "convert to guarded return", |edit| {
  79         let if_indent_level = IndentLevel::from_node(&if_expr.syntax());
  80         let new_if_expr =
  81             if_indent_level.increase_indent(make::if_expression(&expr, early_expression));
  82         let then_block_items = IndentLevel::from(1).decrease_indent(then_block.clone());
  83         let end_of_then = then_block_items.syntax().last_child_or_token().unwrap();
  84         let end_of_then =
  85             if end_of_then.prev_sibling_or_token().map(|n| n.kind()) == Some(WHITESPACE) {
  86                 end_of_then.prev_sibling_or_token().unwrap()
  87             } else {
  88                 end_of_then
  89             };
  90         let mut new_if_and_then_statements = new_if_expr.syntax().children_with_tokens().chain(
  91             then_block_items
  92                 .syntax()
  93                 .children_with_tokens()
  94                 .skip(1)
  95                 .take_while(|i| *i != end_of_then),
  96         );
  97         let new_block = replace_children(
  98             &parent_block.syntax(),
  99             RangeInclusive::new(
 100                 if_expr.clone().syntax().clone().into(),
 101                 if_expr.syntax().clone().into(),
 102             ),
 103             &mut new_if_and_then_statements,
 104         );
 105         edit.target(if_expr.syntax().text_range());
 106         edit.replace_ast(parent_block, ast::Block::cast(new_block).unwrap());
 107         edit.set_cursor(cursor_position);
 108     });
 109     ctx.build()
 110 }
 111
 112 #[cfg(test)]
 113 mod tests {
 114     use super::*;
 115     use crate::helpers::{check_assist, check_assist_not_applicable};
 116
 117     #[test]
 118     fn convert_inside_fn() {
 119         check_assist(
 120             convert_to_guarded_return,
 121             r#"
 122             fn main() {
 123                 bar();
 124                 if<|> true {
 125                     foo();
 126
 127                     //comment
 128                     bar();
 129                 }
 130             }
 131             "#,
 132             r#"
 133             fn main() {
 134                 bar();
 135                 if<|> !true {
 136                     return;
 137                 }
 138                 foo();
 139
 140                 //comment
 141                 bar();
 142             }
 143             "#,
 144         );
 145     }
 146
 147     #[test]
 148     fn convert_inside_while() {
 149         check_assist(
 150             convert_to_guarded_return,
 151             r#"
 152             fn main() {
 153                 while true {
 154                     if<|> true {
 155                         foo();
 156                         bar();
 157                     }
 158                 }
 159             }
 160             "#,
 161             r#"
 162             fn main() {
 163                 while true {
 164                     if<|> !true {
 165                         continue;
 166                     }
 167                     foo();
 168                     bar();
 169                 }
 170             }
 171             "#,
 172         );
 173     }
 174
 175     #[test]
 176     fn convert_inside_loop() {
 177         check_assist(
 178             convert_to_guarded_return,
 179             r#"
 180             fn main() {
 181                 loop {
 182                     if<|> true {
 183                         foo();
 184                         bar();
 185                     }
 186                 }
 187             }
 188             "#,
 189             r#"
 190             fn main() {
 191                 loop {
 192                     if<|> !true {
 193                         continue;
 194                     }
 195                     foo();
 196                     bar();
 197                 }
 198             }
 199             "#,
 200         );
 201     }
 202
 203     #[test]
 204     fn ignore_already_converted_if() {
 205         check_assist_not_applicable(
 206             convert_to_guarded_return,
 207             r#"
 208             fn main() {
 209                 if<|> true {
 210                     return;
 211                 }
 212             }
 213             "#,
 214         );
 215     }
 216
 217     #[test]
 218     fn ignore_already_converted_loop() {
 219         check_assist_not_applicable(
 220             convert_to_guarded_return,
 221             r#"
 222             fn main() {
 223                 loop {
 224                     if<|> true {
 225                         continue;
 226                     }
 227                 }
 228             }
 229             "#,
 230         );
 231     }
 232
 233     #[test]
 234     fn ignore_return() {
 235         check_assist_not_applicable(
 236             convert_to_guarded_return,
 237             r#"
 238             fn main() {
 239                 if<|> true {
 240                     return
 241                 }
 242             }
 243             "#,
 244         );
 245     }
 246
 247     #[test]
 248     fn ignore_else_branch() {
 249         check_assist_not_applicable(
 250             convert_to_guarded_return,
 251             r#"
 252             fn main() {
 253                 if<|> true {
 254                     foo();
 255                 } else {
 256                     bar()
 257                 }
 258             }
 259             "#,
 260         );
 261     }
 262
 263     #[test]
 264     fn ignore_statements_aftert_if() {
 265         check_assist_not_applicable(
 266             convert_to_guarded_return,
 267             r#"
 268             fn main() {
 269                 if<|> true {
 270                     foo();
 271                 }
 272                 bar();
 273             }
 274             "#,
 275         );
 276     }
 277
 278     #[test]
 279     fn ignore_statements_inside_if() {
 280         check_assist_not_applicable(
 281             convert_to_guarded_return,
 282             r#"
 283             fn main() {
 284                 if false {
 285                     if<|> true {
 286                         foo();
 287                     }
 288                 }
 289             }
 290             "#,
 291         );
 292     }
 293 }