feat(data_structures): add first and first_mut methods to stack types

Author: overlookmotel

crates/oxc_data_structures/src/stack/non_empty.rs

@@ -216,6 +216,22 @@ impl<T> NonEmptyStack<T> {
         Self { cursor: start, start, end }
     }
 
+    /// Get reference to first value on stack.
+    #[inline]
+    pub fn first(&self) -> &T {
+        // SAFETY: All methods ensure stack is never empty, so `start` always points to
+        // a valid initialized `T`. `start` is always aligned for `T`.
+        unsafe { self.start.as_ref() }
+    }
+
+    /// Get mutable reference to first value on stack.
+    #[inline]
+    pub fn first_mut(&mut self) -> &mut T {
+        // SAFETY: All methods ensure stack is never empty, so `start` always points to
+        // a valid initialized `T`. `start` is always aligned for `T`.
+        unsafe { self.start.as_mut() }
+    }
+
     /// Get reference to last value on stack.
     #[inline]
     pub fn last(&self) -> &T {
@@ -545,6 +561,64 @@ mod tests {
         stack.pop();
     }
 
+    #[test]
+    fn first() {
+        let mut stack = NonEmptyStack::new(10u64);
+        assert_len_cap_last!(stack, 1, 4, &10);
+        assert_eq!(stack.first(), &10);
+
+        // Make stack grow
+        stack.push(20);
+        stack.push(30);
+        stack.push(40);
+        stack.push(50);
+        assert_len_cap_last!(stack, 5, 8, &50);
+        assert_eq!(stack.first(), &10);
+
+        // Shrink stack back to just 1 entry
+        stack.pop();
+        stack.pop();
+        stack.pop();
+        stack.pop();
+        assert_len_cap_last!(stack, 1, 8, &10);
+        assert_eq!(stack.first(), &10);
+    }
+
+    #[test]
+    fn first_mut() {
+        let mut stack = NonEmptyStack::new(10u64);
+        assert_len_cap_last!(stack, 1, 4, &10);
+        assert_eq!(stack.first_mut(), &mut 10);
+
+        *stack.first_mut() = 11;
+        assert_eq!(stack[0], 11);
+        assert_eq!(stack.first_mut(), &mut 11);
+
+        // Make stack grow
+        stack.push(20);
+        stack.push(30);
+        stack.push(40);
+        stack.push(50);
+        assert_len_cap_last!(stack, 5, 8, &50);
+        assert_eq!(stack.first_mut(), &mut 11);
+
+        *stack.first_mut() = 12;
+        assert_eq!(stack[0], 12);
+        assert_eq!(stack.first_mut(), &mut 12);
+
+        // Shrink stack back to just 1 entry
+        stack.pop();
+        stack.pop();
+        stack.pop();
+        stack.pop();
+        assert_len_cap_last!(stack, 1, 8, &12);
+        assert_eq!(stack.first_mut(), &mut 12);
+
+        *stack.first_mut() = 13;
+        assert_eq!(stack[0], 13);
+        assert_eq!(stack.first_mut(), &mut 13);
+    }
+
     #[test]
     fn last_mut() {
         let mut stack = NonEmptyStack::new(10u64);

crates/oxc_data_structures/src/stack/sparse.rs

@@ -124,6 +124,36 @@ impl<T> SparseStack<T> {
         }
     }
 
+    /// Get reference to value of first entry on the stack.
+    #[inline]
+    pub fn first(&self) -> Option<&T> {
+        let has_value = *self.has_values.first();
+        if has_value {
+            debug_assert!(!self.values.is_empty());
+            // SAFETY: First `self.has_values` is only `true` if there's a corresponding value in `self.values`.
+            // This invariant is maintained in `push`, `pop`, `take_last`, `last_or_init`, and `last_mut_or_init`.
+            let value = unsafe { self.values.get_unchecked(0) };
+            Some(value)
+        } else {
+            None
+        }
+    }
+
+    /// Get mutable reference to value of first entry on the stack.
+    #[inline]
+    pub fn first_mut(&mut self) -> Option<&mut T> {
+        let has_value = *self.has_values.first();
+        if has_value {
+            debug_assert!(!self.values.is_empty());
+            // SAFETY: First `self.has_values` is only `true` if there's a corresponding value in `self.values`.
+            // This invariant is maintained in `push`, `pop`, `take_last`, `last_or_init`, and `last_mut_or_init`.
+            let value = unsafe { self.values.get_unchecked_mut(0) };
+            Some(value)
+        } else {
+            None
+        }
+    }
+
     /// Get reference to value of last entry on the stack.
     #[inline]
     pub fn last(&self) -> Option<&T> {

crates/oxc_data_structures/src/stack/standard.rs

@@ -184,6 +184,12 @@ impl<T> Stack<T> {
         Self { cursor: start, start, end }
     }
 
+    // Note: There is no need to implement `first` and `first_mut` methods.
+    // `NonEmptyStack` can make those methods infallible, but `Stack` can't because `Stack` can be empty.
+    // `std`'s `first` and `first_mut` methods available via `Deref` / `DerefMut` to a `&[T]` / `&mut[T]`
+    // are just as efficient as a hand-written version.
+    // https://godbolt.org/z/rjb1dzob1
+
     /// Get reference to last value on stack.
     #[inline]
     pub fn last(&self) -> Option<&T> {

crates/oxc_transformer/src/common/arrow_function_converter.rs

@@ -194,18 +194,18 @@ impl<'a> Traverse<'a> for ArrowFunctionConverter<'a> {
         );
 
         debug_assert!(self.this_var_stack.len() == 1);
-        debug_assert!(self.this_var_stack.last().is_none());
+        debug_assert!(self.this_var_stack.first().is_none());
         debug_assert!(self.arguments_var_stack.len() == 1);
-        debug_assert!(self.arguments_var_stack.last().is_none());
+        debug_assert!(self.arguments_var_stack.first().is_none());
         debug_assert!(self.constructor_super_stack.len() == 1);
         // TODO: This assertion currently failing because we don't handle `super` in arrow functions
         // in class static properties correctly.
         // e.g. `class C { static f = () => super.prop; }`
-        // debug_assert!(self.constructor_super_stack.last() == &false);
+        // debug_assert!(self.constructor_super_stack.first() == &false);
         debug_assert!(self.super_methods_stack.len() == 1);
-        debug_assert!(self.super_methods_stack.last().is_empty());
+        debug_assert!(self.super_methods_stack.first().is_empty());
         debug_assert!(self.super_needs_transform_stack.len() == 1);
-        debug_assert!(self.super_needs_transform_stack.last() == &false);
+        debug_assert!(self.super_needs_transform_stack.first() == &false);
     }
 
     fn enter_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {

crates/oxc_transformer/src/typescript/enum.rs

@@ -574,9 +574,7 @@ impl IdentifierReferenceRename<'_, '_> {
         //   }
         // }
         // ```
-        //
-        // `NonEmptyStack` guarantees that the stack is not empty.
-        *self.scope_stack.first().unwrap() == symbol_scope_id
+        *self.scope_stack.first() == symbol_scope_id
             // The resolved symbol is declared outside the enum,
             // and we have checked that the name exists in previous_enum_members:
             //
@@ -586,7 +584,7 @@ impl IdentifierReferenceRename<'_, '_> {
             // enum Foo { B = A }
             //                ^ This should be renamed to Foo.A
             // ```
-        || !self.scope_stack.contains(&symbol_scope_id)
+            || !self.scope_stack.contains(&symbol_scope_id)
     }
 }