perf(mangler): use BitSet for keep_names symbols set

crates/oxc_mangler/src/keep_names.rs

@@ -1,4 +1,4 @@
-use rustc_hash::FxHashSet;
+use oxc_allocator::{Allocator, BitSet};
 
 use oxc_ast::{AstKind, ast::*};
 use oxc_semantic::{AstNode, AstNodes, ReferenceId, Scoping, SymbolId};
@@ -32,36 +32,36 @@ impl From<bool> for MangleOptionsKeepNames {
     }
 }
 
-pub fn collect_name_symbols(
+pub fn collect_name_symbols<'a>(
     options: MangleOptionsKeepNames,
+    allocator: &'a Allocator,
     scoping: &Scoping,
     ast_nodes: &AstNodes,
-) -> FxHashSet<SymbolId> {
-    let collector = NameSymbolCollector::new(options, scoping, ast_nodes);
+) -> BitSet<'a> {
+    let collector = NameSymbolCollector::new(options, allocator, scoping, ast_nodes);
     collector.collect()
 }
 
 /// Collects symbols that are used to set `name` properties of functions and classes.
-struct NameSymbolCollector<'a, 'b> {
+struct NameSymbolCollector<'a, 'b, 't> {
     options: MangleOptionsKeepNames,
     scoping: &'b Scoping,
     ast_nodes: &'b AstNodes<'a>,
+    allocator: &'t Allocator,
 }
 
-impl<'a, 'b: 'a> NameSymbolCollector<'a, 'b> {
+impl<'a, 'b: 'a, 't> NameSymbolCollector<'a, 'b, 't> {
     fn new(
         options: MangleOptionsKeepNames,
+        allocator: &'t Allocator,
         scoping: &'b Scoping,
         ast_nodes: &'b AstNodes<'a>,
     ) -> Self {
-        Self { options, scoping, ast_nodes }
+        Self { options, scoping, ast_nodes, allocator }
     }
 
-    fn collect(self) -> FxHashSet<SymbolId> {
-        if !self.options.function && !self.options.class {
-            return FxHashSet::default();
-        }
-
+    fn collect(self) -> BitSet<'t> {
+        let mut symbol_ids = BitSet::new_in(self.scoping.symbols_len(), self.allocator);
         self.scoping
             .symbol_ids()
             .filter(|symbol_id| {
@@ -70,7 +70,8 @@ impl<'a, 'b: 'a> NameSymbolCollector<'a, 'b> {
                 self.is_name_set_declare_node(decl_node, *symbol_id)
                     || self.has_name_set_reference_node(*symbol_id)
             })
-            .collect()
+            .for_each(|symbol_id| symbol_ids.set_bit(symbol_id.index()));
+        symbol_ids
     }
 
     fn has_name_set_reference_node(&self, symbol_id: SymbolId) -> bool {
@@ -246,7 +247,7 @@ impl<'a, 'b: 'a> NameSymbolCollector<'a, 'b> {
 mod test {
     use oxc_allocator::Allocator;
     use oxc_parser::Parser;
-    use oxc_semantic::SemanticBuilder;
+    use oxc_semantic::{SemanticBuilder, SymbolId};
     use oxc_span::SourceType;
     use rustc_hash::FxHashSet;
 
@@ -260,10 +261,12 @@ mod test {
         let ret = SemanticBuilder::new().build(&ret.program);
         assert!(ret.errors.is_empty(), "{source_text}");
         let semantic = ret.semantic;
-        let symbols = collect_name_symbols(opts, semantic.scoping(), semantic.nodes());
+        let symbols = collect_name_symbols(opts, &allocator, semantic.scoping(), semantic.nodes());
         symbols
-            .into_iter()
-            .map(|symbol_id| semantic.scoping().symbol_name(symbol_id).to_string())
+            .ones()
+            .map(|symbol_id| {
+                semantic.scoping().symbol_name(SymbolId::from_usize(symbol_id)).to_string()
+            })
             .collect()
     }
 

crates/oxc_mangler/src/lib.rs

@@ -296,15 +296,19 @@ impl<'t> Mangler<'t> {
     ) {
         let (scoping, ast_nodes) = semantic.scoping_mut_and_nodes();
 
+        let temp_allocator = self.temp_allocator.as_ref();
+
         let (exported_names, exported_symbols) = if self.options.top_level {
             Mangler::collect_exported_symbols(program)
         } else {
             Default::default()
         };
-        let (keep_name_names, keep_name_symbols) =
-            Mangler::collect_keep_name_symbols(self.options.keep_names, scoping, ast_nodes);
-
-        let temp_allocator = self.temp_allocator.as_ref();
+        let (keep_name_names, keep_name_symbols) = Mangler::collect_keep_name_symbols(
+            self.options.keep_names,
+            temp_allocator,
+            scoping,
+            ast_nodes,
+        );
 
         // All symbols with their assigned slots. Keyed by symbol id.
         let mut slots = Vec::from_iter_in(iter::repeat_n(0, scoping.symbols_len()), temp_allocator);
@@ -341,9 +345,11 @@ impl<'t> Mangler<'t> {
 
             // Sort `bindings` in declaration order.
             tmp_bindings.clear();
-            tmp_bindings.extend(
-                bindings.values().copied().filter(|binding| !keep_name_symbols.contains(binding)),
-            );
+            tmp_bindings.extend(bindings.values().copied().filter(|binding| {
+                !keep_name_symbols
+                    .as_ref()
+                    .is_some_and(|keep_name_symbols| keep_name_symbols.has_bit(binding.index()))
+            }));
             if tmp_bindings.is_empty() {
                 continue;
             }
@@ -444,7 +450,7 @@ impl<'t> Mangler<'t> {
         let frequencies = self.tally_slot_frequencies(
             scoping,
             &exported_symbols,
-            &keep_name_symbols,
+            keep_name_symbols.as_ref(),
             total_number_of_slots,
             &slots,
         );
@@ -537,7 +543,7 @@ impl<'t> Mangler<'t> {
         &'a self,
         scoping: &Scoping,
         exported_symbols: &FxHashSet<SymbolId>,
-        keep_name_symbols: &FxHashSet<SymbolId>,
+        keep_name_symbols: Option<&BitSet<'a>>,
         total_number_of_slots: usize,
         slots: &[Slot],
     ) -> Vec<'a, SlotFrequency<'a>> {
@@ -562,7 +568,9 @@ impl<'t> Mangler<'t> {
             if is_special_name(scoping.symbol_name(symbol_id)) {
                 continue;
             }
-            if keep_name_symbols.contains(&symbol_id) {
+            if keep_name_symbols
+                .is_some_and(|keep_name_symbols| keep_name_symbols.has_bit(symbol_id.index()))
+            {
                 continue;
             }
             let index = slot as usize;
@@ -604,11 +612,15 @@ impl<'t> Mangler<'t> {
 
     fn collect_keep_name_symbols<'a>(
         keep_names: MangleOptionsKeepNames,
+        temp_allocator: &'t Allocator,
         scoping: &'a Scoping,
         nodes: &AstNodes,
-    ) -> (FxHashSet<&'a str>, FxHashSet<SymbolId>) {
-        let ids = collect_name_symbols(keep_names, scoping, nodes);
-        (ids.iter().map(|id| scoping.symbol_name(*id)).collect(), ids)
+    ) -> (FxHashSet<&'a str>, Option<BitSet<'t>>) {
+        if !keep_names.function && !keep_names.class {
+            return (FxHashSet::default(), None);
+        }
+        let ids = collect_name_symbols(keep_names, temp_allocator, scoping, nodes);
+        (ids.ones().map(|id| scoping.symbol_name(SymbolId::from_usize(id))).collect(), Some(ids))
     }
 
     /// Collects and generates mangled names for private members using semantic information