Memoize format_expr

Author: HKalbasi

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(),
@@ -991,6 +995,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),