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 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 unsafe { core::ptr::write(ptr.as_ptr(), Self::new(begin).ok_or(kerr!(EINVAL))?) };
57
58 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 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 alloc.free(begin - super::super::PAGE_SIZE, 1);
142 }
143}