diff --git a/crates/oxc_transformer/src/es2015/arrow_functions.rs b/crates/oxc_transformer/src/es2015/arrow_functions.rs index 010723e29a486..4489dc263527e 100644 --- a/crates/oxc_transformer/src/es2015/arrow_functions.rs +++ b/crates/oxc_transformer/src/es2015/arrow_functions.rs @@ -82,7 +82,10 @@ use oxc_syntax::{ use oxc_traverse::{Ancestor, Traverse, TraverseCtx}; use serde::Deserialize; -use crate::{context::Ctx, helpers::bindings::BoundIdentifier}; +use crate::{ + context::Ctx, + helpers::{bindings::BoundIdentifier, stack::SparseStack}, +}; #[derive(Debug, Default, Clone, Deserialize)] pub struct ArrowFunctionsOptions { @@ -97,17 +100,16 @@ pub struct ArrowFunctionsOptions { pub struct ArrowFunctions<'a> { ctx: Ctx<'a>, _options: ArrowFunctionsOptions, - this_var_stack: std::vec::Vec>>, + this_var_stack: SparseStack>, } impl<'a> ArrowFunctions<'a> { pub fn new(options: ArrowFunctionsOptions, ctx: Ctx<'a>) -> Self { - Self { - ctx, - _options: options, - // Initial entry for `Program` scope - this_var_stack: vec![None], - } + // Init stack with empty entry for `Program` (instead of pushing entry in `enter_program`) + let mut this_var_stack = SparseStack::new(); + this_var_stack.push(None); + + Self { ctx, _options: options, this_var_stack } } } @@ -118,10 +120,10 @@ impl<'a> Traverse<'a> for ArrowFunctions<'a> { /// Insert `var _this = this;` for the global scope. fn exit_program(&mut self, program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) { assert!(self.this_var_stack.len() == 1); - let this_var = self.this_var_stack.pop().unwrap(); - if let Some(this_var) = this_var { + if let Some(this_var) = self.this_var_stack.take() { self.insert_this_var_statement_at_the_top_of_statements(&mut program.body, &this_var); } + debug_assert!(self.this_var_stack.len() == 1); } fn enter_function(&mut self, _func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) { @@ -140,8 +142,7 @@ impl<'a> Traverse<'a> for ArrowFunctions<'a> { /// ``` /// Insert the var _this = this; statement outside the arrow function fn exit_function(&mut self, func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) { - let this_var = self.this_var_stack.pop().unwrap(); - if let Some(this_var) = this_var { + if let Some(this_var) = self.this_var_stack.pop() { let Some(body) = &mut func.body else { unreachable!() }; self.insert_this_var_statement_at_the_top_of_statements( @@ -156,8 +157,7 @@ impl<'a> Traverse<'a> for ArrowFunctions<'a> { } fn exit_static_block(&mut self, block: &mut StaticBlock<'a>, _ctx: &mut TraverseCtx<'a>) { - let this_var = self.this_var_stack.pop().unwrap(); - if let Some(this_var) = this_var { + if let Some(this_var) = self.this_var_stack.pop() { self.insert_this_var_statement_at_the_top_of_statements(&mut block.body, &this_var); } } @@ -221,8 +221,7 @@ impl<'a> ArrowFunctions<'a> { // `this` can be in scope at a time. We could create a single `_this` UID and reuse it in each // scope. But this does not match output for some of Babel's test cases. // - let this_var = self.this_var_stack.last_mut().unwrap(); - if this_var.is_none() { + let this_var = self.this_var_stack.get_or_init(|| { let target_scope_id = ctx .scopes() .ancestors(arrow_scope_id) @@ -236,14 +235,13 @@ impl<'a> ArrowFunctions<'a> { }) .unwrap(); - this_var.replace(BoundIdentifier::new_uid( + BoundIdentifier::new_uid( "this", target_scope_id, SymbolFlags::FunctionScopedVariable, ctx, - )); - } - let this_var = this_var.as_ref().unwrap(); + ) + }); Some(this_var.create_spanned_read_reference(span, ctx)) } diff --git a/crates/oxc_transformer/src/helpers/stack.rs b/crates/oxc_transformer/src/helpers/stack.rs new file mode 100644 index 0000000000000..53ca9b039398c --- /dev/null +++ b/crates/oxc_transformer/src/helpers/stack.rs @@ -0,0 +1,115 @@ +/// Stack which is sparsely filled. +/// +/// Functionally equivalent to a stack implemented as `Vec>`, but more memory-efficient +/// in cases where majority of entries in the stack will be empty (`None`). +/// +/// The stack is stored as 2 arrays: +/// 1. `has_values` - Records whether an entry on the stack has a value or not (`Some` or `None`). +/// 2. `values` - Where the stack entry *does* have a value, it's stored in this array. +/// +/// Memory is only consumed for values where values exist. +/// +/// Where value (`T`) is large, and most entries have no value, this will be a significant memory saving. +/// +/// e.g. if `T` is 24 bytes, and 90% of stack entries have no values: +/// * `Vec>` is 24 bytes per entry (or 32 bytes if `T` has no niche). +/// * `SparseStack` is 4 bytes per entry. +/// +/// When the stack grows and reallocates, `SparseStack` has less memory to copy, which is a performance +/// win too. +pub struct SparseStack { + has_values: Vec, + values: Vec, +} + +impl SparseStack { + /// Create new `SparseStack`. + pub fn new() -> Self { + Self { has_values: vec![], values: vec![] } + } + + /// Push an entry to the stack. + #[inline] + pub fn push(&mut self, value: Option) { + let has_value = if let Some(value) = value { + self.values.push(value); + true + } else { + false + }; + self.has_values.push(has_value); + } + + /// Pop last entry from the stack. + /// + /// # Panics + /// Panics if the stack is empty. + pub fn pop(&mut self) -> Option { + let has_value = self.has_values.pop().unwrap(); + if has_value { + debug_assert!(!self.values.is_empty()); + // SAFETY: Last `self.has_values` is only `true` if there's a corresponding value in `self.values`. + // This invariant is maintained in `push`, `take`, and `get_or_init`. + // We maintain it here too because we just popped from `self.has_values`, so that `true` + // has been consumed at the same time we consume its corresponding value from `self.values`. + let value = unsafe { self.values.pop().unwrap_unchecked() }; + Some(value) + } else { + None + } + } + + /// Take value from last entry on the stack. + /// + /// # Panics + /// Panics if the stack is empty. + pub fn take(&mut self) -> Option { + let has_value = self.has_values.last_mut().unwrap(); + if *has_value { + *has_value = false; + + debug_assert!(!self.values.is_empty()); + // SAFETY: Last `self.has_values` is only `true` if there's a corresponding value in `self.values`. + // This invariant is maintained in `push`, `pop`, and `get_or_init`. + // We maintain it here too because we just set last `self.has_values` to `false` + // at the same time as we consume the corresponding value from `self.values`. + let value = unsafe { self.values.pop().unwrap_unchecked() }; + Some(value) + } else { + None + } + } + + /// Initialize the value for top entry on the stack, if it has no value already. + /// Return reference to value. + /// + /// # Panics + /// Panics if the stack is empty. + pub fn get_or_init T>(&mut self, init: I) -> &T { + let has_value = self.has_values.last_mut().unwrap(); + if !*has_value { + *has_value = true; + self.values.push(init()); + } + + debug_assert!(!self.values.is_empty()); + // SAFETY: Last `self.has_values` is only `true` if there's a corresponding value in `self.values`. + // This invariant is maintained in `push`, `pop`, and `take`. + // Here either last `self.has_values` was already `true`, or it's just been set to `true` + // and a value pushed to `self.values` above. + unsafe { self.values.last().unwrap_unchecked() } + } + + /// Get number of entries on the stack. + #[inline] + pub fn len(&self) -> usize { + self.has_values.len() + } + + /// Returns `true` if stack is empty. + #[inline] + #[expect(dead_code)] + pub fn is_empty(&self) -> bool { + self.has_values.is_empty() + } +} diff --git a/crates/oxc_transformer/src/lib.rs b/crates/oxc_transformer/src/lib.rs index 6f191a6efaad9..f995ae719c6f0 100644 --- a/crates/oxc_transformer/src/lib.rs +++ b/crates/oxc_transformer/src/lib.rs @@ -27,6 +27,7 @@ mod typescript; mod helpers { pub mod bindings; pub mod module_imports; + pub mod stack; } use std::{path::Path, rc::Rc};