use std::alloc::{Layout, alloc_zeroed, dealloc, handle_alloc_error};
use std::ptr::NonNull;
pub struct AlignedBytes {
ptr: NonNull<u8>,
len: usize,
align: usize,
}
unsafe impl Send for AlignedBytes {}
unsafe impl Sync for AlignedBytes {}
impl AlignedBytes {
pub fn new_zeroed(len: usize, align: usize) -> Self {
assert!(
align.is_power_of_two(),
"AlignedBytes::new_zeroed: align ({align}) must be a power of two",
);
if len == 0 {
let ptr = NonNull::new(align as *mut u8).unwrap_or(NonNull::dangling());
return Self { ptr, len: 0, align };
}
let layout =
Layout::from_size_align(len, align).expect("AlignedBytes layout: size+align overflow");
let raw = unsafe { alloc_zeroed(layout) };
let ptr = match NonNull::new(raw) {
Some(p) => p,
None => handle_alloc_error(layout),
};
Self { ptr, len, align }
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn align(&self) -> usize {
self.align
}
pub fn as_slice(&self) -> &[u8] {
if self.len == 0 {
return &[];
}
unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
}
pub fn as_slice_mut(&mut self) -> &mut [u8] {
if self.len == 0 {
return &mut [];
}
unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
}
}
impl Drop for AlignedBytes {
fn drop(&mut self) {
if self.len == 0 {
return;
}
let layout = Layout::from_size_align(self.len, self.align)
.expect("AlignedBytes Drop: layout was valid at construction");
unsafe { dealloc(self.ptr.as_ptr(), layout) };
}
}
impl Clone for AlignedBytes {
fn clone(&self) -> Self {
let mut new = Self::new_zeroed(self.len, self.align);
if self.len > 0 {
new.as_slice_mut().copy_from_slice(self.as_slice());
}
new
}
}
impl std::fmt::Debug for AlignedBytes {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AlignedBytes")
.field("len", &self.len)
.field("align", &self.align)
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn allocates_aligned() {
let buf = AlignedBytes::new_zeroed(128, 64);
assert_eq!(buf.len(), 128);
assert_eq!(buf.align(), 64);
let addr = buf.as_slice().as_ptr() as usize;
assert_eq!(addr % 64, 0, "buffer base must be 64-byte aligned");
}
#[test]
fn zero_length_is_safe() {
let buf = AlignedBytes::new_zeroed(0, 64);
assert_eq!(buf.len(), 0);
assert!(buf.as_slice().is_empty());
}
#[test]
fn clone_deep_copies() {
let mut a = AlignedBytes::new_zeroed(16, 64);
a.as_slice_mut()[0] = 42;
let b = a.clone();
a.as_slice_mut()[0] = 99;
assert_eq!(b.as_slice()[0], 42, "clone must be a deep copy");
assert_eq!(a.as_slice()[0], 99);
}
#[test]
fn read_write_round_trip() {
let mut buf = AlignedBytes::new_zeroed(8, 16);
for (i, byte) in buf.as_slice_mut().iter_mut().enumerate() {
*byte = i as u8;
}
assert_eq!(buf.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7]);
}
#[test]
fn allocates_zero_initialized() {
let buf = AlignedBytes::new_zeroed(64, 64);
assert!(buf.as_slice().iter().all(|&b| b == 0));
}
#[test]
#[should_panic(expected = "must be a power of two")]
fn rejects_non_power_of_two_align() {
let _ = AlignedBytes::new_zeroed(8, 7);
}
}