2 files changed, 278 insertions, 0 deletions

diff --git a/crates/ra_assists/src/assists/early_return.rs b/crates/ra_assists/src/assists/early_return.rs
new file mode 100644
index 000000000..8c975714c
--- /dev/null
+++ b/crates/ra_assists/src/assists/early_return.rs
@@ -0,0 +1,276 @@
+//! FIXME: write short doc here
+
+use crate::{
+    assist_ctx::{Assist, AssistCtx},
+    AssistId,
+};
+use hir::db::HirDatabase;
+use ra_syntax::{
+    algo::replace_children,
+    ast::edit::IndentLevel,
+    ast::make,
+    ast::Block,
+    ast::ContinueExpr,
+    ast::IfExpr,
+    ast::ReturnExpr,
+    AstNode,
+    SyntaxKind::{FN_DEF, LOOP_EXPR, L_CURLY, R_CURLY, WHILE_EXPR, WHITESPACE},
+};
+use std::ops::RangeInclusive;
+
+pub(crate) fn convert_to_guarded_return(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
+    let if_expr: IfExpr = ctx.node_at_offset()?;
+    let expr = if_expr.condition()?.expr()?;
+    let then_block = if_expr.then_branch()?.block()?;
+    if if_expr.else_branch().is_some() {
+        return None;
+    }
+
+    let parent_block = if_expr.syntax().parent()?.ancestors().find_map(Block::cast)?;
+
+    if parent_block.expr()? != if_expr.clone().into() {
+        return None;
+    }
+
+    // check for early return and continue
+    let first_in_then_block = then_block.syntax().first_child()?.clone();
+    if ReturnExpr::can_cast(first_in_then_block.kind())
+        || ContinueExpr::can_cast(first_in_then_block.kind())
+        || first_in_then_block
+            .children()
+            .any(|x| ReturnExpr::can_cast(x.kind()) || ContinueExpr::can_cast(x.kind()))
+    {
+        return None;
+    }
+
+    let parent_container = parent_block.syntax().parent()?.parent()?;
+
+    let early_expression = match parent_container.kind() {
+        WHILE_EXPR | LOOP_EXPR => Some("continue;"),
+        FN_DEF => Some("return;"),
+        _ => None,
+    }?;
+
+    if then_block.syntax().first_child_or_token().map(|t| t.kind() == L_CURLY).is_none() {
+        return None;
+    }
+
+    then_block.syntax().last_child_or_token().filter(|t| t.kind() == R_CURLY)?;
+    let cursor_position = ctx.frange.range.start();
+
+    ctx.add_action(AssistId("convert_to_guarded_return"), "convert to guarded return", |edit| {
+        let if_indent_level = IndentLevel::from_node(&if_expr.syntax());
+        let new_if_expr =
+            if_indent_level.increase_indent(make::if_expression(&expr, early_expression));
+        let then_block_items = IndentLevel::from(1).decrease_indent(then_block.clone());
+        let end_of_then = then_block_items.syntax().last_child_or_token().unwrap();
+        let end_of_then =
+            if end_of_then.prev_sibling_or_token().map(|n| n.kind()) == Some(WHITESPACE) {
+                end_of_then.prev_sibling_or_token().unwrap()
+            } else {
+                end_of_then
+            };
+        let mut new_if_and_then_statements = new_if_expr.syntax().children_with_tokens().chain(
+            then_block_items
+                .syntax()
+                .children_with_tokens()
+                .skip(1)
+                .take_while(|i| *i != end_of_then),
+        );
+        let new_block = replace_children(
+            &parent_block.syntax(),
+            RangeInclusive::new(
+                if_expr.clone().syntax().clone().into(),
+                if_expr.syntax().clone().into(),
+            ),
+            &mut new_if_and_then_statements,
+        );
+        edit.target(if_expr.syntax().text_range());
+        edit.replace_ast(parent_block, Block::cast(new_block).unwrap());
+        edit.set_cursor(cursor_position);
+    });
+    ctx.build()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::helpers::{check_assist, check_assist_not_applicable};
+
+    #[test]
+    fn convert_inside_fn() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                bar();
+                if<|> true {
+                    foo();
+
+                    //comment
+                    bar();
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                bar();
+                if<|> !true {
+                    return;
+                }
+                foo();
+
+                //comment
+                bar();
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn convert_inside_while() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                while true {
+                    if<|> true {
+                        foo();
+                        bar();
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                while true {
+                    if<|> !true {
+                        continue;
+                    }
+                    foo();
+                    bar();
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn convert_inside_loop() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                loop {
+                    if<|> true {
+                        foo();
+                        bar();
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                loop {
+                    if<|> !true {
+                        continue;
+                    }
+                    foo();
+                    bar();
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn ignore_already_converted_if() {
+        check_assist_not_applicable(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                if<|> true {
+                    return;
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn ignore_already_converted_loop() {
+        check_assist_not_applicable(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                loop {
+                    if<|> true {
+                        continue;
+                    }
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn ignore_return() {
+        check_assist_not_applicable(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                if<|> true {
+                    return
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn ignore_else_branch() {
+        check_assist_not_applicable(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                if<|> true {
+                    foo();
+                } else {
+                    bar()
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn ignore_statements_aftert_if() {
+        check_assist_not_applicable(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                if<|> true {
+                    foo();
+                }
+                bar();
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn ignore_statements_inside_if() {
+        check_assist_not_applicable(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                if false {
+                    if<|> true {
+                        foo();
+                    }
+                }
+            }
+            "#,
+        );
+    }
+}
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index d2376c475..ab77b46a9 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -108,6 +108,7 @@ mod assists {
     mod add_missing_impl_members;
     mod move_guard;
     mod move_bounds;
+    mod early_return;
 
     pub(crate) fn all<DB: HirDatabase>() -> &'static [fn(AssistCtx<DB>) -> Option<Assist>] {
         &[
@@ -135,6 +136,7 @@ mod assists {
             raw_string::make_raw_string,
             raw_string::make_usual_string,
             raw_string::remove_hash,
+            early_return::convert_to_guarded_return,
         ]
     }
 }