perf(mangler): allocate data in arena (#8471)

crates/oxc_mangler/Cargo.toml

@@ -21,9 +21,11 @@ test = false
 doctest = false
 
 [dependencies]
-itertools = { workspace = true }
+oxc_allocator = { workspace = true }
 oxc_ast = { workspace = true }
 oxc_index = { workspace = true }
 oxc_semantic = { workspace = true }
 oxc_span = { workspace = true }
+
+itertools = { workspace = true }
 rustc-hash = { workspace = true }

crates/oxc_mangler/src/lib.rs

@@ -1,18 +1,20 @@
 use itertools::Itertools;
-use oxc_ast::ast::{Declaration, Program, Statement};
-use oxc_index::{index_vec, Idx, IndexVec};
-use oxc_semantic::{ReferenceId, ScopeTree, SemanticBuilder, SymbolId, SymbolTable};
-use oxc_span::CompactStr;
 use rustc_hash::FxHashSet;
 
-type Slot = usize;
+use oxc_allocator::{Allocator, Vec};
+use oxc_ast::ast::{Declaration, Program, Statement};
+use oxc_index::Idx;
+use oxc_semantic::{ReferenceId, ScopeTree, SemanticBuilder, SymbolId, SymbolTable};
+use oxc_span::{Atom, CompactStr};
 
 #[derive(Default, Debug, Clone, Copy)]
 pub struct MangleOptions {
     pub top_level: bool,
     pub debug: bool,
 }
 
+type Slot = usize;
+
 /// # Name Mangler / Symbol Minification
 ///
 /// See:
@@ -102,18 +104,26 @@ impl Mangler {
             Default::default()
         };
 
+        let allocator = Allocator::default();
+
         // Mangle the symbol table by computing slots from the scope tree.
         // A slot is the occurrence index of a binding identifier inside a scope.
         let mut symbol_table = symbol_table;
 
         // Total number of slots for all scopes
         let mut total_number_of_slots: Slot = 0;
 
-        // All symbols with their assigned slots
-        let mut slots: IndexVec<SymbolId, Slot> = index_vec![0; symbol_table.len()];
+        // All symbols with their assigned slots. Keyed by symbol id.
+        let mut slots: Vec<'_, Slot> = Vec::with_capacity_in(symbol_table.len(), &allocator);
+        for _ in 0..symbol_table.len() {
+            slots.push(0);
+        }
 
         // Keep track of the maximum slot number for each scope
-        let mut max_slot_for_scope = vec![0; scope_tree.len()];
+        let mut max_slot_for_scope = Vec::with_capacity_in(scope_tree.len(), &allocator);
+        for _ in 0..scope_tree.len() {
+            max_slot_for_scope.push(0);
+        }
 
         // Walk the scope tree and compute the slot number for each scope
         for scope_id in scope_tree.descendants_from_root() {
@@ -128,10 +138,10 @@ impl Mangler {
 
             if !bindings.is_empty() {
                 // Sort `bindings` in declaration order.
-                let mut bindings = bindings.values().copied().collect::<Vec<_>>();
+                let mut bindings = bindings.values().copied().collect::<std::vec::Vec<_>>();
                 bindings.sort_unstable();
                 for symbol_id in bindings {
-                    slots[symbol_id] = slot;
+                    slots[symbol_id.index()] = slot;
                     slot += 1;
                 }
             }
@@ -149,12 +159,13 @@ impl Mangler {
             scope_tree,
             total_number_of_slots,
             &slots,
+            &allocator,
         );
 
         let root_unresolved_references = scope_tree.root_unresolved_references();
         let root_bindings = scope_tree.get_bindings(scope_tree.root_scope_id());
 
-        let mut reserved_names = Vec::with_capacity(total_number_of_slots);
+        let mut reserved_names = Vec::with_capacity_in(total_number_of_slots, &allocator);
 
         let generate_name = if self.options.debug { debug_name } else { base54 };
         let mut count = 0;
@@ -199,8 +210,10 @@ impl Mangler {
             // (freq_iter is sorted by frequency from highest to lowest,
             //  so taking means take the N most frequent symbols remaining)
             let slice_of_same_len_strings = slice_of_same_len_strings_group.collect_vec();
-            let mut symbols_renamed_in_this_batch =
-                freq_iter.by_ref().take(slice_of_same_len_strings.len()).collect::<Vec<_>>();
+            let mut symbols_renamed_in_this_batch = freq_iter
+                .by_ref()
+                .take(slice_of_same_len_strings.len())
+                .collect::<std::vec::Vec<_>>();
 
             debug_assert!(symbols_renamed_in_this_batch.len() == slice_of_same_len_strings.len());
 
@@ -226,17 +239,23 @@ impl Mangler {
         self
     }
 
-    fn tally_slot_frequencies(
-        &self,
+    fn tally_slot_frequencies<'a>(
+        &'a self,
         symbol_table: &SymbolTable,
         exported_symbols: &FxHashSet<SymbolId>,
         scope_tree: &ScopeTree,
         total_number_of_slots: usize,
-        slots: &IndexVec<SymbolId, Slot>,
-    ) -> Vec<SlotFrequency> {
+        slots: &[Slot],
+        allocator: &'a Allocator,
+    ) -> Vec<'a, SlotFrequency<'a>> {
         let root_scope_id = scope_tree.root_scope_id();
-        let mut frequencies = vec![SlotFrequency::default(); total_number_of_slots];
-        for (symbol_id, slot) in slots.iter_enumerated() {
+        let mut frequencies = Vec::with_capacity_in(total_number_of_slots, allocator);
+        for _ in 0..total_number_of_slots {
+            frequencies.push(SlotFrequency::new(allocator));
+        }
+
+        for (symbol_id, slot) in slots.iter().copied().enumerate() {
+            let symbol_id = SymbolId::from_usize(symbol_id);
             if symbol_table.get_scope_id(symbol_id) == root_scope_id
                 && (!self.options.top_level || exported_symbols.contains(&symbol_id))
             {
@@ -245,8 +264,8 @@ impl Mangler {
             if is_special_name(symbol_table.get_name(symbol_id)) {
                 continue;
             }
-            let index = *slot;
-            frequencies[index].slot = *slot;
+            let index = slot;
+            frequencies[index].slot = slot;
             frequencies[index].frequency +=
                 symbol_table.get_resolved_reference_ids(symbol_id).len();
             frequencies[index].symbol_ids.push(symbol_id);
@@ -255,7 +274,9 @@ impl Mangler {
         frequencies
     }
 
-    fn collect_exported_symbols(program: &Program) -> (FxHashSet<CompactStr>, FxHashSet<SymbolId>) {
+    fn collect_exported_symbols<'a>(
+        program: &Program<'a>,
+    ) -> (FxHashSet<Atom<'a>>, FxHashSet<SymbolId>) {
         program
             .body
             .iter()
@@ -274,7 +295,7 @@ impl Mangler {
                     itertools::Either::Right(decl.id().into_iter())
                 }
             })
-            .map(|id| (id.name.to_compact_str(), id.symbol_id()))
+            .map(|id| (id.name.clone(), id.symbol_id()))
             .collect()
     }
 }
@@ -283,11 +304,17 @@ fn is_special_name(name: &str) -> bool {
     matches!(name, "exports" | "arguments")
 }
 
-#[derive(Debug, Default, Clone)]
-struct SlotFrequency {
+#[derive(Debug)]
+struct SlotFrequency<'a> {
     pub slot: Slot,
     pub frequency: usize,
-    pub symbol_ids: Vec<SymbolId>,
+    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) }
+    }
 }
 
 #[rustfmt::skip]