Skip to content

Commit df0295f

Browse files
borsbors
committed
Auto merge of #110353 - the8472:in-place-flatten-chunks, r=cuviper
Expand in-place iteration specialization to Flatten, FlatMap and ArrayChunks This enables the following cases to collect in-place: ```rust let v = vec![[0u8; 4]; 1024] let v: Vec<_> = v.into_iter().flatten().collect(); let v: Vec<Option<NonZeroUsize>> = vec![NonZeroUsize::new(0); 1024]; let v: Vec<_> = v.into_iter().flatten().collect(); let v = vec![u8; 4096]; let v: Vec<_> = v.into_iter().array_chunks::<4>().collect(); ``` Especially the nicheful-option-flattening should be useful in real code.
2 parents 46a24ed + 072b51c commit df0295f

25 files changed

+488
-71
lines changed

‎library/alloc/src/collections/binary_heap/mod.rs

+9-2
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@
145145

146146
use core::alloc::Allocator;
147147
use core::fmt;
148-
use core::iter::{FusedIterator, InPlaceIterable, SourceIter, TrustedLen};
148+
use core::iter::{FusedIterator, InPlaceIterable, SourceIter, TrustedFused, TrustedLen};
149149
use core::mem::{self, swap, ManuallyDrop};
150150
use core::num::NonZeroUsize;
151151
use core::ops::{Deref, DerefMut};
@@ -1542,6 +1542,10 @@ impl<T, A: Allocator> ExactSizeIterator for IntoIter<T, A> {
15421542
#[stable(feature = "fused", since = "1.26.0")]
15431543
impl<T, A: Allocator> FusedIterator for IntoIter<T, A> {}
15441544

1545+
#[doc(hidden)]
1546+
#[unstable(issue = "none", feature = "trusted_fused")]
1547+
unsafe impl<T, A: Allocator> TrustedFused for IntoIter<T, A> {}
1548+
15451549
#[stable(feature = "default_iters", since = "1.70.0")]
15461550
impl<T> Default for IntoIter<T> {
15471551
/// Creates an empty `binary_heap::IntoIter`.
@@ -1571,7 +1575,10 @@ unsafe impl<T, A: Allocator> SourceIter for IntoIter<T, A> {
15711575

15721576
#[unstable(issue = "none", feature = "inplace_iteration")]
15731577
#[doc(hidden)]
1574-
unsafe impl<I, A: Allocator> InPlaceIterable for IntoIter<I, A> {}
1578+
unsafe impl<I, A: Allocator> InPlaceIterable for IntoIter<I, A> {
1579+
const EXPAND_BY: Option<NonZeroUsize> = NonZeroUsize::new(1);
1580+
const MERGE_BY: Option<NonZeroUsize> = NonZeroUsize::new(1);
1581+
}
15751582

15761583
unsafe impl<I> AsVecIntoIter for IntoIter<I> {
15771584
type Item = I;

‎library/alloc/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@
154154
#![feature(std_internals)]
155155
#![feature(str_internals)]
156156
#![feature(strict_provenance)]
157+
#![feature(trusted_fused)]
157158
#![feature(trusted_len)]
158159
#![feature(trusted_random_access)]
159160
#![feature(try_trait_v2)]

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

+108-25
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
//! The specialization in this module applies to iterators in the shape of
77
//! `source.adapter().adapter().adapter().collect::<Vec<U>>()`
88
//! where `source` is an owning iterator obtained from [`Vec<T>`], [`Box<[T]>`][box] (by conversion to `Vec`)
9-
//! or [`BinaryHeap<T>`], the adapters each consume one or more items per step
10-
//! (represented by [`InPlaceIterable`]), provide transitive access to `source` (via [`SourceIter`])
11-
//! and thus the underlying allocation. And finally the layouts of `T` and `U` must
12-
//! have the same size and alignment, this is currently ensured via const eval instead of trait bounds
13-
//! in the specialized [`SpecFromIter`] implementation.
9+
//! or [`BinaryHeap<T>`], the adapters guarantee to consume enough items per step to make room
10+
//! for the results (represented by [`InPlaceIterable`]), provide transitive access to `source`
11+
//! (via [`SourceIter`]) and thus the underlying allocation.
12+
//! And finally there are alignment and size constriants to consider, this is currently ensured via
13+
//! const eval instead of trait bounds in the specialized [`SpecFromIter`] implementation.
1414
//!
1515
//! [`BinaryHeap<T>`]: crate::collections::BinaryHeap
1616
//! [box]: crate::boxed::Box
@@ -35,11 +35,28 @@
3535
//! the step of reading a value and getting a reference to write to. Instead raw pointers must be
3636
//! used on the reader and writer side.
3737
//!
38-
//! That writes never clobber a yet-to-be-read item is ensured by the [`InPlaceIterable`] requirements.
38+
//! That writes never clobber a yet-to-be-read items is ensured by the [`InPlaceIterable`] requirements.
3939
//!
4040
//! # Layout constraints
4141
//!
42-
//! [`Allocator`] requires that `allocate()` and `deallocate()` have matching alignment and size.
42+
//! When recycling an allocation between different types we must uphold the [`Allocator`] contract
43+
//! which means that the input and output Layouts have to "fit".
44+
//!
45+
//! To complicate things further `InPlaceIterable` supports splitting or merging items into smaller/
46+
//! larger ones to enable (de)aggregation of arrays.
47+
//!
48+
//! Ultimately each step of the iterator must free up enough *bytes* in the source to make room
49+
//! for the next output item.
50+
//! If `T` and `U` have the same size no fixup is needed.
51+
//! If `T`'s size is a multiple of `U`'s we can compensate by multiplying the capacity accordingly.
52+
//! Otherwise the input capacity (and thus layout) in bytes may not be representable by the output
53+
//! `Vec<U>`. In that case `alloc.shrink()` is used to update the allocation's layout.
54+
//!
55+
//! Alignments of `T` must be the same or larger than `U`. Since alignments are always a power
56+
//! of two _larger_ implies _is a multiple of_.
57+
//!
58+
//! See `in_place_collectible()` for the current conditions.
59+
//!
4360
//! Additionally this specialization doesn't make sense for ZSTs as there is no reallocation to
4461
//! avoid and it would make pointer arithmetic more difficult.
4562
//!
@@ -137,44 +154,73 @@
137154
//! }
138155
//! vec.truncate(write_idx);
139156
//! ```
157+
use crate::alloc::{handle_alloc_error, Global};
158+
use core::alloc::Allocator;
159+
use core::alloc::Layout;
140160
use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
141161
use core::mem::{self, ManuallyDrop, SizedTypeProperties};
142-
use core::ptr::{self};
162+
use core::num::NonZeroUsize;
163+
use core::ptr::{self, NonNull};
143164

144165
use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec};
145166

146-
/// Specialization marker for collecting an iterator pipeline into a Vec while reusing the
147-
/// source allocation, i.e. executing the pipeline in place.
148-
#[rustc_unsafe_specialization_marker]
149-
pub(super) trait InPlaceIterableMarker {}
167+
const fn in_place_collectible<DEST, SRC>(
168+
step_merge: Option<NonZeroUsize>,
169+
step_expand: Option<NonZeroUsize>,
170+
) -> bool {
171+
if DEST::IS_ZST || mem::align_of::<SRC>() < mem::align_of::<DEST>() {
172+
return false;
173+
}
174+
175+
match (step_merge, step_expand) {
176+
(Some(step_merge), Some(step_expand)) => {
177+
// At least N merged source items -> at most M expanded destination items
178+
// e.g.
179+
// - 1 x [u8; 4] -> 4x u8, via flatten
180+
// - 4 x u8 -> 1x [u8; 4], via array_chunks
181+
mem::size_of::<SRC>() * step_merge.get() >= mem::size_of::<DEST>() * step_expand.get()
182+
}
183+
// Fall back to other from_iter impls if an overflow occurred in the step merge/expansion
184+
// tracking.
185+
_ => false,
186+
}
187+
}
150188

151-
impl<T> InPlaceIterableMarker for T where T: InPlaceIterable {}
189+
/// This provides a shorthand for the source type since local type aliases aren't a thing.
190+
#[rustc_specialization_trait]
191+
trait InPlaceCollect: SourceIter<Source: AsVecIntoIter> + InPlaceIterable {
192+
type Src;
193+
}
194+
195+
impl<T> InPlaceCollect for T
196+
where
197+
T: SourceIter<Source: AsVecIntoIter> + InPlaceIterable,
198+
{
199+
type Src = <<T as SourceIter>::Source as AsVecIntoIter>::Item;
200+
}
152201

153202
impl<T, I> SpecFromIter<T, I> for Vec<T>
154203
where
155-
I: Iterator<Item = T> + SourceIter<Source: AsVecIntoIter> + InPlaceIterableMarker,
204+
I: Iterator<Item = T> + InPlaceCollect,
205+
<I as SourceIter>::Source: AsVecIntoIter,
156206
{
157207
default fn from_iter(mut iterator: I) -> Self {
158208
// See "Layout constraints" section in the module documentation. We rely on const
159209
// optimization here since these conditions currently cannot be expressed as trait bounds
160-
if T::IS_ZST
161-
|| mem::size_of::<T>()
162-
!= mem::size_of::<<<I as SourceIter>::Source as AsVecIntoIter>::Item>()
163-
|| mem::align_of::<T>()
164-
!= mem::align_of::<<<I as SourceIter>::Source as AsVecIntoIter>::Item>()
165-
{
210+
if const { !in_place_collectible::<T, I::Src>(I::MERGE_BY, I::EXPAND_BY) } {
166211
// fallback to more generic implementations
167212
return SpecFromIterNested::from_iter(iterator);
168213
}
169214

170-
let (src_buf, src_ptr, dst_buf, dst_end, cap) = unsafe {
215+
let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
171216
let inner = iterator.as_inner().as_into_iter();
172217
(
173218
inner.buf.as_ptr(),
174219
inner.ptr,
220+
inner.cap,
175221
inner.buf.as_ptr() as *mut T,
176222
inner.end as *const T,
177-
inner.cap,
223+
inner.cap * mem::size_of::<I::Src>() / mem::size_of::<T>(),
178224
)
179225
};
180226

@@ -196,18 +242,55 @@ where
196242
}
197243

198244
// The ownership of the allocation and the new `T` values is temporarily moved into `dst_guard`.
199-
// This is safe because `forget_allocation_drop_remaining` immediately forgets the allocation
245+
// This is safe because
246+
// * `forget_allocation_drop_remaining` immediately forgets the allocation
200247
// before any panic can occur in order to avoid any double free, and then proceeds to drop
201248
// any remaining values at the tail of the source.
249+
// * the shrink either panics without invalidating the allocation, aborts or
250+
// succeeds. In the last case we disarm the guard.
202251
//
203252
// Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
204253
// contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
205254
// module documentation why this is ok anyway.
206-
let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap };
255+
let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap: dst_cap };
207256
src.forget_allocation_drop_remaining();
257+
258+
// Adjust the allocation if the alignment didn't match or the source had a capacity in bytes
259+
// that wasn't a multiple of the destination type size.
260+
// Since the discrepancy should generally be small this should only result in some
261+
// bookkeeping updates and no memmove.
262+
if (const {
263+
let src_sz = mem::size_of::<I::Src>();
264+
src_sz > 0 && mem::size_of::<T>() % src_sz != 0
265+
} && src_cap * mem::size_of::<I::Src>() != dst_cap * mem::size_of::<T>())
266+
|| const { mem::align_of::<T>() != mem::align_of::<I::Src>() }
267+
{
268+
let alloc = Global;
269+
unsafe {
270+
// The old allocation exists, therefore it must have a valid layout.
271+
let src_align = mem::align_of::<I::Src>();
272+
let src_size = mem::size_of::<I::Src>().unchecked_mul(src_cap);
273+
let old_layout = Layout::from_size_align_unchecked(src_size, src_align);
274+
275+
// The must be equal or smaller for in-place iteration to be possible
276+
// therefore the new layout must be ≤ the old one and therefore valid.
277+
let dst_align = mem::align_of::<T>();
278+
let dst_size = mem::size_of::<T>().unchecked_mul(dst_cap);
279+
let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align);
280+
281+
let result = alloc.shrink(
282+
NonNull::new_unchecked(dst_buf as *mut u8),
283+
old_layout,
284+
new_layout,
285+
);
286+
let Ok(reallocated) = result else { handle_alloc_error(new_layout) };
287+
dst_buf = reallocated.as_ptr() as *mut T;
288+
}
289+
}
290+
208291
mem::forget(dst_guard);
209292

210-
let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) };
293+
let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) };
211294

212295
vec
213296
}

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

+10-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@ use crate::raw_vec::RawVec;
77
use core::array;
88
use core::fmt;
99
use core::iter::{
10-
FusedIterator, InPlaceIterable, SourceIter, TrustedLen, TrustedRandomAccessNoCoerce,
10+
FusedIterator, InPlaceIterable, SourceIter, TrustedFused, TrustedLen,
11+
TrustedRandomAccessNoCoerce,
1112
};
1213
use core::marker::PhantomData;
1314
use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties};
@@ -337,6 +338,10 @@ impl<T, A: Allocator> ExactSizeIterator for IntoIter<T, A> {
337338
#[stable(feature = "fused", since = "1.26.0")]
338339
impl<T, A: Allocator> FusedIterator for IntoIter<T, A> {}
339340

341+
#[doc(hidden)]
342+
#[unstable(issue = "none", feature = "trusted_fused")]
343+
unsafe impl<T, A: Allocator> TrustedFused for IntoIter<T, A> {}
344+
340345
#[unstable(feature = "trusted_len", issue = "37572")]
341346
unsafe impl<T, A: Allocator> TrustedLen for IntoIter<T, A> {}
342347

@@ -421,7 +426,10 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for IntoIter<T, A> {
421426
// also refer to the vec::in_place_collect module documentation to get an overview
422427
#[unstable(issue = "none", feature = "inplace_iteration")]
423428
#[doc(hidden)]
424-
unsafe impl<T, A: Allocator> InPlaceIterable for IntoIter<T, A> {}
429+
unsafe impl<T, A: Allocator> InPlaceIterable for IntoIter<T, A> {
430+
const EXPAND_BY: Option<NonZeroUsize> = NonZeroUsize::new(1);
431+
const MERGE_BY: Option<NonZeroUsize> = NonZeroUsize::new(1);
432+
}
425433

426434
#[unstable(issue = "none", feature = "inplace_iteration")]
427435
#[doc(hidden)]

‎library/alloc/tests/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#![feature(allocator_api)]
22
#![feature(alloc_layout_extra)]
3+
#![feature(iter_array_chunks)]
34
#![feature(assert_matches)]
45
#![feature(btree_extract_if)]
56
#![feature(cow_is_borrowed)]

‎library/alloc/tests/vec.rs

+43-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
1+
use alloc::vec::Vec;
12
use core::alloc::{Allocator, Layout};
2-
use core::assert_eq;
3-
use core::iter::IntoIterator;
3+
use core::{assert_eq, assert_ne};
4+
use core::iter::{IntoIterator, Iterator};
45
use core::num::NonZeroUsize;
56
use core::ptr::NonNull;
67
use std::alloc::System;
@@ -1184,6 +1185,46 @@ fn test_from_iter_specialization_with_iterator_adapters() {
11841185
assert_eq!(srcptr, sinkptr as *const usize);
11851186
}
11861187

1188+
#[test]
1189+
fn test_in_place_specialization_step_up_down() {
1190+
fn assert_in_place_trait<T: InPlaceIterable>(_: &T) {}
1191+
let src = vec![[0u8; 4]; 256];
1192+
let srcptr = src.as_ptr();
1193+
let src_cap = src.capacity();
1194+
let iter = src.into_iter().flatten();
1195+
assert_in_place_trait(&iter);
1196+
let sink = iter.collect::<Vec<_>>();
1197+
let sinkptr = sink.as_ptr();
1198+
assert_eq!(srcptr as *const u8, sinkptr);
1199+
assert_eq!(src_cap * 4, sink.capacity());
1200+
1201+
let iter = sink.into_iter().array_chunks::<4>();
1202+
assert_in_place_trait(&iter);
1203+
let sink = iter.collect::<Vec<_>>();
1204+
let sinkptr = sink.as_ptr();
1205+
assert_eq!(srcptr, sinkptr);
1206+
assert_eq!(src_cap, sink.capacity());
1207+
1208+
let mut src: Vec<u8> = Vec::with_capacity(17);
1209+
let src_bytes = src.capacity();
1210+
src.resize(8, 0u8);
1211+
let sink: Vec<[u8; 4]> = src.into_iter().array_chunks::<4>().collect();
1212+
let sink_bytes = sink.capacity() * 4;
1213+
assert_ne!(src_bytes, sink_bytes);
1214+
assert_eq!(sink.len(), 2);
1215+
1216+
let src = vec![[0u8; 4]; 256];
1217+
let srcptr = src.as_ptr();
1218+
let iter = src
1219+
.into_iter()
1220+
.flat_map(|a| {
1221+
a.into_iter().map(|b| b.wrapping_add(1))
1222+
});
1223+
assert_in_place_trait(&iter);
1224+
let sink = iter.collect::<Vec<_>>();
1225+
assert_eq!(srcptr as *const u8, sink.as_ptr());
1226+
}
1227+
11871228
#[test]
11881229
fn test_from_iter_specialization_head_tail_drop() {
11891230
let drop_count: Vec<_> = (0..=2).map(|_| Rc::new(())).collect();

‎library/core/src/iter/adapters/array_chunks.rs

+33-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
use crate::array;
2-
use crate::iter::{ByRefSized, FusedIterator, Iterator, TrustedRandomAccessNoCoerce};
2+
use crate::iter::adapters::SourceIter;
3+
use crate::iter::{
4+
ByRefSized, FusedIterator, InPlaceIterable, Iterator, TrustedFused, TrustedRandomAccessNoCoerce,
5+
};
6+
use crate::num::NonZeroUsize;
37
use crate::ops::{ControlFlow, NeverShortCircuit, Try};
48

59
/// An iterator over `N` elements of the iterator at a time.
@@ -159,6 +163,9 @@ where
159163
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
160164
impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
161165

166+
#[unstable(issue = "none", feature = "trusted_fused")]
167+
unsafe impl<I, const N: usize> TrustedFused for ArrayChunks<I, N> where I: TrustedFused + Iterator {}
168+
162169
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
163170
impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
164171
where
@@ -229,3 +236,28 @@ where
229236
accum
230237
}
231238
}
239+
240+
#[unstable(issue = "none", feature = "inplace_iteration")]
241+
unsafe impl<I, const N: usize> SourceIter for ArrayChunks<I, N>
242+
where
243+
I: SourceIter + Iterator,
244+
{
245+
type Source = I::Source;
246+
247+
#[inline]
248+
unsafe fn as_inner(&mut self) -> &mut I::Source {
249+
// SAFETY: unsafe function forwarding to unsafe function with the same requirements
250+
unsafe { SourceIter::as_inner(&mut self.iter) }
251+
}
252+
}
253+
254+
#[unstable(issue = "none", feature = "inplace_iteration")]
255+
unsafe impl<I: InPlaceIterable + Iterator, const N: usize> InPlaceIterable for ArrayChunks<I, N> {
256+
const EXPAND_BY: Option<NonZeroUsize> = I::EXPAND_BY;
257+
const MERGE_BY: Option<NonZeroUsize> = const {
258+
match (I::MERGE_BY, NonZeroUsize::new(N)) {
259+
(Some(m), Some(n)) => m.checked_mul(n),
260+
_ => None,
261+
}
262+
};
263+
}

0 commit comments

Comments
 (0)