feat(data_structures): add CodeBuffer::print_strs_array (#19760)

Add a method `CodeBuffer::print_strs_array`. This works much like `StringBuilder::from_strs_array_in` and `Atom::from_strs_array_in`. It takes an array of `&str`s, computes the total length required to print them all, reserves space for the total number of bytes for all together, and then writes all the strings one after another.

The result is that it's much faster than a series of `print_str` calls, which incur the cost of a bounds check and branch for each string. Instead you only do 1 bounds check for the whole batch.

Like `StringBuilder::from_strs_array_in` there have to be additional checks when computing the total length to ensure the sum of lengths does not exceed `usize::MAX` and wrap around. But when some of the strings are statically known (e.g. `buffer.print_strs_array(["foo", dynamic, "bar"])`) and so it's impossible for overflow to occur, these checks can usually be removed by compiler.
This commit is contained in:
overlookmotel committed 2026-02-26 23:53:31 +00:00
1 parent 845da351a5
commit 46994986e9
1 file changed
+158 -1
+158 -1
View File
@@ -1,6 +1,6 @@
//! A string builder for constructing source code.
use std::iter;
use std::{iter, ptr};
use crate::assert_unchecked;
@@ -449,6 +449,138 @@ impl CodeBuffer {
self.buf.extend_from_slice(bytes);
}
/// Print a series of strings into the buffer.
///
/// This is more efficient than making multiple [`print_str`] calls, because it only does one bounds check,
/// instead of one per string.
///
/// # Panics
///
/// Panics if the sum of length of all strings exceeds `usize::MAX`.
///
/// # Example
/// ```
/// # use oxc_data_structures::code_buffer::CodeBuffer;
/// let mut json = CodeBuffer::new();
/// json.print_str("[");
///
/// const TYPE: &str = "Thing";
/// let names = ["foo", "bar", "baz"];
///
/// for (index, name) in names.into_iter().enumerate() {
/// if index == 0 {
/// json.print_strs_array([r#"{"type":""#, TYPE, r#"","name":""#, name, r#""}"#]);
/// } else {
/// json.print_strs_array([",", r#"{"type":""#, TYPE, r#"","name":""#, name, r#""}"#]);
/// }
/// }
///
/// json.print_str("]");
///
/// let json = json.into_string();
/// assert_eq!(
/// json,
/// r#"[{"type":"Thing","name":"foo"},{"type":"Thing","name":"bar"},{"type":"Thing","name":"baz"}]"#
/// );
/// ```
///
/// [`print_str`]: Self::print_str
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn print_strs_array<const N: usize>(&mut self, strings: [&str; N]) {
// Calculate total length of all the strings concatenated.
//
// We have to use `checked_add` here to guard against additions wrapping around
// if some of the input `&str`s are very long, or there's many of them.
//
// However, `&str`s have max length of `isize::MAX`.
// https://users.rust-lang.org/t/does-str-reliably-have-length-isize-max/126777
// Use `assert_unchecked!` to communicate this invariant to compiler, which allows it to
// optimize out the overflow checks where some of `strings` are static, so their size is known.
//
// e.g. `buf.print_strs_array(["\"", name, "\"", ":"])`, for example,
// requires no checks at all, because the static parts have total length of 3 bytes,
// and `name` has max length of `isize::MAX`. `isize::MAX as usize + 3` cannot overflow `usize`.
// Compiler can see that, and removes the overflow check.
// https://godbolt.org/z/MGh44Yz5d
#[expect(clippy::checked_conversions)]
let total_strings_len = strings.iter().fold(0usize, |total_len, s| {
let len = s.len();
// SAFETY: `&str`s have maximum length of `isize::MAX`
unsafe { assert_unchecked!(len <= (isize::MAX as usize)) };
total_len.checked_add(len).unwrap()
});
// Do actual pushing of the strings into `buf` in a separate function, to ensure that `print_strs_array`
// is inlined, so that compiler has knowledge to remove the overflow checks above.
// When some of `strings` are static, this function is usually only a few instructions.
// Compiler can choose whether or not to inline `print_strs_array_with_total_len`.
// SAFETY: `total_strings_len` has been calculated correctly above.
unsafe { self.print_strs_array_with_total_len(strings, total_strings_len) }
}
/// Print a series of strings into the buffer, with provided `total_strings_len`.
///
/// # SAFETY
/// `total_strings_len` must be the total length of all `strings` concatenated.
unsafe fn print_strs_array_with_total_len<const N: usize>(
&mut self,
strings: [&str; N],
total_strings_len: usize,
) {
// Reserve `total_strings_len` bytes for the strings to be written into
self.reserve(total_strings_len);
// Write each string into `buf`, without bounds checks
let current_len = self.buf.len();
// SAFETY: `buf.ptr + buf.len` is within bounds of `buf`'s allocation.
let start_ptr = unsafe { self.buf.as_mut_ptr().add(current_len) };
let mut dst_ptr = start_ptr;
// Get new length of `buf` after all `strings` are written.
// `current_len + total_strings_len` must be `<= isize::MAX` otherwise `self.reserve(total_strings_len)`
// would have panicked. Therefore this addition cannot wrap around.
let new_len = current_len + total_strings_len;
#[cfg(not(debug_assertions))]
// SAFETY: `new_len` is original length of `buf` plus the total length of all strings concatenated.
// We're going to write those strings into `buf`, so all bytes in `buf` up to `new_len` will be initialized.
// In release mode, nothing in the code in the rest of this function can panic, so we can write the new length
// early, which should be a little bit cheaper than writing it after, as the register containing `new_len`
// is now free, and can be re-used in the code below.
unsafe {
self.buf.set_len(new_len);
}
for str in strings {
let src_ptr = str.as_ptr();
let len = str.len();
// SAFETY:
// `src` is obtained from a `&str` with length `len`, so is valid for reading `len` bytes.
// `dst_ptr` must be within bounds of `buf`'s allocation, because we reserved sufficient space.
// for all the strings. So is `dst_ptr + len`.
// `u8` has no alignment requirements, so `src_ptr` and `dst_ptr` are sufficiently aligned.
// No overlapping, because we're copying from an existing `&str` to newly reserved space in `buf`.
unsafe { ptr::copy_nonoverlapping(src_ptr, dst_ptr, len) };
// SAFETY: We allocated sufficient capacity for all the strings concatenated.
// So `dst_ptr.add(len)` cannot go out of bounds.
dst_ptr = unsafe { dst_ptr.add(len) };
}
debug_assert_eq!(dst_ptr as usize - start_ptr as usize, total_strings_len);
#[cfg(debug_assertions)]
// SAFETY: `new_len` is original length of `buf` plus the total length of all strings concatenated.
// We've written those strings into `buf`, so all bytes in `buf` up to `new_len` are now initialized.
// In debug mode, we need to write the new length last, because the `debug_assert!` above could panic.
unsafe {
self.buf.set_len(new_len);
}
}
/// Print a sequence of bytes, without checking that the buffer still contains a valid UTF-8 string.
///
/// # SAFETY
@@ -673,6 +805,31 @@ mod test {
assert_eq!(String::from(code), "foo");
}
#[test]
fn print_strs_array() {
// Standard operation
let mut code = CodeBuffer::new();
let name = "foo";
let value = "123";
code.print_strs_array(["const ", name, " = ", value, ";"]);
assert_eq!(String::from(code), "const foo = 123;");
// Empty array
let mut code = CodeBuffer::new();
code.print_strs_array([]);
assert_eq!(String::from(code), "");
// Array of empty strings
let mut code = CodeBuffer::new();
code.print_strs_array(["", "", ""]);
assert_eq!(String::from(code), "");
// Array containing empty and non-empty strings
let mut code = CodeBuffer::new();
code.print_strs_array(["", "foo", "", "bar", "", "qux", ""]);
assert_eq!(String::from(code), "foobarqux");
}
#[test]
#[expect(clippy::byte_char_slices)]
fn print_bytes_iter_unchecked() {