perf(formatter): use ArenaVec and ArenaBox (#15420)

This PR aims to replace the Formatter's standard `Vec` and `Box` usages with `ArenaVec` and `ArenaBox` to gain performance, which could reduce the `drop` overhead.

There are three key changes:

1. Change the Document's element from `Vec<FormatElement<'a>>`
```diff
pub struct Document<'a> {
-    elements: Vec<FormatElement<'a>>,
+    elements: &'a [FormatElement<'a>],
}
```

2. Change the `FormatElement::interned`'s argument from `Rc<[FormatElement<'a>]>` to `&'a [FormatElement<'a>]`
```diff
- pub struct Interned<'a>(Rc<[FormatElement<'a>]>);
+ pub struct Interned<'a>(&'a [FormatElement<'a>]);
```

3. Change the `FormatElement::variants`'s variants from `Box<[Box<[FormatElement<'a>]>]>` to `&'a [&'a [FormatElement<'a>]]`
```diff
pub struct BestFittingElement<'a> {
-    variants: Box<[Box<[FormatElement<'a>]>]>,
+    variants: &'a [&'a [FormatElement<'a>]],
}
```

All other changes stem from those mentioned above.

Replacing the `Rc` with `&a` reference approach was suggested and guided by @overlookmotel. Thank you!

Author: Dunqing

crates/oxc_formatter/src/formatter/buffer.rs

@@ -7,6 +7,8 @@ use std::{
 
 use rustc_hash::FxHashMap;
 
+use oxc_allocator::{Allocator, TakeIn, Vec as ArenaVec};
+
 use super::{
     Arguments, Format, FormatElement, FormatResult, FormatState,
     format_element::Interned,
@@ -175,34 +177,35 @@ impl<'ast, W: Buffer<'ast> + ?Sized> Buffer<'ast> for &mut W {
 #[derive(Debug)]
 pub struct VecBuffer<'buf, 'ast> {
     state: &'buf mut FormatState<'ast>,
-    elements: Vec<FormatElement<'ast>>,
+    elements: ArenaVec<'ast, FormatElement<'ast>>,
 }
 
 impl<'buf, 'ast> VecBuffer<'buf, 'ast> {
     pub fn new(state: &'buf mut FormatState<'ast>) -> Self {
-        Self::new_with_vec(state, Vec::new())
+        Self::new_with_vec(state, ArenaVec::new_in(state.context().allocator()))
     }
 
     pub fn new_with_vec(
         state: &'buf mut FormatState<'ast>,
-        elements: Vec<FormatElement<'ast>>,
+        elements: ArenaVec<'ast, FormatElement<'ast>>,
     ) -> Self {
         Self { state, elements }
     }
 
     /// Creates a buffer with the specified capacity
     pub fn with_capacity(capacity: usize, state: &'buf mut FormatState<'ast>) -> Self {
-        Self { state, elements: Vec::with_capacity(capacity) }
+        let elements = ArenaVec::with_capacity_in(capacity, state.context().allocator());
+        Self { state, elements }
     }
 
     /// Consumes the buffer and returns the written [`FormatElement]`s as a vector.
-    pub fn into_vec(self) -> Vec<FormatElement<'ast>> {
+    pub fn into_vec(self) -> ArenaVec<'ast, FormatElement<'ast>> {
         self.elements
     }
 
     /// Takes the elements without consuming self
-    pub fn take_vec(&mut self) -> Vec<FormatElement<'ast>> {
-        std::mem::take(&mut self.elements)
+    pub fn take_vec(&mut self) -> ArenaVec<'ast, FormatElement<'ast>> {
+        self.elements.take_in(self.state.context().allocator())
     }
 }
 
@@ -493,8 +496,13 @@ impl<'buf, 'ast> RemoveSoftLinesBuffer<'buf, 'ast> {
     }
 
     /// Removes the soft line breaks from an interned element.
-    fn clean_interned(&mut self, interned: &Interned<'ast>) -> Interned<'ast> {
-        clean_interned(interned, &mut self.interned_cache, &mut self.conditional_content_stack)
+    fn clean_interned(&mut self, interned: Interned<'ast>) -> Interned<'ast> {
+        clean_interned(
+            interned,
+            &mut self.interned_cache,
+            &mut self.conditional_content_stack,
+            self.inner.state().context().allocator(),
+        )
     }
 
     /// Marker for whether a `StartConditionalContent(mode: Expanded)` has been
@@ -509,11 +517,12 @@ impl<'buf, 'ast> RemoveSoftLinesBuffer<'buf, 'ast> {
 
 // Extracted to function to avoid monomorphization
 fn clean_interned<'ast>(
-    interned: &Interned<'ast>,
+    interned: Interned<'ast>,
     interned_cache: &mut FxHashMap<Interned<'ast>, Interned<'ast>>,
     condition_content_stack: &mut Vec<Condition>,
+    allocator: &'ast Allocator,
 ) -> Interned<'ast> {
-    if let Some(cleaned) = interned_cache.get(interned) {
+    if let Some(cleaned) = interned_cache.get(&interned) {
         cleaned.clone()
     } else {
         // Find the first soft line break element, interned element, or conditional expanded
@@ -522,18 +531,23 @@ fn clean_interned<'ast>(
             FormatElement::Line(LineMode::Soft | LineMode::SoftOrSpace)
             | FormatElement::Tag(Tag::StartConditionalContent(_) | Tag::EndConditionalContent)
             | FormatElement::BestFitting(_) => {
-                let mut cleaned = Vec::new();
-                cleaned.extend_from_slice(&interned[..index]);
+                let mut cleaned =
+                    ArenaVec::from_iter_in(interned[..index].iter().cloned(), allocator);
                 Some((cleaned, &interned[index..]))
             }
             FormatElement::Interned(inner) => {
-                let cleaned_inner = clean_interned(inner, interned_cache, condition_content_stack);
+                let cleaned_inner = clean_interned(
+                    inner.clone(),
+                    interned_cache,
+                    condition_content_stack,
+                    allocator,
+                );
 
                 if &cleaned_inner == inner {
                     None
                 } else {
-                    let mut cleaned = Vec::with_capacity(interned.len());
-                    cleaned.extend_from_slice(&interned[..index]);
+                    let mut cleaned = ArenaVec::with_capacity_in(interned.len(), allocator);
+                    cleaned.extend(interned[..index].iter().cloned());
                     cleaned.push(FormatElement::Interned(cleaned_inner));
                     Some((cleaned, &interned[index + 1..]))
                 }
@@ -567,9 +581,10 @@ fn clean_interned<'ast>(
 
                         FormatElement::Interned(interned) => {
                             cleaned.push(FormatElement::Interned(clean_interned(
-                                interned,
+                                interned.clone(),
                                 interned_cache,
                                 condition_content_stack,
+                                allocator,
                             )));
                         }
                         // Since this buffer aims to simulate infinite print width, we don't need to retain the best fitting.
@@ -588,15 +603,14 @@ fn clean_interned<'ast>(
             None => interned.clone(),
         };
 
-        interned_cache.insert(interned.clone(), result.clone());
+        interned_cache.insert(interned, result.clone());
         result
     }
 }
 
 impl<'ast> Buffer<'ast> for RemoveSoftLinesBuffer<'_, 'ast> {
     fn write_element(&mut self, element: FormatElement<'ast>) -> FormatResult<()> {
-        let mut element_stack = Vec::new();
-        element_stack.push(element);
+        let mut element_stack = Vec::from_iter([element]);
         while let Some(element) = element_stack.pop() {
             match element {
                 FormatElement::Tag(Tag::StartConditionalContent(condition)) => {
@@ -614,14 +628,14 @@ impl<'ast> Buffer<'ast> for RemoveSoftLinesBuffer<'_, 'ast> {
                     self.inner.write_element(FormatElement::Space)?;
                 }
                 FormatElement::Interned(interned) => {
-                    let cleaned = self.clean_interned(&interned);
+                    let cleaned = self.clean_interned(interned);
                     self.inner.write_element(FormatElement::Interned(cleaned))?;
                 }
                 // Since this buffer aims to simulate infinite print width, we don't need to retain the best fitting.
                 // Just extract the flattest variant and then handle elements within it.
                 FormatElement::BestFitting(best_fitting) => {
                     let most_flat = best_fitting.most_flat();
-                    most_flat.iter().rev().for_each(|element| element_stack.push(element.clone()));
+                    element_stack.extend(most_flat.iter().rev().cloned());
                 }
                 element => self.inner.write_element(element)?,
             }

crates/oxc_formatter/src/formatter/builders.rs

@@ -6,6 +6,7 @@ use Tag::{
     StartDedent, StartEntry, StartFill, StartGroup, StartIndent, StartIndentIfGroupBreaks,
     StartLabelled, StartLineSuffix,
 };
+use oxc_allocator::Vec as ArenaVec;
 use oxc_span::{GetSpan, Span};
 use oxc_syntax::identifier::{is_identifier_name, is_line_terminator, is_white_space_single_line};
 
@@ -2549,9 +2550,12 @@ impl<'ast> Format<'ast> for BestFitting<'_, 'ast> {
             buffer.write_fmt(Arguments::from(variant))?;
             buffer.write_element(FormatElement::Tag(EndEntry))?;
 
-            formatted_variants.push(buffer.take_vec().into_boxed_slice());
+            formatted_variants.push(buffer.take_vec().into_bump_slice());
         }
 
+        let formatted_variants =
+            ArenaVec::from_iter_in(formatted_variants, f.context().allocator());
+
         // SAFETY: The constructor guarantees that there are always at least two variants. It's, therefore,
         // safe to call into the unsafe `from_vec_unchecked` function
         let element = unsafe {

crates/oxc_formatter/src/formatter/format_element/document.rs

@@ -2,7 +2,7 @@
 use cow_utils::CowUtils;
 use std::ops::Deref;
 
-use oxc_allocator::Allocator;
+use oxc_allocator::{Allocator, Vec as ArenaVec};
 use oxc_ast::Comment;
 use oxc_span::SourceType;
 use rustc_hash::FxHashMap;
@@ -21,7 +21,7 @@ use crate::{format, write};
 /// A formatted document.
 #[derive(Debug, Clone, Eq, PartialEq, Default)]
 pub struct Document<'a> {
-    elements: Vec<FormatElement<'a>>,
+    elements: &'a [FormatElement<'a>],
 }
 
 impl Document<'_> {
@@ -140,26 +140,26 @@ impl Document<'_> {
     }
 }
 
-impl<'a> From<Vec<FormatElement<'a>>> for Document<'a> {
-    fn from(elements: Vec<FormatElement<'a>>) -> Self {
-        Self { elements }
+impl<'a> From<ArenaVec<'a, FormatElement<'a>>> for Document<'a> {
+    fn from(elements: ArenaVec<'a, FormatElement<'a>>) -> Self {
+        Self { elements: elements.into_bump_slice() }
     }
 }
 
 impl<'a> Deref for Document<'a> {
     type Target = [FormatElement<'a>];
 
     fn deref(&self) -> &Self::Target {
-        self.elements.as_slice()
+        self.elements
     }
 }
 
 impl std::fmt::Display for Document<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let allocator = Allocator::default();
         let context = FormatContext::dummy(&allocator);
-        let formatted = format!(context, [self.elements.as_slice()])
-            .expect("Formatting not to throw any FormatErrors");
+        let formatted =
+            format!(context, [self.elements]).expect("Formatting not to throw any FormatErrors");
 
         f.write_str(formatted.print().expect("Expected a valid document").as_code())
     }

crates/oxc_formatter/src/formatter/format_element/mod.rs

@@ -5,10 +5,13 @@ pub mod tag;
 // use biome_rowan::static_assert;
 use std::hash::{Hash, Hasher};
 use std::num::NonZeroU32;
+use std::ptr;
 use std::{borrow::Cow, ops::Deref, rc::Rc};
 
 use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
 
+use oxc_allocator::{Address, Box as ArenaBox, Vec as ArenaVec};
+
 use crate::{IndentWidth, TabWidth};
 
 use super::{
@@ -146,28 +149,25 @@ impl PrintMode {
 }
 
 #[derive(Clone)]
-pub struct Interned<'a>(Rc<[FormatElement<'a>]>);
+pub struct Interned<'a>(&'a [FormatElement<'a>]);
 
 impl<'a> Interned<'a> {
-    pub(super) fn new(content: Vec<FormatElement<'a>>) -> Self {
-        Self(content.into())
+    pub(super) fn new(content: ArenaVec<'a, FormatElement<'a>>) -> Self {
+        Self(content.into_bump_slice())
     }
 }
 
 impl PartialEq for Interned<'_> {
     fn eq(&self, other: &Interned<'_>) -> bool {
-        Rc::ptr_eq(&self.0, &other.0)
+        ptr::eq(self.0, other.0)
     }
 }
 
 impl Eq for Interned<'_> {}
 
 impl Hash for Interned<'_> {
-    fn hash<H>(&self, hasher: &mut H)
-    where
-        H: Hasher,
-    {
-        Rc::as_ptr(&self.0).hash(hasher);
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0.as_ptr().addr().hash(state);
     }
 }
 
@@ -181,7 +181,7 @@ impl<'a> Deref for Interned<'a> {
     type Target = [FormatElement<'a>];
 
     fn deref(&self) -> &Self::Target {
-        &self.0
+        self.0
     }
 }
 
@@ -305,7 +305,7 @@ pub struct BestFittingElement<'a> {
     /// The different variants for this element.
     /// The first element is the one that takes up the most space horizontally (the most flat),
     /// The last element takes up the least space horizontally (but most horizontal space).
-    variants: Box<[Box<[FormatElement<'a>]>]>,
+    variants: &'a [&'a [FormatElement<'a>]],
 }
 
 impl<'a> BestFittingElement<'a> {
@@ -318,13 +318,13 @@ impl<'a> BestFittingElement<'a> {
     /// ## Safety
     /// The slice must contain at least two variants.
     #[doc(hidden)]
-    pub unsafe fn from_vec_unchecked(variants: Vec<Box<[FormatElement<'a>]>>) -> Self {
+    pub unsafe fn from_vec_unchecked(variants: ArenaVec<'a, &'a [FormatElement<'a>]>) -> Self {
         debug_assert!(
             variants.len() >= 2,
             "Requires at least the least expanded and most expanded variants"
         );
 
-        Self { variants: variants.into_boxed_slice() }
+        Self { variants: variants.into_bump_slice() }
     }
 
     /// Returns the most expanded variant
@@ -334,8 +334,8 @@ impl<'a> BestFittingElement<'a> {
         )
     }
 
-    pub fn variants(&self) -> &[Box<[FormatElement<'a>]>] {
-        &self.variants
+    pub fn variants(&self) -> &[&'a [FormatElement<'a>]] {
+        self.variants
     }
 
     /// Returns the least expanded variant
@@ -348,7 +348,7 @@ impl<'a> BestFittingElement<'a> {
 
 impl std::fmt::Debug for BestFittingElement<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_list().entries(&*self.variants).finish()
+        f.debug_list().entries(self.variants).finish()
     }
 }
 

crates/oxc_formatter/src/formatter/formatter.rs

@@ -1,6 +1,6 @@
 #![allow(clippy::module_inception)]
 
-use oxc_allocator::{Address, Allocator};
+use oxc_allocator::{Address, Allocator, Vec as ArenaVec};
 use oxc_ast::AstKind;
 
 use crate::options::FormatOptions;
@@ -27,7 +27,7 @@ impl<'buf, 'ast> Formatter<'buf, 'ast> {
         Self { buffer }
     }
 
-    pub fn allocator(&self) -> &Allocator {
+    pub fn allocator(&self) -> &'ast Allocator {
         self.context().allocator()
     }
 
@@ -244,7 +244,7 @@ impl<'buf, 'ast> Formatter<'buf, 'ast> {
 
     pub fn intern_vec(
         &mut self,
-        mut elements: Vec<FormatElement<'ast>>,
+        mut elements: ArenaVec<'ast, FormatElement<'ast>>,
     ) -> Option<FormatElement<'ast>> {
         match elements.len() {
             0 => None,

crates/oxc_formatter/src/ir_transform/sort_imports/mod.rs

@@ -2,6 +2,10 @@ mod import_unit;
 mod partitioned_chunk;
 mod source_line;
 
+use std::mem;
+
+use oxc_allocator::{Allocator, Vec as ArenaVec};
+
 use crate::{
     formatter::format_element::{FormatElement, LineMode, document::Document},
     options,
@@ -26,7 +30,7 @@ impl SortImportsTransform {
     // It means that:
     // - There is no redundant spaces, no consecutive line breaks, etc...
     // - Last element is always `FormatElement::Line(Hard)`.
-    pub fn transform<'a>(&self, document: &Document<'a>) -> Document<'a> {
+    pub fn transform<'a>(&self, document: &Document<'a>, allocator: &'a Allocator) -> Document<'a> {
         // Early return for empty files
         if document.len() == 1 && matches!(document[0], FormatElement::Line(LineMode::Hard)) {
             return document.clone();
@@ -136,7 +140,7 @@ impl SortImportsTransform {
 
         // Finally, sort import lines within each chunk.
         // After sorting, flatten everything back to `FormatElement`s.
-        let mut next_elements = Vec::with_capacity(prev_elements.len());
+        let mut next_elements = ArenaVec::with_capacity_in(prev_elements.len(), allocator);
 
         let mut chunks_iter = chunks.into_iter().enumerate().peekable();
         while let Some((idx, chunk)) = chunks_iter.next() {