Turn closure bodies from exprs into blocks if there are comments (#4396)

* Added test cases

* Kind of working, need to refactor.

* Final code

* Update closures.rs

* Update comments.

* Use argument span instead of calculating between_span.

* Add test cases for #[attr]s.

* Use first attr instead of body span, if it exists.

* Use rewrite_missing_comment so comments are correctly folded into blocks.

* Update closures.rs

* Use ? not unwrap()

Author: Solumin

src/formatting/closures.rs

@@ -4,6 +4,7 @@ use rustc_span::Span;
 use crate::config::{lists::*, IndentStyle, SeparatorTactic};
 use crate::formatting::{
     attr::get_attrs_from_stmt,
+    comment::{contains_comment, rewrite_missing_comment},
     expr::{
         block_contains_comment, is_simple_block, is_unsafe_block, rewrite_block_with_visitor,
         rewrite_cond,
@@ -34,6 +35,7 @@ pub(crate) fn rewrite_closure(
     fn_decl: &ast::FnDecl,
     body: &ast::Expr,
     span: Span,
+    arg_span: Span,
     context: &RewriteContext<'_>,
     shape: Shape,
 ) -> Option<String> {
@@ -65,6 +67,34 @@ pub(crate) fn rewrite_closure(
             rewrite_closure_block(block, &prefix, context, body_shape)
         })
     } else {
+        // If there are comments between the fn decl and the body, the body (+ comments) need to be
+        // wrapped in a block. Since there's no return type annotation on closures with expr
+        // bodies, look for comments after the argument block.
+        // Since attrs come before the body, check up to the first attr if there is one.
+        let first_span = body
+            .attrs
+            .first()
+            .map(|attr| attr.span)
+            .unwrap_or(body.span);
+        let between_span = Span::between(arg_span, first_span);
+        if contains_comment(context.snippet(between_span)) {
+            return rewrite_closure_with_block(body, &prefix, context, body_shape).and_then(|rw| {
+                let mut parts = rw.splitn(2, "\n");
+                let head = parts.next()?;
+                let rest = parts.next()?;
+                let block_shape = shape.block_indent(context.config.tab_spaces());
+                let indent = block_shape.indent.to_string_with_newline(context.config);
+                let missing_comment = rewrite_missing_comment(between_span, block_shape, context)?;
+                Some(format!(
+                    "{}{}{}{}{}",
+                    head,
+                    indent,
+                    missing_comment,
+                    indent,
+                    rest.trim()
+                ))
+            });
+        }
         rewrite_closure_expr(body, &prefix, context, body_shape).or_else(|| {
             // The closure originally had a non-block expression, but we can't fit on
             // one line, so we'll insert a block.

src/formatting/expr.rs

@@ -189,11 +189,16 @@ pub(crate) fn format_expr(
                 Some("yield".to_string())
             }
         }
-        ast::ExprKind::Closure(capture, ref is_async, movability, ref fn_decl, ref body, _) => {
-            closures::rewrite_closure(
-                capture, is_async, movability, fn_decl, body, expr.span, context, shape,
-            )
-        }
+        ast::ExprKind::Closure(
+            capture,
+            ref is_async,
+            movability,
+            ref fn_decl,
+            ref body,
+            arg_span,
+        ) => closures::rewrite_closure(
+            capture, is_async, movability, fn_decl, body, expr.span, arg_span, context, shape,
+        ),
         ast::ExprKind::Try(..) | ast::ExprKind::Field(..) | ast::ExprKind::MethodCall(..) => {
             rewrite_chain(expr, context, shape)
         }

tests/source/issue-4384.rs

@@ -0,0 +1,39 @@
+fn main() {
+    let single_comment = ||
+    // this should be indented.
+        1;
+
+    let block_comment = ||
+        /* This is a long,
+         * explanatory comment.
+         * Put this all in a block.
+         */
+        1;
+
+    let nested_closure = ||
+    // indent + wrap in a block.
+        || 1;
+
+    let nested_comment_closure = ||
+    // if you see this code in real life, run.
+        ||
+            // nested closures don't need blocks,
+            // but comments do.
+        || 1;
+
+    let attrb = ||
+        // This block has an attr.
+        #[foo]
+        1;
+
+    let after_attr = ||
+        // There's a comment before...
+        #[attr]
+        // and one after!
+        1;
+
+    let one_line = || /* one line */ 1;
+
+    let one_comment_line = || // put this comment in body
+        1;
+}

tests/target/issue-4384.rs

@@ -0,0 +1,51 @@
+fn main() {
+    let single_comment = || {
+        // this should be indented.
+        1
+    };
+
+    let block_comment = || {
+        /* This is a long,
+         * explanatory comment.
+         * Put this all in a block.
+         */
+        1
+    };
+
+    let nested_closure = || {
+        // indent + wrap in a block.
+        || 1
+    };
+
+    let nested_comment_closure = || {
+        // if you see this code in real life, run.
+        || {
+            // nested closures don't need blocks,
+            // but comments do.
+            || 1
+        }
+    };
+
+    let attrb = || {
+        // This block has an attr.
+        #[foo]
+        1
+    };
+
+    let after_attr = || {
+        // There's a comment before...
+        #[attr]
+        // and one after!
+        1
+    };
+
+    let one_line = || {
+        /* one line */
+        1
+    };
+
+    let one_comment_line = || {
+        // put this comment in body
+        1
+    };
+}