Adds "replace with guarded return" assist

4 files changed, 314 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,
         ]
     }
 }
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index ea92284b8..47bdbb81a 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -284,6 +284,34 @@ impl IndentLevel {
             .collect();
         algo::replace_descendants(&node, &replacements)
     }
+
+    pub fn decrease_indent<N: AstNode>(self, node: N) -> N {
+        N::cast(self._decrease_indent(node.syntax().clone())).unwrap()
+    }
+
+    fn _decrease_indent(self, node: SyntaxNode) -> SyntaxNode {
+        let replacements: FxHashMap<SyntaxElement, SyntaxElement> = node
+            .descendants_with_tokens()
+            .filter_map(|el| el.into_token())
+            .filter_map(ast::Whitespace::cast)
+            .filter(|ws| {
+                let text = ws.syntax().text();
+                text.contains('\n')
+            })
+            .map(|ws| {
+                (
+                    ws.syntax().clone().into(),
+                    make::tokens::whitespace(
+                        &ws.syntax()
+                            .text()
+                            .replace(&format!("\n{:1$}", "", self.0 as usize * 4), "\n"),
+                    )
+                    .into(),
+                )
+            })
+            .collect();
+        algo::replace_descendants(&node, &replacements)
+    }
 }
 
 // FIXME: replace usages with IndentLevel above
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 143835172..00422ea91 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -128,6 +128,14 @@ pub fn where_clause(preds: impl Iterator<Item = ast::WherePred>) -> ast::WhereCl
     }
 }
 
+pub fn if_expression(condition: &ast::Expr, statement: &str) -> ast::IfExpr {
+    return ast_from_text(&format!(
+        "fn f() {{ if !{} {{\n    {}\n}}\n}}",
+        condition.syntax().text(),
+        statement
+    ));
+}
+
 fn ast_from_text<N: AstNode>(text: &str) -> N {
     let parse = SourceFile::parse(text);
     let res = parse.tree().syntax().descendants().find_map(N::cast).unwrap();