Don't remove block label from closure body (#6468)

Author: anatawa12

src/closures.rs

@@ -1,4 +1,4 @@
-use rustc_ast::{ast, ptr};
+use rustc_ast::{Label, ast, ptr};
 use rustc_span::Span;
 use thin_vec::thin_vec;
 use tracing::debug;
@@ -72,7 +72,7 @@ pub(crate) fn rewrite_closure(
 
         result.or_else(|_| {
             // Either we require a block, or tried without and failed.
-            rewrite_closure_block(block, &prefix, context, body_shape)
+            rewrite_closure_block(body, &prefix, context, body_shape)
         })
     } else {
         rewrite_closure_expr(body, &prefix, context, body_shape).or_else(|_| {
@@ -104,8 +104,8 @@ fn get_inner_expr<'a>(
     prefix: &str,
     context: &RewriteContext<'_>,
 ) -> &'a ast::Expr {
-    if let ast::ExprKind::Block(ref block, _) = expr.kind {
-        if !needs_block(block, prefix, context) {
+    if let ast::ExprKind::Block(ref block, ref label) = expr.kind {
+        if !needs_block(block, label, prefix, context) {
             // block.stmts.len() == 1 except with `|| {{}}`;
             // https://github.com/rust-lang/rustfmt/issues/3844
             if let Some(expr) = block.stmts.first().and_then(stmt_expr) {
@@ -118,7 +118,12 @@ fn get_inner_expr<'a>(
 }
 
 // Figure out if a block is necessary.
-fn needs_block(block: &ast::Block, prefix: &str, context: &RewriteContext<'_>) -> bool {
+fn needs_block(
+    block: &ast::Block,
+    label: &Option<Label>,
+    prefix: &str,
+    context: &RewriteContext<'_>,
+) -> bool {
     let has_attributes = block.stmts.first().map_or(false, |first_stmt| {
         !get_attrs_from_stmt(first_stmt).is_empty()
     });
@@ -128,6 +133,7 @@ fn needs_block(block: &ast::Block, prefix: &str, context: &RewriteContext<'_>) -
         || has_attributes
         || block_contains_comment(context, block)
         || prefix.contains('\n')
+        || label.is_some()
 }
 
 fn veto_block(e: &ast::Expr) -> bool {
@@ -230,11 +236,16 @@ fn rewrite_closure_expr(
 
 // Rewrite closure whose body is block.
 fn rewrite_closure_block(
-    block: &ast::Block,
+    block: &ast::Expr,
     prefix: &str,
     context: &RewriteContext<'_>,
     shape: Shape,
 ) -> RewriteResult {
+    debug_assert!(
+        matches!(block.kind, ast::ExprKind::Block(..)),
+        "expected a block expression"
+    );
+
     Ok(format!(
         "{} {}",
         prefix,
@@ -353,6 +364,8 @@ pub(crate) fn rewrite_last_closure(
     expr: &ast::Expr,
     shape: Shape,
 ) -> RewriteResult {
+    debug!("rewrite_last_closure {:?}", expr);
+
     if let ast::ExprKind::Closure(ref closure) = expr.kind {
         let ast::Closure {
             ref binder,
@@ -366,10 +379,11 @@ pub(crate) fn rewrite_last_closure(
             fn_arg_span: _,
         } = **closure;
         let body = match body.kind {
-            ast::ExprKind::Block(ref block, _)
+            ast::ExprKind::Block(ref block, ref label)
                 if !is_unsafe_block(block)
                     && !context.inside_macro()
-                    && is_simple_block(context, block, Some(&body.attrs)) =>
+                    && is_simple_block(context, block, Some(&body.attrs))
+                    && label.is_none() =>
             {
                 stmt_expr(&block.stmts[0]).unwrap_or(body)
             }

tests/source/closure-block-labels.rs

@@ -0,0 +1,192 @@
+// _0: in-macro
+// _1: last argument in function invocation
+// _2: non-last argument in function invocation
+// _3: simple expression
+
+// test0: the cause reported in issue: label is used, and there is usage, multiple statements
+pub fn rustfmt_test0_0(condition: bool) {
+    test_macro!(|transaction| 'block: {
+        if condition {
+            break 'block 0;
+        }
+
+        todo!()
+    });
+}
+
+pub fn rustfmt_test0_1(condition: bool) {
+    test_func(|transaction| 'block: {
+        if condition {
+            break 'block 0;
+        }
+
+        todo!()
+    });
+}
+
+pub fn rustfmt_test0_2(condition: bool) {
+    test_func2(|transaction| 'block: {
+        if condition {
+            break 'block 0;
+        }
+
+        todo!()
+    }, 0);
+}
+
+pub fn rustfmt_test0_3(condition: bool) {
+    let x = |transaction| 'block: {
+        if condition {
+            break 'block 0;
+        }
+
+        todo!()
+    };
+}
+
+
+// test1: label is unused, and there is usage, multiple statements
+pub fn rustfmt_test1_0(condition: bool) {
+    test_macro!(|transaction| 'block: {
+        if condition {
+            todo!("");
+        }
+
+        todo!()
+    });
+}
+
+pub fn rustfmt_test1_1(condition: bool) {
+    test_func(|transaction| 'block: {
+        if condition {
+            todo!("");
+        }
+
+        todo!()
+    });
+}
+
+pub fn rustfmt_test1_2(condition: bool) {
+    test_func2(|transaction| 'block: {
+        if condition {
+            todo!("");
+        }
+
+        todo!()
+    }, 0);
+}
+
+pub fn rustfmt_test1_3(condition: bool) {
+    let x = |transaction| 'block: {
+        if condition {
+            todo!("");
+        }
+
+        todo!()
+    };
+}
+
+
+
+// test2: label is used, single expression
+pub fn rustfmt_test2_0(condition: bool) {
+    test_macro!(|transaction| 'block: {
+        break 'block 0;
+    });
+}
+
+pub fn rustfmt_test2_1(condition: bool) {
+    test_func(|transaction| 'block: {
+        break 'block 0;
+    });
+}
+
+pub fn rustfmt_test2_2(condition: bool) {
+    test_func2(|transaction| 'block: {
+        break 'block 0;
+    }, 0);
+}
+
+pub fn rustfmt_test2_3(condition: bool) {
+    let x = |transaction| 'block: {
+        break 'block 0;
+    };
+}
+
+// test3: label is unused, single general multi-line expression
+pub fn rustfmt_test3_0(condition: bool) {
+    test_macro!(|transaction| 'block: {
+        vec![
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0,
+        ]
+    });
+}
+
+pub fn rustfmt_test3_1(condition: bool) {
+    test_func(|transaction| 'block: {
+        vec![
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0,
+        ]
+    });
+}
+
+pub fn rustfmt_test3_2(condition: bool) {
+    test_func2(|transaction| 'block: {
+        vec![
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0,
+        ]
+    }, 0);
+}
+
+pub fn rustfmt_test3_3(condition: bool) {
+    let x = |transaction| 'block: {
+        vec![
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0,
+        ]
+    };
+}
+
+// test4: label is unused, single block statement-expression
+pub fn rustfmt_test4_0(condition: bool) {
+    test_macro!(|transaction| 'block: {
+        if condition {
+            break 'block 1;
+        } else {
+            break 'block 0;
+        }
+    });
+}
+
+pub fn rustfmt_test4_1(condition: bool) {
+    test_func(|transaction| 'block: {
+        if condition {
+            break 'block 1;
+        } else {
+            break 'block 0;
+        }
+    });
+}
+
+pub fn rustfmt_test4_2(condition: bool) {
+    test_func2(|transaction| 'block: {
+        if condition {
+            break 'block 1;
+        } else {
+            break 'block 0;
+        }
+    }, 1);
+}
+
+pub fn rustfmt_test4_3(condition: bool) {
+    let x = |transaction| 'block: {
+        if condition {
+            break 'block 1;
+        } else {
+            break 'block 0;
+        }
+    };
+}