From 7b90d794d101feea3cee2db774d6bb7f18b76589 Mon Sep 17 00:00:00 2001 From: overlookmotel <557937+overlookmotel@users.noreply.github.com> Date: Mon, 23 Sep 2024 07:52:47 +0000 Subject: [PATCH] perf(transformer): `SparseStack` always keep minimum 1 entry (#5962) Optimize `SparseStack` (which was introduced in #5940). Initialize it with a single entry, and ensure the stack is never emptied. This makes `take`, `get_or_init` and `get_mut_or_init` methods infallible, since there is always an entry on the stack to read. --- .../src/es2015/arrow_functions.rs | 7 +- .../src/es2016/exponentiation_operator.rs | 2 +- .../src/es2020/nullish_coalescing_operator.rs | 2 +- .../es2021/logical_assignment_operators.rs | 2 +- crates/oxc_transformer/src/helpers/stack.rs | 67 +++++++++++++------ 5 files changed, 51 insertions(+), 29 deletions(-) diff --git a/crates/oxc_transformer/src/es2015/arrow_functions.rs b/crates/oxc_transformer/src/es2015/arrow_functions.rs index 4489dc263527e..5ef3bedcff400 100644 --- a/crates/oxc_transformer/src/es2015/arrow_functions.rs +++ b/crates/oxc_transformer/src/es2015/arrow_functions.rs @@ -105,11 +105,8 @@ pub struct ArrowFunctions<'a> { impl<'a> ArrowFunctions<'a> { pub fn new(options: ArrowFunctionsOptions, ctx: Ctx<'a>) -> Self { - // 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 } + // `SparseStack` is created with 1 empty entry, for `Program` + Self { ctx, _options: options, this_var_stack: SparseStack::new() } } } diff --git a/crates/oxc_transformer/src/es2016/exponentiation_operator.rs b/crates/oxc_transformer/src/es2016/exponentiation_operator.rs index cc23ffac005be..5545cc8acf486 100644 --- a/crates/oxc_transformer/src/es2016/exponentiation_operator.rs +++ b/crates/oxc_transformer/src/es2016/exponentiation_operator.rs @@ -65,7 +65,7 @@ impl<'a> ExponentiationOperator<'a> { impl<'a> Traverse<'a> for ExponentiationOperator<'a> { #[inline] // Inline because it's no-op in release mode fn exit_program(&mut self, _program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) { - debug_assert!(self.var_declarations.is_empty()); + debug_assert!(self.var_declarations.len() == 1); } fn enter_statements( diff --git a/crates/oxc_transformer/src/es2020/nullish_coalescing_operator.rs b/crates/oxc_transformer/src/es2020/nullish_coalescing_operator.rs index 1133b326c1c99..d182d635999e1 100644 --- a/crates/oxc_transformer/src/es2020/nullish_coalescing_operator.rs +++ b/crates/oxc_transformer/src/es2020/nullish_coalescing_operator.rs @@ -51,7 +51,7 @@ impl<'a> NullishCoalescingOperator<'a> { impl<'a> Traverse<'a> for NullishCoalescingOperator<'a> { #[inline] // Inline because it's no-op in release mode fn exit_program(&mut self, _program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) { - debug_assert!(self.var_declarations.is_empty()); + debug_assert!(self.var_declarations.len() == 1); } fn enter_statements( diff --git a/crates/oxc_transformer/src/es2021/logical_assignment_operators.rs b/crates/oxc_transformer/src/es2021/logical_assignment_operators.rs index 5ec9f86ebe9a0..1ff16508d36c5 100644 --- a/crates/oxc_transformer/src/es2021/logical_assignment_operators.rs +++ b/crates/oxc_transformer/src/es2021/logical_assignment_operators.rs @@ -76,7 +76,7 @@ impl<'a> LogicalAssignmentOperators<'a> { impl<'a> Traverse<'a> for LogicalAssignmentOperators<'a> { #[inline] // Inline because it's no-op in release mode fn exit_program(&mut self, _program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) { - debug_assert!(self.var_declarations.is_empty()); + debug_assert!(self.var_declarations.len() == 1); } fn enter_statements( diff --git a/crates/oxc_transformer/src/helpers/stack.rs b/crates/oxc_transformer/src/helpers/stack.rs index 1417b311e0ef2..c68fb2744798b 100644 --- a/crates/oxc_transformer/src/helpers/stack.rs +++ b/crates/oxc_transformer/src/helpers/stack.rs @@ -3,6 +3,10 @@ /// 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`). /// +/// Stack is initialized with a single entry which can never be popped off. +/// If `Program` has a entry on the stack, can use this initial entry for it. Get value for `Program` +/// in `exit_program` visitor with `SparseStack::take` instead of `SparseStack::pop`. +/// /// 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. @@ -25,7 +29,10 @@ pub struct SparseStack { impl SparseStack { /// Create new `SparseStack`. pub fn new() -> Self { - Self { has_values: vec![], values: vec![] } + // `has_values` starts with a single empty entry, which will never be popped off. + // This means `take`, `get_or_init`, and `get_mut_or_init` can all be infallible, + // as there's always an entry on the stack to read. + Self { has_values: vec![false], values: vec![] } } /// Push an entry to the stack. @@ -43,9 +50,28 @@ impl SparseStack { /// Pop last entry from the stack. /// /// # Panics - /// Panics if the stack is empty. + /// Panics if the stack has only 1 entry on it. + #[inline] pub fn pop(&mut self) -> Option { - let has_value = self.has_values.pop().unwrap(); + // SAFETY: `self.has_values` starts with 1 entry. Only `pop` removes entries from it. + // We check that popping an entry does not leave the stack empty before performing the pop. + // So `self.has_values` can never be left in an empty state. + // + // This would be equivalent: + // ``` + // assert!(self.has_values.len() > 1); + // self.has_values.pop().unwrap() + // ``` + // But checking `original_len > 1` is 1 more CPU op than decrementing length first, + // and then checking for `new_len > 0`. https://godbolt.org/z/eqx385E5K + let has_value = unsafe { + let new_len = self.has_values.len() - 1; + assert!(new_len > 0); + let has_value = *self.has_values.get_unchecked(new_len); + self.has_values.set_len(new_len); + has_value + }; + 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`. @@ -60,11 +86,12 @@ impl SparseStack { } /// Take value from last entry on the stack. - /// - /// # Panics - /// Panics if the stack is empty. + #[inline] pub fn take(&mut self) -> Option { - let has_value = self.has_values.last_mut().unwrap(); + debug_assert!(!self.has_values.is_empty()); + // SAFETY: `self.has_values` starts with 1 entry. Only `pop` removes entries from it, + // and it ensures `self.has_values` always has at least one entry. + let has_value = unsafe { self.has_values.last_mut().unwrap_unchecked() }; if *has_value { *has_value = false; @@ -82,11 +109,12 @@ impl SparseStack { /// 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. + #[inline] pub fn get_or_init T>(&mut self, init: I) -> &T { - let has_value = self.has_values.last_mut().unwrap(); + debug_assert!(!self.has_values.is_empty()); + // SAFETY: `self.has_values` starts with 1 entry. Only `pop` removes entries from it, + // and it ensures `self.has_values` always has at least one entry. + let has_value = unsafe { self.has_values.last_mut().unwrap_unchecked() }; if !*has_value { *has_value = true; self.values.push(init()); @@ -102,11 +130,12 @@ impl SparseStack { /// Initialize the value for top entry on the stack, if it has no value already. /// Return mutable reference to value. - /// - /// # Panics - /// Panics if the stack is empty. + #[inline] pub fn get_mut_or_init T>(&mut self, init: I) -> &mut T { - let has_value = self.has_values.last_mut().unwrap(); + debug_assert!(!self.has_values.is_empty()); + // SAFETY: `self.has_values` starts with 1 entry. Only `pop` removes entries from it, + // and it ensures `self.has_values` always has at least one entry. + let has_value = unsafe { self.has_values.last_mut().unwrap_unchecked() }; if !*has_value { *has_value = true; self.values.push(init()); @@ -120,14 +149,10 @@ impl SparseStack { } /// Get number of entries on the stack. + /// + /// Number of entries is always at least 1. Stack is never empty. #[inline] pub fn len(&self) -> usize { self.has_values.len() } - - /// Returns `true` if stack is empty. - #[inline] - pub fn is_empty(&self) -> bool { - self.has_values.is_empty() - } }