From c8fd23ca6822477199542d6d1f153da0d41fa629 Mon Sep 17 00:00:00 2001
From: Marcus Klaas <mail@marcusklaas.nl>
Date: Thu, 20 Aug 2015 23:05:41 +0200
Subject: [PATCH] Refactor closure formatting routine

---
 src/expr.rs             | 32 ++++++++++----------------------
 src/visitor.rs          | 15 ++++++++++++---
 tests/source/closure.rs |  4 ++++
 tests/target/closure.rs |  8 ++++++++
 4 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/src/expr.rs b/src/expr.rs
index af6e9ce1e35b..6ed8470fbea5 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -160,7 +160,8 @@ fn rewrite_closure(capture: ast::CaptureClause,
     let prefix = format!("{}|{}|", mover, write_list(&arg_items.collect::<Vec<_>>(), &fmt));
     let block_indent = closure_block_indent(context, offset);
 
-    let body_rewrite = if body.stmts.is_empty() {
+    // Try to format closure body as a single line expression without braces.
+    if body.stmts.is_empty() {
         let expr = body.expr.as_ref().unwrap();
         // All closure bodies are blocks in the eyes of the AST, but we may not
         // want to unwrap them when they only contain a single expression.
@@ -179,29 +180,16 @@ fn rewrite_closure(capture: ast::CaptureClause,
         let accept_rewrite = rewrite.as_ref().map(|result| !result.contains('\n')).unwrap_or(false);
 
         if accept_rewrite {
-            rewrite
-        } else {
-            if let ast::Expr_::ExprBlock(ref inner_body) = expr.node {
-                // Closure body is a proper block, with braces and all.
-                let inner_context = &RewriteContext { block_indent: block_indent, ..*context };
-                inner_body.rewrite(inner_context, 0, 0)
-            } else {
-                // Closure body is an expression, but not a block. It does not
-                // fit a single line, so we format it as it were wrapped in a
-                // block.
-                let inner_context = &RewriteContext {
-                    block_indent: block_indent + context.config.tab_spaces,
-                    ..*context
-                };
-                let indent = offset + context.config.tab_spaces;
-                expr.rewrite(inner_context, context.config.max_width - indent, indent)
-                    .map(|result| {
-                        format!("{{\n{}{}\n{}}}", make_indent(indent), result, make_indent(offset))
-                    })
-            }
+            return Some(format!("{} {}", prefix, rewrite.unwrap()));
         }
+    }
+
+    // We couldn't format the closure body as a single line expression; fall
+    // back to block formatting.
+    let inner_context = &RewriteContext { block_indent: block_indent, ..*context };
+    let body_rewrite = if let ast::Expr_::ExprBlock(ref inner) = body.expr.as_ref().unwrap().node {
+        inner.rewrite(inner_context, 0, 0)
     } else {
-        let inner_context = &RewriteContext { block_indent: block_indent, ..*context };
         body.rewrite(inner_context, 0, 0)
     };
 
diff --git a/src/visitor.rs b/src/visitor.rs
index 86b1197822b4..2489a8c1a9d5 100644
--- a/src/visitor.rs
+++ b/src/visitor.rs
@@ -68,9 +68,18 @@ impl<'a, 'v> visit::Visitor<'v> for FmtVisitor<'a> {
                self.codemap.lookup_char_pos(b.span.lo),
                self.codemap.lookup_char_pos(b.span.hi));
 
-        self.buffer.push_str("{");
-        self.last_pos = self.last_pos + BytePos(1);
+        // Check if this block has braces.
+        let snippet = self.snippet(b.span);
+        let has_braces = snippet.chars().next().unwrap() == '{' || &snippet[..6] == "unsafe";
+        let brace_compensation = if has_braces {
+            BytePos(1)
+        } else {
+            BytePos(0)
+        };
+
+        self.last_pos = self.last_pos + brace_compensation;
         self.block_indent += self.config.tab_spaces;
+        self.buffer.push_str("{");
 
         for stmt in &b.stmts {
             self.visit_stmt(&stmt)
@@ -86,7 +95,7 @@ impl<'a, 'v> visit::Visitor<'v> for FmtVisitor<'a> {
 
         self.block_indent -= self.config.tab_spaces;
         // TODO we should compress any newlines here to just one
-        self.format_missing_with_indent(b.span.hi - BytePos(1));
+        self.format_missing_with_indent(b.span.hi - brace_compensation);
         self.buffer.push_str("}");
         self.last_pos = b.span.hi;
     }
diff --git a/tests/source/closure.rs b/tests/source/closure.rs
index e7d54716e4c2..06b107580118 100644
--- a/tests/source/closure.rs
+++ b/tests/source/closure.rs
@@ -30,6 +30,10 @@ fn main() {
 
     let test = |  | { do_something(); do_something_else(); };
 
+    let arg_test = |big_argument_name, test123| looooooooooooooooooong_function_naaaaaaaaaaaaaaaaame();
+
+    let arg_test = |big_argument_name, test123| {looooooooooooooooooong_function_naaaaaaaaaaaaaaaaame()};
+
     |arg1, arg2, _, _, arg3, arg4| { let temp = arg4 + arg3;
                                      arg2 * arg1 - temp }
 }
diff --git a/tests/target/closure.rs b/tests/target/closure.rs
index e25549b10798..5042f1d037ae 100644
--- a/tests/target/closure.rs
+++ b/tests/target/closure.rs
@@ -51,6 +51,14 @@ fn main() {
                    do_something_else();
                };
 
+    let arg_test = |big_argument_name, test123| {
+                       looooooooooooooooooong_function_naaaaaaaaaaaaaaaaame()
+                   };
+
+    let arg_test = |big_argument_name, test123| {
+                       looooooooooooooooooong_function_naaaaaaaaaaaaaaaaame()
+                   };
+
     |arg1, arg2, _, _, arg3, arg4| {
         let temp = arg4 + arg3;
         arg2 * arg1 - temp