3 files changed, 370 insertions, 0 deletions

diff --git a/crates/ra_assists/src/doc_tests/generated.rs b/crates/ra_assists/src/doc_tests/generated.rs
index d6a34b609..6696cc832 100644
--- a/crates/ra_assists/src/doc_tests/generated.rs
+++ b/crates/ra_assists/src/doc_tests/generated.rs
@@ -728,3 +728,22 @@ use std::{collections::HashMap};
 "#####,
     )
 }
+
+#[test]
+fn doctest_unwrap_block() {
+    check(
+        "unwrap_block",
+        r#####"
+fn foo() {
+    if true {<|>
+        println!("foo");
+    }
+}
+"#####,
+        r#####"
+fn foo() {
+    println!("foo");
+}
+"#####,
+    )
+}
diff --git a/crates/ra_assists/src/handlers/unwrap_block.rs b/crates/ra_assists/src/handlers/unwrap_block.rs
new file mode 100644
index 000000000..8912ce645
--- /dev/null
+++ b/crates/ra_assists/src/handlers/unwrap_block.rs
@@ -0,0 +1,349 @@
+use crate::{Assist, AssistCtx, AssistId};
+
+use ast::{BlockExpr, Expr, ForExpr, IfExpr, LoopBodyOwner, LoopExpr, WhileExpr};
+use ra_fmt::unwrap_trivial_block;
+use ra_syntax::{ast, AstNode, TextRange, T};
+
+// Assist: unwrap_block
+//
+// This assist removes if...else, for, while and loop control statements to just keep the body.
+//
+// ```
+// fn foo() {
+//     if true {<|>
+//         println!("foo");
+//     }
+// }
+// ```
+// ->
+// ```
+// fn foo() {
+//     println!("foo");
+// }
+// ```
+pub(crate) fn unwrap_block(ctx: AssistCtx) -> Option<Assist> {
+    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
+
+    let res = if let Some(if_expr) = l_curly_token.ancestors().find_map(IfExpr::cast) {
+        // if expression
+        let expr_to_unwrap = if_expr.blocks().find_map(|expr| extract_expr(ctx.frange.range, expr));
+        let expr_to_unwrap = expr_to_unwrap?;
+        // Find if we are in a else if block
+        let ancestor = if_expr.syntax().ancestors().skip(1).find_map(ast::IfExpr::cast);
+
+        if let Some(ancestor) = ancestor {
+            Some((ast::Expr::IfExpr(ancestor), expr_to_unwrap))
+        } else {
+            Some((ast::Expr::IfExpr(if_expr), expr_to_unwrap))
+        }
+    } else if let Some(for_expr) = l_curly_token.ancestors().find_map(ForExpr::cast) {
+        // for expression
+        let block_expr = for_expr.loop_body()?;
+        extract_expr(ctx.frange.range, block_expr)
+            .map(|expr_to_unwrap| (ast::Expr::ForExpr(for_expr), expr_to_unwrap))
+    } else if let Some(while_expr) = l_curly_token.ancestors().find_map(WhileExpr::cast) {
+        // while expression
+        let block_expr = while_expr.loop_body()?;
+        extract_expr(ctx.frange.range, block_expr)
+            .map(|expr_to_unwrap| (ast::Expr::WhileExpr(while_expr), expr_to_unwrap))
+    } else if let Some(loop_expr) = l_curly_token.ancestors().find_map(LoopExpr::cast) {
+        // loop expression
+        let block_expr = loop_expr.loop_body()?;
+        extract_expr(ctx.frange.range, block_expr)
+            .map(|expr_to_unwrap| (ast::Expr::LoopExpr(loop_expr), expr_to_unwrap))
+    } else {
+        None
+    };
+
+    let (expr, expr_to_unwrap) = res?;
+    ctx.add_assist(AssistId("unwrap_block"), "Unwrap block", |edit| {
+        edit.set_cursor(expr.syntax().text_range().start());
+        edit.target(expr_to_unwrap.syntax().text_range());
+
+        let pat_start: &[_] = &[' ', '{', '\n'];
+        let expr_to_unwrap = expr_to_unwrap.to_string();
+        let expr_string = expr_to_unwrap.trim_start_matches(pat_start);
+        let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();
+        expr_string_lines.pop(); // Delete last line
+
+        let expr_string = expr_string_lines
+            .into_iter()
+            .map(|line| line.replacen("    ", "", 1)) // Delete indentation
+            .collect::<Vec<String>>()
+            .join("\n");
+
+        edit.replace(expr.syntax().text_range(), expr_string);
+    })
+}
+
+fn extract_expr(cursor_range: TextRange, block_expr: BlockExpr) -> Option<Expr> {
+    let block = block_expr.block()?;
+    let cursor_in_range = block.l_curly_token()?.text_range().contains_range(cursor_range);
+
+    if cursor_in_range {
+        Some(unwrap_trivial_block(block_expr))
+    } else {
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::helpers::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn simple_if() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                bar();
+                if true {<|>
+                    foo();
+
+                    //comment
+                    bar();
+                } else {
+                    println!("bar");
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                bar();
+                <|>foo();
+
+                //comment
+                bar();
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_else() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                bar();
+                if true {
+                    foo();
+
+                    //comment
+                    bar();
+                } else {<|>
+                    println!("bar");
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                bar();
+                <|>println!("bar");
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_else_if() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {<|>
+                    println!("bar");
+                } else {
+                    println!("foo");
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                //bar();
+                <|>println!("bar");
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_bad_cursor_position() {
+        check_assist_not_applicable(
+            unwrap_block,
+            r#"
+            fn main() {
+                bar();<|>
+                if true {
+                    foo();
+
+                    //comment
+                    bar();
+                } else {
+                    println!("bar");
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_for() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                for i in 0..5 {<|>
+                    if true {
+                        foo();
+
+                        //comment
+                        bar();
+                    } else {
+                        println!("bar");
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                <|>if true {
+                    foo();
+
+                    //comment
+                    bar();
+                } else {
+                    println!("bar");
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_in_for() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                for i in 0..5 {
+                    if true {<|>
+                        foo();
+
+                        //comment
+                        bar();
+                    } else {
+                        println!("bar");
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                for i in 0..5 {
+                    <|>foo();
+
+                    //comment
+                    bar();
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_loop() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                loop {<|>
+                    if true {
+                        foo();
+
+                        //comment
+                        bar();
+                    } else {
+                        println!("bar");
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                <|>if true {
+                    foo();
+
+                    //comment
+                    bar();
+                } else {
+                    println!("bar");
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_while() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                while true {<|>
+                    if true {
+                        foo();
+
+                        //comment
+                        bar();
+                    } else {
+                        println!("bar");
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                <|>if true {
+                    foo();
+
+                    //comment
+                    bar();
+                } else {
+                    println!("bar");
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_in_while_bad_cursor_position() {
+        check_assist_not_applicable(
+            unwrap_block,
+            r#"
+            fn main() {
+                while true {
+                    if true {
+                        foo();<|>
+
+                        //comment
+                        bar();
+                    } else {
+                        println!("bar");
+                    }
+                }
+            }
+            "#,
+        );
+    }
+}
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index 64bd87afb..c5df86600 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -143,6 +143,7 @@ mod handlers {
     mod split_import;
     mod add_from_impl_for_enum;
     mod reorder_fields;
+    mod unwrap_block;
 
     pub(crate) fn all() -> &'static [AssistHandler] {
         &[
@@ -181,6 +182,7 @@ mod handlers {
             replace_unwrap_with_match::replace_unwrap_with_match,
             split_import::split_import,
             add_from_impl_for_enum::add_from_impl_for_enum,
+            unwrap_block::unwrap_block,
             // These are manually sorted for better priorities
             add_missing_impl_members::add_missing_impl_members,
             add_missing_impl_members::add_missing_default_members,