1use std::alloc::{Layout, alloc, dealloc};
2use std::marker::PhantomData;
3
4pub const PAGE_SIZE: usize = 4096;
5
6#[derive(Debug)]
7pub struct AlignedBuf {
8 ptr: *mut u8,
9 layout: Layout,
10 _marker: PhantomData<[u8]>,
11}
12
13impl AlignedBuf {
14 pub fn new(size: usize, align: usize) -> Result<Self, std::io::Error> {
15 let layout = Layout::from_size_align(size, align)
16 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))?;
17
18 let ptr = unsafe { alloc(layout) };
19 if ptr.is_null() {
20 return Err(std::io::Error::new(
21 std::io::ErrorKind::OutOfMemory,
22 "allocation failed",
23 ));
24 }
25
26 Ok(Self {
27 ptr,
28 layout,
29 _marker: PhantomData,
30 })
31 }
32
33 pub fn new_page_aligned(size: usize) -> Result<Self, std::io::Error> {
34 Self::new(size, PAGE_SIZE)
35 }
36
37 pub fn from_slice(data: &[u8], align: usize) -> Result<Self, std::io::Error> {
38 let mut buf = Self::new(data.len(), align)?;
39 buf.as_mut_slice().copy_from_slice(data);
40 Ok(buf)
41 }
42
43 pub fn from_slice_page_aligned(data: &[u8]) -> Result<Self, std::io::Error> {
44 Self::from_slice(data, PAGE_SIZE)
45 }
46
47 pub fn as_slice(&self) -> &[u8] {
48 unsafe { std::slice::from_raw_parts(self.ptr, self.layout.size()) }
49 }
50
51 pub fn as_mut_slice(&mut self) -> &mut [u8] {
52 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.layout.size()) }
53 }
54
55 pub fn as_mut_ptr(&mut self) -> *mut u8 {
56 self.ptr
57 }
58
59 pub fn as_ptr(&self) -> *const u8 {
60 self.ptr
61 }
62
63 pub fn len(&self) -> usize {
64 self.layout.size()
65 }
66
67 pub fn is_empty(&self) -> bool {
68 self.layout.size() == 0
69 }
70}
71
72impl Drop for AlignedBuf {
73 fn drop(&mut self) {
74 unsafe { dealloc(self.ptr, self.layout) }
75 }
76}
77
78unsafe impl Send for AlignedBuf {}
79unsafe impl Sync for AlignedBuf {}
80
81pub fn align_down(n: usize, align: usize) -> usize {
82 n & !(align - 1)
83}
84
85pub fn align_up(n: usize, align: usize) -> usize {
86 (n + align - 1) & !(align - 1)
87}
88
89pub fn is_aligned(n: usize, align: usize) -> bool {
90 n & (align - 1) == 0
91}
92
93#[cfg(test)]
94mod tests {
95 use super::*;
96
97 #[test]
98 fn test_aligned_buf() {
99 let mut buf = AlignedBuf::new(8192, 4096).unwrap();
100 assert_eq!(buf.len(), 8192);
101 buf.as_mut_slice()[0] = 42;
102 assert_eq!(buf.as_slice()[0], 42);
103 }
104
105 #[test]
106 fn test_from_slice() {
107 let data = b"hello world";
108 let buf = AlignedBuf::from_slice(data, 4096).unwrap();
109 assert_eq!(buf.as_slice(), data);
110 }
111
112 #[test]
113 fn test_align_functions() {
114 assert_eq!(align_down(100, 16), 96);
115 assert_eq!(align_up(100, 16), 112);
116 assert!(is_aligned(16, 16));
117 assert!(!is_aligned(17, 16));
118 }
119}