Skip to content

Commit

Permalink
Fix in-place collect not reallocating when necessary
Browse files Browse the repository at this point in the history
  • Loading branch information
the8472 committed Nov 29, 2023
1 parent df0295f commit 4268a02
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 7 deletions.
37 changes: 30 additions & 7 deletions library/alloc/src/vec/in_place_collect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ const fn in_place_collectible<DEST, SRC>(
step_merge: Option<NonZeroUsize>,
step_expand: Option<NonZeroUsize>,
) -> 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;
}

Expand All @@ -186,6 +186,32 @@ 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 true;
}

if const {
// examples that may require reallocs unless src/dest capacities turn out to be multiples at runtime
// Vec<u8> -> Vec<[u8; 4]>
let dst_larger = mem::size_of::<SRC>() < mem::size_of::<DEST>();

// Vec<[u8; 3]> -> Vec<[u8; 2]>
// dest_sz can't actually be 0 since in_place_collectible() checks for that but const eval doesn't know that
let dest_sz = mem::size_of::<DEST>();
let src_not_multiple_of_dest =
dest_sz == 0 || mem::size_of::<SRC>() % mem::size_of::<DEST>() != 0;

dst_larger || src_not_multiple_of_dest
} {
return src_cap * mem::size_of::<SRC>() != dst_cap * mem::size_of::<DEST>();
}

// Equal size + alignment won't need a realloc.
// src size being an integer multiple of the dest size works too
return false;
}

/// This provides a shorthand for the source type since local type aliases aren't a thing.
#[rustc_specialization_trait]
trait InPlaceCollect: SourceIter<Source: AsVecIntoIter> + InPlaceIterable {
Expand Down Expand Up @@ -259,12 +285,7 @@ where
// that wasn't a multiple of the destination type size.
// Since the discrepancy should generally be small this should only result in some
// bookkeeping updates and no memmove.
if (const {
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>() }
{
if needs_realloc::<I::Src, T>(src_cap, dst_cap) {
let alloc = Global;
unsafe {
// The old allocation exists, therefore it must have a valid layout.
Expand All @@ -286,6 +307,8 @@ where
let Ok(reallocated) = result else { handle_alloc_error(new_layout) };
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);
Expand Down
8 changes: 8 additions & 0 deletions library/alloc/tests/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1213,6 +1213,14 @@ fn test_in_place_specialization_step_up_down() {
assert_ne!(src_bytes, sink_bytes);
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 srcptr = src.as_ptr();
let iter = src
Expand Down

0 comments on commit 4268a02

Please sign in to comment.