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,44 @@ pub(crate) fn format_expr(
     expr_type: ExprType,
     context: &RewriteContext<'_>,
     shape: Shape,
+) -> Option<String> {
+    if context.inside_macro() || context.is_macro_def {
+        return format_expr_inner(expr, expr_type, context, shape);
+    }
+    let mut map = context.memoize.borrow_mut();
+    let clean;
+    let query_id = QueryId {
+        shape,
+        span: expr.span.data(),
+    };
+    if let Some(map) = &*map {
+        if let Some(r) = map.get(&query_id) {
+            return r.clone();
+        }
+        clean = false;
+    } else {
+        *map = Some(HashMap::default());
+        clean = true;
+    }
+    drop(map);
+    let r = format_expr_inner(expr, expr_type, context, shape);
+    let mut map = context.memoize.borrow_mut();
+    if clean {
+        *map = None;
+    } else {
+        if let Some(map) = &mut *map {
+            map.insert(query_id, r.clone());
+        }
+    }
+    drop(map);
+    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;
@@ -188,6 +189,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/lib.rs

@@ -1,4 +1,5 @@
 #![feature(rustc_private)]
+#![feature(once_cell)]
 #![deny(rust_2018_idioms)]
 #![warn(unreachable_pub)]
 #![recursion_limit = "256"]

src/rewrite.rs

@@ -1,6 +1,7 @@
 // A generic trait to abstract the rewriting of an element (of the AST).
 
 use std::cell::{Cell, RefCell};
+use std::collections::HashMap;
 use std::rc::Rc;
 
 use rustc_ast::ptr;
@@ -12,6 +13,7 @@ use crate::skip::SkipContext;
 use crate::syntux::session::ParseSess;
 use crate::visitor::SnippetProvider;
 use crate::FormatReport;
+use rustc_span::SpanData;
 
 pub(crate) trait Rewrite {
     /// Rewrite self into shape.
@@ -24,10 +26,19 @@ impl<T: Rewrite> Rewrite for ptr::P<T> {
     }
 }
 
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub(crate) struct QueryId {
+    pub(crate) shape: Shape,
+    pub(crate) span: SpanData,
+}
+
+pub(crate) type Memoize = Option<HashMap<QueryId, Option<String>>>;
+
 #[derive(Clone)]
 pub(crate) struct RewriteContext<'a> {
     pub(crate) parse_sess: &'a ParseSess,
     pub(crate) config: &'a Config,
+    pub(crate) memoize: Rc<RefCell<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

@@ -16,7 +16,7 @@ use crate::items::{
 };
 use crate::macros::{macro_style, rewrite_macro, rewrite_macro_def, MacroPosition};
 use crate::modules::Module;
-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: Rc<RefCell<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: Rc<RefCell<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),