src/pairs.rs

   1 use rustc_ast::ast;
   2
   3 use crate::config::lists::*;
   4 use crate::config::IndentStyle;
   5 use crate::rewrite::{Rewrite, RewriteContext};
   6 use crate::shape::Shape;
   7 use crate::utils::{
   8     first_line_width, is_single_line, last_line_width, trimmed_last_line_width, wrap_str,
   9 };
  10
  11 /// Sigils that decorate a binop pair.
  12 #[derive(new, Clone, Copy)]
  13 pub(crate) struct PairParts<'a> {
  14     prefix: &'a str,
  15     infix: &'a str,
  16     suffix: &'a str,
  17 }
  18
  19 impl<'a> PairParts<'a> {
  20     pub(crate) fn infix(infix: &'a str) -> PairParts<'a> {
  21         PairParts {
  22             prefix: "",
  23             infix,
  24             suffix: "",
  25         }
  26     }
  27 }
  28
  29 // Flattens a tree of pairs into a list and tries to rewrite them all at once.
  30 // FIXME would be nice to reuse the lists API for this, but because each separator
  31 // can be different, we can't.
  32 pub(crate) fn rewrite_all_pairs(
  33     expr: &ast::Expr,
  34     shape: Shape,
  35     context: &RewriteContext<'_>,
  36 ) -> Option<String> {
  37     expr.flatten(context, shape).and_then(|list| {
  38         // First we try formatting on one line.
  39         rewrite_pairs_one_line(&list, shape, context)
  40             .or_else(|| rewrite_pairs_multiline(&list, shape, context))
  41     })
  42 }
  43
  44 // This may return a multi-line result since we allow the last expression to go
  45 // multiline in a 'single line' formatting.
  46 fn rewrite_pairs_one_line<T: Rewrite>(
  47     list: &PairList<'_, '_, T>,
  48     shape: Shape,
  49     context: &RewriteContext<'_>,
  50 ) -> Option<String> {
  51     assert!(list.list.len() >= 2, "Not a pair?");
  52
  53     let mut result = String::new();
  54     let base_shape = shape.block();
  55
  56     for ((_, rewrite), s) in list.list.iter().zip(list.separators.iter()) {
  57         if let Some(rewrite) = rewrite {
  58             if !is_single_line(rewrite) || result.len() > shape.width {
  59                 return None;
  60             }
  61
  62             result.push_str(rewrite);
  63             result.push(' ');
  64             result.push_str(s);
  65             result.push(' ');
  66         } else {
  67             return None;
  68         }
  69     }
  70
  71     let prefix_len = result.len();
  72     let last = list.list.last()?.0;
  73     let cur_shape = base_shape.offset_left(last_line_width(&result))?;
  74     let last_rewrite = last.rewrite(context, cur_shape)?;
  75     result.push_str(&last_rewrite);
  76
  77     if first_line_width(&result) > shape.width {
  78         return None;
  79     }
  80
  81     // Check the last expression in the list. We sometimes let this expression
  82     // go over multiple lines, but we check for some ugly conditions.
  83     if !(is_single_line(&result) || last_rewrite.starts_with('{'))
  84         && (last_rewrite.starts_with('(') || prefix_len > context.config.tab_spaces())
  85     {
  86         return None;
  87     }
  88
  89     wrap_str(result, context.config.max_width(), shape)
  90 }
  91
  92 fn rewrite_pairs_multiline<T: Rewrite>(
  93     list: &PairList<'_, '_, T>,
  94     shape: Shape,
  95     context: &RewriteContext<'_>,
  96 ) -> Option<String> {
  97     let rhs_offset = shape.rhs_overhead(context.config);
  98     let nested_shape = (match context.config.indent_style() {
  99         IndentStyle::Visual => shape.visual_indent(0),
 100         IndentStyle::Block => shape.block_indent(context.config.tab_spaces()),
 101     })
 102     .with_max_width(context.config)
 103     .sub_width(rhs_offset)?;
 104
 105     let indent_str = nested_shape.indent.to_string_with_newline(context.config);
 106     let mut result = String::new();
 107
 108     result.push_str(list.list[0].1.as_ref()?);
 109
 110     for ((e, default_rw), s) in list.list[1..].iter().zip(list.separators.iter()) {
 111         // The following test checks if we should keep two subexprs on the same
 112         // line. We do this if not doing so would create an orphan and there is
 113         // enough space to do so.
 114         let offset = if result.contains('\n') {
 115             0
 116         } else {
 117             shape.used_width()
 118         };
 119         if last_line_width(&result) + offset <= nested_shape.used_width() {
 120             // We must snuggle the next line onto the previous line to avoid an orphan.
 121             if let Some(line_shape) =
 122                 shape.offset_left(s.len() + 2 + trimmed_last_line_width(&result))
 123             {
 124                 if let Some(rewrite) = e.rewrite(context, line_shape) {
 125                     result.push(' ');
 126                     result.push_str(s);
 127                     result.push(' ');
 128                     result.push_str(&rewrite);
 129                     continue;
 130                 }
 131             }
 132         }
 133
 134         match context.config.binop_separator() {
 135             SeparatorPlace::Back => {
 136                 result.push(' ');
 137                 result.push_str(s);
 138                 result.push_str(&indent_str);
 139             }
 140             SeparatorPlace::Front => {
 141                 result.push_str(&indent_str);
 142                 result.push_str(s);
 143                 result.push(' ');
 144             }
 145         }
 146
 147         result.push_str(default_rw.as_ref()?);
 148     }
 149     Some(result)
 150 }
 151
 152 // Rewrites a single pair.
 153 pub(crate) fn rewrite_pair<LHS, RHS>(
 154     lhs: &LHS,
 155     rhs: &RHS,
 156     pp: PairParts<'_>,
 157     context: &RewriteContext<'_>,
 158     shape: Shape,
 159     separator_place: SeparatorPlace,
 160 ) -> Option<String>
 161 where
 162     LHS: Rewrite,
 163     RHS: Rewrite,
 164 {
 165     let tab_spaces = context.config.tab_spaces();
 166     let lhs_overhead = match separator_place {
 167         SeparatorPlace::Back => shape.used_width() + pp.prefix.len() + pp.infix.trim_end().len(),
 168         SeparatorPlace::Front => shape.used_width(),
 169     };
 170     let lhs_shape = Shape {
 171         width: context.budget(lhs_overhead),
 172         ..shape
 173     };
 174     let lhs_result = lhs
 175         .rewrite(context, lhs_shape)
 176         .map(|lhs_str| format!("{}{}", pp.prefix, lhs_str))?;
 177
 178     // Try to put both lhs and rhs on the same line.
 179     let rhs_orig_result = shape
 180         .offset_left(last_line_width(&lhs_result) + pp.infix.len())
 181         .and_then(|s| s.sub_width(pp.suffix.len()))
 182         .and_then(|rhs_shape| rhs.rewrite(context, rhs_shape));
 183     if let Some(ref rhs_result) = rhs_orig_result {
 184         // If the length of the lhs is equal to or shorter than the tab width or
 185         // the rhs looks like block expression, we put the rhs on the same
 186         // line with the lhs even if the rhs is multi-lined.
 187         let allow_same_line = lhs_result.len() <= tab_spaces
 188             || rhs_result
 189                 .lines()
 190                 .next()
 191                 .map(|first_line| first_line.ends_with('{'))
 192                 .unwrap_or(false);
 193         if !rhs_result.contains('\n') || allow_same_line {
 194             let one_line_width = last_line_width(&lhs_result)
 195                 + pp.infix.len()
 196                 + first_line_width(rhs_result)
 197                 + pp.suffix.len();
 198             if one_line_width <= shape.width {
 199                 return Some(format!(
 200                     "{}{}{}{}",
 201                     lhs_result, pp.infix, rhs_result, pp.suffix
 202                 ));
 203             }
 204         }
 205     }
 206
 207     // We have to use multiple lines.
 208     // Re-evaluate the rhs because we have more space now:
 209     let mut rhs_shape = match context.config.indent_style() {
 210         IndentStyle::Visual => shape
 211             .sub_width(pp.suffix.len() + pp.prefix.len())?
 212             .visual_indent(pp.prefix.len()),
 213         IndentStyle::Block => {
 214             // Try to calculate the initial constraint on the right hand side.
 215             let rhs_overhead = shape.rhs_overhead(context.config);
 216             Shape::indented(shape.indent.block_indent(context.config), context.config)
 217                 .sub_width(rhs_overhead)?
 218         }
 219     };
 220     let infix = match separator_place {
 221         SeparatorPlace::Back => pp.infix.trim_end(),
 222         SeparatorPlace::Front => pp.infix.trim_start(),
 223     };
 224     if separator_place == SeparatorPlace::Front {
 225         rhs_shape = rhs_shape.offset_left(infix.len())?;
 226     }
 227     let rhs_result = rhs.rewrite(context, rhs_shape)?;
 228     let indent_str = rhs_shape.indent.to_string_with_newline(context.config);
 229     let infix_with_sep = match separator_place {
 230         SeparatorPlace::Back => format!("{}{}", infix, indent_str),
 231         SeparatorPlace::Front => format!("{}{}", indent_str, infix),
 232     };
 233     Some(format!(
 234         "{}{}{}{}",
 235         lhs_result, infix_with_sep, rhs_result, pp.suffix
 236     ))
 237 }
 238
 239 // A pair which forms a tree and can be flattened (e.g., binops).
 240 trait FlattenPair: Rewrite + Sized {
 241     fn flatten(&self, _: &RewriteContext<'_>, _: Shape) -> Option<PairList<'_, '_, Self>> {
 242         None
 243     }
 244 }
 245
 246 struct PairList<'a, 'b, T: Rewrite> {
 247     list: Vec<(&'b T, Option<String>)>,
 248     separators: Vec<&'a str>,
 249 }
 250
 251 impl FlattenPair for ast::Expr {
 252     fn flatten(
 253         &self,
 254         context: &RewriteContext<'_>,
 255         shape: Shape,
 256     ) -> Option<PairList<'_, '_, ast::Expr>> {
 257         let top_op = match self.kind {
 258             ast::ExprKind::Binary(op, _, _) => op.node,
 259             _ => return None,
 260         };
 261
 262         let default_rewrite = |node: &ast::Expr, sep: usize, is_first: bool| {
 263             if is_first {
 264                 return node.rewrite(context, shape);
 265             }
 266             let nested_overhead = sep + 1;
 267             let rhs_offset = shape.rhs_overhead(context.config);
 268             let nested_shape = (match context.config.indent_style() {
 269                 IndentStyle::Visual => shape.visual_indent(0),
 270                 IndentStyle::Block => shape.block_indent(context.config.tab_spaces()),
 271             })
 272             .with_max_width(context.config)
 273             .sub_width(rhs_offset)?;
 274             let default_shape = match context.config.binop_separator() {
 275                 SeparatorPlace::Back => nested_shape.sub_width(nested_overhead)?,
 276                 SeparatorPlace::Front => nested_shape.offset_left(nested_overhead)?,
 277             };
 278             node.rewrite(context, default_shape)
 279         };
 280
 281         // Turn a tree of binop expressions into a list using a depth-first,
 282         // in-order traversal.
 283         let mut stack = vec![];
 284         let mut list = vec![];
 285         let mut separators = vec![];
 286         let mut node = self;
 287         loop {
 288             match node.kind {
 289                 ast::ExprKind::Binary(op, ref lhs, _) if op.node == top_op => {
 290                     stack.push(node);
 291                     node = lhs;
 292                 }
 293                 _ => {
 294                     let op_len = separators.last().map_or(0, |s: &&str| s.len());
 295                     let rw = default_rewrite(node, op_len, list.is_empty());
 296                     list.push((node, rw));
 297                     if let Some(pop) = stack.pop() {
 298                         match pop.kind {
 299                             ast::ExprKind::Binary(op, _, ref rhs) => {
 300                                 separators.push(op.node.to_string());
 301                                 node = rhs;
 302                             }
 303                             _ => unreachable!(),
 304                         }
 305                     } else {
 306                         break;
 307                     }
 308                 }
 309             }
 310         }
 311
 312         assert_eq!(list.len() - 1, separators.len());
 313         Some(PairList { list, separators })
 314     }
 315 }
 316
 317 impl FlattenPair for ast::Ty {}
 318 impl FlattenPair for ast::Pat {}