Skip to main content

pulith_fs/
align.rs

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}