diff --git a/src/atom.rs b/src/atom.rs index 5a8aa7f..c11b3aa 100644 --- a/src/atom.rs +++ b/src/atom.rs @@ -336,14 +336,21 @@ impl Ord for Atom { // over the one from &str. impl Atom { fn from_mutated_str(s: &str, f: F) -> Self { - let mut buffer = mem::MaybeUninit::<[u8; 64]>::uninit(); - let buffer = unsafe { &mut *buffer.as_mut_ptr() }; - - if let Some(buffer_prefix) = buffer.get_mut(..s.len()) { - buffer_prefix.copy_from_slice(s.as_bytes()); - let as_str = unsafe { ::std::str::from_utf8_unchecked_mut(buffer_prefix) }; - f(as_str); - Atom::from(&*as_str) + const MAX_STACK_SIZE: usize = 64; + let mut buffer = mem::MaybeUninit::<[u8; MAX_STACK_SIZE]>::uninit(); + + if s.len() <= MAX_STACK_SIZE { + let buffer_ptr = buffer.as_mut_ptr() as *mut u8; + // SAFETY: `buffer_ptr` points to the `MaybeUninit` array. + // We use `copy_nonoverlapping` to write valid data into it, + // and then create a slice covering ONLY the initialized portion. + unsafe { + std::ptr::copy_nonoverlapping(s.as_ptr(), buffer_ptr, s.len()); + let buffer_slice = slice::from_raw_parts_mut(buffer_ptr, s.len()); + let as_str = ::std::str::from_utf8_unchecked_mut(buffer_slice); + f(as_str); + Atom::from(&*as_str) + } } else { let mut string = s.to_owned(); f(&mut string); diff --git a/tests/ub_test.rs b/tests/ub_test.rs new file mode 100644 index 0000000..64c47bf --- /dev/null +++ b/tests/ub_test.rs @@ -0,0 +1,8 @@ +use string_cache::DefaultAtom; + +#[test] +fn test_to_ascii_uppercase() { + let s = DefaultAtom::from("hello"); + let upper = s.to_ascii_uppercase(); + assert_eq!(&*upper, "HELLO"); +}