Memoize format_expr

1 files changed, 50 insertions, 1 deletions

diff --git a/src/expr.rs b/src/expr.rs
index 4f333cd27ce..cfecc9b9d89 100644
--- a/src/expr.rs
+++ b/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;
@@ -54,6 +55,54 @@ pub(crate) fn format_expr(
     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);
 
     if contains_skip(&*expr.attrs) {