perf(data_structures): copy short strings without a memcpy call (#27953)

Fast path for short strings.

---

According to local benchmark, it seemed to be effective for the formatter as well, but on CodSpeed, it appears to be effective only for codegen... 🤔

crates/oxc_codegen/src/lib.rs

@@ -354,7 +354,7 @@ impl<'a> Codegen<'a> {
     /// Push str into the buffer
     #[inline]
     pub fn print_str(&mut self, s: &str) {
-        self.code.print_str(s);
+        self.code.print_str_short(s);
     }
 
     /// Print an identifier name (binding, reference, label, private name, property key),
@@ -362,7 +362,7 @@ impl<'a> Codegen<'a> {
     #[inline]
     pub fn print_name(&mut self, s: &str) {
         if !self.options.ascii_only || s.is_ascii() {
-            self.code.print_str(s);
+            self.code.print_str_short(s);
         } else {
             self.print_non_ascii_escaped(s, NonAsciiEscape::Identifier);
             self.need_space_before_identifier = self.code.len();

crates/oxc_data_structures/src/code_buffer.rs

@@ -395,6 +395,100 @@ impl CodeBuffer {
         self.buf.extend_from_slice(s.as_ref().as_bytes());
     }
 
+    /// Push a string into the buffer, like [`print_str`],
+    /// with a fast path for a string of at most 16 bytes (most tokens and identifiers).
+    ///
+    /// The fast path copies with at most two overlapping loads and stores of the same width,
+    /// instead of a `memcpy` call.
+    ///
+    /// # Example
+    /// ```
+    /// # use oxc_data_structures::code_buffer::CodeBuffer;
+    /// let mut code = CodeBuffer::with_capacity(16);
+    /// code.print_str_short("let");
+    /// code.print_str_short(" identifier");
+    /// assert_eq!(code.as_str(), "let identifier");
+    /// ```
+    ///
+    /// [`print_str`]: CodeBuffer::print_str
+    #[inline]
+    pub fn print_str_short(&mut self, s: &str) {
+        // SAFETY: `s` is a `str`, so its bytes are valid UTF-8
+        unsafe { self.print_bytes_short_unchecked(s.as_bytes()) };
+    }
+
+    /// Push a slice of bytes into the buffer, like [`print_bytes_unchecked`],
+    /// with the fast path of [`print_str_short`].
+    ///
+    /// # SAFETY
+    /// Same as [`print_bytes_unchecked`]: the buffer must stay valid UTF-8
+    /// when it is read with [`as_str`] or converted with [`into_string`].
+    ///
+    /// # Example
+    /// ```
+    /// # use oxc_data_structures::code_buffer::CodeBuffer;
+    /// let mut code = CodeBuffer::new();
+    ///
+    /// // SAFETY: All bytes in this slice are ASCII
+    /// unsafe {
+    ///     code.print_bytes_short_unchecked("abcd".as_bytes());
+    /// }
+    /// ```
+    ///
+    /// [`print_bytes_unchecked`]: CodeBuffer::print_bytes_unchecked
+    /// [`print_str_short`]: CodeBuffer::print_str_short
+    /// [`as_str`]: CodeBuffer::as_str
+    /// [`into_string`]: CodeBuffer::into_string
+    #[inline]
+    pub unsafe fn print_bytes_short_unchecked(&mut self, bytes: &[u8]) {
+        #[cold]
+        #[inline(never)]
+        fn write_slow(code_buffer: &mut CodeBuffer, bytes: &[u8]) {
+            code_buffer.buf.extend_from_slice(bytes);
+        }
+
+        /// Copies `size_of::<T>()` bytes at `0` and at `n - size_of::<T>()`, which covers all `n` bytes.
+        ///
+        /// # SAFETY
+        /// `size_of::<T>() <= n`, `src` is readable and `dst` is writable for `n` bytes.
+        #[inline]
+        unsafe fn copy_overlapping<T: Copy>(src: *const u8, dst: *mut u8, n: usize) {
+            let tail = n - size_of::<T>();
+            // SAFETY: Both copies are within the first `n` bytes of `src` and `dst`
+            unsafe {
+                dst.cast::<T>().write_unaligned(src.cast::<T>().read_unaligned());
+                dst.add(tail)
+                    .cast::<T>()
+                    .write_unaligned(src.add(tail).cast::<T>().read_unaligned());
+            }
+        }
+
+        let n = bytes.len();
+        let len = self.len();
+        if n > 16 || self.capacity() - len < n {
+            write_slow(self, bytes);
+            return;
+        }
+
+        let src = bytes.as_ptr();
+        // SAFETY: There are at least `n` bytes spare capacity after `len`,
+        // and each copy is no wider than `n`.
+        // The caller guarantees that the buffer stays valid UTF-8.
+        unsafe {
+            let dst = self.buf.as_mut_ptr().add(len);
+            if n >= 8 {
+                copy_overlapping::<u64>(src, dst, n);
+            } else if n >= 4 {
+                copy_overlapping::<u32>(src, dst, n);
+            } else if n >= 2 {
+                copy_overlapping::<u16>(src, dst, n);
+            } else if n == 1 {
+                *dst = *src;
+            }
+            self.buf.set_len(len + n);
+        }
+    }
+
     /// Push a sequence of ASCII characters into the buffer.
     ///
     /// # Panics
@@ -828,6 +922,21 @@ mod test {
         assert_eq!(String::from(code), "foo");
     }
 
+    #[test]
+    fn print_str_short() {
+        let text = "abcdefghijklmnopqrstu";
+        for n in 0..=text.len() {
+            // Enough spare capacity, and none
+            for capacity in [32, 1] {
+                let mut code = CodeBuffer::with_capacity(capacity);
+                code.print_str("x");
+                code.print_str_short(&text[..n]);
+                code.print_str_short("é");
+                assert_eq!(code.as_str(), format!("x{}é", &text[..n]));
+            }
+        }
+    }
+
     #[test]
     fn print_strs_array() {
         // Standard operation

crates/oxc_estree/src/serialize/strings.rs

@@ -289,7 +289,7 @@ fn write_str<T: EscapeTable>(s: &str, buffer: &mut CodeBuffer) {
         unsafe {
             let len = current_ptr.offset_from_unsigned(chunk_start_ptr);
             let chunk = slice::from_raw_parts(chunk_start_ptr, len);
-            buffer.print_bytes_unchecked(chunk);
+            buffer.print_bytes_short_unchecked(chunk);
         }
 
         write_char_escape(escape, byte, buffer);
@@ -310,7 +310,7 @@ fn write_str<T: EscapeTable>(s: &str, buffer: &mut CodeBuffer) {
     unsafe {
         let len = iter.end_ptr().offset_from_unsigned(chunk_start_ptr);
         let chunk = slice::from_raw_parts(chunk_start_ptr, len);
-        buffer.print_bytes_unchecked(chunk);
+        buffer.print_bytes_short_unchecked(chunk);
     }
 
     buffer.print_ascii_byte(b'"');

crates/oxc_formatter_core/src/printer/mod.rs

@@ -831,24 +831,21 @@ impl<'a> Printer<'a> {
 
         match text {
             Text::Token(text) => {
-                // SAFETY: `text` is a ASCII-only string
-                unsafe {
-                    self.state.buffer.print_bytes_unchecked(text.as_bytes());
-                }
+                self.state.buffer.print_str_short(text);
                 self.state.line_width += text.len();
             }
             Text::Text { text, width } => {
                 if width.is_multiline() {
                     let line_break_position = text.find('\n').unwrap_or(text.len());
                     let (first_line, remaining) = text.split_at(line_break_position);
-                    self.state.buffer.print_str(first_line);
+                    self.state.buffer.print_str_short(first_line);
                     self.state.line_width += width.value() as usize;
                     // Print the remaining lines
                     for char in remaining.chars() {
                         self.print_char(char);
                     }
                 } else {
-                    self.state.buffer.print_str(text);
+                    self.state.buffer.print_str_short(text);
                     self.state.line_width += width.value() as usize;
                 }
             }