perf(formatter_core): stage IR buffers on the heap to reduce arena memory (#24582)

Part of #24558

Adds `HeapVecBuffer` to `oxc_formatter_core`: a staging twin of `VecBuffer` that grows on a pooled heap vector and moves the finished sequence into the arena as one exactly-sized allocation.

Converted JS formatter call sites: `Formatter::intern`, `BestFitting` variant construction, and the three grouped-call-argument variant builders.

The other formatters would automatically benefit from this as well, and there was no need to make any individual changes.

---

TIL: `antd.js` is a bundled file and is not intended for general formatting.

crates/oxc_formatter/src/formatter/mod.rs

@@ -47,7 +47,7 @@ pub mod buffer {
 
 pub use oxc_formatter_core::{
     Argument, Arguments, Buffer, BufferExtensions, Format, FormatElement, FormatOptions,
-    FormatState, Formatted, Formatter, GroupId, MemoizeFormat, Memoized, SourceText,
+    FormatState, Formatted, Formatter, GroupId, HeapVecBuffer, MemoizeFormat, Memoized, SourceText,
     UniqueGroupIdBuilder, VecBuffer,
 };
 

crates/oxc_formatter/src/print/call_like_expression/arguments.rs

@@ -7,8 +7,8 @@ use crate::{
     ast_nodes::{AstNode, AstNodes},
     format_args,
     formatter::{
-        Comments, FormatElement, JoinBuilderJsExt as _, JsFormatContext, JsFormatter,
-        JsFormatterExt as _, SourceText, VecBuffer,
+        Comments, FormatElement, HeapVecBuffer, JoinBuilderJsExt as _, JsFormatContext,
+        JsFormatter, JsFormatterExt as _, SourceText,
         buffer::RemoveSoftLinesBuffer,
         format_element,
         prelude::{
@@ -737,14 +737,14 @@ fn write_grouped_arguments<'a>(
 
     // First write the most expanded variant because it needs `arguments`.
     let most_expanded = {
-        let mut buffer = VecBuffer::new(f.state_mut());
+        let mut buffer = HeapVecBuffer::new(f.state_mut());
         buffer.write_element(FormatElement::Tag(Tag::StartEntry));
 
         format_all_elements_broken_out(node, elements.iter().cloned(), true, &mut buffer);
 
         buffer.write_element(FormatElement::Tag(Tag::EndEntry));
 
-        buffer.into_vec().into_arena_slice()
+        buffer.take_into_arena_slice()
     };
 
     // Now reformat the first or last argument if they happen to be a function or arrow function expression.
@@ -825,7 +825,7 @@ fn write_grouped_arguments<'a>(
 
     // Write the second variant that forces the group of the first/last argument to expand.
     let middle_variant = {
-        let mut buffer = VecBuffer::new(f.state_mut());
+        let mut buffer = HeapVecBuffer::new(f.state_mut());
 
         buffer.write_element(FormatElement::Tag(Tag::StartEntry));
 
@@ -863,7 +863,7 @@ fn write_grouped_arguments<'a>(
 
         buffer.write_element(FormatElement::Tag(Tag::EndEntry));
 
-        buffer.into_vec().into_arena_slice()
+        buffer.take_into_arena_slice()
     };
 
     // If the grouped content breaks, then we can skip the most_flat variant,
@@ -874,7 +874,7 @@ fn write_grouped_arguments<'a>(
     } else {
         // Write the most flat variant with the first or last argument grouped.
         let most_flat = {
-            let mut buffer = VecBuffer::new(f.state_mut());
+            let mut buffer = HeapVecBuffer::new(f.state_mut());
             buffer.write_element(FormatElement::Tag(Tag::StartEntry));
 
             write!(
@@ -898,7 +898,7 @@ fn write_grouped_arguments<'a>(
 
             buffer.write_element(FormatElement::Tag(Tag::EndEntry));
 
-            buffer.into_vec().into_arena_slice()
+            buffer.take_into_arena_slice()
         };
 
         ArenaVec::from_array_in([most_flat, middle_variant, most_expanded], f)

crates/oxc_formatter_core/AGENTS.md

@@ -26,6 +26,19 @@ Prettier doc primitives are ported on demand; still missing:
 - `hardlineWithoutBreakParent` (markdown tables)
 - and the `trim` doc
 
+### Choosing a staging buffer
+
+The arena is a bump allocator and never reclaims, so a vector grown in it strands every grown-out-of allocation for the rest of the format run.
+Pick by what you're building:
+
+- Root document (feeds `Document::new` / `EmbeddedIr`): `VecBuffer` (arena)
+  - it moves into the `Document` for free, and heap-staging it costs an extra copy for no benefit
+- Unknown-length staging that ends interned/sliced: `HeapVecBuffer`
+  - grows on a pooled heap vector, the arena receives one exactly-sized copy (see its rustdoc for the full rationale)
+- Known-length sequences: build exact-sized directly (e.g. `ArenaVec::from_iter_in`)
+
+`Formatter::intern` and `BestFitting` already stage on the heap; consumer crates get this for free.
+
 ### Generic context design
 
 The core is parameterized over a consumer-supplied context so it stays language-agnostic:

crates/oxc_formatter_core/src/buffer.rs

@@ -7,7 +7,7 @@ use std::{
 
 use rustc_hash::FxHashMap;
 
-use oxc_allocator::{Allocator, ArenaVec, TakeIn};
+use oxc_allocator::{Allocator, ArenaVec};
 
 use crate::{
     Argument, Arguments, Format, FormatElement, FormatState,
@@ -44,11 +44,12 @@ pub trait Buffer<'ast, C> {
 
     /// Replaces the elements starting at `start` with `replacement`.
     ///
-    /// Used by streaming IR transforms (currently `SortImportsTransform`) to splice a reordered
-    /// chunk back into the buffer. Only `VecBuffer` supports this; the wrapper buffers
-    /// (`PreambleBuffer`, `Inspect`, `RemoveSoftLinesBuffer`) are only ever active inside
-    /// inner-expression contexts, never on the call stack while a streaming chunk is being
-    /// flushed, so they implement this as `unreachable!()`.
+    /// Used by streaming IR transforms (currently `SortImportsTransform`)
+    /// to splice a reordered chunk back into the buffer.
+    /// Only the vector-backed buffers (`VecBuffer`, `HeapVecBuffer`) support this;
+    /// the wrapper buffers (`PreambleBuffer`, `Inspect`, `RemoveSoftLinesBuffer`) are only ever active
+    /// inside inner-expression contexts, never on the call stack while a streaming chunk is being flushed,
+    /// so they implement this as `unreachable!()`.
     fn replace_end(&mut self, start: usize, replacement: &[FormatElement<'ast>]);
 }
 
@@ -110,11 +111,6 @@ impl<'buf, 'ast, C> VecBuffer<'buf, 'ast, C> {
     pub fn into_vec(self) -> ArenaVec<'ast, FormatElement<'ast>> {
         self.elements
     }
-
-    /// Takes the elements without consuming self
-    pub fn take_vec(&mut self) -> ArenaVec<'ast, FormatElement<'ast>> {
-        self.elements.take_in(self.state)
-    }
 }
 
 impl<'ast, C> Deref for VecBuffer<'_, 'ast, C> {
@@ -153,6 +149,93 @@ impl<'ast, C> Buffer<'ast, C> for VecBuffer<'_, 'ast, C> {
     }
 }
 
+/// Heap-backed [`Buffer`], the staging twin of [`VecBuffer`].
+///
+/// Use it to build an element sequence whose final length isn't known up front
+/// (interned content, best-fitting variants), then move the result into the arena
+/// exactly-sized via [`HeapVecBuffer::take_into_arena_vec`] / [`HeapVecBuffer::take_into_arena_slice`].
+///
+/// Growing a vector directly in the arena strands every grown-out-of allocation for the rest of the format run
+/// (the arena is a bump allocator and never reclaims;
+/// growth can reuse the allocation only when nothing else was bumped in between, which practically never holds while formatting).
+/// Heap growth is reclaimed by the system allocator, and the arena receives only the final, exactly-sized copy.
+///
+/// The backing vector is borrowed from the [`FormatState`] scratch pool and returned on drop,
+/// so nested/repeated staging allocates nothing in the steady state.
+pub struct HeapVecBuffer<'buf, 'ast, C> {
+    state: &'buf mut FormatState<'ast, C>,
+    elements: Vec<FormatElement<'ast>>,
+}
+
+impl<'buf, 'ast, C> HeapVecBuffer<'buf, 'ast, C> {
+    pub fn new(state: &'buf mut FormatState<'ast, C>) -> Self {
+        let elements = state.take_scratch_buffer();
+        Self { state, elements }
+    }
+
+    /// Removes the last element written to this buffer, if any.
+    pub fn pop(&mut self) -> Option<FormatElement<'ast>> {
+        self.elements.pop()
+    }
+
+    /// Moves the buffered elements into an exactly-sized arena vector, leaving the buffer empty.
+    pub fn take_into_arena_vec(&mut self) -> ArenaVec<'ast, FormatElement<'ast>> {
+        let len = self.elements.len();
+        let allocator = self.state.allocator();
+        let mut vec = ArenaVec::with_capacity_in(len, &allocator);
+        // SAFETY: `vec` has capacity for `len` elements, and the heap and arena buffers cannot overlap.
+        // `FormatElement` is not `Drop` (`ArenaVec` const-asserts `!needs_drop` on construction),
+        // so the bitwise move duplicates no owning state and the `clear()` below merely forgets the moved-out elements.
+        unsafe {
+            std::ptr::copy_nonoverlapping(self.elements.as_ptr(), vec.as_mut_ptr(), len);
+            vec.set_len(len);
+        }
+        self.elements.clear();
+        vec
+    }
+
+    /// Moves the buffered elements into an exactly-sized arena slice, leaving the buffer empty.
+    pub fn take_into_arena_slice(&mut self) -> &'ast [FormatElement<'ast>] {
+        self.take_into_arena_vec().into_arena_slice()
+    }
+}
+
+impl<C> Drop for HeapVecBuffer<'_, '_, C> {
+    fn drop(&mut self) {
+        self.state.return_scratch_buffer(std::mem::take(&mut self.elements));
+    }
+}
+
+impl<'ast, C> Deref for HeapVecBuffer<'_, 'ast, C> {
+    type Target = [FormatElement<'ast>];
+
+    fn deref(&self) -> &Self::Target {
+        &self.elements
+    }
+}
+
+impl<'ast, C> Buffer<'ast, C> for HeapVecBuffer<'_, 'ast, C> {
+    fn write_element(&mut self, element: FormatElement<'ast>) {
+        self.elements.push(element);
+    }
+
+    fn elements(&self) -> &[FormatElement<'ast>] {
+        self
+    }
+
+    fn state(&self) -> &FormatState<'ast, C> {
+        self.state
+    }
+
+    fn state_mut(&mut self) -> &mut FormatState<'ast, C> {
+        self.state
+    }
+
+    fn replace_end(&mut self, start: usize, replacement: &[FormatElement<'ast>]) {
+        self.elements.splice(start.., replacement.iter().cloned());
+    }
+}
+
 /// This struct wraps an existing buffer and emits a preamble text when the first text is written.
 ///
 /// This can be useful if you, for example, want to write some content if what gets written next isn't empty.

crates/oxc_formatter_core/src/builders.rs

@@ -9,7 +9,7 @@ use oxc_allocator::ArenaVec;
 
 use crate::{
     Argument, Arguments, Buffer, Format, FormatContext, FormatElement, FormatOptions, Formatter,
-    GroupId, VecBuffer,
+    GroupId, HeapVecBuffer,
     format::write,
     format_element::{
         self, LineMode, PrintMode, TextWidth,
@@ -883,7 +883,9 @@ impl<'fmt, 'ast, C> BestFitting<'fmt, 'ast, C> {
 
 impl<'ast, C> Format<'ast, C> for BestFitting<'_, 'ast, C> {
     fn fmt(&self, f: &mut Formatter<'_, 'ast, C>) {
-        let mut buffer = VecBuffer::new(f.state_mut());
+        // Stage each variant on the heap,
+        // so only the exactly-sized variant slices land in the arena (see `HeapVecBuffer`).
+        let mut buffer = HeapVecBuffer::new(f.state_mut());
         let variants = self.variants.items();
 
         let mut formatted_variants = Vec::with_capacity(variants.len());
@@ -893,8 +895,9 @@ impl<'ast, C> Format<'ast, C> for BestFitting<'_, 'ast, C> {
             buffer.write_fmt(Arguments::from(variant));
             buffer.write_element(FormatElement::Tag(EndEntry));
 
-            formatted_variants.push(buffer.take_vec().into_arena_slice());
+            formatted_variants.push(buffer.take_into_arena_slice());
         }
+        drop(buffer);
 
         let formatted_variants = ArenaVec::from_iter_in(formatted_variants, f);
 

crates/oxc_formatter_core/src/formatter.rs

@@ -5,7 +5,8 @@ use std::borrow::Cow;
 use oxc_allocator::{Allocator, ArenaVec, GetAllocator};
 
 use crate::{
-    Argument, Arguments, Buffer, FormatContext, FormatElement, FormatState, VecBuffer,
+    Argument, Arguments, Buffer, FormatContext, FormatElement, FormatState,
+    buffer::HeapVecBuffer,
     builders::{FillBuilder, JoinBuilder},
     format::{Format, write},
     format_element::Interned,
@@ -60,12 +61,20 @@ impl<'buf, 'ast, C> Formatter<'buf, 'ast, C> {
     }
 
     /// Formats `content` into an interned element without writing it to the formatter's buffer.
+    /// The content is staged in a [`HeapVecBuffer`] and lands in the arena exactly-sized,
+    /// so the arena only ever holds the final interned slice, not the staging growth.
     pub fn intern(&mut self, content: &dyn Format<'ast, C>) -> Option<FormatElement<'ast>> {
-        let mut buffer = VecBuffer::new(self.state_mut());
+        let mut buffer = HeapVecBuffer::new(self.state_mut());
         write(&mut buffer, Arguments::new(&[Argument::new(&content)]));
-        let elements = buffer.into_vec();
 
-        self.intern_vec(elements)
+        // Intentionally not delegated to `intern_vec`:
+        // dispatching on the heap buffer lets the 0/1-element cases return without ever allocating an `ArenaVec`.
+        match buffer.len() {
+            0 => None,
+            // Doesn't get cheaper than calling clone, use the element directly
+            1 => buffer.pop(),
+            _ => Some(FormatElement::Interned(Interned::new(buffer.take_into_arena_vec()))),
+        }
     }
 
     #[expect(clippy::unused_self)] // Keep `self` the same as the original source

crates/oxc_formatter_core/src/lib.rs

@@ -39,8 +39,8 @@ pub mod test_support;
 
 pub use arguments::{Argument, Arguments};
 pub use buffer::{
-    Buffer, BufferExtensions, Inspect, PreambleBuffer, Recorded, Recording, RemoveSoftLinesBuffer,
-    VecBuffer,
+    Buffer, BufferExtensions, HeapVecBuffer, Inspect, PreambleBuffer, Recorded, Recording,
+    RemoveSoftLinesBuffer, VecBuffer,
 };
 pub use diagnostics::{ActualStart, FormatError, InvalidDocumentError, PrintError};
 pub use embedded::{

crates/oxc_formatter_core/src/state.rs

@@ -1,7 +1,7 @@
 use oxc_allocator::{Allocator, GetAllocator};
 use rustc_hash::FxHashMap;
 
-use crate::{GroupId, UniqueGroupIdBuilder, format_element::Interned};
+use crate::{FormatElement, GroupId, UniqueGroupIdBuilder, format_element::Interned};
 
 /// This structure stores the state that is relevant for the formatting of the whole document.
 ///
@@ -15,6 +15,11 @@ pub struct FormatState<'ast, C> {
     // For the document IR printing process
     /// The interned elements that have been printed to this point
     printed_interned_elements: FxHashMap<Interned<'ast>, usize>,
+    /// Reusable heap staging buffers for [`crate::HeapVecBuffer`], LIFO;
+    /// see there for why staging is heap-backed.
+    /// The pool holds one buffer per nesting level, each retaining its high-water capacity,
+    /// so staging costs no allocation at all in the steady state.
+    scratch_buffers: Vec<Vec<FormatElement<'ast>>>,
 }
 
 impl<C: std::fmt::Debug> std::fmt::Debug for FormatState<'_, C> {
@@ -31,6 +36,25 @@ impl<'ast, C> FormatState<'ast, C> {
             allocator,
             group_id_builder: UniqueGroupIdBuilder::default(),
             printed_interned_elements: FxHashMap::default(),
+            scratch_buffers: Vec::new(),
+        }
+    }
+
+    /// Takes a heap staging buffer from the pool (empty, capacity retained), or creates one.
+    ///
+    /// Used by [`crate::HeapVecBuffer`]; see [`FormatState::scratch_buffers`] for why staging happens on the heap.
+    /// Return it with [`FormatState::return_scratch_buffer`] once done.
+    pub(crate) fn take_scratch_buffer(&mut self) -> Vec<FormatElement<'ast>> {
+        self.scratch_buffers.pop().unwrap_or_default()
+    }
+
+    /// Returns a staging buffer to the pool so its capacity is reused.
+    pub(crate) fn return_scratch_buffer(&mut self, mut buffer: Vec<FormatElement<'ast>>) {
+        // A capacity-less buffer has nothing worth pooling
+        // (a buffer that never grew beyond its `Vec::default` state).
+        if buffer.capacity() > 0 {
+            buffer.clear();
+            self.scratch_buffers.push(buffer);
         }
     }
 

tasks/track_memory_allocations/allocs_formatter.yaml

@@ -1,77 +1,77 @@
 checker.ts:
   file size: 2922154 # 2.92 MB
-  sys allocs: 11013
-  sys reallocs: 492
-  sys deallocs: 11012
-  sys alloc bytes: 5918285 # 5.92 MB
+  sys allocs: 11027
+  sys reallocs: 560
+  sys deallocs: 11026
+  sys alloc bytes: 6131149 # 6.13 MB
   sys peak growth: 2990058 # 2.99 MB
-  arena allocs: 1081815
-  arena reallocs: 118690
-  arena size: 138513064 # 138.51 MB
+  arena allocs: 1079319
+  arena reallocs: 3951
+  arena size: 116637808 # 116.64 MB
 
 App.tsx:
   file size: 415342 # 415.34 kB
-  sys allocs: 2532
-  sys reallocs: 80
-  sys deallocs: 2531
-  sys alloc bytes: 985182 # 985.18 kB
-  sys peak growth: 443550 # 443.55 kB
-  arena allocs: 212668
-  arena reallocs: 25343
-  arena size: 29608600 # 29.61 MB
+  sys allocs: 2544
+  sys reallocs: 149
+  sys deallocs: 2543
+  sys alloc bytes: 1542526 # 1.54 MB
+  sys peak growth: 582900 # 582.90 kB
+  arena allocs: 211998
+  arena reallocs: 1679
+  arena size: 22302080 # 22.30 MB
 
 RadixUIAdoptionSection.jsx:
   file size: 2520 # 2.52 kB
-  sys allocs: 54
-  sys reallocs: 26
-  sys deallocs: 53
-  sys alloc bytes: 10116 # 10.12 kB
-  sys peak growth: 2856 # 2.86 kB
+  sys allocs: 57
+  sys reallocs: 31
+  sys deallocs: 56
+  sys alloc bytes: 15492 # 15.49 kB
+  sys peak growth: 7208 # 7.21 kB
   arena allocs: 1155
-  arena reallocs: 157
-  arena size: 210048 # 210.05 kB
+  arena reallocs: 69
+  arena size: 169128 # 169.13 kB
 
 pdf.mjs:
   file size: 567296 # 567.30 kB
-  sys allocs: 4019
-  sys reallocs: 172
-  sys deallocs: 4018
-  sys alloc bytes: 2050032 # 2.05 MB
+  sys allocs: 4031
+  sys reallocs: 250
+  sys deallocs: 4030
+  sys alloc bytes: 3057776 # 3.06 MB
   sys peak growth: 1164624 # 1.16 MB
-  arena allocs: 437639
-  arena reallocs: 43489
-  arena size: 44872224 # 44.87 MB
+  arena allocs: 434612
+  arena reallocs: 5696
+  arena size: 36857064 # 36.86 MB
 
 antd.js:
   file size: 6686316 # 6.69 MB
-  sys allocs: 34121
-  sys reallocs: 3994
-  sys deallocs: 34120
-  sys alloc bytes: 20850396 # 20.85 MB
-  sys peak growth: 13841112 # 13.84 MB
-  arena allocs: 2586207
-  arena reallocs: 302459
-  arena size: 525725120 # 525.73 MB
+  sys allocs: 34143
+  sys reallocs: 4138
+  sys deallocs: 34142
+  sys alloc bytes: 105935484 # 105.94 MB
+  sys peak growth: 85678576 # 85.68 MB
+  arena allocs: 2559844
+  arena reallocs: 14941
+  arena size: 332361480 # 332.36 MB
 
 binder.ts:
   file size: 193077 # 193.08 kB
-  sys allocs: 457
-  sys reallocs: 35
-  sys deallocs: 456
-  sys alloc bytes: 291781 # 291.78 kB
+  sys allocs: 464
+  sys reallocs: 61
+  sys deallocs: 463
+  sys alloc bytes: 315013 # 315.01 kB
   sys peak growth: 195661 # 195.66 kB
-  arena allocs: 64083
-  arena reallocs: 6932
-  arena size: 8574752 # 8.57 MB
+  arena allocs: 64023
+  arena reallocs: 244
+  arena size: 7370072 # 7.37 MB
 
 kitchen-sink.tsx:
   file size: 732898 # 732.90 kB
-  sys allocs: 9565
-  sys reallocs: 2815
-  sys deallocs: 9564
-  sys alloc bytes: 3180244 # 3.18 MB
+  sys allocs: 9578
+  sys reallocs: 2876
+  sys deallocs: 9577
+  sys alloc bytes: 3308148 # 3.31 MB
   sys peak growth: 1499988 # 1.50 MB
-  arena allocs: 569903
-  arena reallocs: 60147
-  arena size: 57887072 # 57.89 MB
+  arena allocs: 567769
+  arena reallocs: 5992
+  arena size: 48529696 # 48.53 MB