diff --git a/library/alloc/src/str.rs b/library/alloc/src/str.rs index 8c64bc0801259..89e6eee925e10 100644 --- a/library/alloc/src/str.rs +++ b/library/alloc/src/str.rs @@ -90,8 +90,8 @@ impl> Join<&str> for [S] { } } -macro_rules! spezialize_for_lengths { - ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => { +macro_rules! specialize_for_lengths { + ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => {{ let mut target = $target; let iter = $iter; let sep_bytes = $separator; @@ -102,7 +102,8 @@ macro_rules! spezialize_for_lengths { $num => { for s in iter { copy_slice_and_advance!(target, sep_bytes); - copy_slice_and_advance!(target, s.borrow().as_ref()); + let content_bytes = s.borrow().as_ref(); + copy_slice_and_advance!(target, content_bytes); } }, )* @@ -110,11 +111,13 @@ macro_rules! spezialize_for_lengths { // arbitrary non-zero size fallback for s in iter { copy_slice_and_advance!(target, sep_bytes); - copy_slice_and_advance!(target, s.borrow().as_ref()); + let content_bytes = s.borrow().as_ref(); + copy_slice_and_advance!(target, content_bytes); } } } - }; + target + }} } macro_rules! copy_slice_and_advance { @@ -153,7 +156,7 @@ where // if the `len` calculation overflows, we'll panic // we would have run out of memory anyway and the rest of the function requires // the entire Vec pre-allocated for safety - let len = sep_len + let reserved_len = sep_len .checked_mul(iter.len()) .and_then(|n| { slice.iter().map(|s| s.borrow().as_ref().len()).try_fold(n, usize::checked_add) @@ -161,22 +164,25 @@ where .expect("attempt to join into collection with len > usize::MAX"); // crucial for safety - let mut result = Vec::with_capacity(len); - assert!(result.capacity() >= len); + let mut result = Vec::with_capacity(reserved_len); + debug_assert!(result.capacity() >= reserved_len); result.extend_from_slice(first.borrow().as_ref()); unsafe { - { - let pos = result.len(); - let target = result.get_unchecked_mut(pos..len); - - // copy separator and slices over without bounds checks - // generate loops with hardcoded offsets for small separators - // massive improvements possible (~ x2) - spezialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4); - } - result.set_len(len); + let pos = result.len(); + let target = result.get_unchecked_mut(pos..reserved_len); + + // copy separator and slices over without bounds checks + // generate loops with hardcoded offsets for small separators + // massive improvements possible (~ x2) + let remain = specialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4); + + // issue #80335: A weird borrow implementation can return different + // slices for the length calculation and the actual copy, so + // `remain.len()` might be non-zero. + let result_len = reserved_len - remain.len(); + result.set_len(result_len); } result } diff --git a/library/alloc/tests/str.rs b/library/alloc/tests/str.rs index 604835e6cc4a6..6df8d8c2f354f 100644 --- a/library/alloc/tests/str.rs +++ b/library/alloc/tests/str.rs @@ -160,6 +160,36 @@ fn test_join_for_different_lengths_with_long_separator() { test_join!("~~~~~a~~~~~bc", ["", "a", "bc"], "~~~~~"); } +#[test] +fn test_join_isue_80335() { + use core::{borrow::Borrow, cell::Cell}; + + struct WeirdBorrow { + state: Cell, + } + + impl Default for WeirdBorrow { + fn default() -> Self { + WeirdBorrow { state: Cell::new(false) } + } + } + + impl Borrow for WeirdBorrow { + fn borrow(&self) -> &str { + let state = self.state.get(); + if state { + "0" + } else { + self.state.set(true); + "123456" + } + } + } + + let arr: [WeirdBorrow; 3] = Default::default(); + test_join!("0-0-0", arr, "-"); +} + #[test] #[cfg_attr(miri, ignore)] // Miri is too slow fn test_unsafe_slice() {