Fix exponential execution time by memoizing format_expr

src/expr.rs

@@ -1,5 +1,6 @@
 use std::borrow::Cow;
 use std::cmp::min;
+use std::collections::HashMap;
 
 use itertools::Itertools;
 use rustc_ast::token::{DelimToken, LitKind};
@@ -22,7 +23,7 @@ use crate::macros::{rewrite_macro, MacroPosition};
 use crate::matches::rewrite_match;
 use crate::overflow::{self, IntoOverflowableItem, OverflowableItem};
 use crate::pairs::{rewrite_all_pairs, rewrite_pair, PairParts};
-use crate::rewrite::{Rewrite, RewriteContext};
+use crate::rewrite::{QueryId, Rewrite, RewriteContext};
 use crate::shape::{Indent, Shape};
 use crate::source_map::{LineRangeUtils, SpanUtils};
 use crate::spanned::Spanned;
@@ -53,6 +54,54 @@ pub(crate) fn format_expr(
  expr_type: ExprType,
  context: &RewriteContext<'_>,
  shape: Shape,
+) -> Option<String> {
+ // when max_width is tight, we should check all possible formattings, in order to find
+ // if we can fit expression in the limit. Doing it recursively takes exponential time
+ // relative to input size, and people hit it with rustfmt takes minutes in #4476 #4867 #5128
+ // By memoization of format_expr function, we format each pair of expression and shape
+ // only once, so worst case execution time becomes O(n*max_width^3).
+ if context.inside_macro() || context.is_macro_def {
+ // span ids are not unique in macros, so we don't memoize result of them.
+ return format_expr_inner(expr, expr_type, context, shape);
+ }
+ let clean;
+ let query_id = QueryId {
+ shape,
+ span: expr.span,
+ };
+ if let Some(map) = context.memoize.take() {
+ if let Some(r) = map.get(&query_id) {
+ let r = r.clone();
+ context.memoize.set(Some(map)); // restore map in the memoize cell for other users
+ return r;
+ }
+ context.memoize.set(Some(map));
+ clean = false;
+ } else {
+ context.memoize.set(Some(HashMap::default()));
+ clean = true; // We got None, so we are the top level called function. When
+ // this function finishes, no one is interested in what is in the map, because
+ // all of them are sub expressions of this top level expression, and this is
+ // done. So we should clean up memoize map to save some memory.
+ }
+
+ let r = format_expr_inner(expr, expr_type, context, shape);
+ if clean {
+ context.memoize.set(None);
+ } else {
+ if let Some(mut map) = context.memoize.take() {
+ map.insert(query_id, r.clone()); // insert the result in the memoize map
+ context.memoize.set(Some(map)); // so it won't be computed again
+ }
+ }
+ r
+}
+
+fn format_expr_inner(
+ expr: &ast::Expr,
+ expr_type: ExprType,
+ context: &RewriteContext<'_>,
+ shape: Shape,
 ) -> Option<String> {
  skip_out_of_file_lines_range!(context, expr.span);
 

src/formatting.rs

@@ -2,6 +2,7 @@
 
 use std::collections::HashMap;
 use std::io::{self, Write};
+use std::rc::Rc;
 use std::time::{Duration, Instant};
 
 use rustc_ast::ast;
@@ -190,6 +191,7 @@ impl<'a, T: FormatHandler + 'a> FormatContext<'a, T> {
  self.config,
  &snippet_provider,
  self.report.clone(),
+ Rc::default(),
  );
  visitor.skip_context.update_with_attrs(&self.krate.attrs);
  visitor.is_macro_def = is_macro_def;

src/rewrite.rs

@@ -12,6 +12,7 @@ use crate::shape::Shape;
 use crate::skip::SkipContext;
 use crate::visitor::SnippetProvider;
 use crate::FormatReport;
+use rustc_data_structures::stable_map::FxHashMap;
 
 pub(crate) trait Rewrite {
  /// Rewrite self into shape.
@@ -24,10 +25,22 @@ impl<T: Rewrite> Rewrite for ptr::P<T> {
  }
 }
 
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub(crate) struct QueryId {
+ pub(crate) shape: Shape,
+ pub(crate) span: Span,
+}
+
+// We use Option<HashMap> instead of HashMap, because in case of `None`
+// the function clean the memoize map, but it doesn't clean when
+// there is `Some(empty)`, so they are different.
+pub(crate) type Memoize = Rc<Cell<Option<FxHashMap<QueryId, Option<String>>>>>;
+
 #[derive(Clone)]
 pub(crate) struct RewriteContext<'a> {
  pub(crate) parse_sess: &'a ParseSess,
  pub(crate) config: &'a Config,
+ pub(crate) memoize: Memoize,
  pub(crate) inside_macro: Rc<Cell<bool>>,
  // Force block indent style even if we are using visual indent style.
  pub(crate) use_block: Cell<bool>,

src/shape.rs

@@ -4,7 +4,7 @@ use std::ops::{Add, Sub};
 
 use crate::Config;
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) struct Indent {
  // Width of the block indent, in characters. Must be a multiple of
  // Config::tab_spaces.
@@ -139,7 +139,7 @@ impl Sub<usize> for Indent {
 // 8096 is close enough to infinite for rustfmt.
 const INFINITE_SHAPE_WIDTH: usize = 8096;
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) struct Shape {
  pub(crate) width: usize,
  // The current indentation of code.

src/visitor.rs

@@ -17,7 +17,7 @@ use crate::items::{
 use crate::macros::{macro_style, rewrite_macro, rewrite_macro_def, MacroPosition};
 use crate::modules::Module;
 use crate::parse::session::ParseSess;
-use crate::rewrite::{Rewrite, RewriteContext};
+use crate::rewrite::{Memoize, Rewrite, RewriteContext};
 use crate::shape::{Indent, Shape};
 use crate::skip::{is_skip_attr, SkipContext};
 use crate::source_map::{LineRangeUtils, SpanUtils};
@@ -71,6 +71,7 @@ impl SnippetProvider {
 
 pub(crate) struct FmtVisitor<'a> {
  parent_context: Option<&'a RewriteContext<'a>>,
+ pub(crate) memoize: Memoize,
  pub(crate) parse_sess: &'a ParseSess,
  pub(crate) buffer: String,
  pub(crate) last_pos: BytePos,
@@ -754,6 +755,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
  ctx.config,
  ctx.snippet_provider,
  ctx.report.clone(),
+ ctx.memoize.clone(),
  );
  visitor.skip_context.update(ctx.skip_context.clone());
  visitor.set_parent_context(ctx);
@@ -765,10 +767,12 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
  config: &'a Config,
  snippet_provider: &'a SnippetProvider,
  report: FormatReport,
+ memoize: Memoize,
  ) -> FmtVisitor<'a> {
  FmtVisitor {
  parent_context: None,
  parse_sess: parse_session,
+ memoize,
  buffer: String::with_capacity(snippet_provider.big_snippet.len() * 2),
  last_pos: BytePos(0),
  block_indent: Indent::empty(),
@@ -987,6 +991,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
  RewriteContext {
  parse_sess: self.parse_sess,
  config: self.config,
+ memoize: self.memoize.clone(),
  inside_macro: Rc::new(Cell::new(false)),
  use_block: Cell::new(false),
  is_if_else_block: Cell::new(false),