Improve formatting of series of binop expressions

This commit changes the handling of binops (and potentially other pairs), where
the expressions are logically a list, e.g., `a + b + c`. It makes the single
line vs multi-line approaches explicit and introduces a lowering step.

This improves formatting in a number of places, mostly improving consistency of
formatting with very short sub-expressions, but also some weird indentation.

Closes #2802

2 files changed, 221 insertions, 88 deletions

diff --git a/src/expr.rs b/src/expr.rs
index 8f167650507..7f39e93b267 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -31,7 +31,7 @@ use lists::{
 use macros::{rewrite_macro, MacroArg, MacroPosition};
 use matches::rewrite_match;
 use overflow;
-use pairs::{rewrite_pair, PairParts};
+use pairs::{rewrite_all_pairs, rewrite_pair, PairParts};
 use patterns::{can_be_overflowed_pat, is_short_pattern, TuplePatField};
 use rewrite::{Rewrite, RewriteContext};
 use shape::{Indent, Shape};
@@ -89,7 +89,7 @@ pub fn format_expr(
         ast::ExprKind::Paren(ref subexpr) => rewrite_paren(context, subexpr, shape, expr.span),
         ast::ExprKind::Binary(op, ref lhs, ref rhs) => {
             // FIXME: format comments between operands and operator
-            rewrite_simple_binaries(context, expr, shape, op).or_else(|| {
+            rewrite_all_pairs(expr, shape, context).or_else(|| {
                 rewrite_pair(
                     &**lhs,
                     &**rhs,
@@ -362,80 +362,6 @@ pub fn format_expr(
         })
 }
 
-/// Collect operands that appears in the given binary operator in the opposite order.
-/// e.g. `collect_binary_items(e, ||)` for `a && b || c || d` returns `[d, c, a && b]`.
-fn collect_binary_items<'a>(mut expr: &'a ast::Expr, binop: ast::BinOp) -> Vec<&'a ast::Expr> {
-    let mut result = vec![];
-    let mut prev_lhs = None;
-    loop {
-        match expr.node {
-            ast::ExprKind::Binary(inner_binop, ref lhs, ref rhs)
-                if inner_binop.node == binop.node =>
-            {
-                result.push(&**rhs);
-                expr = lhs;
-                prev_lhs = Some(lhs);
-            }
-            _ => {
-                if let Some(lhs) = prev_lhs {
-                    result.push(lhs);
-                }
-                break;
-            }
-        }
-    }
-    result
-}
-
-/// Rewrites a binary expression whose operands fits within a single line.
-fn rewrite_simple_binaries(
-    context: &RewriteContext,
-    expr: &ast::Expr,
-    shape: Shape,
-    op: ast::BinOp,
-) -> Option<String> {
-    let op_str = context.snippet(op.span);
-
-    // 2 = spaces around a binary operator.
-    let sep_overhead = op_str.len() + 2;
-    let nested_overhead = sep_overhead - 1;
-
-    let nested_shape = (match context.config.indent_style() {
-        IndentStyle::Visual => shape.visual_indent(0),
-        IndentStyle::Block => shape.block_indent(context.config.tab_spaces()),
-    }).with_max_width(context.config);
-    let nested_shape = match context.config.binop_separator() {
-        SeparatorPlace::Back => nested_shape.sub_width(nested_overhead)?,
-        SeparatorPlace::Front => nested_shape.offset_left(nested_overhead)?,
-    };
-
-    let opt_rewrites: Option<Vec<_>> = collect_binary_items(expr, op)
-        .iter()
-        .rev()
-        .map(|e| e.rewrite(context, nested_shape))
-        .collect();
-    if let Some(rewrites) = opt_rewrites {
-        if rewrites.iter().all(|e| ::utils::is_single_line(e)) {
-            let total_width = rewrites.iter().map(|s| s.len()).sum::<usize>()
-                + sep_overhead * (rewrites.len() - 1);
-
-            let sep_str = if total_width <= shape.width {
-                format!(" {} ", op_str)
-            } else {
-                let indent_str = nested_shape.indent.to_string_with_newline(context.config);
-                match context.config.binop_separator() {
-                    SeparatorPlace::Back => format!(" {}{}", op_str.trim_right(), indent_str),
-                    SeparatorPlace::Front => format!("{}{} ", indent_str, op_str.trim_left()),
-                }
-            };
-
-            return wrap_str(rewrites.join(&sep_str), context.config.max_width(), shape);
-        }
-    }
-
-    None
-}
-
 pub fn rewrite_array<T: Rewrite + Spanned + ToExpr>(
     name: &str,
     exprs: &[&T],
@@ -2004,7 +1930,7 @@ fn choose_rhs<R: Rewrite>(
                 }
                 (None, Some(ref new_rhs)) => Some(format!("{}{}", new_indent_str, new_rhs)),
                 (None, None) => None,
-                (Some(ref orig_rhs), _) => Some(format!(" {}", orig_rhs)),
+                (Some(orig_rhs), _) => Some(format!(" {}", orig_rhs)),
             }
         }
     }
diff --git a/src/pairs.rs b/src/pairs.rs
index 55c78b641a0..3cd044d606e 100644
--- a/src/pairs.rs
+++ b/src/pairs.rs
@@ -8,23 +8,24 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use config::lists::*;
+use syntax::ast;
 
+use config::lists::*;
 use config::IndentStyle;
 use rewrite::{Rewrite, RewriteContext};
 use shape::Shape;
-use utils::{first_line_width, last_line_width};
+use utils::{first_line_width, is_single_line, last_line_width, trimmed_last_line_width, wrap_str};
 
 /// Sigils that decorate a binop pair.
 #[derive(new, Clone, Copy)]
-pub struct PairParts<'a> {
+pub(crate) struct PairParts<'a> {
     prefix: &'a str,
     infix: &'a str,
     suffix: &'a str,
 }
 
 impl<'a> PairParts<'a> {
-    pub fn infix(infix: &'a str) -> PairParts<'a> {
+    pub(crate) fn infix(infix: &'a str) -> PairParts<'a> {
         PairParts {
             prefix: "",
             infix,
@@ -33,7 +34,148 @@ impl<'a> PairParts<'a> {
     }
 }
 
-pub fn rewrite_pair<LHS, RHS>(
+// Flattens a tree of pairs into a list and tries to rewrite them all at once.
+// FIXME would be nice to reuse the lists API for this, but because each separator
+// can be different, we can't.
+pub(crate) fn rewrite_all_pairs(
+    expr: &ast::Expr,
+    shape: Shape,
+    context: &RewriteContext,
+) -> Option<String> {
+    // First we try formatting on one line.
+    if let Some(list) = expr.flatten(context, false) {
+        if let Some(r) = rewrite_pairs_one_line(&list, shape, context) {
+            return Some(r);
+        }
+    }
+
+    // We can't format on line, so try many. When we flatten here we make sure
+    // to only flatten pairs with the same operator, that way we don't
+    // necessarily need one line per sub-expression, but we don't do anything
+    // too funny wrt precedence.
+    expr.flatten(context, true)
+        .and_then(|list| rewrite_pairs_multiline(list, shape, context))
+}
+
+// This may return a multi-line result since we allow the last expression to go
+// multiline in a 'single line' formatting.
+fn rewrite_pairs_one_line<T: Rewrite>(
+    list: &PairList<T>,
+    shape: Shape,
+    context: &RewriteContext,
+) -> Option<String> {
+    assert!(list.list.len() >= 2, "Not a pair?");
+
+    let mut result = String::new();
+    let base_shape = shape.block();
+
+    for (e, s) in list.list[..list.list.len()]
+        .iter()
+        .zip(list.separators.iter())
+    {
+        let cur_shape = base_shape.offset_left(last_line_width(&result))?;
+        let rewrite = e.rewrite(context, cur_shape)?;
+
+        if !is_single_line(&rewrite) || result.len() > shape.width {
+            return None;
+        }
+
+        result.push_str(&rewrite);
+        result.push(' ');
+        result.push_str(s);
+        result.push(' ');
+    }
+
+    let last = list.list.last().unwrap();
+    let cur_shape = base_shape.offset_left(last_line_width(&result))?;
+    let rewrite = last.rewrite(context, cur_shape)?;
+    result.push_str(&rewrite);
+
+    if first_line_width(&result) > shape.width {
+        return None;
+    }
+
+    // Check the last expression in the list. We let this expression go over
+    // multiple lines, but we check that if this is necessary, then we can't
+    // do better using multi-line formatting.
+    if !is_single_line(&result) {
+        let multiline_shape = shape.offset_left(list.separators.last().unwrap().len() + 1)?;
+        let multiline_list: PairList<T> = PairList {
+            list: vec![last],
+            separators: vec![],
+            separator_place: list.separator_place,
+        };
+        // Format as if we were multi-line.
+        if let Some(rewrite) = rewrite_pairs_multiline(multiline_list, multiline_shape, context) {
+            // Also, don't let expressions surrounded by parens go multi-line,
+            // this looks really bad.
+            if rewrite.starts_with('(') || is_single_line(&rewrite) {
+                return None;
+            }
+        }
+    }
+
+    wrap_str(result, context.config.max_width(), shape)
+}
+
+fn rewrite_pairs_multiline<T: Rewrite>(
+    list: PairList<T>,
+    shape: Shape,
+    context: &RewriteContext,
+) -> Option<String> {
+    let rhs_offset = shape.rhs_overhead(&context.config);
+    let nested_shape = (match context.config.indent_style() {
+        IndentStyle::Visual => shape.visual_indent(0),
+        IndentStyle::Block => shape.block_indent(context.config.tab_spaces()),
+    }).with_max_width(&context.config)
+        .sub_width(rhs_offset)?;
+
+    let indent_str = nested_shape.indent.to_string_with_newline(context.config);
+    let mut result = String::new();
+
+    let rewrite = list.list[0].rewrite(context, shape)?;
+    result.push_str(&rewrite);
+
+    for (e, s) in list.list[1..].iter().zip(list.separators.iter()) {
+        if trimmed_last_line_width(&result) <= context.config.tab_spaces() {
+            // We must snuggle the next line onto the previous line to avoid an orphan.
+            if let Some(line_shape) =
+                shape.offset_left(s.len() + 2 + trimmed_last_line_width(&result))
+            {
+                if let Some(rewrite) = e.rewrite(context, line_shape) {
+                    result.push(' ');
+                    result.push_str(s);
+                    result.push(' ');
+                    result.push_str(&rewrite);
+                    continue;
+                }
+            }
+        }
+
+        let nested_overhead = s.len() + 1;
+        let line_shape = match context.config.binop_separator() {
+            SeparatorPlace::Back => {
+                result.push(' ');
+                result.push_str(s);
+                result.push_str(&indent_str);
+                nested_shape.sub_width(nested_overhead)?
+            }
+            SeparatorPlace::Front => {
+                result.push_str(&indent_str);
+                result.push_str(s);
+                result.push(' ');
+                nested_shape.offset_left(nested_overhead)?
+            }
+        };
+
+        let rewrite = e.rewrite(context, line_shape)?;
+        result.push_str(&rewrite);
+    }
+    Some(result)
+}
+
+// Rewrites a single pair.
+pub(crate) fn rewrite_pair<LHS, RHS>(
     lhs: &LHS,
     rhs: &RHS,
     pp: PairParts,
@@ -45,6 +187,7 @@ where
     LHS: Rewrite,
     RHS: Rewrite,
 {
+    let tab_spaces = context.config.tab_spaces();
     let lhs_overhead = match separator_place {
         SeparatorPlace::Back => shape.used_width() + pp.prefix.len() + pp.infix.trim_right().len(),
         SeparatorPlace::Front => shape.used_width(),
@@ -66,12 +209,11 @@ where
         // If the length of the lhs is equal to or shorter than the tab width or
         // the rhs looks like block expression, we put the rhs on the same
         // line with the lhs even if the rhs is multi-lined.
-        let allow_same_line = lhs_result.len() <= context.config.tab_spaces()
-            || rhs_result
-                .lines()
-                .next()
-                .map(|first_line| first_line.ends_with('{'))
-                .unwrap_or(false);
+        let allow_same_line = lhs_result.len() <= tab_spaces || rhs_result
+            .lines()
+            .next()
+            .map(|first_line| first_line.ends_with('{'))
+            .unwrap_or(false);
         if !rhs_result.contains('\n') || allow_same_line {
             let one_line_width = last_line_width(&lhs_result)
                 + pp.infix.len()
@@ -117,3 +259,68 @@ where
         lhs_result, infix_with_sep, rhs_result, pp.suffix
     ))
 }
+
+// A pair which forms a tree and can be flattened (e.g., binops).
+trait FlattenPair: Rewrite + Sized {
+    // If `_same_op` is `true`, then we only combine binops with the same
+    // operator into the list. E.g,, if the source is `a * b + c`, if `_same_op`
+    // is true, we make `[(a * b), c]` if `_same_op` is false, we make
+    // `[a, b, c]`
+    fn flatten(&self, _context: &RewriteContext, _same_op: bool) -> Option<PairList<Self>> {
+        None
+    }
+}
+
+struct PairList<'a, 'b, T: Rewrite + 'b> {
+    list: Vec<&'b T>,
+    separators: Vec<&'a str>,
+    separator_place: SeparatorPlace,
+}
+
+impl FlattenPair for ast::Expr {
+    fn flatten(&self, context: &RewriteContext, same_op: bool) -> Option<PairList<ast::Expr>> {
+        let top_op = match self.node {
+            ast::ExprKind::Binary(op, _, _) => op.node,
+            _ => return None,
+        };
+
+        // Turn a tree of binop expressions into a list using a depth-first,
+        // in-order traversal.
+        let mut stack = vec![];
+        let mut list = vec![];
+        let mut separators = vec![];
+        let mut node = self;
+        loop {
+            match node.node {
+                ast::ExprKind::Binary(op, ref lhs, _) if !same_op || op.node == top_op => {
+                    stack.push(node);
+                    node = lhs;
+                }
+                _ => {
+                    list.push(node);
+                    if let Some(pop) = stack.pop() {
+                        match pop.node {
+                            ast::ExprKind::Binary(op, _, ref rhs) => {
+                                separators.push(op.node.to_string());
+                                node = rhs;
+                            }
+                            _ => unreachable!(),
+                        }
+                    } else {
+                        break;
+                    }
+                }
+            }
+        }
+
+        assert_eq!(list.len() - 1, separators.len());
+        Some(PairList {
+            list,
+            separators,
+            separator_place: context.config.binop_separator(),
+        })
+    }
+}
+
+impl FlattenPair for ast::Ty {}
+impl FlattenPair for ast::Pat {}