Skip to content

Commit 2734112

Browse files
committed
Add Vec::flatten (as unstable, of course)
This has been a common question on IRC for years -- with people wanting to flatten Vecs of `[u8; 4]` colours or `[f32; 3]` linalg vectors. I'm irrationally excited that const generics has a reached a state where this works now. (Named after `Iterator::flatten`, which is the same conceptual operation, just slower and doesn't currently work right with array `Item`s.)
1 parent 6312b89 commit 2734112

File tree

2 files changed

+80
-0
lines changed

2 files changed

+80
-0
lines changed

src/liballoc/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
#![feature(box_syntax)]
7878
#![feature(cfg_target_has_atomic)]
7979
#![feature(coerce_unsized)]
80+
#![cfg_attr(not(bootstrap), feature(const_generics))]
8081
#![feature(dispatch_from_dyn)]
8182
#![feature(core_intrinsics)]
8283
#![feature(custom_attribute)]

src/liballoc/vec.rs

+79
Original file line numberDiff line numberDiff line change
@@ -1369,6 +1369,85 @@ impl<T> Vec<T> {
13691369
}
13701370
}
13711371

1372+
#[cfg(not(bootstrap))]
1373+
impl<T, const N: usize> Vec<[T; N]> {
1374+
/// Flatten a `Vec<[T; N]>` into an N-times longer `Vec<T>`.
1375+
///
1376+
/// This is O(1) and non-allocating.
1377+
///
1378+
/// # Panics
1379+
///
1380+
/// If the length of the result would be too large to store in a `usize`,
1381+
/// which can only happen if `T` is a zero-sized type.
1382+
///
1383+
/// # Examples
1384+
///
1385+
/// ```
1386+
/// #![feature(vec_flatten)]
1387+
///
1388+
/// let v = vec![ [1, 2, 3], [7, 8, 9] ];
1389+
/// assert_eq!(v.flatten(), [1, 2, 3, 7, 8, 9]);
1390+
///
1391+
/// let v = vec![[0; 7]; 13];
1392+
/// assert_eq!(v.flatten().len(), 7 * 13);
1393+
///
1394+
/// let v = vec![[(); 500]; 2_000];
1395+
/// assert_eq!(v.flatten().len(), 1_000_000);
1396+
///
1397+
/// enum Never {}
1398+
/// let v: Vec<[Never; 0]> = vec![[], [], []];
1399+
/// assert_eq!(v.flatten().len(), 0);
1400+
/// ```
1401+
///
1402+
/// ```should_panic
1403+
/// #![feature(vec_flatten)]
1404+
///
1405+
/// let v = vec![[(); std::usize::MAX]; 2];
1406+
/// v.flatten(); // panics for length overflow
1407+
/// ```
1408+
#[inline]
1409+
#[unstable(feature = "vec_flatten", issue = "88888888",
1410+
reason = "new API, and needs const generics stabilized first")]
1411+
pub fn flatten(mut self) -> Vec<T> {
1412+
if N == 0 {
1413+
// The allocator-related safety implications of switching pointers
1414+
// between ZSTs and other types are non-obvious, so just do the
1415+
// simple thing -- this check is compile-time anyway.
1416+
return Vec::new();
1417+
}
1418+
1419+
let ptr = self.as_mut_ptr() as *mut T;
1420+
let (len, cap) =
1421+
if mem::size_of::<T>() == 0 {
1422+
// Since ZSTs are limited only by the size of the counters,
1423+
// it's completely possible to overflow here. Capacity just
1424+
// saturates because it usually comes in as `usize::MAX` so
1425+
// checking the multiplication would make it almost always panic.
1426+
(
1427+
self.len().checked_mul(N).expect("length overflow"),
1428+
self.capacity().saturating_mul(N),
1429+
)
1430+
} else {
1431+
// But for types with an actual size these multiplications
1432+
// cannot overflow as the memory is allocated in self, so don't
1433+
// codegen a reference to panic machinery.
1434+
(self.len() * N, self.capacity() * N)
1435+
};
1436+
mem::forget(self);
1437+
1438+
// SAFETY:
1439+
// - The pointer came from the same allocator with which it's being used.
1440+
// - The layout of the allocation hasn't changed, because the alignment
1441+
// of an array matches that of its elements and we increased the length
1442+
// by the appropriate amount to counteract the decrease in type size.
1443+
// - Length is less than or equal to capacity because we increased both
1444+
// proportionately or set capacity to the largest possible value.
1445+
unsafe {
1446+
Vec::from_raw_parts(ptr, len, cap)
1447+
}
1448+
}
1449+
}
1450+
13721451
impl<T: Clone> Vec<T> {
13731452
/// Resizes the `Vec` in-place so that `len` is equal to `new_len`.
13741453
///

0 commit comments

Comments
 (0)