Auto merge of #118460 - the8472:fix-vec-realloc, r=saethlin

Fix in-place collect not reallocating when necessary

Regression introduced in https://github.com/rust-lang/rust/pull/110353.
This was [caught by miri](404764617)

r? `@saethlin`
This commit is contained in:
bors 2023-12-06 08:45:11 +00:00
commit 15bb3e204a
2 changed files with 35 additions and 7 deletions

View file

@ -168,7 +168,7 @@ const fn in_place_collectible<DEST, SRC>(
step_merge: Option<NonZeroUsize>, step_merge: Option<NonZeroUsize>,
step_expand: Option<NonZeroUsize>, step_expand: Option<NonZeroUsize>,
) -> bool { ) -> bool {
if DEST::IS_ZST || mem::align_of::<SRC>() < mem::align_of::<DEST>() { if const { SRC::IS_ZST || DEST::IS_ZST || mem::align_of::<SRC>() < mem::align_of::<DEST>() } {
return false; return false;
} }
@ -186,6 +186,27 @@ const fn in_place_collectible<DEST, SRC>(
} }
} }
const fn needs_realloc<SRC, DEST>(src_cap: usize, dst_cap: usize) -> bool {
if const { mem::align_of::<SRC>() != mem::align_of::<DEST>() } {
return src_cap > 0;
}
// If src type size is an integer multiple of the destination type size then
// the caller will have calculated a `dst_cap` that is an integer multiple of
// `src_cap` without remainder.
if const {
let src_sz = mem::size_of::<SRC>();
let dest_sz = mem::size_of::<DEST>();
dest_sz != 0 && src_sz % dest_sz == 0
} {
return false;
}
// type layouts don't guarantee a fit, so do a runtime check to see if
// the allocations happen to match
return src_cap > 0 && src_cap * mem::size_of::<SRC>() != dst_cap * mem::size_of::<DEST>();
}
/// This provides a shorthand for the source type since local type aliases aren't a thing. /// This provides a shorthand for the source type since local type aliases aren't a thing.
#[rustc_specialization_trait] #[rustc_specialization_trait]
trait InPlaceCollect: SourceIter<Source: AsVecIntoIter> + InPlaceIterable { trait InPlaceCollect: SourceIter<Source: AsVecIntoIter> + InPlaceIterable {
@ -259,13 +280,10 @@ where
// that wasn't a multiple of the destination type size. // that wasn't a multiple of the destination type size.
// Since the discrepancy should generally be small this should only result in some // Since the discrepancy should generally be small this should only result in some
// bookkeeping updates and no memmove. // bookkeeping updates and no memmove.
if (const { if needs_realloc::<I::Src, T>(src_cap, dst_cap) {
let src_sz = mem::size_of::<I::Src>();
src_sz > 0 && mem::size_of::<T>() % src_sz != 0
} && src_cap * mem::size_of::<I::Src>() != dst_cap * mem::size_of::<T>())
|| const { mem::align_of::<T>() != mem::align_of::<I::Src>() }
{
let alloc = Global; let alloc = Global;
debug_assert_ne!(src_cap, 0);
debug_assert_ne!(dst_cap, 0);
unsafe { unsafe {
// The old allocation exists, therefore it must have a valid layout. // The old allocation exists, therefore it must have a valid layout.
let src_align = mem::align_of::<I::Src>(); let src_align = mem::align_of::<I::Src>();
@ -286,6 +304,8 @@ where
let Ok(reallocated) = result else { handle_alloc_error(new_layout) }; let Ok(reallocated) = result else { handle_alloc_error(new_layout) };
dst_buf = reallocated.as_ptr() as *mut T; dst_buf = reallocated.as_ptr() as *mut T;
} }
} else {
debug_assert_eq!(src_cap * mem::size_of::<I::Src>(), dst_cap * mem::size_of::<T>());
} }
mem::forget(dst_guard); mem::forget(dst_guard);

View file

@ -1213,6 +1213,14 @@ fn test_in_place_specialization_step_up_down() {
assert_ne!(src_bytes, sink_bytes); assert_ne!(src_bytes, sink_bytes);
assert_eq!(sink.len(), 2); assert_eq!(sink.len(), 2);
let mut src: Vec<[u8; 3]> = Vec::with_capacity(17);
src.resize( 8, [0; 3]);
let iter = src.into_iter().map(|[a, b, _]| [a, b]);
assert_in_place_trait(&iter);
let sink: Vec<[u8; 2]> = iter.collect();
assert_eq!(sink.len(), 8);
assert!(sink.capacity() <= 25);
let src = vec![[0u8; 4]; 256]; let src = vec![[0u8; 4]; 256];
let srcptr = src.as_ptr(); let srcptr = src.as_ptr();
let iter = src let iter = src