use std::collections::BTreeMap;
use std::mem;
use oxc_allocator::Vec as ArenaVec;
use oxc_ast::ast::*;
use oxc_span::{GetSpan, SPAN, Span};
use oxc_traverse::TraverseCtx;
use crate::pragma::IgnoreType;
use super::counters::{
ClassFieldHoist, CounterKind, CounterType, PendingInsertion, build_counter_stmt,
prepend_counter,
};
use super::coverage_map::is_synthetic_span;
use super::ignore::mark_ignored_declarator_fn;
use super::names::{declarator_function_name, property_key_to_name};
use super::{CoverageState, CoverageTransform};
impl<'arena> CoverageTransform<'_, 'arena> {
pub(super) fn register_statement_counter(
&mut self,
stmt: &Statement<'arena>,
ctx: &TraverseCtx<'arena, CoverageState>,
) {
let span = stmt.span();
let parent_ignored = self.in_ignored_subtree();
let is_injected = is_synthetic_span(span);
let has_ignore_next = !is_injected
&& (ctx.state.pragmas.get(span.start) == Some(IgnoreType::Next)
|| self.is_in_ignored_if_arm(span));
self.ignored_stmt_stack.push(has_ignore_next);
if is_injected || is_container_statement(stmt) || parent_ignored {
return;
}
if has_ignore_next {
self.skip_next = true;
return;
}
if self.skip_next {
self.skip_next = false;
return;
}
if let Some(stmt_id) = self.add_statement(span) {
self.pending_insertions.push(PendingInsertion {
target_start: span.start,
counter_id: stmt_id,
counter_type: CounterType::Statement,
});
}
}
pub(super) fn finish_statement(&mut self) {
self.skip_next = false;
self.ignored_stmt_stack.pop();
}
pub(super) fn insert_pending_statement_counters(
&mut self,
stmts: &mut ArenaVec<'arena, Statement<'arena>>,
ctx: &TraverseCtx<'arena, CoverageState>,
) {
if self.pending_insertions.is_empty() {
return;
}
let cov_fn = self.cov_fn_name;
let mut insertions_by_target = BTreeMap::new();
for stmt in stmts.iter() {
let span = stmt.span();
if is_synthetic_span(span) {
continue;
}
insertions_by_target.entry(span.start).or_insert_with(Vec::new);
}
let mut unmatched = Vec::with_capacity(self.pending_insertions.len());
let mut matched_count = 0;
for pending in mem::take(&mut self.pending_insertions) {
if let Some(target) = insertions_by_target.get_mut(&pending.target_start) {
target.push(pending);
matched_count += 1;
} else {
unmatched.push(pending);
}
}
self.pending_insertions = unmatched;
if matched_count == 0 {
return;
}
let original = mem::replace(stmts, ArenaVec::new_in(ctx));
let mut rebuilt = ArenaVec::with_capacity_in(original.len() + matched_count, ctx);
for stmt in original {
let span = stmt.span();
if let Some(insertions) = insertions_by_target.remove(&span.start) {
for pending in insertions {
rebuilt
.push(build_counter_stmt(CounterKind::from_pending(cov_fn, &pending), ctx));
}
}
rebuilt.push(stmt);
}
*stmts = rebuilt;
}
pub(super) fn instrument_variable_declarator(
&mut self,
decl: &mut VariableDeclarator<'arena>,
ctx: &TraverseCtx<'arena, CoverageState>,
) {
if self.skip_current_var_decl {
mark_ignored_declarator_fn(decl, &mut self.skip_next);
return;
}
if let Some(name) = declarator_function_name(decl) {
self.pending_name = Some(name);
}
let Some(init) = decl.init.as_mut() else { return };
if self.in_ignored_subtree() {
return;
}
let init_span = init.span();
if is_synthetic_span(init_span) {
return;
}
let is_named_initializer = matches!(
init,
Expression::FunctionExpression(_)
| Expression::ArrowFunctionExpression(_)
| Expression::ClassExpression(_)
);
if is_named_initializer
&& let Some(hoist_target_start) = enclosing_var_decl_hoist_target(ctx)
{
self.try_hoist_named_initializer_counter(init_span, hoist_target_start);
return;
}
if let Some(stmt_id) = self.add_statement(init_span) {
prepend_counter(init, CounterKind::stmt(self.cov_fn_name, stmt_id), ctx);
}
}
pub(super) fn try_hoist_named_initializer_counter(
&mut self,
init_span: Span,
hoist_target_start: u32,
) {
if let Some(stmt_id) = self.add_statement(init_span) {
self.pending_insertions.push(PendingInsertion {
target_start: hoist_target_start,
counter_id: stmt_id,
counter_type: CounterType::Statement,
});
}
}
pub(super) fn try_hoist_named_property_initializer(
&mut self,
prop: &PropertyDefinition<'_>,
span: Span,
) {
if let Some(name) = property_key_to_name(&prop.key) {
self.pending_name = Some(name);
}
if let Some(stmt_id) = self.add_statement(span) {
let target_start = prop.span.start;
let is_static = prop.r#static;
if let Some(top) = self.pending_class_field_hoists.last_mut() {
top.push(ClassFieldHoist { target_start, counter_id: stmt_id, is_static });
}
}
}
fn drain_pending_insertions_for_target(
&mut self,
target_start: u32,
) -> impl Iterator<Item = PendingInsertion> + '_ {
self.pending_insertions.extract_if(.., move |p| p.target_start == target_start)
}
pub(super) fn retarget_pending_insertions(&mut self, from_start: u32, to_start: u32) {
for pending in &mut self.pending_insertions {
if pending.target_start == from_start {
pending.target_start = to_start;
}
}
}
pub(super) fn inject_pending_counters_into_statement_child(
&mut self,
body: &mut Statement<'arena>,
ctx: &mut TraverseCtx<'arena, CoverageState>,
) {
if matches!(body, Statement::BlockStatement(_)) {
return;
}
let span = body.span();
if is_synthetic_span(span) {
return;
}
let pending: Vec<_> = self.drain_pending_insertions_for_target(span.start).collect();
if pending.is_empty() {
return;
}
let cov_fn = self.cov_fn_name;
let scope_id = ctx.create_child_scope_of_current(oxc_syntax::scope::ScopeFlags::empty());
let original = mem::replace(body, Statement::new_empty_statement(SPAN, ctx));
let mut stmts = ArenaVec::new_in(ctx);
for insertion in pending {
stmts.push(build_counter_stmt(CounterKind::from_pending(cov_fn, &insertion), ctx));
}
stmts.push(original);
*body = Statement::new_block_statement_with_scope_id(SPAN, stmts, scope_id, ctx);
}
}
pub(super) fn is_container_statement(stmt: &Statement<'_>) -> bool {
matches!(
stmt,
Statement::BlockStatement(_)
| Statement::EmptyStatement(_)
| Statement::FunctionDeclaration(_)
| Statement::ClassDeclaration(_)
| Statement::VariableDeclaration(_)
| Statement::ImportDeclaration(_)
| Statement::ExportNamedDeclaration(_)
| Statement::ExportDefaultDeclaration(_)
| Statement::ExportAllDeclaration(_)
| Statement::TSTypeAliasDeclaration(_)
| Statement::TSInterfaceDeclaration(_)
| Statement::TSEnumDeclaration(_)
| Statement::TSModuleDeclaration(_)
| Statement::TSImportEqualsDeclaration(_)
| Statement::TSExportAssignment(_)
| Statement::TSNamespaceExportDeclaration(_)
)
}
pub(super) fn enclosing_var_decl_hoist_target(ctx: &TraverseCtx<'_, CoverageState>) -> Option<u32> {
use oxc_traverse::Ancestor;
let mut iter = ctx.ancestors();
let var_decl_span = match iter.next()? {
Ancestor::VariableDeclarationDeclarations(a) => *a.span(),
_ => return None,
};
match iter.next()? {
Ancestor::ForStatementInit(_)
| Ancestor::ForInStatementLeft(_)
| Ancestor::ForOfStatementLeft(_) => None,
Ancestor::ExportNamedDeclarationDeclaration(a) => Some(a.span().start),
_ => Some(var_decl_span.start),
}
}