Skip to content

Commit 0ad927c

Browse files
borsbors
committed
Auto merge of #122785 - the8472:const-select-specialized-impl, r=Amanieu
select Vec::from_iter impls in a const block to optimize compile times Ignoring whitespace diffs should make this easier to review. This relies on the trick from #122301 Split out from #120682
2 parents a0569fa + 3ff1e44 commit 0ad927c

File tree

1 file changed

+89
-79
lines changed

1 file changed

+89
-79
lines changed

‎library/alloc/src/vec/in_place_collect.rs‎

Lines changed: 89 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -229,96 +229,106 @@ where
229229
I: Iterator<Item = T> + InPlaceCollect,
230230
<I as SourceIter>::Source: AsVecIntoIter,
231231
{
232-
default fn from_iter(mut iterator: I) -> Self {
233-
// See "Layout constraints" section in the module documentation. We rely on const
234-
// optimization here since these conditions currently cannot be expressed as trait bounds
235-
if const { !in_place_collectible::<T, I::Src>(I::MERGE_BY, I::EXPAND_BY) } {
236-
// fallback to more generic implementations
237-
return SpecFromIterNested::from_iter(iterator);
238-
}
239-
240-
let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
241-
let inner = iterator.as_inner().as_into_iter();
242-
(
243-
inner.buf.as_ptr(),
244-
inner.ptr,
245-
inner.cap,
246-
inner.buf.as_ptr() as *mut T,
247-
inner.end as *const T,
248-
inner.cap * mem::size_of::<I::Src>() / mem::size_of::<T>(),
249-
)
232+
default fn from_iter(iterator: I) -> Self {
233+
// Select the implementation in const eval to avoid codegen of the dead branch to improve compile times.
234+
let fun: fn(I) -> Vec<T> = const {
235+
// See "Layout constraints" section in the module documentation. We use const conditions here
236+
// since these conditions currently cannot be expressed as trait bounds
237+
if in_place_collectible::<T, I::Src>(I::MERGE_BY, I::EXPAND_BY) {
238+
from_iter_in_place
239+
} else {
240+
// fallback
241+
SpecFromIterNested::<T, I>::from_iter
242+
}
250243
};
251244

252-
// SAFETY: `dst_buf` and `dst_end` are the start and end of the buffer.
253-
let len = unsafe { SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf, dst_end) };
245+
fun(iterator)
246+
}
247+
}
254248

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

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

284-
// Adjust the allocation if the source had a capacity in bytes that wasn't a multiple
285-
// of the destination type size.
286-
// Since the discrepancy should generally be small this should only result in some
287-
// bookkeeping updates and no memmove.
288-
if needs_realloc::<I::Src, T>(src_cap, dst_cap) {
289-
let alloc = Global;
290-
debug_assert_ne!(src_cap, 0);
291-
debug_assert_ne!(dst_cap, 0);
292-
unsafe {
293-
// The old allocation exists, therefore it must have a valid layout.
294-
let src_align = mem::align_of::<I::Src>();
295-
let src_size = mem::size_of::<I::Src>().unchecked_mul(src_cap);
296-
let old_layout = Layout::from_size_align_unchecked(src_size, src_align);
269+
let src = unsafe { iterator.as_inner().as_into_iter() };
270+
// check if SourceIter contract was upheld
271+
// caveat: if they weren't we might not even make it to this point
272+
debug_assert_eq!(src_buf, src.buf.as_ptr());
273+
// check InPlaceIterable contract. This is only possible if the iterator advanced the
274+
// source pointer at all. If it uses unchecked access via TrustedRandomAccess
275+
// then the source pointer will stay in its initial position and we can't use it as reference
276+
if src.ptr != src_ptr {
277+
debug_assert!(
278+
unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
279+
"InPlaceIterable contract violation, write pointer advanced beyond read pointer"
280+
);
281+
}
297282

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

304-
let result = alloc.shrink(
305-
NonNull::new_unchecked(dst_buf as *mut u8),
306-
old_layout,
307-
new_layout,
308-
);
309-
let Ok(reallocated) = result else { handle_alloc_error(new_layout) };
310-
dst_buf = reallocated.as_ptr() as *mut T;
311-
}
312-
} else {
313-
debug_assert_eq!(src_cap * mem::size_of::<I::Src>(), dst_cap * mem::size_of::<T>());
298+
// Adjust the allocation if the source had a capacity in bytes that wasn't a multiple
299+
// of the destination type size.
300+
// Since the discrepancy should generally be small this should only result in some
301+
// bookkeeping updates and no memmove.
302+
if needs_realloc::<I::Src, T>(src_cap, dst_cap) {
303+
let alloc = Global;
304+
debug_assert_ne!(src_cap, 0);
305+
debug_assert_ne!(dst_cap, 0);
306+
unsafe {
307+
// The old allocation exists, therefore it must have a valid layout.
308+
let src_align = mem::align_of::<I::Src>();
309+
let src_size = mem::size_of::<I::Src>().unchecked_mul(src_cap);
310+
let old_layout = Layout::from_size_align_unchecked(src_size, src_align);
311+
312+
// The allocation must be equal or smaller for in-place iteration to be possible
313+
// therefore the new layout must be ≤ the old one and therefore valid.
314+
let dst_align = mem::align_of::<T>();
315+
let dst_size = mem::size_of::<T>().unchecked_mul(dst_cap);
316+
let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align);
317+
318+
let result =
319+
alloc.shrink(NonNull::new_unchecked(dst_buf as *mut u8), old_layout, new_layout);
320+
let Ok(reallocated) = result else { handle_alloc_error(new_layout) };
321+
dst_buf = reallocated.as_ptr() as *mut T;
314322
}
323+
} else {
324+
debug_assert_eq!(src_cap * mem::size_of::<I::Src>(), dst_cap * mem::size_of::<T>());
325+
}
315326

316-
mem::forget(dst_guard);
327+
mem::forget(dst_guard);
317328

318-
let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) };
329+
let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) };
319330

320-
vec
321-
}
331+
vec
322332
}
323333

324334
fn write_in_place_with_drop<T>(

0 commit comments

Comments
 (0)