Force block closure with a single control flow body

Author: topecongiro

src/expr.rs

@@ -621,8 +621,7 @@ fn rewrite_closure(
         }
 
         // Figure out if the block is necessary.
-        let needs_block = block.rules != ast::BlockCheckMode::Default || block.stmts.len() > 1 ||
-            context.inside_macro ||
+        let needs_block = is_unsafe_block(block) || block.stmts.len() > 1 || context.inside_macro ||
             block_contains_comment(block, context.codemap) ||
             prefix.contains('\n');
 
@@ -632,9 +631,13 @@ fn rewrite_closure(
             false
         };
         if no_return_type && !needs_block {
-            // lock.stmts.len() == 1
+            // block.stmts.len() == 1
             if let Some(ref expr) = stmt_expr(&block.stmts[0]) {
-                if let Some(rw) = rewrite_closure_expr(expr, &prefix, context, body_shape) {
+                if let Some(rw) = if is_block_closure_forced(expr) {
+                    rewrite_closure_with_block(context, shape, &prefix, expr)
+                } else {
+                    rewrite_closure_expr(expr, &prefix, context, body_shape)
+                } {
                     return Some(rw);
                 }
             }
@@ -2287,7 +2290,23 @@ fn rewrite_last_closure(
         if prefix.contains('\n') {
             return None;
         }
+
         let body_shape = try_opt!(shape.offset_left(extra_offset));
+
+        // We force to use block for the body of the closure for certain kinds of expressions.
+        if is_block_closure_forced(body) {
+            return rewrite_closure_with_block(context, body_shape, &prefix, body).and_then(
+                |body_str| {
+                    // If the expression fits in a single line, we need not force block closure.
+                    if !body_str.contains('\n') {
+                        rewrite_closure_expr(body, &prefix, context, shape)
+                    } else {
+                        Some(body_str)
+                    }
+                },
+            );
+        }
+
         // When overflowing the closure which consists of a single control flow expression,
         // force to use block if its condition uses multi line.
         let is_multi_lined_cond = rewrite_cond(context, body, body_shape)
@@ -2303,6 +2322,23 @@ fn rewrite_last_closure(
     None
 }
 
+fn is_block_closure_forced(expr: &ast::Expr) -> bool {
+    match expr.node {
+        ast::ExprKind::If(..) |
+        ast::ExprKind::IfLet(..) |
+        ast::ExprKind::Loop(..) |
+        ast::ExprKind::While(..) |
+        ast::ExprKind::WhileLet(..) |
+        ast::ExprKind::ForLoop(..) => true,
+        ast::ExprKind::AddrOf(_, ref expr) |
+        ast::ExprKind::Box(ref expr) |
+        ast::ExprKind::Try(ref expr) |
+        ast::ExprKind::Unary(_, ref expr) |
+        ast::ExprKind::Cast(ref expr, _) => is_block_closure_forced(expr),
+        _ => false,
+    }
+}
+
 fn rewrite_last_arg_with_overflow<'a, T>(
     context: &RewriteContext,
     args: &[&T],