1 files changed, 85 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..1fb13f481 100644
--- a/crates/ra_assists/src/handlers/unwrap_block.rs
+++ b/crates/ra_assists/src/handlers/unwrap_block.rs
@@ -1,7 +1,10 @@
 use ra_fmt::unwrap_trivial_block;
 use ra_syntax::{
-    ast::{self, ElseBranch, Expr, LoopBodyOwner},
-    match_ast, AstNode, TextRange, T,
+    ast::{
+        self,
+        edit::{AstNodeEdit, IndentLevel},
+    },
+    AstNode, TextRange, T,
 };
 
 use crate::{AssistContext, AssistId, Assists};
@@ -24,94 +27,73 @@ use crate::{AssistContext, AssistId, Assists};
 // }
 // ```
 pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
-    let block = ast::BlockExpr::cast(l_curly_token.parent())?;
-    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 l_curly_token = ctx.find_token_at_offset(T!['{'])?;
+    let mut block = ast::BlockExpr::cast(l_curly_token.parent())?;
+    let mut parent = block.syntax().parent()?;
+    if ast::MatchArm::can_cast(parent.kind()) {
+        parent = parent.ancestors().find(|it| ast::MatchExpr::can_cast(it.kind()))?
+    }
 
-                resp?
-            },
-            _ => return None,
+    let parent = ast::Expr::cast(parent)?;
+
+    match parent.clone() {
+        ast::Expr::ForExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::LoopExpr(_) => (),
+        ast::Expr::MatchExpr(_) => block = block.dedent(IndentLevel(1)),
+        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(), &[' ', '{']),
+                        );
+                    });
+                }
+            } 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();
@@ -490,6 +472,30 @@ mod tests {
     }
 
     #[test]
+    fn unwrap_match_arm() {
+        check_assist(
+            unwrap_block,
+            r#"
+fn main() {
+    match rel_path {
+        Ok(rel_path) => {<|>
+            let rel_path = RelativePathBuf::from_path(rel_path).ok()?;
+            Some((*id, rel_path))
+        }
+        Err(_) => None,
+    }
+}
+"#,
+            r#"
+fn main() {
+    let rel_path = RelativePathBuf::from_path(rel_path).ok()?;
+    Some((*id, rel_path))
+}
+"#,
+        );
+    }
+
+    #[test]
     fn simple_if_in_while_bad_cursor_position() {
         check_assist_not_applicable(
             unwrap_block,