Skip to content

Commit 680031b

Browse files
nagisanagisa
committed
Implement [T]::align_to
1 parent d453782 commit 680031b

File tree

5 files changed

+287
-114
lines changed

5 files changed

+287
-114
lines changed

‎src/libcore/slice/mod.rs

Lines changed: 162 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1697,27 +1697,169 @@ impl<T> [T] {
16971697
}
16981698
}
16991699

1700-
// #[unstable(feature = "slice_align_to", issue = "44488")]
1701-
// pub fn align_to<U>(&self) -> (&[T], &[U], &[T]) {
1702-
// // First, find at what point do we split between the first and 2nd slice.
1703-
// let x = self.as_ptr();
1704-
// let offset = x.align_offset(::mem::align_of::<U>());
1705-
// if offset > x * ::mem::size_of::<T>() {
1706-
// return (self, [], []);
1707-
// }
1708-
1709-
// }
1710-
1711-
// #[unstable(feature = "slice_align_to", issue = "44488")]
1712-
// pub fn align_to_mut<U>(&mut self) -> (&mut [T], &mut [U], &mut [T]) {
1713-
// }
1714-
}}
1700+
/// Function to calculate lenghts of the middle and trailing slice for `align_to{,_mut}`.
1701+
fn align_to_offsets<U>(&self) -> (usize, usize) {
1702+
// What we gonna do about `rest` is figure out what multiple of `U`s we can put in a
1703+
// lowest number of `T`s. And how many `T`s we need for each such "multiple".
1704+
//
1705+
// Consider for example T=u8 U=u16. Then we can put 1 U in 2 Ts. Simple. Now, consider
1706+
// for example a case where size_of::<T> = 16, size_of::<U> = 24. We can put 2 Us in
1707+
// place of every 3 Ts in the `rest` slice. A bit more complicated.
1708+
//
1709+
// Formula to calculate this is:
1710+
//
1711+
// Us = lcm(size_of::<T>, size_of::<U>) / size_of::<U>
1712+
// Ts = lcm(size_of::<T>, size_of::<U>) / size_of::<T>
1713+
//
1714+
// Expanded and simplified:
1715+
//
1716+
// Us = size_of::<T> / gcd(size_of::<T>, size_of::<U>)
1717+
// Ts = size_of::<U> / gcd(size_of::<T>, size_of::<U>)
1718+
//
1719+
// Luckily since all this is constant-evaluated... performance here matters not!
1720+
#[inline]
1721+
fn gcd(a: usize, b: usize) -> usize {
1722+
// iterative stein’s algorithm
1723+
// We should still make this `const fn` (and revert to recursive algorithm if we do)
1724+
// because relying on llvm to consteval all this is… well, it makes me
1725+
let (ctz_a, mut ctz_b) = unsafe {
1726+
if a == 0 { return b; }
1727+
if b == 0 { return a; }
1728+
(::intrinsics::cttz_nonzero(a), ::intrinsics::cttz_nonzero(b))
1729+
};
1730+
let k = ctz_a.min(ctz_b);
1731+
let mut a = a >> ctz_a;
1732+
let mut b = b;
1733+
loop {
1734+
// remove all factors of 2 from b
1735+
b >>= ctz_b;
1736+
if a > b {
1737+
::mem::swap(&mut a, &mut b);
1738+
}
1739+
b = b - a;
1740+
unsafe {
1741+
if b == 0 {
1742+
break;
1743+
}
1744+
ctz_b = ::intrinsics::cttz_nonzero(b);
1745+
}
1746+
}
1747+
return a << k;
1748+
}
1749+
let gcd: usize = gcd(::mem::size_of::<T>(), ::mem::size_of::<U>());
1750+
let ts: usize = ::mem::size_of::<U>() / gcd;
1751+
let us: usize = ::mem::size_of::<T>() / gcd;
17151752

1716-
#[lang = "slice"]
1717-
#[cfg(not(test))]
1718-
#[cfg(not(stage0))]
1719-
impl<T> [T] {
1720-
slice_core_methods!();
1753+
// Armed with this knowledge, we can find how many `U`s we can fit!
1754+
let us_len = self.len() / ts * us;
1755+
// And how many `T`s will be in the trailing slice!
1756+
let ts_len = self.len() % ts;
1757+
return (us_len, ts_len);
1758+
}
1759+
1760+
/// Transmute the slice to a slice of another type, ensuring aligment of the types is
1761+
/// maintained.
1762+
///
1763+
/// This method splits the slice into three distinct slices: prefix, correctly aligned middle
1764+
/// slice of a new type, and the suffix slice. The middle slice will have the greatest length
1765+
/// possible for a given type and input slice.
1766+
///
1767+
/// This method has no purpose when either input element `T` or output element `U` are
1768+
/// zero-sized and will return the original slice without splitting anything.
1769+
///
1770+
/// # Unsafety
1771+
///
1772+
/// This method is essentially a `transmute` with respect to the elements in the returned
1773+
/// middle slice, so all the usual caveats pertaining to `transmute::<T, U>` also apply here.
1774+
///
1775+
/// # Examples
1776+
///
1777+
/// Basic usage:
1778+
///
1779+
/// ```
1780+
/// # #![feature(slice_align_to)]
1781+
/// unsafe {
1782+
/// let bytes: [u8; 7] = [1, 2, 3, 4, 5, 6, 7];
1783+
/// let (prefix, shorts, suffix) = bytes.align_to::<u16>();
1784+
/// // less_efficient_algorithm_for_bytes(prefix);
1785+
/// // more_efficient_algorithm_for_aligned_shorts(shorts);
1786+
/// // less_efficient_algorithm_for_bytes(suffix);
1787+
/// }
1788+
/// ```
1789+
#[unstable(feature = "slice_align_to", issue = "44488")]
1790+
#[cfg(not(stage0))]
1791+
pub unsafe fn align_to<U>(&self) -> (&[T], &[U], &[T]) {
1792+
// Note that most of this function will be constant-evaluated,
1793+
if ::mem::size_of::<U>() == 0 || ::mem::size_of::<T>() == 0 {
1794+
// handle ZSTs specially, which is – don't handle them at all.
1795+
return (self, &[], &[]);
1796+
}
1797+
let ptr = self.as_ptr();
1798+
let offset = ::intrinsics::align_offset(ptr, ::mem::align_of::<U>());
1799+
if offset > self.len() {
1800+
return (self, &[], &[]);
1801+
} else {
1802+
let (left, rest) = self.split_at(offset);
1803+
let (us_len, ts_len) = rest.align_to_offsets::<U>();
1804+
return (left,
1805+
from_raw_parts(rest.as_ptr() as *const U, us_len),
1806+
from_raw_parts(rest.as_ptr().offset((rest.len() - ts_len) as isize), ts_len))
1807+
}
1808+
}
1809+
1810+
/// Transmute the slice to a slice of another type, ensuring aligment of the types is
1811+
/// maintained.
1812+
///
1813+
/// This method splits the slice into three distinct slices: prefix, correctly aligned middle
1814+
/// slice of a new type, and the suffix slice. The middle slice will have the greatest length
1815+
/// possible for a given type and input slice.
1816+
///
1817+
/// This method has no purpose when either input element `T` or output element `U` are
1818+
/// zero-sized and will return the original slice without splitting anything.
1819+
///
1820+
/// # Unsafety
1821+
///
1822+
/// This method is essentially a `transmute` with respect to the elements in the returned
1823+
/// middle slice, so all the usual caveats pertaining to `transmute::<T, U>` also apply here.
1824+
///
1825+
/// # Examples
1826+
///
1827+
/// Basic usage:
1828+
///
1829+
/// ```
1830+
/// # #![feature(slice_align_to)]
1831+
/// unsafe {
1832+
/// let mut bytes: [u8; 7] = [1, 2, 3, 4, 5, 6, 7];
1833+
/// let (prefix, shorts, suffix) = bytes.align_to_mut::<u16>();
1834+
/// // less_efficient_algorithm_for_bytes(prefix);
1835+
/// // more_efficient_algorithm_for_aligned_shorts(shorts);
1836+
/// // less_efficient_algorithm_for_bytes(suffix);
1837+
/// }
1838+
/// ```
1839+
#[unstable(feature = "slice_align_to", issue = "44488")]
1840+
#[cfg(not(stage0))]
1841+
pub unsafe fn align_to_mut<U>(&mut self) -> (&mut [T], &mut [U], &mut [T]) {
1842+
// Note that most of this function will be constant-evaluated,
1843+
if ::mem::size_of::<U>() == 0 || ::mem::size_of::<T>() == 0 {
1844+
// handle ZSTs specially, which is – don't handle them at all.
1845+
return (self, &mut [], &mut []);
1846+
}
1847+
1848+
// First, find at what point do we split between the first and 2nd slice. Easy with
1849+
// ptr.align_offset.
1850+
let ptr = self.as_ptr();
1851+
let offset = ::intrinsics::align_offset(ptr, ::mem::align_of::<U>());
1852+
if offset > self.len() {
1853+
return (self, &mut [], &mut []);
1854+
} else {
1855+
let (left, rest) = self.split_at_mut(offset);
1856+
let (us_len, ts_len) = rest.align_to_offsets::<U>();
1857+
let mut_ptr = rest.as_mut_ptr();
1858+
return (left,
1859+
from_raw_parts_mut(mut_ptr as *mut U, us_len),
1860+
from_raw_parts_mut(mut_ptr.offset((rest.len() - ts_len) as isize), ts_len))
1861+
}
1862+
}
17211863
}
17221864

17231865
#[lang = "slice_u8"]

‎src/libcore/tests/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141
#![feature(try_from)]
4242
#![feature(try_trait)]
4343
#![feature(exact_chunks)]
44+
#![feature(slice_align_to)]
45+
#![feature(align_offset)]
4446
#![feature(reverse_bits)]
4547
#![feature(inclusive_range_methods)]
4648
#![feature(iterator_find_map)]

‎src/libcore/tests/ptr.rs

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,3 +296,92 @@ fn write_unaligned_drop() {
296296
}
297297
DROPS.with(|d| assert_eq!(*d.borrow(), [0]));
298298
}
299+
300+
#[test]
301+
fn align_offset_zst() {
302+
// For pointers of stride = 0, the pointer is already aligned or it cannot be aligned at
303+
// all, because no amount of elements will align the pointer.
304+
let mut p = 1;
305+
while p < 1024 {
306+
assert_eq!((p as *const ()).align_offset(p), 0);
307+
if p != 1 {
308+
assert_eq!(((p + 1) as *const ()).align_offset(p), !0);
309+
}
310+
p = (p + 1).next_power_of_two();
311+
}
312+
}
313+
314+
#[test]
315+
fn align_offset_stride1() {
316+
// For pointers of stride = 1, the pointer can always be aligned. The offset is equal to
317+
// number of bytes.
318+
let mut align = 1;
319+
while align < 1024 {
320+
for ptr in 1..2*align {
321+
let expected = ptr % align;
322+
let offset = if expected == 0 { 0 } else { align - expected };
323+
assert_eq!((ptr as *const u8).align_offset(align), offset,
324+
"ptr = {}, align = {}, size = 1", ptr, align);
325+
align = (align + 1).next_power_of_two();
326+
}
327+
}
328+
}
329+
330+
#[test]
331+
fn align_offset_weird_strides() {
332+
#[repr(packed)]
333+
struct A3(u16, u8);
334+
struct A4(u32);
335+
#[repr(packed)]
336+
struct A5(u32, u8);
337+
#[repr(packed)]
338+
struct A6(u32, u16);
339+
#[repr(packed)]
340+
struct A7(u32, u16, u8);
341+
#[repr(packed)]
342+
struct A8(u32, u32);
343+
#[repr(packed)]
344+
struct A9(u32, u32, u8);
345+
#[repr(packed)]
346+
struct A10(u32, u32, u16);
347+
348+
unsafe fn test_weird_stride<T>(ptr: *const T, align: usize) -> bool {
349+
let numptr = ptr as usize;
350+
let mut expected = usize::max_value();
351+
// Naive but definitely correct way to find the *first* aligned element of stride::<T>.
352+
for el in 0..align {
353+
if (numptr + el * ::std::mem::size_of::<T>()) % align == 0 {
354+
expected = el;
355+
break;
356+
}
357+
}
358+
let got = ptr.align_offset(align);
359+
if got != expected {
360+
eprintln!("aligning {:p} (with stride of {}) to {}, expected {}, got {}", ptr,
361+
::std::mem::size_of::<T>(), align, expected, got);
362+
return true;
363+
}
364+
return false;
365+
}
366+
367+
// For pointers of stride != 1, we verify the algorithm against the naivest possible
368+
// implementation
369+
let mut align = 1;
370+
let mut x = false;
371+
while align < 1024 {
372+
for ptr in 1usize..4*align {
373+
unsafe {
374+
x |= test_weird_stride::<A3>(ptr as *const A3, align);
375+
x |= test_weird_stride::<A4>(ptr as *const A4, align);
376+
x |= test_weird_stride::<A5>(ptr as *const A5, align);
377+
x |= test_weird_stride::<A6>(ptr as *const A6, align);
378+
x |= test_weird_stride::<A7>(ptr as *const A7, align);
379+
x |= test_weird_stride::<A8>(ptr as *const A8, align);
380+
x |= test_weird_stride::<A9>(ptr as *const A9, align);
381+
x |= test_weird_stride::<A10>(ptr as *const A10, align);
382+
}
383+
}
384+
align = (align + 1).next_power_of_two();
385+
}
386+
assert!(!x);
387+
}

‎src/libcore/tests/slice.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -812,3 +812,37 @@ pub mod memchr {
812812
}
813813
}
814814
}
815+
816+
#[test]
817+
fn test_align_to_simple() {
818+
let bytes = [1u8, 2, 3, 4, 5, 6, 7];
819+
let (prefix, aligned, suffix) = unsafe { bytes.align_to::<u16>() };
820+
assert_eq!(aligned.len(), 3);
821+
assert!(prefix == [1] || suffix == [7]);
822+
let expect1 = [1 << 8 | 2, 3 << 8 | 4, 5 << 8 | 6];
823+
let expect2 = [1 | 2 << 8, 3 | 4 << 8, 5 | 6 << 8];
824+
let expect3 = [2 | 3 << 8, 4 | 5 << 8, 6 | 7 << 8];
825+
let expect4 = [2 | 3 << 8, 4 | 5 << 8, 6 | 7 << 8];
826+
assert!(aligned == expect1 || aligned == expect2 || aligned == expect3 || aligned == expect4,
827+
"aligned={:?} expected={:?} || {:?} || {:?} || {:?}",
828+
aligned, expect1, expect2, expect3, expect4);
829+
}
830+
831+
#[test]
832+
fn test_align_to_zst() {
833+
let bytes = [1, 2, 3, 4, 5, 6, 7];
834+
let (prefix, aligned, suffix) = unsafe { bytes.align_to::<()>() };
835+
assert_eq!(aligned.len(), 0);
836+
assert!(prefix == [1, 2, 3, 4, 5, 6, 7] || suffix == [1, 2, 3, 4, 5, 6, 7]);
837+
}
838+
839+
#[test]
840+
fn test_align_to_non_trivial() {
841+
#[repr(align(8))] struct U64(u64, u64);
842+
#[repr(align(8))] struct U64U64U32(u64, u64, u32);
843+
let data = [U64(1, 2), U64(3, 4), U64(5, 6), U64(7, 8), U64(9, 10), U64(11, 12), U64(13, 14),
844+
U64(15, 16)];
845+
let (prefix, aligned, suffix) = unsafe { data.align_to::<U64U64U32>() };
846+
assert_eq!(aligned.len(), 4);
847+
assert_eq!(prefix.len() + suffix.len(), 2);
848+
}

0 commit comments

Comments
 (0)