diff --git a/crates/oxc_ast/src/ast_builder_impl.rs b/crates/oxc_ast/src/ast_builder_impl.rs index b2b4b8edfd086..695ac064c2d58 100644 --- a/crates/oxc_ast/src/ast_builder_impl.rs +++ b/crates/oxc_ast/src/ast_builder_impl.rs @@ -1,4 +1,4 @@ -use std::borrow::Cow; +use std::{alloc::Layout, borrow::Cow, mem::MaybeUninit, slice, str}; use oxc_allocator::{Allocator, AllocatorAccessor, Box, FromIn, IntoIn, Vec}; use oxc_span::{SPAN, Span}; @@ -332,44 +332,60 @@ fn escape_template_element_raw<'a>(raw: &str, ast: AstBuilder<'a>) -> Str<'a> { return ast.str(raw); } - // Allocate directly in arena + // Allocate directly in arena. + // It's impossible for this addition to overflow, because max length of a `&str` is `isize::MAX` + // and we've at most doubled the length, which cannot overflow `usize::MAX`. let len = bytes.len() + extra_bytes; - let layout = std::alloc::Layout::array::(len).unwrap(); + let layout = Layout::array::(len).unwrap(); let ptr = ast.allocator.alloc_layout(layout); - // SAFETY: `ptr` points to `len` bytes of valid memory allocated by the arena. - // Input is valid UTF-8, we only escape ASCII bytes, so output is also valid UTF-8. - unsafe { - let escaped = std::slice::from_raw_parts_mut(ptr.as_ptr(), len); - let mut j = 0; - for i in 0..bytes.len() { + // SAFETY: `ptr` points to `len` bytes of memory allocated by the arena. + // `MaybeUninit` has the same layout as `u8` and does not require its contents to be initialized, + // so it's sound to form a `&mut [MaybeUninit]` over this uninitialized memory. + let dest = unsafe { slice::from_raw_parts_mut(ptr.as_ptr().cast::>(), len) }; + + let mut j = 0; + for i in 0..bytes.len() { + // SAFETY: For each input byte we write either 1 or 2 bytes, and `len` was sized to fit + // exactly that many bytes, so `j` and `j + 1` are always in bounds. + // Note: Compiler merges each pair of writes into a single 2-byte write. + unsafe { match bytes[i] { b'\\' => { - *escaped.get_unchecked_mut(j) = b'\\'; - *escaped.get_unchecked_mut(j + 1) = b'\\'; + dest.get_unchecked_mut(j).write(b'\\'); + dest.get_unchecked_mut(j + 1).write(b'\\'); j += 2; } b'`' => { - *escaped.get_unchecked_mut(j) = b'\\'; - *escaped.get_unchecked_mut(j + 1) = b'`'; + dest.get_unchecked_mut(j).write(b'\\'); + dest.get_unchecked_mut(j + 1).write(b'`'); j += 2; } b'$' if bytes.get(i + 1) == Some(&b'{') => { - *escaped.get_unchecked_mut(j) = b'\\'; - *escaped.get_unchecked_mut(j + 1) = b'$'; + dest.get_unchecked_mut(j).write(b'\\'); + dest.get_unchecked_mut(j + 1).write(b'$'); j += 2; } b'\r' => { - *escaped.get_unchecked_mut(j) = b'\\'; - *escaped.get_unchecked_mut(j + 1) = b'r'; + dest.get_unchecked_mut(j).write(b'\\'); + dest.get_unchecked_mut(j + 1).write(b'r'); j += 2; } b => { - *escaped.get_unchecked_mut(j) = b; + dest.get_unchecked_mut(j).write(b); j += 1; } } } - Str::from(std::str::from_utf8_unchecked(escaped)) } + + debug_assert_eq!(j, len); + + // SAFETY: The loop above initialized all `len` bytes of `dest`. + // `MaybeUninit` has the same layout as `u8`, so it's sound to read those bytes back as `&[u8]` + // via a pointer cast. `MaybeUninit::slice_assume_init_ref` would express this directly, but it is unstable. + let bytes = unsafe { slice::from_raw_parts(dest.as_ptr().cast::(), len) }; + // SAFETY: Input is valid UTF-8 and we only insert ASCII bytes replacing existing ASCII, so output is valid UTF-8 + let escaped = unsafe { str::from_utf8_unchecked(bytes) }; + Str::from(escaped) }