refactor(mangler): use external allocator

Author: Cameron McHenry

crates/oxc_mangler/src/lib.rs

@@ -37,6 +37,26 @@ pub struct MangleOptions {
 
 type Slot = usize;
 
+/// Enum to handle both owned and borrowed allocators. This is not `Cow` because that type
+/// requires `ToOwned`/`Clone`, which is not implemented for `Allocator`. Although this does
+/// incur some pointer indirect on each reference to the allocator, it allows the API to be
+/// more ergonomic by either accepting an existing allocator, or allowing an internal one to
+/// be created and used temporarily automatically.
+enum TempAllocator<'t> {
+    Owned(Allocator),
+    Borrowed(&'t Allocator),
+}
+
+impl TempAllocator<'_> {
+    /// Get a reference to the allocator, regardless of whether it's owned or borrowed
+    fn as_ref(&self) -> &Allocator {
+        match self {
+            TempAllocator::Owned(allocator) => allocator,
+            TempAllocator::Borrowed(allocator) => allocator,
+        }
+    }
+}
+
 /// # Name Mangler / Symbol Minification
 ///
 /// ## Example
@@ -155,17 +175,80 @@ type Slot = usize;
 /// - slot 1: `top_level_b`, `foo_a`, `bar_a`
 /// - slot 2: `foo`
 /// - slot 3: `bar`
-#[derive(Default)]
-pub struct Mangler {
+pub struct Mangler<'t> {
     options: MangleOptions,
+    /// An allocator meant to be used for temporary allocations during mangling.
+    /// It can be cleared after mangling is done, to free up memory for subsequent
+    /// files or other operations.
+    temp_allocator: TempAllocator<'t>,
 }
 
-impl Mangler {
+impl Default for Mangler<'_> {
+    fn default() -> Self {
+        Self {
+            options: MangleOptions::default(),
+            temp_allocator: TempAllocator::Owned(Allocator::default()),
+        }
+    }
+}
+
+impl<'t> Mangler<'t> {
     #[must_use]
     pub fn new() -> Self {
         Self::default()
     }
 
+    /// Creates a new `Mangler` using an existing temporary allocator. This is an allocator
+    /// that can be reset after mangling and is only used for temporary allocations during
+    /// the mangling process. This makes processing multiple files at once much more efficient,
+    /// because the same memory can be used for mangling each file.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use oxc_allocator::Allocator;
+    /// use oxc_mangler::Mangler;
+    /// use oxc_parser::Parser;
+    /// use oxc_span::SourceType;
+    ///
+    /// let allocator = Allocator::default();
+    /// let temp_allocator = Allocator::default();
+    /// let source = "function myFunction(param) { return param + 1; }";
+    ///
+    /// let parsed = Parser::new(&allocator, source, SourceType::mjs()).parse();
+    /// let mangled_symbols = Mangler::new_with_temp_allocator(&temp_allocator)
+    ///     .build(&parsed.program);
+    ///
+    /// // Reset the allocator to free temporary memory
+    /// temp_allocator.reset();
+    /// ```
+    ///
+    /// Processing multiple files:
+    ///
+    /// ```rust
+    /// # use oxc_allocator::Allocator;
+    /// # use oxc_mangler::Mangler;
+    /// # use oxc_parser::Parser;
+    /// # use oxc_span::SourceType;
+    /// let allocator = Allocator::default();
+    /// let temp_allocator = Allocator::default();
+    /// let files = ["function foo() {}", "function bar() {}"];
+    ///
+    /// for source in files {
+    ///     let parsed = Parser::new(&allocator, source, SourceType::mjs()).parse();
+    ///     let mangled_symbols = Mangler::new_with_temp_allocator(&temp_allocator)
+    ///         .build(&parsed.program);
+    ///     temp_allocator.reset(); // Free memory between files
+    /// }
+    /// ```
+    #[must_use]
+    pub fn new_with_temp_allocator(temp_allocator: &'t Allocator) -> Self {
+        Self {
+            options: MangleOptions::default(),
+            temp_allocator: TempAllocator::Borrowed(temp_allocator),
+        }
+    }
+
     #[must_use]
     pub fn with_options(mut self, options: MangleOptions) -> Self {
         self.options = options;
@@ -216,10 +299,11 @@ impl Mangler {
         let (keep_name_names, keep_name_symbols) =
             Mangler::collect_keep_name_symbols(self.options.keep_names, scoping, ast_nodes);
 
-        let allocator = Allocator::default();
-
         // All symbols with their assigned slots. Keyed by symbol id.
-        let mut slots = Vec::from_iter_in(iter::repeat_n(0, scoping.symbols_len()), &allocator);
+        let mut slots = Vec::from_iter_in(
+            iter::repeat_n(0, scoping.symbols_len()),
+            self.temp_allocator.as_ref(),
+        );
 
         // Stores the lived scope ids for each slot. Keyed by slot number.
         let mut slot_liveness: std::vec::Vec<FixedBitSet> = vec![];
@@ -300,13 +384,13 @@ impl Mangler {
             &keep_name_symbols,
             total_number_of_slots,
             &slots,
-            &allocator,
         );
 
         let root_unresolved_references = scoping.root_unresolved_references();
         let root_bindings = scoping.get_bindings(scoping.root_scope_id());
 
-        let mut reserved_names = Vec::with_capacity_in(total_number_of_slots, &allocator);
+        let mut reserved_names =
+            Vec::with_capacity_in(total_number_of_slots, self.temp_allocator.as_ref());
 
         let mut count = 0;
         for _ in 0..total_number_of_slots {
@@ -387,12 +471,12 @@ impl Mangler {
         keep_name_symbols: &FxHashSet<SymbolId>,
         total_number_of_slots: usize,
         slots: &[Slot],
-        allocator: &'a Allocator,
     ) -> Vec<'a, SlotFrequency<'a>> {
         let root_scope_id = scoping.root_scope_id();
         let mut frequencies = Vec::from_iter_in(
-            repeat_with(|| SlotFrequency::new(allocator)).take(total_number_of_slots),
-            allocator,
+            repeat_with(|| SlotFrequency::new(self.temp_allocator.as_ref()))
+                .take(total_number_of_slots),
+            self.temp_allocator.as_ref(),
         );
 
         for (symbol_id, slot) in slots.iter().copied().enumerate() {
@@ -463,9 +547,9 @@ struct SlotFrequency<'a> {
     pub symbol_ids: Vec<'a, SymbolId>,
 }
 
-impl<'a> SlotFrequency<'a> {
-    fn new(allocator: &'a Allocator) -> Self {
-        Self { slot: 0, frequency: 0, symbol_ids: Vec::new_in(allocator) }
+impl<'t> SlotFrequency<'t> {
+    fn new(temp_allocator: &'t Allocator) -> Self {
+        Self { slot: 0, frequency: 0, symbol_ids: Vec::new_in(temp_allocator) }
     }
 }
 

tasks/benchmark/benches/minifier.rs

@@ -57,14 +57,17 @@ fn bench_mangler(criterion: &mut Criterion) {
         let source_type = SourceType::from_path(&file.file_name).unwrap();
         let source_text = file.source_text.as_str();
         let mut allocator = Allocator::default();
+        let mut temp_allocator = Allocator::default();
         group.bench_function(id, |b| {
             b.iter_with_setup_wrapper(|runner| {
                 allocator.reset();
+                temp_allocator.reset();
                 let program = Parser::new(&allocator, source_text, source_type).parse().program;
                 let mut semantic =
                     SemanticBuilder::new().with_scope_tree_child_ids(true).build(&program).semantic;
                 runner.run(|| {
-                    Mangler::new().build_with_semantic(&mut semantic, &program);
+                    Mangler::new_with_temp_allocator(&temp_allocator)
+                        .build_with_semantic(&mut semantic, &program);
                 });
             });
         });