Merge pull request #288 from marcusklaas/closure-return

Format closures' return types

2 files changed, 61 insertions, 22 deletions

diff --git a/src/expr.rs b/src/expr.rs
index 2923aa45097..ab20dc1738d 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -163,9 +163,12 @@ fn rewrite_closure(capture: ast::CaptureClause,
 
     // 4 = "|| {".len(), which is overconservative when the closure consists of
     // a single expression.
-    let argument_budget = try_opt!(width.checked_sub(4 + mover.len()));
+    let budget = try_opt!(width.checked_sub(4 + mover.len()));
     // 1 = |
     let argument_offset = offset + 1;
+    let ret_str = try_opt!(fn_decl.output.rewrite(context, budget, argument_offset));
+    // 1 = space between arguments and return type.
+    let horizontal_budget = budget.checked_sub(ret_str.len() + 1).unwrap_or(0);
 
     let arg_items = itemize_list(context.codemap,
                                  fn_decl.inputs.iter(),
@@ -176,13 +179,37 @@ fn rewrite_closure(capture: ast::CaptureClause,
                                  span_after(span, "|", context.codemap),
                                  body.span.lo);
 
-    let fmt = ListFormatting::for_fn(argument_budget, argument_offset);
+    let fmt = ListFormatting {
+        tactic: ListTactic::HorizontalVertical,
+        separator: ",",
+        trailing_separator: SeparatorTactic::Never,
+        indent: argument_offset,
+        h_width: horizontal_budget,
+        v_width: budget,
+        ends_with_newline: false,
+    };
     let list_str = try_opt!(write_list(&arg_items.collect::<Vec<_>>(), &fmt));
-    let prefix = format!("{}|{}|", mover, list_str);
+    let mut prefix = format!("{}|{}|", mover, list_str);
+
+    if !ret_str.is_empty() {
+        if prefix.contains('\n') {
+            prefix.push('\n');
+            prefix.push_str(&make_indent(argument_offset));
+        } else {
+            prefix.push(' ');
+        }
+        prefix.push_str(&ret_str);
+    }
+
     let closure_indent = closure_indent(context, offset);
 
     // Try to format closure body as a single line expression without braces.
-    if body.stmts.is_empty() {
+    if is_simple_block(body, context.codemap) && !prefix.contains('\n') {
+        let (spacer, closer) = if ret_str.is_empty() {
+            (" ", "")
+        } else {
+            (" { ", " }")
+        };
         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.
@@ -192,28 +219,31 @@ fn rewrite_closure(capture: ast::CaptureClause,
             }
             _ => expr,
         };
-
-        // 1 = the separating space between arguments and the body.
-        let extra_offset = extra_offset(&prefix, offset) + 1;
-        let budget = try_opt!(width.checked_sub(extra_offset));
+        let extra_offset = extra_offset(&prefix, offset) + spacer.len();
+        let budget = try_opt!(width.checked_sub(extra_offset + closer.len()));
         let rewrite = inner_expr.rewrite(context, budget, offset + extra_offset);
 
         // Checks if rewrite succeeded and fits on a single line.
         let accept_rewrite = rewrite.as_ref().map(|result| !result.contains('\n')).unwrap_or(false);
 
         if accept_rewrite {
-            return Some(format!("{} {}", prefix, rewrite.unwrap()));
+            return Some(format!("{}{}{}{}", prefix, spacer, rewrite.unwrap(), closer));
         }
     }
 
     // We couldn't format the closure body as a single line expression; fall
     // back to block formatting.
     let inner_context = context.overflow_context(closure_indent - context.block_indent);
-    let body_rewrite = if let ast::Expr_::ExprBlock(ref inner) = body.expr.as_ref().unwrap().node {
-        inner.rewrite(&inner_context, 0, 0)
-    } else {
-        body.rewrite(&inner_context, 0, 0)
-    };
+    let body_rewrite = body.expr
+                           .as_ref()
+                           .and_then(|body_expr| {
+                               if let ast::Expr_::ExprBlock(ref inner) = body_expr.node {
+                                   Some(inner.rewrite(&inner_context, 2, 0))
+                               } else {
+                                   None
+                               }
+                           })
+                           .unwrap_or_else(|| body.rewrite(&inner_context, 2, 0));
 
     Some(format!("{} {}", prefix, try_opt!(body_rewrite)))
 }
diff --git a/src/items.rs b/src/items.rs
index e2224ce4eff..7555ea01199 100644
--- a/src/items.rs
+++ b/src/items.rs
@@ -17,7 +17,7 @@ use lists::{write_list, itemize_list, ListItem, ListFormatting, SeparatorTactic,
 use expr::rewrite_assign_rhs;
 use comment::FindUncommented;
 use visitor::FmtVisitor;
-use rewrite::Rewrite;
+use rewrite::{Rewrite, RewriteContext};
 use config::{Config, BlockIndentStyle, Density};
 
 use syntax::{ast, abi};
@@ -207,7 +207,8 @@ impl<'a> FmtVisitor<'a> {
                                                           generics_span));
         result.push_str(&generics_str);
 
-        let ret_str = self.rewrite_return(&fd.output, indent);
+        let context = self.get_context();
+        let ret_str = fd.output.rewrite(&context, self.config.max_width - indent, indent).unwrap();
 
         // Args.
         let (one_line_budget, multi_line_budget, mut arg_indent) =
@@ -902,14 +903,22 @@ impl<'a> FmtVisitor<'a> {
             Some(format!(" where {}", preds_str))
         }
     }
+}
 
-    fn rewrite_return(&self, ret: &ast::FunctionRetTy, indent: usize) -> String {
-        match *ret {
-            ast::FunctionRetTy::DefaultReturn(_) => String::new(),
-            ast::FunctionRetTy::NoReturn(_) => "-> !".to_owned(),
+impl Rewrite for ast::FunctionRetTy {
+    fn rewrite(&self, context: &RewriteContext, width: usize, offset: usize) -> Option<String> {
+        match *self {
+            ast::FunctionRetTy::DefaultReturn(_) => Some(String::new()),
+            ast::FunctionRetTy::NoReturn(_) => {
+                if width >= 4 {
+                    Some("-> !".to_owned())
+                } else {
+                    None
+                }
+            }
             ast::FunctionRetTy::Return(ref ty) => {
-                let ctxt = &self.get_context();
-                format!("-> {}", ty.rewrite(ctxt, ctxt.config.max_width, indent).unwrap())
+                let inner_width = try_opt!(width.checked_sub(3));
+                ty.rewrite(context, inner_width, offset + 3).map(|r| format!("-> {}", r))
             }
         }
     }