do not remove then block when you unwrap else block #4361

Signed-off-by: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com>

crates/ra_assists/src/handlers/unwrap_block.rs

@@ -1,6 +1,6 @@
 use crate::{AssistContext, AssistId, Assists};
 
-use ast::LoopBodyOwner;
+use ast::{ElseBranch, Expr, LoopBodyOwner};
 use ra_fmt::unwrap_trivial_block;
 use ra_syntax::{ast, match_ast, AstNode, TextRange, T};
 
@@ -25,19 +25,11 @@ 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::IfExpr(if_expr) => {
-                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().parent().and_then(ast::IfExpr::cast);
-
-                match ancestor {
-                    None => (ast::Expr::IfExpr(if_expr), expr_to_unwrap),
-                    Some(ancestor) => (ast::Expr::IfExpr(ancestor), expr_to_unwrap),
-                }
-            },
             ast::ForExpr(for_expr) => {
                 let block_expr = for_expr.loop_body()?;
                 let expr_to_unwrap = extract_expr(ctx.frange.range, block_expr)?;
@@ -53,27 +45,62 @@ pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
                 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.set_cursor(ancestor_then_branch.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.set_cursor(then_branch.syntax().text_range().end());
+                                    edit.delete(range_to_del);
+                                    edit.replace(target, update_expr_string(else_block.to_string(), &[' ', '{']));
+                                });
+                            }
+                        },
+                        ElseBranch::IfExpr(_) => {},
+                    }
+                }
+
+                resp?
+            },
             _ => return None,
         }
     };
 
     let target = expr_to_unwrap.syntax().text_range();
-    acc.add(AssistId("unwrap_block"), "Unwrap block", target, |edit| {
+    acc.add(assist_id, assist_label, target, |edit| {
         edit.set_cursor(expr.syntax().text_range().start());
 
-        let pat_start: &[_] = &[' ', '{', '\n'];
+        edit.replace(
-        let expr_to_unwrap = expr_to_unwrap.to_string();
+            expr.syntax().text_range(),
-        let expr_string = expr_to_unwrap.trim_start_matches(pat_start);
+            update_expr_string(expr_to_unwrap.to_string(), &[' ', '{', '\n']),
-        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);
     })
 }
 
@@ -87,6 +114,18 @@ fn extract_expr(cursor_range: TextRange, block: ast::BlockExpr) -> Option<ast::E
     }
 }
 
+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();
+    expr_string_lines.pop(); // Delete last line
+
+    expr_string_lines
+        .into_iter()
+        .map(|line| line.replacen("    ", "", 1)) // Delete indentation
+        .collect::<Vec<String>>()
+        .join("\n")
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -142,7 +181,13 @@ mod tests {
             r#"
             fn main() {
                 bar();
-                <|>println!("bar");
+                if true {
+                    foo();
+
+                    //comment
+                    bar();
+                }<|>
+                println!("bar");
             }
             "#,
         );
@@ -170,7 +215,127 @@ mod tests {
             r#"
             fn main() {
                 //bar();
-                <|>println!("bar");
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                }<|>
+                println!("bar");
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_else_if_nested() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {
+                    println!("bar");
+                } else if true {<|>
+                    println!("foo");
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {
+                    println!("bar");
+                }<|>
+                println!("foo");
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_else_if_nested_else() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {
+                    println!("bar");
+                } else if true {
+                    println!("foo");
+                } else {<|>
+                    println!("else");
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {
+                    println!("bar");
+                } else if true {
+                    println!("foo");
+                }<|>
+                println!("else");
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn simple_if_else_if_nested_middle() {
+        check_assist(
+            unwrap_block,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {
+                    println!("bar");
+                } else if true {<|>
+                    println!("foo");
+                } else {
+                    println!("else");
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                //bar();
+                if true {
+                    println!("true");
+
+                    //comment
+                    //bar();
+                } else if false {
+                    println!("bar");
+                }<|>
+                println!("foo");
             }
             "#,
         );