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::SourceType;
use oxc_str::{CompactStr, Ident, Str};
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;
const SLOT_UNASSIGNED: Slot = Slot::MAX;
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()
.with_build_nodes(true)
.with_class_table(true)
.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 allocator = self.temp_allocator.as_ref();
let constraints =
Constraints::collect(allocator, scoping, ast_nodes, program, self.options);
let slots = SlotAssignment::compute(allocator, scoping, ast_nodes, &constraints);
let ranking = SlotRanking::tally(allocator, scoping, &constraints, &slots);
let names =
NameTable::generate(allocator, scoping, &constraints, &ranking, &slots, generate_name);
names.apply(allocator, scoping, &ranking);
}
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) }
}
}
struct Constraints<'a, 's> {
top_level: bool,
exported_names: HashSet<'a, Str<'a>>,
exported_symbols: Option<BitSet<'a>>,
keep_name_names: FxHashSet<&'s str>,
keep_name_symbols: Option<BitSet<'a>>,
}
struct SlotAssignment<'a, 's> {
slots: Vec<'a, Slot>,
total_slots: usize,
eval_reserved_names: FxHashSet<&'s str>,
}
struct SlotRanking<'a> {
frequencies: Vec<'a, SlotFrequency<'a>>,
}
struct NameTable<'a, const CAPACITY: usize> {
names: Vec<'a, InlineString<CAPACITY, u8>>,
}
impl<'a, 's> Constraints<'a, 's> {
fn collect(
allocator: &'a Allocator,
scoping: &'s Scoping,
ast_nodes: &AstNodes,
program: &'a Program<'a>,
options: MangleOptions,
) -> Self {
let top_level = options.top_level(program.source_type);
let (exported_names, exported_symbols) = if top_level && program.source_type.is_module() {
collect_exported_symbols(program, allocator, scoping.symbols_len())
} else {
(HashSet::new_in(allocator), None)
};
let (keep_name_names, keep_name_symbols) =
collect_keep_name_symbols(options.keep_names, allocator, scoping, ast_nodes);
Self { top_level, exported_names, exported_symbols, keep_name_names, keep_name_symbols }
}
}
impl<'a, 's> SlotAssignment<'a, 's> {
fn compute(
allocator: &'a Allocator,
scoping: &'s Scoping,
ast_nodes: &AstNodes,
constraints: &Constraints,
) -> Self {
let keep_name_symbols = constraints.keep_name_symbols.as_ref();
let mut eval_reserved_names: FxHashSet<&'s str> = FxHashSet::default();
let mut slots =
Vec::from_iter_in(iter::repeat_n(SLOT_UNASSIGNED, scoping.symbols_len()), allocator);
let mut slot_liveness: Vec<BitSet> =
Vec::with_capacity_in(scoping.symbols_len(), allocator);
let mut tmp_bindings = Vec::with_capacity_in(100, allocator);
let mut reusable_slots = Vec::new_in(allocator);
let mut ancestor_set = BitSet::new_in(scoping.scopes_len(), allocator);
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.is_some_and(|keep| keep.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(), 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);
}
}
}
if scoping.scope_flags(scope_id).is_function()
&& let Some(func) = ast_nodes.kind(scoping.get_node_id(scope_id)).as_function()
&& func.is_expression()
&& let Some(id) = &func.id
&& let Some(&shadower) = bindings.get(&id.name)
&& shadower != id.symbol_id()
&& slots[shadower.index()] != SLOT_UNASSIGNED
{
slots[id.symbol_id().index()] = slots[shadower.index()];
}
}
let total_slots = slot_liveness.len();
Self { slots, total_slots, eval_reserved_names }
}
}
impl<'a> SlotRanking<'a> {
fn tally(
allocator: &'a Allocator,
scoping: &Scoping,
constraints: &Constraints,
slots: &SlotAssignment,
) -> Self {
let exported_symbols = constraints.exported_symbols.as_ref();
let keep_name_symbols = constraints.keep_name_symbols.as_ref();
let root_scope_id = scoping.root_scope_id();
let mut frequencies = Vec::from_iter_in(
repeat_with(|| SlotFrequency::new(allocator)).take(slots.total_slots),
allocator,
);
for (symbol_id, &slot) in slots.slots.iter().enumerate() {
if slot == SLOT_UNASSIGNED {
continue;
}
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
&& (!constraints.top_level
|| exported_symbols.is_some_and(|e| e.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| keep.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));
Self { frequencies }
}
}
impl<'a, const CAPACITY: usize> NameTable<'a, CAPACITY> {
fn generate(
allocator: &'a Allocator,
scoping: &Scoping,
constraints: &Constraints,
ranking: &SlotRanking,
slots: &SlotAssignment,
generate_name: impl Fn(u32) -> InlineString<CAPACITY, u8>,
) -> Self {
let root_unresolved_references = scoping.root_unresolved_references();
let root_bindings = scoping.get_bindings(scoping.root_scope_id());
let is_reserved = |name: &str| {
oxc_syntax::keyword::is_reserved_keyword(name)
|| is_special_name(name)
|| root_unresolved_references.contains_key(name)
|| (root_bindings.contains_key(name)
&& (!constraints.top_level || constraints.exported_names.contains(name)))
|| constraints.keep_name_names.contains(name)
|| slots.eval_reserved_names.contains(name)
};
let count = ranking.frequencies.len();
let mut names = Vec::with_capacity_in(count, allocator);
let mut candidate = 0;
for _ in 0..count {
let name = loop {
let name = generate_name(candidate);
candidate += 1;
if !is_reserved(name.as_str()) {
break name;
}
};
names.push(name);
}
Self { names }
}
fn apply(self, allocator: &'a Allocator, scoping: &mut Scoping, ranking: &SlotRanking) {
let mut freq_iter = ranking.frequencies.iter();
let mut symbols_renamed_in_this_batch = Vec::with_capacity_in(100, allocator);
let mut slice_of_same_len_strings = Vec::with_capacity_in(100, allocator);
for (_, group) in &self.names.into_iter().chunk_by(InlineString::len) {
slice_of_same_len_strings.clear();
slice_of_same_len_strings.extend(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: &&SlotFrequency| a.slot);
for (symbol_to_rename, new_name) in
symbols_renamed_in_this_batch.iter().zip(slice_of_same_len_strings.iter())
{
for &symbol_id in &symbol_to_rename.symbol_ids {
scoping.set_symbol_name(symbol_id, Ident::from(new_name.as_str()));
}
}
}
}
}
fn collect_exported_symbols<'a>(
program: &Program<'a>,
allocator: &'a Allocator,
symbols_len: usize,
) -> (HashSet<'a, Str<'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_arena_str());
exported_symbols.set_bit(id.symbol_id().index());
}
}
} else if let Some(id) = decl.id() {
exported_names.insert(id.name.as_arena_str());
exported_symbols.set_bit(id.symbol_id().index());
}
}
(exported_names, Some(exported_symbols))
}
fn collect_keep_name_symbols<'alloc, 's>(
keep_names: MangleOptionsKeepNames,
allocator: &'alloc Allocator,
scoping: &'s Scoping,
nodes: &AstNodes,
) -> (FxHashSet<&'s str>, Option<BitSet<'alloc>>) {
if !keep_names.function && !keep_names.class {
return (FxHashSet::default(), None);
}
let ids = collect_name_symbols(keep_names, allocator, scoping, nodes);
(ids.ones().map(|id| scoping.symbol_name(SymbolId::from_usize(id))).collect(), Some(ids))
}
fn debug_name(n: u32) -> InlineString<15, u8> {
InlineString::from_str(&format!("slot_{n}"))
}