feat(mangler): support keep_names option (#9898)

Same with #9873, but built on top of #9897.

close #9873
close #9711

Author: sapphi-red

crates/oxc_mangler/src/keep_names.rs

@@ -3,24 +3,65 @@ use oxc_ast::{AstKind, ast::*};
 use oxc_semantic::{AstNode, AstNodes, ReferenceId, Scoping, SymbolId};
 use rustc_hash::FxHashSet;
 
-#[cfg_attr(not(test), expect(dead_code))]
-pub fn collect_name_symbols(scoping: &Scoping, ast_nodes: &AstNodes) -> FxHashSet<SymbolId> {
-    let collector = NameSymbolCollector::new(scoping, ast_nodes);
+#[derive(Debug, Clone, Copy, Default)]
+pub struct MangleOptionsKeepNames {
+    /// Preserve `name` property for functions.
+    ///
+    /// Default `false`
+    pub function: bool,
+
+    /// Preserve `name` property for classes.
+    ///
+    /// Default `false`
+    pub class: bool,
+}
+
+impl MangleOptionsKeepNames {
+    pub fn all_false() -> Self {
+        Self { function: false, class: false }
+    }
+
+    pub fn all_true() -> Self {
+        Self { function: true, class: true }
+    }
+}
+
+impl From<bool> for MangleOptionsKeepNames {
+    fn from(keep_names: bool) -> Self {
+        if keep_names { Self::all_true() } else { Self::all_false() }
+    }
+}
+
+pub fn collect_name_symbols(
+    options: MangleOptionsKeepNames,
+    scoping: &Scoping,
+    ast_nodes: &AstNodes,
+) -> FxHashSet<SymbolId> {
+    let collector = NameSymbolCollector::new(options, scoping, ast_nodes);
     collector.collect()
 }
 
 /// Collects symbols that are used to set `name` properties of functions and classes.
 struct NameSymbolCollector<'a, 'b> {
+    options: MangleOptionsKeepNames,
     scoping: &'b Scoping,
     ast_nodes: &'b AstNodes<'a>,
 }
 
 impl<'a, 'b: 'a> NameSymbolCollector<'a, 'b> {
-    fn new(scoping: &'b Scoping, ast_nodes: &'b AstNodes<'a>) -> Self {
-        Self { scoping, ast_nodes }
+    fn new(
+        options: MangleOptionsKeepNames,
+        scoping: &'b Scoping,
+        ast_nodes: &'b AstNodes<'a>,
+    ) -> Self {
+        Self { options, scoping, ast_nodes }
     }
 
     fn collect(self) -> FxHashSet<SymbolId> {
+        if !self.options.function && !self.options.class {
+            return FxHashSet::default();
+        }
+
         self.scoping
             .symbol_ids()
             .filter(|symbol_id| {
@@ -42,9 +83,12 @@ impl<'a, 'b: 'a> NameSymbolCollector<'a, 'b> {
     fn is_name_set_declare_node(&self, node: &'a AstNode, symbol_id: SymbolId) -> bool {
         match node.kind() {
             AstKind::Function(function) => {
-                function.id.as_ref().is_some_and(|id| id.symbol_id() == symbol_id)
+                self.options.function
+                    && function.id.as_ref().is_some_and(|id| id.symbol_id() == symbol_id)
+            }
+            AstKind::Class(cls) => {
+                self.options.class && cls.id.as_ref().is_some_and(|id| id.symbol_id() == symbol_id)
             }
-            AstKind::Class(cls) => cls.id.as_ref().is_some_and(|id| id.symbol_id() == symbol_id),
             AstKind::VariableDeclarator(decl) => {
                 if let BindingPatternKind::BindingIdentifier(id) = &decl.id.kind {
                     if id.symbol_id() == symbol_id {
@@ -176,9 +220,18 @@ impl<'a, 'b: 'a> NameSymbolCollector<'a, 'b> {
         }
     }
 
-    #[expect(clippy::unused_self)]
     fn is_expression_whose_name_needs_to_be_kept(&self, expr: &Expression) -> bool {
-        expr.is_anonymous_function_definition()
+        let is_anonymous = expr.is_anonymous_function_definition();
+        if !is_anonymous {
+            return false;
+        }
+
+        if self.options.class && self.options.function {
+            return true;
+        }
+
+        let is_class = matches!(expr, Expression::ClassExpression(_));
+        (self.options.class && is_class) || (self.options.function && !is_class)
     }
 }
 
@@ -189,81 +242,83 @@ mod test {
     use oxc_semantic::SemanticBuilder;
     use oxc_span::SourceType;
     use rustc_hash::FxHashSet;
-    use std::iter::once;
 
-    use super::collect_name_symbols;
+    use super::{MangleOptionsKeepNames, collect_name_symbols};
 
-    fn collect(source_text: &str) -> FxHashSet<String> {
+    fn collect(opts: MangleOptionsKeepNames, source_text: &str) -> FxHashSet<String> {
         let allocator = Allocator::default();
         let ret = Parser::new(&allocator, source_text, SourceType::mjs()).parse();
         assert!(!ret.panicked, "{source_text}");
         assert!(ret.errors.is_empty(), "{source_text}");
         let ret = SemanticBuilder::new().build(&ret.program);
         assert!(ret.errors.is_empty(), "{source_text}");
         let semantic = ret.semantic;
-        let symbols = collect_name_symbols(semantic.scoping(), semantic.nodes());
+        let symbols = collect_name_symbols(opts, semantic.scoping(), semantic.nodes());
         symbols
             .into_iter()
             .map(|symbol_id| semantic.scoping().symbol_name(symbol_id).to_string())
             .collect()
     }
 
+    fn data(s: &str) -> FxHashSet<String> {
+        FxHashSet::from_iter([s.to_string()])
+    }
+
+    fn function_only() -> MangleOptionsKeepNames {
+        MangleOptionsKeepNames { function: true, class: false }
+    }
+
+    fn class_only() -> MangleOptionsKeepNames {
+        MangleOptionsKeepNames { function: false, class: true }
+    }
+
     #[test]
     fn test_declarations() {
-        assert_eq!(collect("function foo() {}"), once("foo".to_string()).collect());
-        assert_eq!(collect("class Foo {}"), once("Foo".to_string()).collect());
+        assert_eq!(collect(function_only(), "function foo() {}"), data("foo"));
+        assert_eq!(collect(class_only(), "class Foo {}"), data("Foo"));
     }
 
     #[test]
     fn test_simple_declare_init() {
-        assert_eq!(collect("var foo = function() {}"), once("foo".to_string()).collect());
-        assert_eq!(collect("var foo = () => {}"), once("foo".to_string()).collect());
-        assert_eq!(collect("var Foo = class {}"), once("Foo".to_string()).collect());
+        assert_eq!(collect(function_only(), "var foo = function() {}"), data("foo"));
+        assert_eq!(collect(function_only(), "var foo = () => {}"), data("foo"));
+        assert_eq!(collect(class_only(), "var Foo = class {}"), data("Foo"));
     }
 
     #[test]
     fn test_simple_assign() {
-        assert_eq!(collect("var foo; foo = function() {}"), once("foo".to_string()).collect());
-        assert_eq!(collect("var foo; foo = () => {}"), once("foo".to_string()).collect());
-        assert_eq!(collect("var Foo; Foo = class {}"), once("Foo".to_string()).collect());
+        assert_eq!(collect(function_only(), "var foo; foo = function() {}"), data("foo"));
+        assert_eq!(collect(function_only(), "var foo; foo = () => {}"), data("foo"));
+        assert_eq!(collect(class_only(), "var Foo; Foo = class {}"), data("Foo"));
 
-        assert_eq!(collect("var foo; foo ||= function() {}"), once("foo".to_string()).collect());
-        assert_eq!(
-            collect("var foo = 1; foo &&= function() {}"),
-            once("foo".to_string()).collect()
-        );
-        assert_eq!(collect("var foo; foo ??= function() {}"), once("foo".to_string()).collect());
+        assert_eq!(collect(function_only(), "var foo; foo ||= function() {}"), data("foo"));
+        assert_eq!(collect(function_only(), "var foo = 1; foo &&= function() {}"), data("foo"));
+        assert_eq!(collect(function_only(), "var foo; foo ??= function() {}"), data("foo"));
     }
 
     #[test]
     fn test_default_declarations() {
-        assert_eq!(collect("var [foo = function() {}] = []"), once("foo".to_string()).collect());
-        assert_eq!(collect("var [foo = () => {}] = []"), once("foo".to_string()).collect());
-        assert_eq!(collect("var [Foo = class {}] = []"), once("Foo".to_string()).collect());
-        assert_eq!(collect("var { foo = function() {} } = {}"), once("foo".to_string()).collect());
+        assert_eq!(collect(function_only(), "var [foo = function() {}] = []"), data("foo"));
+        assert_eq!(collect(function_only(), "var [foo = () => {}] = []"), data("foo"));
+        assert_eq!(collect(class_only(), "var [Foo = class {}] = []"), data("Foo"));
+        assert_eq!(collect(function_only(), "var { foo = function() {} } = {}"), data("foo"));
     }
 
     #[test]
     fn test_default_assign() {
+        assert_eq!(collect(function_only(), "var foo; [foo = function() {}] = []"), data("foo"));
+        assert_eq!(collect(function_only(), "var foo; [foo = () => {}] = []"), data("foo"));
+        assert_eq!(collect(class_only(), "var Foo; [Foo = class {}] = []"), data("Foo"));
         assert_eq!(
-            collect("var foo; [foo = function() {}] = []"),
-            once("foo".to_string()).collect()
-        );
-        assert_eq!(collect("var foo; [foo = () => {}] = []"), once("foo".to_string()).collect());
-        assert_eq!(collect("var Foo; [Foo = class {}] = []"), once("Foo".to_string()).collect());
-        assert_eq!(
-            collect("var foo; ({ foo = function() {} } = {})"),
-            once("foo".to_string()).collect()
+            collect(function_only(), "var foo; ({ foo = function() {} } = {})"),
+            data("foo")
         );
     }
 
     #[test]
     fn test_for_in_declaration() {
-        assert_eq!(
-            collect("for (var foo = function() {} in []) {}"),
-            once("foo".to_string()).collect()
-        );
-        assert_eq!(collect("for (var foo = () => {} in []) {}"), once("foo".to_string()).collect());
-        assert_eq!(collect("for (var Foo = class {} in []) {}"), once("Foo".to_string()).collect());
+        assert_eq!(collect(function_only(), "for (var foo = function() {} in []) {}"), data("foo"));
+        assert_eq!(collect(function_only(), "for (var foo = () => {} in []) {}"), data("foo"));
+        assert_eq!(collect(class_only(), "for (var Foo = class {} in []) {}"), data("Foo"));
     }
 }

crates/oxc_mangler/src/lib.rs

@@ -2,26 +2,32 @@ use std::iter::{self, repeat_with};
 
 use fixedbitset::FixedBitSet;
 use itertools::Itertools;
+use keep_names::collect_name_symbols;
 use rustc_hash::FxHashSet;
 
 use base54::base54;
 use oxc_allocator::{Allocator, Vec};
 use oxc_ast::ast::{Declaration, Program, Statement};
 use oxc_data_structures::inline_string::InlineString;
 use oxc_index::Idx;
-use oxc_semantic::{Scoping, Semantic, SemanticBuilder, SymbolId};
+use oxc_semantic::{AstNodes, Scoping, Semantic, SemanticBuilder, SymbolId};
 use oxc_span::Atom;
 
 pub(crate) mod base54;
 mod keep_names;
 
+pub use keep_names::MangleOptionsKeepNames;
+
 #[derive(Default, Debug, Clone, Copy)]
 pub struct MangleOptions {
     /// Pass true to mangle names declared in the top level scope.
     ///
     /// Default: `false`
     pub top_level: bool,
 
+    /// Keep function / class names
+    pub keep_names: MangleOptionsKeepNames,
+
     /// Use more readable mangled names
     /// (e.g. `slot_0`, `slot_1`, `slot_2`, ...) for debugging.
     ///
@@ -207,6 +213,8 @@ impl Mangler {
         } else {
             Default::default()
         };
+        let (keep_name_names, keep_name_symbols) =
+            Mangler::collect_keep_name_symbols(self.options.keep_names, &scoping, &ast_nodes);
 
         let allocator = Allocator::default();
 
@@ -226,6 +234,16 @@ impl Mangler {
                 continue;
             }
 
+            // Sort `bindings` in declaration order.
+            tmp_bindings.clear();
+            tmp_bindings.extend(
+                bindings.values().copied().filter(|binding| !keep_name_symbols.contains(binding)),
+            );
+            tmp_bindings.sort_unstable();
+            if tmp_bindings.is_empty() {
+                continue;
+            }
+
             let mut slot = slot_liveness.len();
 
             reusable_slots.clear();
@@ -236,11 +254,11 @@ impl Mangler {
                     .enumerate()
                     .filter(|(_, slot_liveness)| !slot_liveness.contains(scope_id.index()))
                     .map(|(slot, _)| slot)
-                    .take(bindings.len()),
+                    .take(tmp_bindings.len()),
             );
 
             // The number of new slots that needs to be allocated.
-            let remaining_count = bindings.len() - reusable_slots.len();
+            let remaining_count = tmp_bindings.len() - reusable_slots.len();
             reusable_slots.extend(slot..slot + remaining_count);
 
             slot += remaining_count;
@@ -249,10 +267,6 @@ impl Mangler {
                     .resize_with(slot, || FixedBitSet::with_capacity(scoping.scopes_len()));
             }
 
-            // Sort `bindings` in declaration order.
-            tmp_bindings.clear();
-            tmp_bindings.extend(bindings.values().copied());
-            tmp_bindings.sort_unstable();
             for (&symbol_id, assigned_slot) in
                 tmp_bindings.iter().zip(reusable_slots.iter().copied())
             {
@@ -283,6 +297,7 @@ impl Mangler {
         let frequencies = self.tally_slot_frequencies(
             &scoping,
             &exported_symbols,
+            &keep_name_symbols,
             total_number_of_slots,
             &slots,
             &allocator,
@@ -305,6 +320,8 @@ impl Mangler {
                     && !root_unresolved_references.contains_key(n)
                     && !(root_bindings.contains_key(n)
                         && (!self.options.top_level || exported_names.contains(n)))
+                        // TODO: only skip the names that are kept in the current scope
+                        && !keep_name_names.contains(n)
                 {
                     break name;
                 }
@@ -369,6 +386,7 @@ impl Mangler {
         &'a self,
         scoping: &Scoping,
         exported_symbols: &FxHashSet<SymbolId>,
+        keep_name_symbols: &FxHashSet<SymbolId>,
         total_number_of_slots: usize,
         slots: &[Slot],
         allocator: &'a Allocator,
@@ -389,6 +407,9 @@ impl Mangler {
             if is_special_name(scoping.symbol_name(symbol_id)) {
                 continue;
             }
+            if keep_name_symbols.contains(&symbol_id) {
+                continue;
+            }
             let index = slot;
             frequencies[index].slot = slot;
             frequencies[index].frequency += scoping.get_resolved_reference_ids(symbol_id).len();
@@ -422,6 +443,15 @@ impl Mangler {
             .map(|id| (id.name, id.symbol_id()))
             .collect()
     }
+
+    fn collect_keep_name_symbols<'a>(
+        keep_names: MangleOptionsKeepNames,
+        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)
+    }
 }
 
 fn is_special_name(name: &str) -> bool {