1 files changed, 49 insertions, 79 deletions

diff --git a/crates/ra_assists/src/handlers/unwrap_block.rs b/crates/ra_assists/src/handlers/unwrap_block.rs
index 8440c7d0f..c48ecaae8 100644
--- a/crates/ra_assists/src/handlers/unwrap_block.rs
+++ b/crates/ra_assists/src/handlers/unwrap_block.rs
@@ -1,8 +1,5 @@
 use ra_fmt::unwrap_trivial_block;
-use ra_syntax::{
-    ast::{self, ElseBranch, Expr, LoopBodyOwner},
-    match_ast, AstNode, TextRange, T,
-};
+use ra_syntax::{ast, AstNode, TextRange, T};
 
 use crate::{AssistContext, AssistId, Assists};
 
@@ -29,89 +26,62 @@ pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let parent = block.syntax().parent()?;
     let assist_id = AssistId("unwrap_block");
     let assist_label = "Unwrap block";
-
-    let (expr, expr_to_unwrap) = match_ast! {
-        match parent {
-            ast::ForExpr(for_expr) => {
-                let block_expr = for_expr.loop_body()?;
-                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
-                (ast::Expr::ForExpr(for_expr), expr_to_unwrap)
-            },
-            ast::WhileExpr(while_expr) => {
-                let block_expr = while_expr.loop_body()?;
-                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
-                (ast::Expr::WhileExpr(while_expr), expr_to_unwrap)
-            },
-            ast::LoopExpr(loop_expr) => {
-                let block_expr = loop_expr.loop_body()?;
-                let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
-                (ast::Expr::LoopExpr(loop_expr), expr_to_unwrap)
-            },
-            ast::IfExpr(if_expr) => {
-                let mut resp = None;
-
-                let then_branch = if_expr.then_branch()?;
-                if then_branch.l_curly_token()?.text_range().contains_range(ctx.frange.range) {
-                    if let Some(ancestor) = if_expr.syntax().parent().and_then(ast::IfExpr::cast) {
-                        // For `else if` blocks
-                        let ancestor_then_branch = ancestor.then_branch()?;
-                        let l_curly_token = then_branch.l_curly_token()?;
-
-                        let target = then_branch.syntax().text_range();
-                        return acc.add(assist_id, assist_label, target, |edit| {
-                            let range_to_del_else_if = TextRange::new(ancestor_then_branch.syntax().text_range().end(), l_curly_token.text_range().start());
-                            let range_to_del_rest = TextRange::new(then_branch.syntax().text_range().end(), if_expr.syntax().text_range().end());
-
-                            edit.delete(range_to_del_rest);
-                            edit.delete(range_to_del_else_if);
-                            edit.replace(target, update_expr_string(then_branch.to_string(), &[' ', '{']));
-                        });
-                    } else {
-                        resp = Some((ast::Expr::IfExpr(if_expr.clone()), Expr::BlockExpr(then_branch)));
-                    }
-                } else if let Some(else_branch) = if_expr.else_branch() {
-                    match else_branch {
-                        ElseBranch::Block(else_block) => {
-                            let l_curly_token = else_block.l_curly_token()?;
-                            if l_curly_token.text_range().contains_range(ctx.frange.range) {
-                                let target = else_block.syntax().text_range();
-                                return acc.add(assist_id, assist_label, target, |edit| {
-                                    let range_to_del = TextRange::new(then_branch.syntax().text_range().end(), l_curly_token.text_range().start());
-
-                                    edit.delete(range_to_del);
-                                    edit.replace(target, update_expr_string(else_block.to_string(), &[' ', '{']));
-                                });
-                            }
-                        },
-                        ElseBranch::IfExpr(_) => {},
-                    }
+    let parent = ast::Expr::cast(parent)?;
+
+    match parent.clone() {
+        ast::Expr::ForExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::LoopExpr(_) => (),
+        ast::Expr::IfExpr(if_expr) => {
+            let then_branch = if_expr.then_branch()?;
+            if then_branch == block {
+                if let Some(ancestor) = if_expr.syntax().parent().and_then(ast::IfExpr::cast) {
+                    // For `else if` blocks
+                    let ancestor_then_branch = ancestor.then_branch()?;
+
+                    let target = then_branch.syntax().text_range();
+                    return acc.add(assist_id, assist_label, target, |edit| {
+                        let range_to_del_else_if = TextRange::new(
+                            ancestor_then_branch.syntax().text_range().end(),
+                            l_curly_token.text_range().start(),
+                        );
+                        let range_to_del_rest = TextRange::new(
+                            then_branch.syntax().text_range().end(),
+                            if_expr.syntax().text_range().end(),
+                        );
+
+                        edit.delete(range_to_del_rest);
+                        edit.delete(range_to_del_else_if);
+                        edit.replace(
+                            target,
+                            update_expr_string(then_branch.to_string(), &[' ', '{']),
+                        );
+                    });
                 }
-
-                resp?
-            },
-            _ => return None,
+            } else {
+                let target = block.syntax().text_range();
+                return acc.add(assist_id, assist_label, target, |edit| {
+                    let range_to_del = TextRange::new(
+                        then_branch.syntax().text_range().end(),
+                        l_curly_token.text_range().start(),
+                    );
+
+                    edit.delete(range_to_del);
+                    edit.replace(target, update_expr_string(block.to_string(), &[' ', '{']));
+                });
+            }
         }
+        _ => return None,
     };
 
-    let target = expr_to_unwrap.syntax().text_range();
-    acc.add(assist_id, assist_label, target, |edit| {
-        edit.replace(
-            expr.syntax().text_range(),
-            update_expr_string(expr_to_unwrap.to_string(), &[' ', '{', '\n']),
+    let unwrapped = unwrap_trivial_block(block);
+    let target = unwrapped.syntax().text_range();
+    acc.add(assist_id, assist_label, target, |builder| {
+        builder.replace(
+            parent.syntax().text_range(),
+            update_expr_string(unwrapped.to_string(), &[' ', '{', '\n']),
         );
     })
 }
 
-fn extract_expr(cursor_range: TextRange, block: ast::BlockExpr) -> Option<ast::Expr> {
-    let cursor_in_range = block.l_curly_token()?.text_range().contains_range(cursor_range);
-
-    if cursor_in_range {
-        Some(unwrap_trivial_block(block))
-    } else {
-        None
-    }
-}
-
 fn update_expr_string(expr_str: String, trim_start_pat: &[char]) -> String {
     let expr_string = expr_str.trim_start_matches(trim_start_pat);
     let mut expr_string_lines: Vec<&str> = expr_string.lines().collect();