3 files changed, 43 insertions, 4 deletions

diff --git a/crates/ra_assists/src/handlers/replace_if_let_with_match.rs b/crates/ra_assists/src/handlers/replace_if_let_with_match.rs
index 2442f049b..79097621e 100644
--- a/crates/ra_assists/src/handlers/replace_if_let_with_match.rs
+++ b/crates/ra_assists/src/handlers/replace_if_let_with_match.rs
@@ -1,4 +1,3 @@
-use ra_fmt::unwrap_trivial_block;
 use syntax::{
     ast::{
         self,
@@ -8,7 +7,10 @@ use syntax::{
     AstNode,
 };
 
-use crate::{utils::TryEnum, AssistContext, AssistId, AssistKind, Assists};
+use crate::{
+    utils::{unwrap_trivial_block, TryEnum},
+    AssistContext, AssistId, AssistKind, Assists,
+};
 
 // Assist: replace_if_let_with_match
 //
diff --git a/crates/ra_assists/src/handlers/unwrap_block.rs b/crates/ra_assists/src/handlers/unwrap_block.rs
index 2879090b8..3851aeb3e 100644
--- a/crates/ra_assists/src/handlers/unwrap_block.rs
+++ b/crates/ra_assists/src/handlers/unwrap_block.rs
@@ -1,4 +1,3 @@
-use ra_fmt::unwrap_trivial_block;
 use syntax::{
     ast::{
         self,
@@ -7,7 +6,7 @@ use syntax::{
     AstNode, TextRange, T,
 };
 
-use crate::{AssistContext, AssistId, AssistKind, Assists};
+use crate::{utils::unwrap_trivial_block, AssistContext, AssistId, AssistKind, Assists};
 
 // Assist: unwrap_block
 //
diff --git a/crates/ra_assists/src/utils.rs b/crates/ra_assists/src/utils.rs
index 6d85661c4..a20453dd8 100644
--- a/crates/ra_assists/src/utils.rs
+++ b/crates/ra_assists/src/utils.rs
@@ -4,6 +4,7 @@ pub(crate) mod insert_use;
 use std::{iter, ops};
 
 use hir::{Adt, Crate, Enum, ScopeDef, Semantics, Trait, Type};
+use itertools::Itertools;
 use ra_ide_db::RootDatabase;
 use rustc_hash::FxHashSet;
 use syntax::{
@@ -17,6 +18,43 @@ use crate::assist_config::SnippetCap;
 
 pub(crate) use insert_use::{find_insert_use_container, insert_use_statement};
 
+pub(crate) fn unwrap_trivial_block(block: ast::BlockExpr) -> ast::Expr {
+    extract_trivial_expression(&block)
+        .filter(|expr| !expr.syntax().text().contains_char('\n'))
+        .unwrap_or_else(|| block.into())
+}
+
+pub fn extract_trivial_expression(block: &ast::BlockExpr) -> Option<ast::Expr> {
+    let has_anything_else = |thing: &SyntaxNode| -> bool {
+        let mut non_trivial_children =
+            block.syntax().children_with_tokens().filter(|it| match it.kind() {
+                WHITESPACE | T!['{'] | T!['}'] => false,
+                _ => it.as_node() != Some(thing),
+            });
+        non_trivial_children.next().is_some()
+    };
+
+    if let Some(expr) = block.expr() {
+        if has_anything_else(expr.syntax()) {
+            return None;
+        }
+        return Some(expr);
+    }
+    // Unwrap `{ continue; }`
+    let (stmt,) = block.statements().next_tuple()?;
+    if let ast::Stmt::ExprStmt(expr_stmt) = stmt {
+        if has_anything_else(expr_stmt.syntax()) {
+            return None;
+        }
+        let expr = expr_stmt.expr()?;
+        match expr.syntax().kind() {
+            CONTINUE_EXPR | BREAK_EXPR | RETURN_EXPR => return Some(expr),
+            _ => (),
+        }
+    }
+    None
+}
+
 #[derive(Clone, Copy, Debug)]
 pub(crate) enum Cursor<'a> {
     Replace(&'a SyntaxNode),