Skip to main content

osiris/mem/pfa/
bitset.rs

1use core::pin::Pin;
2use core::ptr::NonNull;
3
4use crate::hal::mem::PhysAddr;
5
6use crate::{
7    error::Result,
8    types::{
9        bitset::BitAlloc,
10        boxed::{self, Box},
11    },
12};
13
14pub struct Allocator<const N: usize> {
15    begin: PhysAddr,
16    bitalloc: BitAlloc<N>,
17}
18
19impl<const WORDS: usize> Allocator<WORDS> {
20    pub fn new(begin: PhysAddr) -> Option<Self> {
21        if !begin.is_multiple_of(super::PAGE_SIZE) {
22            return None;
23        }
24
25        if begin > PhysAddr::MAX - (WORDS * super::PAGE_SIZE * usize::BITS as usize) {
26            return None;
27        }
28
29        Some(Self {
30            begin,
31            bitalloc: BitAlloc::new(WORDS * BitAlloc::<WORDS>::BITS_PER_WORD)?,
32        })
33    }
34}
35
36impl<const WORDS: usize> super::Allocator<WORDS> for Allocator<WORDS> {
37    fn initializer() -> unsafe fn(PhysAddr, usize) -> Result<Pin<Box<Self>>> {
38        |addr: PhysAddr, pcnt: usize| -> Result<Pin<Box<Self>>> {
39            if pcnt > WORDS {
40                todo!("Runtime page frame allocator for more than {} pages", WORDS)
41            }
42
43            if !addr.is_multiple_of(core::mem::align_of::<Self>()) {
44                return Err(kerr!(EINVAL));
45            }
46
47            let ptr = NonNull::new(addr.as_mut_ptr::<Self>()).ok_or(kerr!(EINVAL))?;
48            // Align this up to PAGE_SIZE
49            let begin = addr + size_of::<Self>();
50            let begin = if begin.is_multiple_of(super::PAGE_SIZE) {
51                begin
52            } else {
53                PhysAddr::new((begin.as_usize() + super::PAGE_SIZE - 1) & !(super::PAGE_SIZE - 1))
54            };
55            // TODO: Subtract the needed pages from the available
56            unsafe { core::ptr::write(ptr.as_ptr(), Self::new(begin).ok_or(kerr!(EINVAL))?) };
57
58            // Safety: Ptr is properly aligned and non-null. The validity of the memory at that address is valid by the call contract.
59            Ok(Pin::new(unsafe { boxed::Box::from_raw(ptr) }))
60        }
61    }
62
63    fn alloc(&mut self, page_count: usize) -> Option<PhysAddr> {
64        let idx = self.bitalloc.alloc(page_count)?;
65        Some(self.begin + (idx * super::PAGE_SIZE))
66    }
67
68    fn free(&mut self, addr: PhysAddr, page_count: usize) {
69        bug_on!(
70            !addr.is_multiple_of(super::PAGE_SIZE),
71            "free address {} is not page-aligned",
72            addr
73        );
74        // diff() is absolute, so a sub-begin address would silently map to a
75        // bit elsewhere in the bitmap.
76        bug_on!(
77            addr < self.begin,
78            "free address {} below allocator begin {}",
79            addr,
80            self.begin
81        );
82        let idx = addr.diff(self.begin) / super::PAGE_SIZE;
83        self.bitalloc.free(idx, page_count);
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::super::Allocator as _;
90    use super::*;
91
92    fn test_begin() -> PhysAddr {
93        let layout = std::alloc::Layout::from_size_align(
94            2 * 64 * super::super::PAGE_SIZE,
95            super::super::PAGE_SIZE,
96        )
97        .unwrap();
98        let ptr = unsafe { std::alloc::alloc(layout) };
99        PhysAddr::new(ptr as usize)
100    }
101
102    #[test]
103    fn alloc_free_roundtrip() {
104        let begin = test_begin();
105        let mut alloc = Allocator::<2>::new(begin).unwrap();
106
107        let a = alloc.alloc(1).unwrap();
108        let b = alloc.alloc(1).unwrap();
109        assert_ne!(a, b);
110
111        alloc.free(a, 1);
112        let c = alloc.alloc(1).unwrap();
113        assert_eq!(a, c, "freed page is returned by next alloc");
114    }
115
116    #[test]
117    fn alloc_returns_addresses_in_range() {
118        let begin = test_begin();
119        let mut alloc = Allocator::<1>::new(begin).unwrap();
120        let end = begin + 64 * super::super::PAGE_SIZE;
121
122        while let Some(addr) = alloc.alloc(1) {
123            assert!(
124                addr >= begin && addr < end,
125                "addr {addr} outside [{begin}, {end})"
126            );
127            assert!(
128                addr.is_multiple_of(super::super::PAGE_SIZE),
129                "addr {addr} not page-aligned"
130            );
131        }
132    }
133
134    #[test]
135    #[should_panic(expected = "below allocator begin")]
136    fn free_below_begin_panics() {
137        let begin = test_begin() + super::super::PAGE_SIZE;
138        let mut alloc = Allocator::<2>::new(begin).unwrap();
139        // diff() is absolute, so without the bound check a sub-begin address
140        // would silently clear a bit elsewhere in the bitmap.
141        alloc.free(begin - super::super::PAGE_SIZE, 1);
142    }
143}