pub use crate::offset::{align_offset, AlignedOffset};
#[cfg(feature = "alloc")]
use alloc::{
borrow::{Cow, ToOwned},
boxed::Box,
};
use core::ops::{
Deref, DerefMut, Index, IndexMut, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo,
RangeToInclusive,
};
use core::{
cmp::{max, min},
fmt::Debug,
};
#[cfg(feature = "std")]
use std::io::IoSliceMut;
pub unsafe trait Alignment: Debug {
const ALIGNMENT: usize;
}
pub unsafe trait AlignedTo<A: Alignment>: Alignment {}
#[derive(Debug, Copy, Clone)]
pub struct A1;
unsafe impl Alignment for A1 {
const ALIGNMENT: usize = 1;
}
unsafe impl AlignedTo<A1> for A1 {}
#[repr(align(2))]
#[derive(Debug, Copy, Clone)]
pub struct A2;
unsafe impl Alignment for A2 {
const ALIGNMENT: usize = 2;
}
unsafe impl AlignedTo<A1> for A2 {}
unsafe impl AlignedTo<A2> for A2 {}
#[repr(align(4))]
#[derive(Debug, Copy, Clone)]
pub struct A4;
unsafe impl Alignment for A4 {
const ALIGNMENT: usize = 4;
}
unsafe impl AlignedTo<A1> for A4 {}
unsafe impl AlignedTo<A2> for A4 {}
unsafe impl AlignedTo<A4> for A4 {}
#[repr(align(8))]
#[derive(Debug, Copy, Clone)]
pub struct A8;
unsafe impl Alignment for A8 {
const ALIGNMENT: usize = 8;
}
unsafe impl AlignedTo<A1> for A8 {}
unsafe impl AlignedTo<A2> for A8 {}
unsafe impl AlignedTo<A4> for A8 {}
unsafe impl AlignedTo<A8> for A8 {}
pub trait AsAligned<A: Alignment> {
fn as_aligned(&self) -> &AlignedSlice<A>;
}
pub trait AsAlignedMut<A: Alignment>: AsAligned<A> {
fn as_aligned_mut(&mut self) -> &mut AlignedSlice<A>;
}
pub trait TryAsAligned<A: Alignment> {
fn try_as_aligned(&self) -> Result<&AlignedSlice<A>, Misaligned>;
}
pub trait TryAsAlignedMut<A: Alignment>: TryAsAligned<A> {
fn try_as_aligned_mut(&mut self) -> Result<&mut AlignedSlice<A>, Misaligned>;
}
#[repr(C)]
#[derive(Debug, Eq)]
pub struct AlignedSlice<A: Alignment> {
alignment: A,
data: [u8],
}
impl<A: Alignment> PartialEq for AlignedSlice<A> {
fn eq(&self, other: &Self) -> bool {
self.data == other.data
}
}
#[cfg(feature = "alloc")]
impl<A: Alignment> ToOwned for AlignedSlice<A> {
type Owned = Box<AlignedSlice<A>>;
fn to_owned(&self) -> Self::Owned {
let mut owned = alloc_aligned(self.len());
owned.copy_from_slice(self);
owned
}
}
#[cfg(feature = "alloc")]
pub fn copy_to_align<'a, A: Alignment>(data: &'a [u8]) -> Cow<'a, AlignedSlice<A>> {
if is_aligned_to::<A>(data) {
Cow::Borrowed(data.try_as_aligned().unwrap())
} else {
let mut copy = alloc_aligned::<A>(data.len());
copy.copy_from_slice(data);
Cow::Owned(copy)
}
}
#[cfg(feature = "alloc")]
pub fn alloc_aligned<A: Alignment>(size: usize) -> Box<AlignedSlice<A>> {
let layout = alloc::alloc::Layout::from_size_align(size, A::ALIGNMENT).unwrap();
unsafe {
let bs = core::slice::from_raw_parts_mut(alloc::alloc::alloc_zeroed(layout), size);
Box::from_raw(to_alignedslice_unchecked_mut(bs))
}
}
#[cfg(feature = "std")]
pub fn read_to_slice<A: Alignment, R: std::io::Read>(
mut r: R,
size_hint: Option<usize>,
) -> std::io::Result<Box<AlignedSlice<A>>> {
let mut bytes_read = 0;
let mut eof_byte = [0u8];
let mut buf = alloc_aligned(size_hint.unwrap_or(4000));
loop {
let mut ios = [
IoSliceMut::new(&mut buf.as_mut()[bytes_read..]),
IoSliceMut::new(&mut eof_byte),
];
let this_read = r.read_vectored(&mut ios)?;
bytes_read += this_read;
if this_read == 0 {
return Ok(realloc(buf, bytes_read));
}
if bytes_read > buf.len() {
buf = realloc(buf, max(bytes_read * 2, 4000));
buf[bytes_read - 1] = eof_byte[0];
}
}
}
fn realloc<A: Alignment>(s: Box<AlignedSlice<A>>, new_size: usize) -> Box<AlignedSlice<A>> {
if s.len() == new_size {
s
} else {
let mut new_buf = alloc_aligned(new_size);
let up_to = min(new_size, s.len());
new_buf.as_mut()[..up_to].copy_from_slice(s[..up_to].as_ref());
new_buf
}
}
impl<A: Alignment> AlignedSlice<A> {
pub fn split_at(&self, mid: usize) -> (&AlignedSlice<A>, &[u8]) {
let (before, after) = self.data.split_at(mid);
(unsafe { to_alignedslice_unchecked(before) }, after)
}
}
impl<A: Alignment> Deref for AlignedSlice<A> {
type Target = [u8];
fn deref(&self) -> &Self::Target {
&self.data
}
}
impl<A: Alignment> DerefMut for AlignedSlice<A> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.data
}
}
impl<T: AsRef<[u8]>> AsAligned<A1> for T {
fn as_aligned(&self) -> &AlignedSlice<A1> {
let s = self.as_ref();
s.as_aligned()
}
}
impl<T: AsRef<[u8]> + AsMut<[u8]>> AsAlignedMut<A1> for T {
fn as_aligned_mut(&mut self) -> &mut AlignedSlice<A1> {
let s = self.as_mut();
s.as_aligned_mut()
}
}
impl<A: Alignment, T: AsRef<[u8]>> TryAsAligned<A> for T {
fn try_as_aligned(&self) -> Result<&AlignedSlice<A>, Misaligned> {
self.as_ref().try_as_aligned()
}
}
impl<A: Alignment, T: AsRef<[u8]> + AsMut<[u8]>> TryAsAlignedMut<A> for T {
fn try_as_aligned_mut(&mut self) -> Result<&mut AlignedSlice<A>, Misaligned> {
self.as_mut().try_as_aligned_mut()
}
}
impl AsAligned<A1> for [u8] {
fn as_aligned(&self) -> &AlignedSlice<A1> {
unsafe { to_alignedslice_unchecked(self) }
}
}
impl AsAlignedMut<A1> for [u8] {
fn as_aligned_mut(&mut self) -> &mut AlignedSlice<A1> {
unsafe { to_alignedslice_unchecked_mut(self) }
}
}
impl<A: Alignment> TryAsAligned<A> for [u8] {
fn try_as_aligned(&self) -> Result<&AlignedSlice<A>, Misaligned> {
self.as_aligned().try_as_aligned()
}
}
impl<A: Alignment> TryAsAlignedMut<A> for [u8] {
fn try_as_aligned_mut(&mut self) -> Result<&mut AlignedSlice<A>, Misaligned> {
self.as_aligned_mut().try_as_aligned_mut()
}
}
impl<FromA: Alignment, ToA: Alignment> AsAligned<ToA> for AlignedSlice<FromA>
where
FromA: AlignedTo<ToA>,
{
fn as_aligned(&self) -> &AlignedSlice<ToA> {
let initial_align = core::mem::align_of_val(self);
let out = unsafe { &*(self as *const Self as *const AlignedSlice<ToA>) };
assert!(initial_align >= core::mem::align_of_val(out));
out
}
}
impl<FromA: Alignment, ToA: Alignment> AsAlignedMut<ToA> for AlignedSlice<FromA>
where
FromA: AlignedTo<ToA>,
{
fn as_aligned_mut(&mut self) -> &mut AlignedSlice<ToA> {
let initial_align = core::mem::align_of_val(self);
let out = unsafe { &mut *(self as *mut Self as *mut AlignedSlice<ToA>) };
assert!(initial_align >= core::mem::align_of_val(out));
out
}
}
impl<FromA: Alignment, ToA: Alignment> TryAsAligned<ToA> for AlignedSlice<FromA> {
fn try_as_aligned(&self) -> Result<&AlignedSlice<ToA>, Misaligned> {
if core::mem::align_of::<FromA>() >= core::mem::align_of::<ToA>()
|| is_aligned_to::<ToA>(self)
{
Ok(unsafe { &*(self as *const Self as *const AlignedSlice<ToA>) })
} else {
Err(Misaligned {})
}
}
}
impl<FromA: Alignment, ToA: Alignment> TryAsAlignedMut<ToA> for AlignedSlice<FromA> {
fn try_as_aligned_mut(&mut self) -> Result<&mut AlignedSlice<ToA>, Misaligned> {
if core::mem::align_of::<FromA>() >= core::mem::align_of::<ToA>()
|| is_aligned_to::<ToA>(self)
{
Ok(unsafe { &mut *(self as *mut Self as *mut AlignedSlice<ToA>) })
} else {
Err(Misaligned {})
}
}
}
impl<FromA, ToA: Alignment> AsRef<AlignedSlice<ToA>> for AlignedSlice<FromA>
where
FromA: AlignedTo<ToA>,
{
fn as_ref(&self) -> &AlignedSlice<ToA> {
self.as_aligned()
}
}
impl<FromA, ToA: Alignment> AsMut<AlignedSlice<ToA>> for AlignedSlice<FromA>
where
FromA: AlignedTo<ToA>,
{
fn as_mut(&mut self) -> &mut AlignedSlice<ToA> {
self.as_aligned_mut()
}
}
impl<A: Alignment> Index<usize> for AlignedSlice<A> {
type Output = u8;
fn index(&self, index: usize) -> &Self::Output {
&self.data[index]
}
}
impl<A: Alignment> Index<RangeTo<usize>> for AlignedSlice<A> {
type Output = AlignedSlice<A>;
fn index(&self, index: RangeTo<usize>) -> &Self::Output {
unsafe { to_alignedslice_unchecked(&self.data[index]) }
}
}
impl<A: Alignment> Index<RangeToInclusive<usize>> for AlignedSlice<A> {
type Output = AlignedSlice<A>;
fn index(&self, index: RangeToInclusive<usize>) -> &Self::Output {
unsafe { to_alignedslice_unchecked(&self.data[index]) }
}
}
impl<A: Alignment> Index<Range<usize>> for AlignedSlice<A> {
type Output = AlignedSlice<A1>;
fn index(&self, index: Range<usize>) -> &Self::Output {
&self.data[index].as_aligned()
}
}
impl<A: Alignment> Index<RangeInclusive<usize>> for AlignedSlice<A> {
type Output = AlignedSlice<A1>;
fn index(&self, index: RangeInclusive<usize>) -> &Self::Output {
&self.data[index].as_aligned()
}
}
impl<A: Alignment> Index<RangeFrom<usize>> for AlignedSlice<A> {
type Output = AlignedSlice<A1>;
fn index(&self, index: RangeFrom<usize>) -> &Self::Output {
&self.data[index].as_aligned()
}
}
impl<A: Alignment> Index<RangeFull> for AlignedSlice<A> {
type Output = AlignedSlice<A>;
fn index(&self, _: RangeFull) -> &Self::Output {
&self
}
}
impl<A: Alignment> Index<Range<AlignedOffset<A>>> for AlignedSlice<A> {
type Output = AlignedSlice<A>;
fn index(&self, index: Range<AlignedOffset<A>>) -> &Self::Output {
unsafe {
to_alignedslice_unchecked(&self.data[index.start.to_usize()..index.end.to_usize()])
}
}
}
impl<A: Alignment> Index<RangeInclusive<AlignedOffset<A>>> for AlignedSlice<A> {
type Output = AlignedSlice<A>;
fn index(&self, index: RangeInclusive<AlignedOffset<A>>) -> &Self::Output {
unsafe {
to_alignedslice_unchecked(&self.data[index.start().to_usize()..=index.end().to_usize()])
}
}
}
impl<A: Alignment> Index<RangeFrom<AlignedOffset<A>>> for AlignedSlice<A> {
type Output = AlignedSlice<A>;
fn index(&self, index: RangeFrom<AlignedOffset<A>>) -> &Self::Output {
unsafe { to_alignedslice_unchecked(&self.data[index.start.to_usize()..]) }
}
}
impl<A: Alignment> IndexMut<usize> for AlignedSlice<A> {
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.data[index]
}
}
impl<A: Alignment> IndexMut<RangeTo<usize>> for AlignedSlice<A> {
fn index_mut(&mut self, index: RangeTo<usize>) -> &mut Self::Output {
unsafe { to_alignedslice_unchecked_mut(&mut self.data[index]) }
}
}
impl<A: Alignment> IndexMut<RangeToInclusive<usize>> for AlignedSlice<A> {
fn index_mut(&mut self, index: RangeToInclusive<usize>) -> &mut Self::Output {
unsafe { to_alignedslice_unchecked_mut(&mut self.data[index]) }
}
}
impl<A: Alignment> IndexMut<Range<usize>> for AlignedSlice<A> {
fn index_mut(&mut self, index: Range<usize>) -> &mut Self::Output {
self.data[index].as_aligned_mut()
}
}
impl<A: Alignment> IndexMut<RangeInclusive<usize>> for AlignedSlice<A> {
fn index_mut(&mut self, index: RangeInclusive<usize>) -> &mut Self::Output {
self.data[index].as_aligned_mut()
}
}
impl<A: Alignment> IndexMut<RangeFrom<usize>> for AlignedSlice<A> {
fn index_mut(&mut self, index: RangeFrom<usize>) -> &mut Self::Output {
self.data[index].as_aligned_mut()
}
}
impl<A: Alignment> IndexMut<RangeFull> for AlignedSlice<A> {
fn index_mut(&mut self, _: RangeFull) -> &mut Self::Output {
self
}
}
impl<A: Alignment> IndexMut<Range<AlignedOffset<A>>> for AlignedSlice<A> {
fn index_mut(&mut self, index: Range<AlignedOffset<A>>) -> &mut Self::Output {
unsafe {
to_alignedslice_unchecked_mut(
&mut self.data[index.start.to_usize()..index.end.to_usize()],
)
}
}
}
impl<A: Alignment> IndexMut<RangeInclusive<AlignedOffset<A>>> for AlignedSlice<A> {
fn index_mut(&mut self, index: RangeInclusive<AlignedOffset<A>>) -> &mut Self::Output {
unsafe {
to_alignedslice_unchecked_mut(
&mut self.data[index.start().to_usize()..=index.end().to_usize()],
)
}
}
}
impl<A: Alignment> IndexMut<RangeFrom<AlignedOffset<A>>> for AlignedSlice<A> {
fn index_mut(&mut self, index: RangeFrom<AlignedOffset<A>>) -> &mut Self::Output {
unsafe { to_alignedslice_unchecked_mut(&mut self.data[index.start.to_usize()..]) }
}
}
unsafe fn to_alignedslice_unchecked<'a, A: Alignment>(value: &'a [u8]) -> &'a AlignedSlice<A> {
debug_assert!(is_aligned_to::<A>(value));
#[allow(unused_unsafe)]
unsafe {
&*(value as *const [u8] as *const AlignedSlice<A>)
}
}
unsafe fn to_alignedslice_unchecked_mut<'a, A: Alignment>(
value: &'a mut [u8],
) -> &'a mut AlignedSlice<A> {
debug_assert!(is_aligned_to::<A>(value));
#[allow(unused_unsafe)]
unsafe {
&mut *(value as *mut [u8] as *mut AlignedSlice<A>)
}
}
fn is_aligned_to<A: Alignment>(value: &[u8]) -> bool {
is_aligned(value, A::ALIGNMENT)
}
pub(crate) fn is_aligned(value: &[u8], alignment: usize) -> bool {
value as *const [u8] as *const u8 as usize % alignment == 0
}
#[derive(Debug)]
pub struct Misaligned {}
#[cfg(feature = "std")]
impl std::error::Error for Misaligned {}
impl core::fmt::Display for Misaligned {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "Misaligned")
}
}
impl<'a> core::convert::From<&'a [u8]> for &'a AlignedSlice<A1> {
fn from(value: &'a [u8]) -> Self {
unsafe { to_alignedslice_unchecked(value) }
}
}
impl<'a> core::convert::From<&'a mut [u8]> for &'a mut AlignedSlice<A1> {
fn from(value: &'a mut [u8]) -> Self {
unsafe { to_alignedslice_unchecked_mut(value) }
}
}
impl<A: Alignment> AsRef<[u8]> for AlignedSlice<A> {
fn as_ref(&self) -> &[u8] {
&self.data
}
}
impl<A: Alignment> AsMut<[u8]> for AlignedSlice<A> {
fn as_mut(&mut self) -> &mut [u8] {
&mut self.data
}
}
fn align_bytes<'a, A: Alignment>(value: &'a [u8]) -> &'a AlignedSlice<A> {
let p = value as *const [u8] as *const u8 as usize;
let offset = p.wrapping_neg() & (A::ALIGNMENT - 1);
value[offset..].try_as_aligned().unwrap()
}
pub fn empty_aligned<A: Alignment>() -> &'static AlignedSlice<A> {
&align_bytes(b" ")[..0]
}
#[cfg(test)]
mod test {
use super::{read_to_slice, AlignedSlice, A8};
#[test]
fn test_read_to_slice() {
let mut d: Vec<u8> = vec![0; 16384];
for x in 0..16384 {
d[x] = (x % 256) as u8;
}
for size_hint in &[
Some(11),
None,
Some(0),
Some(1),
Some(7999),
Some(8000),
Some(8001),
] {
let s: Box<AlignedSlice<A8>> = read_to_slice(b"".as_ref(), *size_hint).unwrap();
assert_eq!(**s, *b"");
let s: Box<AlignedSlice<A8>> = read_to_slice(&d[..12], *size_hint).unwrap();
assert_eq!(**s, d[..12]);
let s: Box<AlignedSlice<A8>> = read_to_slice(&d[..8000], *size_hint).unwrap();
assert_eq!(**s, d[..8000]);
let s: Box<AlignedSlice<A8>> = read_to_slice(d.as_slice(), *size_hint).unwrap();
assert_eq!(&**s, d.as_slice());
}
}
}