1 files changed, 44 insertions, 3 deletions
diff --git a/crates/ide_assists/src/handlers/unwrap_block.rs b/crates/ide_assists/src/handlers/unwrap_block.rs
index 9171874e9d4..42c5e6f0883 100644
--- a/crates/ide_assists/src/handlers/unwrap_block.rs
+++ b/crates/ide_assists/src/handlers/unwrap_block.rs
@@ -6,7 +6,7 @@ use syntax::{
     AstNode, SyntaxKind, TextRange, T,
 };
 
-use crate::{utils::unwrap_trivial_block, AssistContext, AssistId, AssistKind, Assists};
+use crate::{AssistContext, AssistId, AssistKind, Assists};
 
 // Assist: unwrap_block
 //
@@ -88,9 +88,8 @@ pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
         _ => return None,
     };
 
-    let unwrapped = unwrap_trivial_block(block);
     acc.add(assist_id, assist_label, target, |builder| {
-        builder.replace(parent.syntax().text_range(), update_expr_string(unwrapped.to_string()));
+        builder.replace(parent.syntax().text_range(), update_expr_string(block.to_string()));
     })
 }
 
@@ -675,4 +674,46 @@ fn main() {
 "#,
         );
     }
+
+    #[test]
+    fn if_single_statement() {
+        check_assist(
+            unwrap_block,
+            r#"
+fn main() {
+    if true {$0
+        return 3;
+    }
+}
+"#,
+            r#"
+fn main() {
+    return 3;
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn multiple_statements() {
+        check_assist(
+            unwrap_block,
+            r#"
+fn main() -> i32 {
+    if 2 > 1 {$0
+        let a = 5;
+        return 3;
+    }
+    5
+}
+"#,
+            r#"
+fn main() -> i32 {
+    let a = 5;
+    return 3;
+    5
+}
+"#,
+        );
+    }
 }