feat(data_structures): add StringExt trait (#24006)

Introduce a `StringExt` trait which adds 2 unsafe methods to `String`:

- `push_unchecked`
- `push_str_unchecked`

Both push to the string without checking there's sufficient capacity for it. The caller guarantees the push cannot cause the `String`'s length to exceed its capacity.

Both methods call the corresponding safe methods of `String` internally, with `assert_unchecked!`s preceding those calls to allow compiler to remove the bounds check. The shape of these assertions matches the code in `String`'s methods exactly.

Examination of generated assembly confirms that the optimizations work.

crates/oxc_data_structures/Cargo.toml

@@ -36,6 +36,7 @@ all = [
   "rope",
   "slice_iter",
   "stack",
+  "string_ext",
 ]
 assert_unchecked = []
 box_macros = []
@@ -46,3 +47,4 @@ non_null = []
 rope = ["dep:ropey"]
 slice_iter = ["assert_unchecked"]
 stack = []
+string_ext = ["assert_unchecked"]

crates/oxc_data_structures/README.md

@@ -16,6 +16,7 @@ This crate provides specialized data structures and utilities that are used thro
 - **Box macros**: Macros for creating boxed arrays / slices (similar to `vec!` macro)
 - **Fieldless enums macro**: Macro for creating enums with a `VARIANTS` constant listing all variants
 - **Non-null pointers**: `NonNullConst<T>` and `NonNullMut<T>` - non-null pointer types with explicit const/mut permissions
+- **`StringExt` trait**: Add methods to `String` to push without capacity checks
 
 ## Architecture
 

crates/oxc_data_structures/src/lib.rs

@@ -26,3 +26,6 @@ pub mod slice_iter;
 
 #[cfg(feature = "stack")]
 pub mod stack;
+
+#[cfg(feature = "string_ext")]
+pub mod string_ext;

crates/oxc_data_structures/src/string_ext.rs

@@ -0,0 +1,109 @@
+//! [`StringExt`] trait - extension methods for [`String`].
+
+use crate::assert_unchecked;
+
+/// Extension trait for [`String`].
+pub trait StringExt {
+    /// Append a [`char`] to the end of this `String`, without checking that the `String`
+    /// has sufficient spare capacity to hold it.
+    ///
+    /// This is a faster version of [`String::push`] which skips the check for whether the
+    /// `String`'s backing buffer needs to grow (and the cold grow/reallocation branch).
+    ///
+    /// # SAFETY
+    /// `String` must have spare capacity for `ch`
+    /// i.e. `self.len() + ch.len_utf8() <= self.capacity()`.
+    unsafe fn push_unchecked(&mut self, ch: char);
+
+    /// Append a `&str` to the end of this `String`, without checking that the `String`
+    /// has sufficient spare capacity to hold it.
+    ///
+    /// This is a faster version of [`String::push_str`] which skips the check for whether the
+    /// `String`'s backing buffer needs to grow (and the cold grow/reallocation branch).
+    ///
+    /// # SAFETY
+    /// `String` must have spare capacity for `s`
+    /// i.e. `self.len() + s.len() <= self.capacity()`.
+    unsafe fn push_str_unchecked(&mut self, s: &str);
+}
+
+impl StringExt for String {
+    #[inline]
+    unsafe fn push_unchecked(&mut self, ch: char) {
+        // SAFETY: Caller guarantees `self.len() + ch.len_utf8() <= self.capacity()`.
+        //
+        // The 2nd assertion is the promise to the compiler that no growth is needed. It is phrased as
+        // `additional <= capacity - len` (rather than `len + additional <= capacity`) to match
+        // the exact shape of the grow check inside `String::push` (`additional > capacity.wrapping_sub(len)`).
+        // LLVM does not derive the difference form from the sum form, so only this phrasing removes
+        // the capacity check and grow branch.
+        // The 1st assertion guarantees `capacity - len` does not wrap.
+        unsafe {
+            assert_unchecked!(self.len() <= self.capacity());
+            assert_unchecked!(ch.len_utf8() <= self.capacity() - self.len());
+        }
+
+        self.push(ch);
+    }
+
+    #[inline]
+    unsafe fn push_str_unchecked(&mut self, s: &str) {
+        // SAFETY: Caller guarantees `self.len() + s.len() <= self.capacity()`.
+        // See `push_unchecked` above for why the assertions are phrased this way.
+        unsafe {
+            assert_unchecked!(self.len() <= self.capacity());
+            assert_unchecked!(s.len() <= self.capacity() - self.len());
+        }
+
+        self.push_str(s);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::StringExt;
+
+    #[test]
+    fn push_unchecked_ascii() {
+        let mut s = String::with_capacity(8);
+        // SAFETY: capacity 8, pushing single bytes well within capacity
+        unsafe {
+            s.push_unchecked('a');
+            s.push_unchecked('$');
+        }
+        assert_eq!(s, "a$");
+    }
+
+    #[test]
+    fn push_unchecked_multibyte() {
+        let mut s = String::with_capacity(8);
+        // SAFETY: capacity 8; 'é' is 2 bytes, '가' is 3 bytes -> 5 bytes total <= 8
+        unsafe {
+            s.push_unchecked('é');
+            s.push_unchecked('가');
+        }
+        assert_eq!(s, "é가");
+    }
+
+    #[test]
+    fn push_str_unchecked_works() {
+        let mut s = String::with_capacity(16);
+        // SAFETY: capacity 16, "foo" + "bar" = 6 bytes <= 16
+        unsafe {
+            s.push_str_unchecked("foo");
+            s.push_str_unchecked("bar");
+        }
+        assert_eq!(s, "foobar");
+    }
+
+    #[test]
+    fn push_up_to_exact_capacity() {
+        let mut s = String::with_capacity(3);
+        // SAFETY: capacity 3; filling it exactly (2 + 1 bytes)
+        unsafe {
+            s.push_str_unchecked("ab");
+            s.push_unchecked('c');
+        }
+        assert_eq!(s, "abc");
+    }
+}