Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

select Vec::from_iter impls in a const block to optimize compile times #122785

Merged
merged 1 commit into from
Mar 21, 2024
Merged

select Vec::from_iter impls in a const block to optimize compile times #122785

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 89 additions & 79 deletions library/alloc/src/vec/in_place_collect.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,96 +229,106 @@ where
I: Iterator<Item = T> + InPlaceCollect,
<I as SourceIter>::Source: AsVecIntoIter,
{
default fn from_iter(mut iterator: I) -> Self {
// See "Layout constraints" section in the module documentation. We rely on const
// optimization here since these conditions currently cannot be expressed as trait bounds
if const { !in_place_collectible::<T, I::Src>(I::MERGE_BY, I::EXPAND_BY) } {
// fallback to more generic implementations
return SpecFromIterNested::from_iter(iterator);
}

let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
let inner = iterator.as_inner().as_into_iter();
(
inner.buf.as_ptr(),
inner.ptr,
inner.cap,
inner.buf.as_ptr() as *mut T,
inner.end as *const T,
inner.cap * mem::size_of::<I::Src>() / mem::size_of::<T>(),
)
default fn from_iter(iterator: I) -> Self {
// Select the implementation in const eval to avoid codegen of the dead branch to improve compile times.
let fun: fn(I) -> Vec<T> = const {
// See "Layout constraints" section in the module documentation. We use const conditions here
// since these conditions currently cannot be expressed as trait bounds
if in_place_collectible::<T, I::Src>(I::MERGE_BY, I::EXPAND_BY) {
from_iter_in_place
} else {
// fallback
SpecFromIterNested::<T, I>::from_iter
}
};

// SAFETY: `dst_buf` and `dst_end` are the start and end of the buffer.
let len = unsafe { SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf, dst_end) };
fun(iterator)
}
}

let src = unsafe { iterator.as_inner().as_into_iter() };
// check if SourceIter contract was upheld
// caveat: if they weren't we might not even make it to this point
debug_assert_eq!(src_buf, src.buf.as_ptr());
// check InPlaceIterable contract. This is only possible if the iterator advanced the
// source pointer at all. If it uses unchecked access via TrustedRandomAccess
// then the source pointer will stay in its initial position and we can't use it as reference
if src.ptr != src_ptr {
debug_assert!(
unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
"InPlaceIterable contract violation, write pointer advanced beyond read pointer"
);
}
fn from_iter_in_place<I, T>(mut iterator: I) -> Vec<T>
where
I: Iterator<Item = T> + InPlaceCollect,
<I as SourceIter>::Source: AsVecIntoIter,
{
let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
let inner = iterator.as_inner().as_into_iter();
(
inner.buf.as_ptr(),
inner.ptr,
inner.cap,
inner.buf.as_ptr() as *mut T,
inner.end as *const T,
inner.cap * mem::size_of::<I::Src>() / mem::size_of::<T>(),
)
};

// The ownership of the source allocation and the new `T` values is temporarily moved into `dst_guard`.
// This is safe because
// * `forget_allocation_drop_remaining` immediately forgets the allocation
// before any panic can occur in order to avoid any double free, and then proceeds to drop
// any remaining values at the tail of the source.
// * the shrink either panics without invalidating the allocation, aborts or
// succeeds. In the last case we disarm the guard.
//
// Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
// contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
// module documentation why this is ok anyway.
let dst_guard =
InPlaceDstDataSrcBufDrop { ptr: dst_buf, len, src_cap, src: PhantomData::<I::Src> };
src.forget_allocation_drop_remaining();
// SAFETY: `dst_buf` and `dst_end` are the start and end of the buffer.
let len = unsafe { SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf, dst_end) };

// Adjust the allocation if the source had a capacity in bytes 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 needs_realloc::<I::Src, T>(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::<I::Src>();
let src_size = mem::size_of::<I::Src>().unchecked_mul(src_cap);
let old_layout = Layout::from_size_align_unchecked(src_size, src_align);
let src = unsafe { iterator.as_inner().as_into_iter() };
// check if SourceIter contract was upheld
// caveat: if they weren't we might not even make it to this point
debug_assert_eq!(src_buf, src.buf.as_ptr());
// check InPlaceIterable contract. This is only possible if the iterator advanced the
// source pointer at all. If it uses unchecked access via TrustedRandomAccess
// then the source pointer will stay in its initial position and we can't use it as reference
if src.ptr != src_ptr {
debug_assert!(
unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
"InPlaceIterable contract violation, write pointer advanced beyond read pointer"
);
}

// The allocation must be equal or smaller for in-place iteration to be possible
// therefore the new layout must be ≤ the old one and therefore valid.
let dst_align = mem::align_of::<T>();
let dst_size = mem::size_of::<T>().unchecked_mul(dst_cap);
let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align);
// The ownership of the source allocation and the new `T` values is temporarily moved into `dst_guard`.
// This is safe because
// * `forget_allocation_drop_remaining` immediately forgets the allocation
// before any panic can occur in order to avoid any double free, and then proceeds to drop
// any remaining values at the tail of the source.
// * the shrink either panics without invalidating the allocation, aborts or
// succeeds. In the last case we disarm the guard.
//
// Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
// contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
// module documentation why this is ok anyway.
let dst_guard =
InPlaceDstDataSrcBufDrop { ptr: dst_buf, len, src_cap, src: PhantomData::<I::Src> };
src.forget_allocation_drop_remaining();

let result = alloc.shrink(
NonNull::new_unchecked(dst_buf as *mut u8),
old_layout,
new_layout,
);
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>());
// Adjust the allocation if the source had a capacity in bytes 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 needs_realloc::<I::Src, T>(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::<I::Src>();
let src_size = mem::size_of::<I::Src>().unchecked_mul(src_cap);
let old_layout = Layout::from_size_align_unchecked(src_size, src_align);

// The allocation must be equal or smaller for in-place iteration to be possible
// therefore the new layout must be ≤ the old one and therefore valid.
let dst_align = mem::align_of::<T>();
let dst_size = mem::size_of::<T>().unchecked_mul(dst_cap);
let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align);

let result =
alloc.shrink(NonNull::new_unchecked(dst_buf as *mut u8), old_layout, new_layout);
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);
mem::forget(dst_guard);

let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) };
let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) };

vec
}
vec
}

fn write_in_place_with_drop<T>(
Expand Down
Loading