SparseStack abstraction

crates/oxc_transformer/src/es2015/arrow_functions.rs

@@ -80,7 +80,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 {
@@ -95,27 +98,16 @@ pub struct ArrowFunctionsOptions {
 pub struct ArrowFunctions<'a> {
     ctx: Ctx<'a>,
     _options: ArrowFunctionsOptions,
-    // Stack to record scopes which require a `var _this;` statement to be added.
-    // Stack is split into 2 parts for memory efficiency.
-    // `needs_this_var_stack` is pushed to for every `var`-holding scope.
-    // Initially `false`, meaning no `var _this;` statement required.
-    // When a `_this` var needs to be added, the top entry of `needs_this_var_stack` is set to `true`
-    // and details of the var to be added are pushed to `this_var_stack`.
-    // As most scopes won't need a `var _this` added, stack contains only a bool for most scopes (1 byte),
-    // rather than an `Option<BoundIdentifier>` (24 bytes).
-    needs_this_var_stack: std::vec::Vec<bool>,
-    this_var_stack: std::vec::Vec<BoundIdentifier<'a>>,
+    this_var_stack: SparseStack<BoundIdentifier<'a>>,
 }
 
 impl<'a> ArrowFunctions<'a> {
     pub fn new(options: ArrowFunctionsOptions, ctx: Ctx<'a>) -> Self {
-        Self {
-            ctx,
-            _options: options,
-            // Initial entry for `Program` scope
-            needs_this_var_stack: vec![false],
-            this_var_stack: vec![],
-        }
+        // 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_empty();
+
+        Self { ctx, _options: options, this_var_stack }
     }
 }
 
@@ -125,16 +117,16 @@ 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.needs_this_var_stack.len() == 1);
-        if *self.needs_this_var_stack.last().unwrap() {
-            self.insert_this_var_statement_at_the_top_of_statements(&mut program.body);
+        assert!(self.this_var_stack.len() == 1);
+        if let Some(this_var) = self.this_var_stack.pop() {
+            self.insert_this_var_statement_at_the_top_of_statements(&mut program.body, &this_var);
         }
         debug_assert!(self.this_var_stack.is_empty());
     }
 
     fn enter_function(&mut self, func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) {
         if func.body.is_some() {
-            self.needs_this_var_stack.push(false);
+            self.this_var_stack.push_empty();
         }
     }
 
@@ -150,22 +142,25 @@ 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 Some(body) = func.body.as_mut() else {
+        let Some(body) = &mut func.body else {
             return;
         };
 
-        if self.needs_this_var_stack.pop().unwrap() {
-            self.insert_this_var_statement_at_the_top_of_statements(&mut body.statements);
+        if let Some(this_var) = self.this_var_stack.pop() {
+            self.insert_this_var_statement_at_the_top_of_statements(
+                &mut body.statements,
+                &this_var,
+            );
         }
     }
 
     fn enter_static_block(&mut self, _block: &mut StaticBlock<'a>, _ctx: &mut TraverseCtx<'a>) {
-        self.needs_this_var_stack.push(false);
+        self.this_var_stack.push_empty();
     }
 
     fn exit_static_block(&mut self, block: &mut StaticBlock<'a>, _ctx: &mut TraverseCtx<'a>) {
-        if self.needs_this_var_stack.pop().unwrap() {
-            self.insert_this_var_statement_at_the_top_of_statements(&mut block.body);
+        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);
         }
     }
 
@@ -228,12 +223,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.
         // <https://github.com/oxc-project/oxc/pull/5840>
-        let needs_this_var = self.needs_this_var_stack.last_mut().unwrap();
-        let this_var = if *needs_this_var {
-            self.this_var_stack.last().unwrap().clone()
-        } else {
-            *needs_this_var = true;
-
+        let this_var = self.this_var_stack.get_or_init(|| {
             let target_scope_id = ctx
                 .scopes()
                 .ancestors(arrow_scope_id)
@@ -247,15 +237,13 @@ impl<'a> ArrowFunctions<'a> {
                 })
                 .unwrap();
 
-            let this_var = BoundIdentifier::new_uid(
+            BoundIdentifier::new_uid(
                 "this",
                 target_scope_id,
                 SymbolFlags::FunctionScopedVariable,
                 ctx,
-            );
-            self.this_var_stack.push(this_var.clone());
-            this_var
-        };
+            )
+        });
         Some(this_var.create_spanned_read_reference(span, ctx))
     }
 
@@ -370,9 +358,8 @@ impl<'a> ArrowFunctions<'a> {
     fn insert_this_var_statement_at_the_top_of_statements(
         &mut self,
         statements: &mut Vec<'a, Statement<'a>>,
+        this_var: &BoundIdentifier<'a>,
     ) {
-        let this_var = self.this_var_stack.pop().unwrap();
-
         let binding_pattern = self.ctx.ast.binding_pattern(
             self.ctx
                 .ast

crates/oxc_transformer/src/helpers/stack.rs

@@ -0,0 +1,94 @@
+/// Stack which is sparsely filled.
+///
+/// Functionally equivalent to a stack implemented as `Vec<Option<T>>`, 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<Option<T>>` is 24 bytes per entry (or 32 bytes if `T` has no niche).
+/// * `SparseStack<T>` 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<T> {
+    has_values: Vec<bool>,
+    values: Vec<T>,
+}
+
+impl<T> SparseStack<T> {
+    /// Create new `SparseStack`.
+    pub fn new() -> Self {
+        Self { has_values: vec![], values: vec![] }
+    }
+
+    /// Push an entry to the stack.
+    #[expect(dead_code)]
+    pub fn push(&mut self, value: Option<T>) {
+        let has_value = if let Some(value) = value {
+            self.values.push(value);
+            true
+        } else {
+            false
+        };
+        self.has_values.push(has_value);
+    }
+
+    /// Push an empty entry to the stack.
+    pub fn push_empty(&mut self) {
+        self.has_values.push(false);
+    }
+
+    /// Pop last entry from the stack.
+    ///
+    /// # Panics
+    /// Panics if the stack is empty.
+    pub fn pop(&mut self) -> Option<T> {
+        let has_value = self.has_values.pop().unwrap();
+        if has_value {
+            // SAFETY: `self.has_values` only contains `true` if there's a corresponding value
+            // in `self.values`. This invariant is maintained in `push` 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 entry = unsafe { self.values.pop().unwrap_unchecked() };
+            Some(entry)
+        } else {
+            None
+        }
+    }
+
+    /// Initialize the value for top entry on the stack, if it has no value already.
+    /// Returns reference to value.
+    ///
+    /// # Panics
+    /// Panics if the stack is empty.
+    pub fn get_or_init<I: FnMut() -> T>(&mut self, mut init: I) -> &T {
+        let has_value = self.has_values.last_mut().unwrap();
+        if !*has_value {
+            *has_value = true;
+            self.values.push(init());
+        }
+
+        // SAFETY: `self.has_values` only contains `true` if there's a corresponding value
+        // in `self.values`. This invariant is maintained in `push` and `pop`.
+        // Here either `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.
+    pub fn len(&self) -> usize {
+        self.has_values.len()
+    }
+
+    /// Returns `true` if stack is empty.
+    pub fn is_empty(&self) -> bool {
+        self.has_values.is_empty()
+    }
+}

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};