perf(formatter_core): share one thread-cached scratch vector across staging buffers (#24583)

Apply #24539, share `ScratchBuffer`'s buffer.

NOTE: checked-in vectors keep their high-water capacity for the thread's lifetime (same policy as #24539). This is ~KB for regular sources, and fine.

crates/oxc_formatter_core/AGENTS.md

@@ -34,7 +34,7 @@ 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)
+  - a watermarked view over a shared, thread-cached scratch 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.

crates/oxc_formatter_core/src/buffer.rs

@@ -160,37 +160,53 @@ impl<'ast, C> Buffer<'ast, C> for VecBuffer<'_, 'ast, C> {
 /// 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.
+/// The buffer is a watermarked view over the format run's single shared staging vector:
+/// its content is the vector's tail past the length recorded at creation, drained on completion (or on drop).
+/// See `ScratchBuffer` in `state.rs` for the sharing scheme and the LIFO discipline it relies on.
 pub struct HeapVecBuffer<'buf, 'ast, C> {
     state: &'buf mut FormatState<'ast, C>,
-    elements: Vec<FormatElement<'ast>>,
+    watermark: usize,
 }
 
 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 }
+        let watermark = state.scratch().len();
+        Self { state, watermark }
+    }
+
+    /// The number of elements written to this buffer.
+    ///
+    /// Shadows the [`Deref`] slice method: computed without the subslice bounds check.
+    pub fn len(&self) -> usize {
+        self.state.scratch().len() - self.watermark
+    }
+
+    /// Returns `true` if nothing has been written to this buffer.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
     }
 
     /// Removes the last element written to this buffer, if any.
     pub fn pop(&mut self) -> Option<FormatElement<'ast>> {
-        self.elements.pop()
+        if self.is_empty() { None } else { self.state.scratch_mut().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 watermark = self.watermark;
+        let scratch = self.state.scratch_mut();
+        let tail = &scratch[watermark..];
+        let len = tail.len();
         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.
+        // so the bitwise move duplicates no owning state and the `truncate()` below merely forgets the moved-out elements.
         unsafe {
-            std::ptr::copy_nonoverlapping(self.elements.as_ptr(), vec.as_mut_ptr(), len);
+            std::ptr::copy_nonoverlapping(tail.as_ptr(), vec.as_mut_ptr(), len);
             vec.set_len(len);
         }
-        self.elements.clear();
+        scratch.truncate(watermark);
         vec
     }
 
@@ -202,7 +218,9 @@ impl<'buf, 'ast, C> HeapVecBuffer<'buf, 'ast, C> {
 
 impl<C> Drop for HeapVecBuffer<'_, '_, C> {
     fn drop(&mut self) {
-        self.state.return_scratch_buffer(std::mem::take(&mut self.elements));
+        // Forget any elements not taken, restoring the shared vector for the enclosing buffer.
+        let watermark = self.watermark;
+        self.state.scratch_mut().truncate(watermark);
     }
 }
 
@@ -210,13 +228,13 @@ impl<'ast, C> Deref for HeapVecBuffer<'_, 'ast, C> {
     type Target = [FormatElement<'ast>];
 
     fn deref(&self) -> &Self::Target {
-        &self.elements
+        &self.state.scratch()[self.watermark..]
     }
 }
 
 impl<'ast, C> Buffer<'ast, C> for HeapVecBuffer<'_, 'ast, C> {
     fn write_element(&mut self, element: FormatElement<'ast>) {
-        self.elements.push(element);
+        self.state.scratch_mut().push(element);
     }
 
     fn elements(&self) -> &[FormatElement<'ast>] {
@@ -232,7 +250,8 @@ impl<'ast, C> Buffer<'ast, C> for HeapVecBuffer<'_, 'ast, C> {
     }
 
     fn replace_end(&mut self, start: usize, replacement: &[FormatElement<'ast>]) {
-        self.elements.splice(start.., replacement.iter().cloned());
+        let start = self.watermark + start;
+        self.state.scratch_mut().splice(start.., replacement.iter().cloned());
     }
 }
 

crates/oxc_formatter_core/src/state.rs

@@ -1,8 +1,60 @@
+use std::{cell::RefCell, mem};
+
 use oxc_allocator::{Allocator, GetAllocator};
 use rustc_hash::FxHashMap;
 
 use crate::{FormatElement, GroupId, UniqueGroupIdBuilder, format_element::Interned};
 
+thread_local! {
+    /// Cache of heap staging vectors, keeping their high-water capacity alive across format runs on the same thread.
+    /// once a thread is warm, staging performs no heap allocation at all.
+    /// A stack because format runs nest (embedded-language formatting creates a child [`FormatState`] on the same thread).
+    static SCRATCH_CACHE: RefCell<Vec<Vec<FormatElement<'static>>>> =
+        const { RefCell::new(Vec::new()) };
+}
+
+/// The heap staging vector backing [`crate::HeapVecBuffer`] (see there for why staging is heap-backed).
+/// One vector per format run, shared by all nesting levels like a stack via watermarks:
+/// each buffer records the length at creation, pushes its elements, and drains its own tail on completion.
+/// Sound because IR staging is strictly LIFO (an inner buffer always completes before its enclosing one resumes).
+///
+/// Checked out of [`SCRATCH_CACHE`] on creation and returned on drop (panics included).
+pub struct ScratchBuffer<'ast>(Vec<FormatElement<'ast>>);
+
+impl<'ast> ScratchBuffer<'ast> {
+    fn checkout() -> Self {
+        let vec = SCRATCH_CACHE.with_borrow_mut(Vec::pop).unwrap_or_default();
+        // SAFETY: cached vectors are always empty (checkin clears them);
+        // an empty vector holds no values, so re-branding its element lifetime is sound.
+        Self(unsafe {
+            mem::transmute::<Vec<FormatElement<'static>>, Vec<FormatElement<'ast>>>(vec)
+        })
+    }
+}
+
+impl Drop for ScratchBuffer<'_> {
+    fn drop(&mut self) {
+        // Every `HeapVecBuffer` drains its own tail on completion;
+        // a leftover means a buffer leaked elements past its watermark.
+        // Checked here because every format run: root, fragment, or embedded — ends by dropping its `FormatState`.
+        // (Skipped mid-unwind: buffers up the stack haven't truncated their tails yet.)
+        debug_assert!(
+            self.0.is_empty() || std::thread::panicking(),
+            "scratch buffer not fully drained at the end of the format run"
+        );
+        let mut vec = mem::take(&mut self.0);
+        // A buffer that never grew has nothing worth caching
+        if vec.capacity() == 0 {
+            return;
+        }
+        vec.clear();
+        // SAFETY: just cleared; see `checkout`
+        let vec =
+            unsafe { mem::transmute::<Vec<FormatElement<'_>>, Vec<FormatElement<'static>>>(vec) };
+        SCRATCH_CACHE.with_borrow_mut(|cache| cache.push(vec));
+    }
+}
+
 /// This structure stores the state that is relevant for the formatting of the whole document.
 ///
 /// This structure is different from [crate::Formatter] in that the formatting infrastructure
@@ -15,11 +67,8 @@ 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>>>,
+    /// Heap staging vector for [`crate::HeapVecBuffer`]; see [`ScratchBuffer`].
+    scratch: ScratchBuffer<'ast>,
 }
 
 impl<C: std::fmt::Debug> std::fmt::Debug for FormatState<'_, C> {
@@ -36,26 +85,18 @@ impl<'ast, C> FormatState<'ast, C> {
             allocator,
             group_id_builder: UniqueGroupIdBuilder::default(),
             printed_interned_elements: FxHashMap::default(),
-            scratch_buffers: Vec::new(),
+            scratch: ScratchBuffer::checkout(),
         }
     }
 
-    /// 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()
+    /// The heap staging vector shared by all [`crate::HeapVecBuffer`]s of this format run.
+    pub(crate) fn scratch(&self) -> &[FormatElement<'ast>] {
+        &self.scratch.0
     }
 
-    /// 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);
-        }
+    /// Mutable access to the heap staging vector; see [`FormatState::scratch`].
+    pub(crate) fn scratch_mut(&mut self) -> &mut Vec<FormatElement<'ast>> {
+        &mut self.scratch.0
     }
 
     /// Returns the allocator used for arena-allocating format elements.

tasks/track_memory_allocations/allocs_formatter.yaml

@@ -1,9 +1,9 @@
 checker.ts:
   file size: 2922154 # 2.92 MB
-  sys allocs: 11027
-  sys reallocs: 560
-  sys deallocs: 11026
-  sys alloc bytes: 6131149 # 6.13 MB
+  sys allocs: 11013
+  sys reallocs: 492
+  sys deallocs: 11012
+  sys alloc bytes: 5918285 # 5.92 MB
   sys peak growth: 2990058 # 2.99 MB
   arena allocs: 1079319
   arena reallocs: 3951
@@ -11,32 +11,32 @@ checker.ts:
 
 App.tsx:
   file size: 415342 # 415.34 kB
-  sys allocs: 2544
-  sys reallocs: 149
-  sys deallocs: 2543
-  sys alloc bytes: 1542526 # 1.54 MB
-  sys peak growth: 582900 # 582.90 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: 211998
   arena reallocs: 1679
   arena size: 22302080 # 22.30 MB
 
 RadixUIAdoptionSection.jsx:
   file size: 2520 # 2.52 kB
-  sys allocs: 57
-  sys reallocs: 31
-  sys deallocs: 56
-  sys alloc bytes: 15492 # 15.49 kB
-  sys peak growth: 7208 # 7.21 kB
+  sys allocs: 54
+  sys reallocs: 26
+  sys deallocs: 53
+  sys alloc bytes: 10116 # 10.12 kB
+  sys peak growth: 2856 # 2.86 kB
   arena allocs: 1155
   arena reallocs: 69
   arena size: 169128 # 169.13 kB
 
 pdf.mjs:
   file size: 567296 # 567.30 kB
-  sys allocs: 4031
-  sys reallocs: 250
-  sys deallocs: 4030
-  sys alloc bytes: 3057776 # 3.06 MB
+  sys allocs: 4019
+  sys reallocs: 172
+  sys deallocs: 4018
+  sys alloc bytes: 2050032 # 2.05 MB
   sys peak growth: 1164624 # 1.16 MB
   arena allocs: 434612
   arena reallocs: 5696
@@ -44,21 +44,21 @@ pdf.mjs:
 
 antd.js:
   file size: 6686316 # 6.69 MB
-  sys allocs: 34143
-  sys reallocs: 4138
-  sys deallocs: 34142
-  sys alloc bytes: 105935484 # 105.94 MB
-  sys peak growth: 85678576 # 85.68 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: 2559844
   arena reallocs: 14941
   arena size: 332361480 # 332.36 MB
 
 binder.ts:
   file size: 193077 # 193.08 kB
-  sys allocs: 464
-  sys reallocs: 61
-  sys deallocs: 463
-  sys alloc bytes: 315013 # 315.01 kB
+  sys allocs: 457
+  sys reallocs: 35
+  sys deallocs: 456
+  sys alloc bytes: 291781 # 291.78 kB
   sys peak growth: 195661 # 195.66 kB
   arena allocs: 64023
   arena reallocs: 244
@@ -66,10 +66,10 @@ binder.ts:
 
 kitchen-sink.tsx:
   file size: 732898 # 732.90 kB
-  sys allocs: 9578
-  sys reallocs: 2876
-  sys deallocs: 9577
-  sys alloc bytes: 3308148 # 3.31 MB
+  sys allocs: 9565
+  sys reallocs: 2815
+  sys deallocs: 9564
+  sys alloc bytes: 3180244 # 3.18 MB
   sys peak growth: 1499988 # 1.50 MB
   arena allocs: 567769
   arena reallocs: 5992