use oxc_allocator::{Allocator, ArenaVec, BitSet, GetAllocator};
use oxc_ast::ast::*;
#[cfg(debug_assertions)]
use oxc_ast_visit::{VisitJs, walk_js::walk_function};
use oxc_ecmascript::BoundNames;
use oxc_semantic::Scoping;
use oxc_span::SourceType;
use oxc_syntax::{reference::ReferenceId, scope::ScopeId, symbol::SymbolId};
use crate::{CompressOptions, CompressOptionsUnused, TraverseCtx};
pub struct SymbolLiveness<'a> {
implicitly_observable: BitSet<'a>,
recursive_functions: Option<FunctionGraph<'a>>,
}
impl<'a> SymbolLiveness<'a> {
pub fn new_if_enabled(
source_type: SourceType,
options: &CompressOptions,
scoping: &Scoping,
allocator: &'a Allocator,
) -> Option<Self> {
let recursive_functions_enabled = options.unused != CompressOptionsUnused::Keep;
if !source_type.is_module() && !recursive_functions_enabled {
return None;
}
Some(Self::new(source_type, scoping, allocator))
}
fn new(source_type: SourceType, scoping: &Scoping, allocator: &'a Allocator) -> Self {
let symbols_len = scoping.symbols_len();
let mut implicitly_observable = BitSet::new_in(symbols_len, allocator);
if source_type.is_script() {
for &symbol_id in scoping.get_bindings(scoping.root_scope_id()).values() {
implicitly_observable.set_bit(symbol_id.index());
}
}
Self { implicitly_observable, recursive_functions: None }
}
#[inline]
pub fn is_implicitly_observable(&self, symbol_id: SymbolId) -> bool {
self.implicitly_observable.contains(symbol_id.index())
}
#[inline]
pub fn function_is_dead(&self, symbol_id: SymbolId) -> bool {
self.recursive_functions.as_ref().is_some_and(|graph| graph.is_dead(symbol_id))
}
fn mark_implicitly_observable(&mut self, symbol_id: SymbolId) {
self.implicitly_observable.set_bit(symbol_id.index());
}
fn mark_bound_names<'b>(&mut self, node: &impl BoundNames<'b>) {
node.bound_names(&mut |ident| {
if let Some(symbol_id) = ident.symbol_id.get() {
self.mark_implicitly_observable(symbol_id);
}
});
}
fn register_function(
&mut self,
function: &Function<'_>,
source_type: SourceType,
scoping: &Scoping,
allocator: &'a Allocator,
) {
let Some(symbol_id) = function.id.as_ref().and_then(|id| id.symbol_id.get()) else {
return;
};
let Some(scope_id) = function.scope_id.get() else { return };
let binding_scope_id = scoping.symbol_scope_id(symbol_id);
let binding_scope_flags = scoping.scope_flags(binding_scope_id);
if !function.r#async
&& !function.generator
&& !binding_scope_flags.is_var()
&& !binding_scope_flags.is_strict_mode()
{
self.mark_implicitly_observable(symbol_id);
return;
}
if source_type.is_script() && binding_scope_id == scoping.root_scope_id() {
return;
}
let graph =
self.recursive_functions.get_or_insert_with(|| FunctionGraph::new(scoping, allocator));
graph.register(scope_id, symbol_id);
}
fn analyze(&mut self, scoping: &Scoping) -> bool {
let Some(graph) = &mut self.recursive_functions else { return false };
graph.analyze(scoping, &self.implicitly_observable)
}
#[cfg(debug_assertions)]
fn dead_functions(&self) -> Option<&BitSet<'a>> {
self.recursive_functions.as_ref().map(|graph| &graph.dead)
}
}
struct FunctionGraph<'a> {
candidates: BitSet<'a>,
function_by_scope: ArenaVec<'a, Option<SymbolId>>,
dead: BitSet<'a>,
scratch: GraphScratch<'a>,
}
impl<'a> FunctionGraph<'a> {
fn new(scoping: &Scoping, allocator: &'a Allocator) -> Self {
let symbols_len = scoping.symbols_len();
let function_by_scope =
ArenaVec::from_iter_in(std::iter::repeat_n(None, scoping.scopes_len()), &allocator);
Self {
candidates: BitSet::new_in(symbols_len, allocator),
function_by_scope,
dead: BitSet::new_in(symbols_len, allocator),
scratch: GraphScratch::new(symbols_len, allocator),
}
}
fn register(&mut self, scope_id: ScopeId, symbol_id: SymbolId) {
self.function_by_scope[scope_id.index()] = Some(symbol_id);
self.candidates.set_bit(symbol_id.index());
}
#[inline]
fn is_dead(&self, symbol_id: SymbolId) -> bool {
self.dead.contains(symbol_id.index())
}
fn owner(&self, scoping: &Scoping, scope_id: ScopeId) -> Option<SymbolId> {
scoping
.scope_ancestors(scope_id)
.find_map(|scope_id| self.function_by_scope.get(scope_id.index()).copied().flatten())
}
fn analyze(&mut self, scoping: &Scoping, implicitly_observable: &BitSet<'_>) -> bool {
if scoping.root_scope_flags().contains_direct_eval() {
debug_assert!(self.dead.is_empty(), "direct eval formed after liveness was published");
self.dead.clear();
return false;
}
self.scratch.reset();
for bit in self.candidates.ones() {
let target = SymbolId::from_usize(bit);
if implicitly_observable.contains(bit) {
self.scratch.mark_live(target);
}
let mut has_registered_function_owner = false;
for &reference_id in scoping.get_resolved_reference_ids(target) {
let reference = scoping.get_reference(reference_id);
let Some(owner) = self.owner(scoping, reference.scope_id()) else {
self.scratch.mark_live(target);
continue;
};
has_registered_function_owner = true;
if implicitly_observable.contains(owner.index()) {
self.scratch.mark_live(target);
continue;
}
if self.candidates.contains(owner.index()) {
self.scratch.owned_references.push((owner, target));
} else if !scoping.symbol_is_unused(owner) {
self.scratch.mark_live(target);
}
}
if !has_registered_function_owner && !self.dead.contains(bit) {
self.scratch.candidates_to_untrack.set_bit(bit);
}
}
for index in 0..self.scratch.owned_references.len() {
let (owner, target) = self.scratch.owned_references[index];
let owner_leaves_graph = self.scratch.candidates_to_untrack.contains(owner.index());
let owner_stays_live_by_count = !scoping.symbol_is_unused(owner);
if owner_leaves_graph && owner_stays_live_by_count {
self.scratch.mark_live(target);
}
}
for bit in self.scratch.candidates_to_untrack.ones() {
self.candidates.unset_bit(bit);
}
#[cfg(debug_assertions)]
for bit in self.dead.ones() {
assert!(
self.candidates.contains(bit),
"dead function symbol {bit} was untracked; dead candidates must remain graph \
candidates",
);
}
self.scratch.propagate_liveness(&self.candidates);
#[cfg(debug_assertions)]
for bit in self.dead.ones() {
assert!(
!self.scratch.live.contains(bit),
"function liveness resurrected dead symbol {bit}; transforms must not create a \
new path to a previously unreachable binding",
);
}
let mut published_new_dead = false;
for bit in self.candidates.ones() {
if !self.scratch.live.contains(bit) && !self.dead.contains(bit) {
self.dead.set_bit(bit);
published_new_dead = true;
}
}
published_new_dead
}
}
struct GraphScratch<'a> {
live: BitSet<'a>,
live_worklist: ArenaVec<'a, SymbolId>,
owned_references: ArenaVec<'a, (SymbolId, SymbolId)>,
candidates_to_untrack: BitSet<'a>,
}
impl<'a> GraphScratch<'a> {
fn new(symbols_len: usize, allocator: &'a Allocator) -> Self {
Self {
live: BitSet::new_in(symbols_len, allocator),
live_worklist: ArenaVec::new_in(&allocator),
owned_references: ArenaVec::new_in(&allocator),
candidates_to_untrack: BitSet::new_in(symbols_len, allocator),
}
}
fn reset(&mut self) {
self.live.clear();
self.live_worklist.clear();
self.owned_references.clear();
self.candidates_to_untrack.clear();
}
fn mark_live(&mut self, symbol_id: SymbolId) {
let bit = symbol_id.index();
if !self.live.contains(bit) {
self.live.set_bit(bit);
self.live_worklist.push(symbol_id);
}
}
fn mark_targets_live(&mut self, owner: SymbolId) {
let mut index = self
.owned_references
.partition_point(|&(reference_owner, _)| reference_owner.index() < owner.index());
while let Some((reference_owner, target)) = self.owned_references.get(index).copied() {
if reference_owner != owner {
break;
}
self.mark_live(target);
index += 1;
}
}
fn propagate_liveness(&mut self, candidates: &BitSet<'_>) {
self.owned_references.retain(|(owner, _)| candidates.contains(owner.index()));
self.owned_references.sort_unstable_by_key(|&(owner, _)| owner.index());
while let Some(owner) = self.live_worklist.pop() {
if !candidates.contains(owner.index()) {
continue;
}
self.mark_targets_live(owner);
}
}
}
pub fn register_function(function: &Function<'_>, ctx: &mut TraverseCtx<'_>) {
if !function.is_declaration() {
return;
}
if ctx.options().unused == CompressOptionsUnused::Keep {
return;
}
let source_type = ctx.source_type();
let allocator = ctx.allocator();
let TraverseCtx { state, scoping, .. } = ctx;
if let Some(liveness) = state.symbols.liveness_mut() {
liveness.register_function(function, source_type, scoping.scoping(), allocator);
}
}
pub fn register_using_declaration(
declaration: &VariableDeclaration<'_>,
ctx: &mut TraverseCtx<'_>,
) {
if !declaration.kind.is_using() {
return;
}
let source_type = ctx.source_type();
let allocator = ctx.allocator();
let TraverseCtx { state, scoping, .. } = ctx;
let liveness = state
.symbols
.ensure_liveness(|| SymbolLiveness::new(source_type, scoping.scoping(), allocator));
liveness.mark_bound_names(declaration);
}
pub fn register_export_declaration(declaration: &ExportDeclaration<'_>, ctx: &mut TraverseCtx<'_>) {
let TraverseCtx { state, .. } = ctx;
let Some(liveness) = state.symbols.liveness_mut() else { return };
if !declaration.export_kind().is_type() {
liveness.mark_bound_names(&declaration.declaration);
}
}
pub fn register_named_export(declaration: &ExportNamedDeclaration<'_>, ctx: &mut TraverseCtx<'_>) {
if declaration.export_kind.is_type() {
return;
}
let TraverseCtx { state, scoping, .. } = ctx;
let Some(liveness) = state.symbols.liveness_mut() else { return };
for specifier in &declaration.specifiers {
if specifier.export_kind.is_type() {
continue;
}
let ModuleExportName::IdentifierReference(local) = &specifier.local else { continue };
let Some(reference_id) = local.reference_id.get() else { continue };
let reference = scoping.scoping().get_reference(reference_id);
if reference.flags().is_type_only() {
continue;
}
if let Some(symbol_id) = reference.symbol_id() {
liveness.mark_implicitly_observable(symbol_id);
}
}
}
pub fn register_default_export(
declaration: &ExportDefaultDeclaration<'_>,
ctx: &mut TraverseCtx<'_>,
) {
let symbol_id = match &declaration.declaration {
ExportDefaultDeclarationKind::FunctionDeclaration(function) => {
function.id.as_ref().and_then(|id| id.symbol_id.get())
}
ExportDefaultDeclarationKind::ClassDeclaration(class) => {
class.id.as_ref().and_then(|id| id.symbol_id.get())
}
_ => None,
};
if let Some(symbol_id) = symbol_id
&& let Some(liveness) = ctx.state.symbols.liveness_mut()
{
liveness.mark_implicitly_observable(symbol_id);
}
}
pub fn dead_references_affect_analysis(ctx: &TraverseCtx<'_>) -> bool {
let Some(liveness) = ctx.state.symbols.liveness() else { return false };
let Some(graph) = &liveness.recursive_functions else { return false };
if ctx.scoping().root_scope_flags().contains_direct_eval() {
return false;
}
ctx.state.pass_changes.removed_references.ones().any(|bit| {
ctx.scoping()
.get_reference(ReferenceId::from_usize(bit))
.symbol_id()
.is_some_and(|symbol_id| graph.candidates.contains(symbol_id.index()))
})
}
pub fn analyze<'a>(program: &Program<'a>, ctx: &mut TraverseCtx<'a>, recompute: bool) -> bool {
#[cfg(not(debug_assertions))]
let _ = program;
#[cfg(debug_assertions)]
if let Some(dead) = ctx.state.symbols.liveness().and_then(SymbolLiveness::dead_functions) {
debug_assert_dead_function_declarations_removed(program, ctx.scoping(), dead);
}
if !recompute {
return false;
}
let TraverseCtx { state, scoping, .. } = ctx;
state.symbols.liveness_mut().is_some_and(|liveness| liveness.analyze(scoping.scoping()))
}
#[cfg(debug_assertions)]
fn debug_assert_dead_function_declarations_removed(
program: &Program<'_>,
scoping: &Scoping,
dead: &BitSet<'_>,
) {
if dead.is_empty() {
return;
}
DeadFunctionSweep { scoping, dead }.visit_program(program);
}
#[cfg(debug_assertions)]
struct DeadFunctionSweep<'s, 'd, 'a> {
scoping: &'s Scoping,
dead: &'d BitSet<'a>,
}
#[cfg(debug_assertions)]
impl<'a> VisitJs<'a> for DeadFunctionSweep<'_, '_, '_> {
fn visit_function(&mut self, function: &Function<'a>, flags: oxc_syntax::scope::ScopeFlags) {
if function.is_declaration()
&& let Some(symbol_id) = function.id.as_ref().and_then(|id| id.symbol_id.get())
{
assert!(
!self.dead.contains(symbol_id.index()),
"dead function `{}` survived the pass after its deadness was published",
self.scoping.symbol_name(symbol_id),
);
}
walk_function(self, function, flags);
}
}