use std::iter::{self, repeat_with};
use itertools::Itertools;
use keep_names::collect_name_symbols;
use oxc_index::IndexVec;
use oxc_syntax::class::ClassId;
use rustc_hash::{FxHashMap, FxHashSet};
use base54::base54;
use oxc_allocator::{Allocator, BitSet, HashSet, Vec};
use oxc_ast::ast::{Declaration, Program, Statement};
use oxc_data_structures::inline_string::InlineString;
use oxc_semantic::{AstNodes, Reference, Scoping, Semantic, SemanticBuilder, SymbolId};
use oxc_span::{Atom, CompactStr, Ident, SourceType};
pub(crate) mod base54;
mod keep_names;
pub use keep_names::MangleOptionsKeepNames;
#[derive(Default, Debug, Clone, Copy)]
pub struct MangleOptions {
pub top_level: Option<bool>,
pub keep_names: MangleOptionsKeepNames,
pub debug: bool,
}
impl MangleOptions {
fn top_level(self, source_type: SourceType) -> bool {
self.top_level.unwrap_or(source_type.is_module() || source_type.is_commonjs())
}
}
type Slot = u32;
enum TempAllocator<'t> {
Owned(Allocator),
Borrowed(&'t Allocator),
}
impl TempAllocator<'_> {
fn as_ref(&self) -> &Allocator {
match self {
TempAllocator::Owned(allocator) => allocator,
TempAllocator::Borrowed(allocator) => allocator,
}
}
}
pub struct ManglerReturn {
pub scoping: Scoping,
pub class_private_mappings: IndexVec<ClassId, FxHashMap<String, CompactStr>>,
}
pub struct Mangler<'t> {
options: MangleOptions,
temp_allocator: TempAllocator<'t>,
}
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()
}
#[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;
self
}
#[must_use]
pub fn build(self, program: &Program<'_>) -> ManglerReturn {
let mut semantic = SemanticBuilder::new().build(program).semantic;
let class_private_mappings = self.build_with_semantic(&mut semantic, program);
ManglerReturn { scoping: semantic.into_scoping(), class_private_mappings }
}
pub fn build_with_semantic(
self,
semantic: &mut Semantic<'_>,
program: &Program<'_>,
) -> IndexVec<ClassId, FxHashMap<String, CompactStr>> {
let class_private_mappings = Self::collect_private_members_from_semantic(semantic);
if self.options.debug {
self.build_with_semantic_impl(semantic, program, debug_name);
} else {
self.build_with_semantic_impl(semantic, program, base54);
}
class_private_mappings
}
fn build_with_semantic_impl<const CAPACITY: usize, G: Fn(u32) -> InlineString<CAPACITY, u8>>(
self,
semantic: &mut Semantic<'_>,
program: &Program<'_>,
generate_name: G,
) {
let (scoping, ast_nodes) = semantic.scoping_mut_and_nodes();
let symbols_len = scoping.symbols_len();
let temp_allocator = self.temp_allocator.as_ref();
let top_level = self.options.top_level(program.source_type);
let (exported_names, exported_symbols) = if top_level && program.source_type.is_module() {
Mangler::collect_exported_symbols(program, temp_allocator, symbols_len)
} else {
(HashSet::new_in(temp_allocator), None)
};
let (keep_name_names, keep_name_symbols) = Mangler::collect_keep_name_symbols(
self.options.keep_names,
temp_allocator,
scoping,
ast_nodes,
);
let mut slots = Vec::from_iter_in(iter::repeat_n(0, scoping.symbols_len()), temp_allocator);
let mut slot_liveness: Vec<BitSet> =
Vec::with_capacity_in(scoping.symbols_len(), temp_allocator);
let mut tmp_bindings = Vec::with_capacity_in(100, temp_allocator);
let mut reusable_slots = Vec::new_in(temp_allocator);
let mut ancestor_set = BitSet::new_in(scoping.scopes_len(), temp_allocator);
let mut eval_reserved_names: FxHashSet<&str> = FxHashSet::default();
for (scope_id, bindings) in scoping.iter_bindings() {
if bindings.is_empty() {
continue;
}
if scoping.scope_flags(scope_id).contains_direct_eval() {
for (name, _) in bindings {
eval_reserved_names.insert(name.as_str());
}
continue;
}
tmp_bindings.clear();
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;
}
tmp_bindings.sort_unstable();
let mut slot = slot_liveness.len();
reusable_slots.clear();
reusable_slots.extend(
slot_liveness
.iter()
.enumerate()
.filter(|(_, slot_liveness)| !slot_liveness.has_bit(scope_id.index()))
.map(
#[expect(clippy::cast_possible_truncation)]
|(slot, _)| slot as Slot,
)
.take(tmp_bindings.len()),
);
let remaining_count = tmp_bindings.len() - reusable_slots.len();
#[expect(clippy::cast_possible_truncation)]
reusable_slots.extend((slot as Slot)..(slot + remaining_count) as Slot);
slot += remaining_count;
if slot_liveness.len() < slot {
slot_liveness.extend(
iter::repeat_with(|| BitSet::new_in(scoping.scopes_len(), temp_allocator))
.take(remaining_count),
);
}
ancestor_set.clear();
for ancestor_id in scoping.scope_ancestors(scope_id).skip(1) {
ancestor_set.set_bit(ancestor_id.index());
}
let scope_id_index = scope_id.index();
for (&symbol_id, &assigned_slot) in tmp_bindings.iter().zip(&reusable_slots) {
slots[symbol_id.index()] = assigned_slot;
let declared_scope_id =
ast_nodes.get_node(scoping.symbol_declaration(symbol_id)).scope_id();
let redeclared_scope_ids = scoping
.symbol_redeclarations(symbol_id)
.iter()
.map(|r| ast_nodes.get_node(r.declaration).scope_id());
let referenced_scope_ids =
scoping.get_resolved_references(symbol_id).map(Reference::scope_id);
let slot_liveness_bitset = &mut slot_liveness[assigned_slot as usize];
for used_scope_id in referenced_scope_ids
.chain(redeclared_scope_ids)
.chain([scope_id, declared_scope_id])
{
for ancestor_id in scoping.scope_ancestors(used_scope_id) {
let ancestor_index = ancestor_id.index();
if ancestor_index == scope_id_index || ancestor_set.has_bit(ancestor_index)
{
break;
}
if slot_liveness_bitset.has_bit(ancestor_index) {
debug_assert!(
scoping.scope_ancestors(ancestor_id).skip(1).all(|a| {
let idx = a.index();
slot_liveness_bitset.has_bit(idx)
|| idx == scope_id_index
|| ancestor_set.has_bit(idx)
}),
"Invariant violated: ancestor chain should be fully marked live"
);
break;
}
slot_liveness_bitset.set_bit(ancestor_index);
}
}
}
}
let total_number_of_slots = slot_liveness.len();
let frequencies = self.tally_slot_frequencies(
scoping,
exported_symbols.as_ref(),
keep_name_symbols.as_ref(),
total_number_of_slots,
&slots,
top_level,
);
let root_unresolved_references = scoping.root_unresolved_references();
let root_bindings = scoping.get_bindings(scoping.root_scope_id());
let names_needed = frequencies.len();
let mut reserved_names = Vec::with_capacity_in(names_needed, temp_allocator);
let mut count = 0;
for _ in 0..names_needed {
let name = loop {
let name = generate_name(count);
count += 1;
let n = name.as_str();
if !oxc_syntax::keyword::is_reserved_keyword(n)
&& !is_special_name(n)
&& !root_unresolved_references.contains_key(n)
&& !(root_bindings.contains_key(n)
&& (!top_level || exported_names.contains(n)))
&& !keep_name_names.contains(n)
&& !eval_reserved_names.contains(n)
{
break name;
}
};
reserved_names.push(name);
}
let mut freq_iter = frequencies.iter();
let mut symbols_renamed_in_this_batch = Vec::with_capacity_in(100, temp_allocator);
let mut slice_of_same_len_strings = Vec::with_capacity_in(100, temp_allocator);
for (_, slice_of_same_len_strings_group) in
&reserved_names.into_iter().chunk_by(InlineString::len)
{
slice_of_same_len_strings.clear();
slice_of_same_len_strings.extend(slice_of_same_len_strings_group);
symbols_renamed_in_this_batch.clear();
symbols_renamed_in_this_batch
.extend(freq_iter.by_ref().take(slice_of_same_len_strings.len()));
debug_assert_eq!(symbols_renamed_in_this_batch.len(), slice_of_same_len_strings.len());
symbols_renamed_in_this_batch.sort_unstable_by_key(|a| a.slot);
let symbols_to_rename_with_new_names =
symbols_renamed_in_this_batch.iter().zip(slice_of_same_len_strings.iter());
for (symbol_to_rename, new_name) in symbols_to_rename_with_new_names {
for &symbol_id in &symbol_to_rename.symbol_ids {
scoping.set_symbol_name(symbol_id, Ident::from(new_name.as_str()));
}
}
}
}
fn tally_slot_frequencies<'a>(
&'a self,
scoping: &Scoping,
exported_symbols: Option<&BitSet<'a>>,
keep_name_symbols: Option<&BitSet<'a>>,
total_number_of_slots: usize,
slots: &[Slot],
top_level: bool,
) -> Vec<'a, SlotFrequency<'a>> {
let root_scope_id = scoping.root_scope_id();
let temp_allocator = self.temp_allocator.as_ref();
let mut frequencies = Vec::from_iter_in(
repeat_with(|| SlotFrequency::new(temp_allocator)).take(total_number_of_slots),
temp_allocator,
);
for (symbol_id, &slot) in slots.iter().enumerate() {
let symbol_id = SymbolId::from_usize(symbol_id);
let symbol_scope_id = scoping.symbol_scope_id(symbol_id);
if symbol_scope_id == root_scope_id
&& (!top_level
|| exported_symbols.is_some_and(|exported_symbols| {
exported_symbols.has_bit(symbol_id.index())
}))
{
continue;
}
if scoping.scope_flags(symbol_scope_id).contains_direct_eval() {
continue;
}
if is_special_name(scoping.symbol_name(symbol_id)) {
continue;
}
if keep_name_symbols
.is_some_and(|keep_name_symbols| keep_name_symbols.has_bit(symbol_id.index()))
{
continue;
}
let index = slot as usize;
frequencies[index].slot = slot;
frequencies[index].frequency += scoping.get_resolved_reference_ids(symbol_id).len();
frequencies[index].symbol_ids.push(symbol_id);
}
frequencies.retain(|x| !x.symbol_ids.is_empty());
frequencies.sort_unstable_by_key(|x| std::cmp::Reverse(x.frequency));
frequencies
}
fn collect_exported_symbols<'a>(
program: &Program<'a>,
allocator: &'a Allocator,
symbols_len: usize,
) -> (HashSet<'a, Atom<'a>>, Option<BitSet<'a>>) {
let mut exported_symbols = BitSet::new_in(symbols_len, allocator);
let mut exported_names = HashSet::new_in(allocator);
for statement in &program.body {
let Statement::ExportNamedDeclaration(v) = statement else { continue };
let Some(decl) = &v.declaration else { continue };
if let Declaration::VariableDeclaration(decl) = decl {
for decl in &decl.declarations {
if let Some(id) = decl.id.get_binding_identifier() {
exported_names.insert(id.name.as_atom());
exported_symbols.set_bit(id.symbol_id().index());
}
}
} else if let Some(id) = decl.id() {
exported_names.insert(id.name.as_atom());
exported_symbols.set_bit(id.symbol_id().index());
}
}
(exported_names, Some(exported_symbols))
}
fn collect_keep_name_symbols<'a>(
keep_names: MangleOptionsKeepNames,
temp_allocator: &'t Allocator,
scoping: &'a Scoping,
nodes: &AstNodes,
) -> (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))
}
fn collect_private_members_from_semantic(
semantic: &Semantic<'_>,
) -> IndexVec<ClassId, FxHashMap<String, CompactStr>> {
let classes = semantic.classes();
let private_member_count: IndexVec<ClassId, usize> = classes
.elements
.iter()
.map(|class_elements| {
class_elements.iter().filter(|element| element.is_private).count()
})
.collect();
let parent_private_member_count: IndexVec<ClassId, usize> = classes
.declarations
.iter_enumerated()
.map(|(class_id, _)| {
classes
.ancestors(class_id)
.skip(1)
.map(|id| private_member_count[id])
.sum::<usize>()
})
.collect();
classes
.elements
.iter_enumerated()
.map(|(class_id, class_elements)| {
let parent_private_member_count = parent_private_member_count[class_id];
assert!(
u32::try_from(class_elements.len() + parent_private_member_count).is_ok(),
"too many class elements"
);
class_elements
.iter()
.filter_map(|element| {
if element.is_private { Some(element.name.to_string()) } else { None }
})
.enumerate()
.map(|(i, name)| {
#[expect(
clippy::cast_possible_truncation,
reason = "checked above with assert"
)]
let mangled = CompactStr::new(
base54((parent_private_member_count + i) as u32).as_str(),
);
(name, mangled)
})
.collect::<FxHashMap<_, _>>()
})
.collect()
}
}
fn is_special_name(name: &str) -> bool {
matches!(name, "arguments")
}
#[derive(Debug)]
struct SlotFrequency<'a> {
pub slot: Slot,
pub frequency: usize,
pub symbol_ids: Vec<'a, SymbolId>,
}
impl<'t> SlotFrequency<'t> {
fn new(temp_allocator: &'t Allocator) -> Self {
Self { slot: 0, frequency: 0, symbol_ids: Vec::new_in(temp_allocator) }
}
}
fn debug_name(n: u32) -> InlineString<15, u8> {
InlineString::from_str(&format!("slot_{n}"))
}