diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index e4f96fd764040..a6cbed092c0ac 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -168,7 +168,7 @@ const fn in_place_collectible( step_merge: Option, step_expand: Option, ) -> bool { - if DEST::IS_ZST || mem::align_of::() < mem::align_of::() { + if const { SRC::IS_ZST || DEST::IS_ZST || mem::align_of::() < mem::align_of::() } { return false; } @@ -186,6 +186,27 @@ const fn in_place_collectible( } } +const fn needs_realloc(src_cap: usize, dst_cap: usize) -> bool { + if const { mem::align_of::() != mem::align_of::() } { + 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::(); + let dest_sz = mem::size_of::(); + 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::() != dst_cap * mem::size_of::(); +} + /// This provides a shorthand for the source type since local type aliases aren't a thing. #[rustc_specialization_trait] trait InPlaceCollect: SourceIter + InPlaceIterable { @@ -259,13 +280,10 @@ 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::(); - src_sz > 0 && mem::size_of::() % src_sz != 0 - } && src_cap * mem::size_of::() != dst_cap * mem::size_of::()) - || const { mem::align_of::() != mem::align_of::() } - { + if needs_realloc::(src_cap, dst_cap) { let alloc = Global; + debug_assert_ne!(src_cap, 0); + debug_assert_ne!(dst_cap, 0); unsafe { // The old allocation exists, therefore it must have a valid layout. let src_align = mem::align_of::(); @@ -286,6 +304,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::(), dst_cap * mem::size_of::()); } mem::forget(dst_guard); diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index 81de7085e097a..3d3175ba3a9ca 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -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