lewis.moe
Nothing to see here, move along
1use crate::error::KernelError;
2
3pub mod bitmap_seal {
4 pub trait Sealed {}
5 impl<const N: usize> Sealed for [u64; N] {}
6}
7
8pub trait BitmapBacking: bitmap_seal::Sealed {
9 fn chunks(&self) -> &[u64];
10 fn chunks_mut(&mut self) -> &mut [u64];
11}
12
13impl<const N: usize> BitmapBacking for [u64; N] {
14 fn chunks(&self) -> &[u64] {
15 self.as_slice()
16 }
17 fn chunks_mut(&mut self) -> &mut [u64] {
18 self.as_mut_slice()
19 }
20}
21
22pub struct BitmapAllocator<B: BitmapBacking> {
23 bits: B,
24 total_items: usize,
25 used_items: usize,
26}
27
28impl<const N: usize> BitmapAllocator<[u64; N]> {
29 pub const fn new() -> Self {
30 Self {
31 bits: [!0u64; N],
32 total_items: 0,
33 used_items: 0,
34 }
35 }
36}
37
38impl<B: BitmapBacking> BitmapAllocator<B> {
39 pub const fn from_backing(backing: B) -> Self {
40 Self {
41 bits: backing,
42 total_items: 0,
43 used_items: 0,
44 }
45 }
46
47 pub fn init(&mut self, total: usize) {
48 let capacity = self.bits.chunks().len() * 64;
49 self.bits
50 .chunks_mut()
51 .iter_mut()
52 .for_each(|chunk| *chunk = !0u64);
53 self.total_items = total.min(capacity);
54 self.used_items = self.total_items;
55 }
56
57 pub fn mark_free(&mut self, idx: usize) {
58 assert!(idx < self.total_items, "mark_free index out of range");
59 let chunk = idx / 64;
60 let bit = idx % 64;
61 let mask = 1u64 << bit;
62 if self.bits.chunks()[chunk] & mask != 0 {
63 self.bits.chunks_mut()[chunk] &= !mask;
64 self.used_items -= 1;
65 }
66 }
67
68 pub fn mark_used(&mut self, idx: usize) {
69 assert!(idx < self.total_items, "mark_used index out of range");
70 let chunk = idx / 64;
71 let bit = idx % 64;
72 let mask = 1u64 << bit;
73 if self.bits.chunks()[chunk] & mask == 0 {
74 self.bits.chunks_mut()[chunk] |= mask;
75 self.used_items += 1;
76 }
77 }
78
79 pub fn mark_range_free(&mut self, range: core::ops::Range<usize>) {
80 let total = self.total_items;
81 range
82 .filter(|&idx| idx < total)
83 .for_each(|idx| self.mark_free(idx));
84 }
85
86 pub fn allocate(&mut self) -> Option<usize> {
87 let chunks_to_scan = self.total_items.div_ceil(64);
88 let result = (0..chunks_to_scan)
89 .find(|&chunk_idx| self.bits.chunks()[chunk_idx] != !0u64)
90 .and_then(|chunk_idx| {
91 let chunk = self.bits.chunks()[chunk_idx];
92 let bit = (!chunk).trailing_zeros() as usize;
93 let idx = chunk_idx * 64 + bit;
94 (idx < self.total_items).then_some((chunk_idx, bit, idx))
95 });
96
97 result.map(|(chunk_idx, bit, idx)| {
98 self.bits.chunks_mut()[chunk_idx] |= 1u64 << bit;
99 self.used_items += 1;
100 idx
101 })
102 }
103
104 pub fn allocate_contiguous(&mut self, count: usize) -> Option<usize> {
105 if count == 0 {
106 return None;
107 }
108 if count == 1 {
109 return self.allocate();
110 }
111
112 let chunks_to_scan = self.total_items.div_ceil(64);
113 let mut run_start = 0usize;
114 let mut run_len = 0usize;
115
116 let mut i = 0usize;
117 let result = loop {
118 if i >= self.total_items {
119 break None;
120 }
121 let chunk_idx = i / 64;
122 if chunk_idx >= chunks_to_scan {
123 break None;
124 }
125 let chunk = self.bits.chunks()[chunk_idx];
126 if chunk == !0u64 {
127 run_start = i + 64 - (i % 64);
128 run_len = 0;
129 i = run_start;
130 continue;
131 }
132 let bit = i % 64;
133 if chunk & (1u64 << bit) != 0 {
134 run_start = i + 1;
135 run_len = 0;
136 i += 1;
137 continue;
138 }
139 if run_len == 0 {
140 run_start = i;
141 }
142 run_len += 1;
143 if run_len >= count {
144 break Some(run_start);
145 }
146 i += 1;
147 };
148
149 result.inspect(|&base| {
150 (base..base + count).for_each(|idx| self.mark_used(idx));
151 })
152 }
153
154 pub fn deallocate(&mut self, idx: usize) -> Result<(), KernelError> {
155 if idx >= self.total_items {
156 return Err(KernelError::InvalidAddress);
157 }
158 let chunk = idx / 64;
159 let bit = idx % 64;
160 let mask = 1u64 << bit;
161 if self.bits.chunks()[chunk] & mask == 0 {
162 return Err(KernelError::InvalidAddress);
163 }
164 self.bits.chunks_mut()[chunk] &= !mask;
165 self.used_items -= 1;
166 Ok(())
167 }
168
169 pub const fn total_items(&self) -> usize {
170 self.total_items
171 }
172
173 pub const fn used_items(&self) -> usize {
174 self.used_items
175 }
176
177 pub const fn free_items(&self) -> usize {
178 self.total_items - self.used_items
179 }
180
181 #[allow(dead_code)]
182 pub fn is_used(&self, idx: usize) -> bool {
183 if idx >= self.total_items {
184 return false;
185 }
186 let chunk = idx / 64;
187 let bit = idx % 64;
188 self.bits.chunks()[chunk] & (1u64 << bit) != 0
189 }
190}
191
192#[cfg(kani)]
193mod kani_proofs {
194 use super::*;
195
196 fn make_bitmap(total: usize) -> BitmapAllocator<[u64; 1]> {
197 let mut bm = BitmapAllocator::<[u64; 1]>::new();
198 bm.init(total);
199 bm.mark_range_free(0..total);
200 bm
201 }
202
203 #[kani::proof]
204 #[kani::unwind(12)]
205 fn invariant_holds_after_op_sequence() {
206 let total: usize = kani::any();
207 kani::assume(total > 0 && total <= 8);
208 let mut bm = make_bitmap(total);
209
210 let ops: [u8; 8] = kani::any();
211 let mut last_allocated: Option<usize> = None;
212
213 ops.iter().for_each(|&op| {
214 match op % 3 {
215 0 => {
216 last_allocated = bm.allocate();
217 }
218 1 => {
219 last_allocated.take().into_iter().for_each(|idx| {
220 let _ = bm.deallocate(idx);
221 });
222 }
223 _ => {}
224 }
225 assert!(bm.used_items() + bm.free_items() == bm.total_items());
226 });
227 }
228
229 #[kani::proof]
230 #[kani::unwind(12)]
231 fn allocate_returns_in_range() {
232 let total: usize = kani::any();
233 kani::assume(total > 0 && total <= 8);
234 let mut bm = make_bitmap(total);
235
236 (0..8).for_each(|_| {
237 bm.allocate().into_iter().for_each(|idx| {
238 assert!(idx < total);
239 });
240 });
241 }
242
243 #[kani::proof]
244 #[kani::unwind(12)]
245 fn allocate_exhaustion() {
246 let total: usize = kani::any();
247 kani::assume(total > 0 && total <= 8);
248 let mut bm = make_bitmap(total);
249
250 (0..total).for_each(|_| {
251 assert!(bm.allocate().is_some());
252 });
253 assert!(bm.allocate().is_none());
254 assert!(bm.free_items() == 0);
255 }
256
257 #[kani::proof]
258 #[kani::unwind(12)]
259 fn double_free_detected() {
260 let total: usize = kani::any();
261 kani::assume(total > 0 && total <= 8);
262 let mut bm = make_bitmap(total);
263
264 let idx = bm.allocate().unwrap();
265 assert!(bm.deallocate(idx).is_ok());
266 assert!(bm.deallocate(idx).is_err());
267 }
268
269 #[kani::proof]
270 #[kani::unwind(12)]
271 fn mark_free_idempotent() {
272 let total: usize = kani::any();
273 kani::assume(total > 0 && total <= 8);
274 let mut bm = make_bitmap(total);
275
276 let idx: usize = kani::any();
277 kani::assume(idx < total);
278
279 let _ = bm.allocate();
280 bm.mark_free(idx);
281 let free_after_first = bm.free_items();
282 bm.mark_free(idx);
283 let free_after_second = bm.free_items();
284
285 assert!(free_after_first == free_after_second);
286 assert!(bm.used_items() + bm.free_items() == bm.total_items());
287 }
288
289 #[kani::proof]
290 #[kani::unwind(12)]
291 fn contiguous_allocation_returns_free_run() {
292 let total: usize = kani::any();
293 kani::assume(total >= 4 && total <= 8);
294
295 let count: usize = kani::any();
296 kani::assume(count >= 2 && count <= 4);
297
298 let mut bm = make_bitmap(total);
299
300 bm.allocate_contiguous(count).into_iter().for_each(|base| {
301 assert!(base + count <= total);
302 (base..base + count).for_each(|i| {
303 assert!(bm.is_used(i));
304 });
305 });
306
307 assert!(bm.used_items() + bm.free_items() == bm.total_items());
308 }
309
310 #[kani::proof]
311 #[kani::unwind(12)]
312 fn deallocate_out_of_range_returns_error() {
313 let total: usize = kani::any();
314 kani::assume(total > 0 && total <= 8);
315 let mut bm = make_bitmap(total);
316
317 let idx: usize = kani::any();
318 kani::assume(idx >= total);
319 assert!(bm.deallocate(idx).is_err());
320 }
321}