Nothing to see here, move along meow
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 pub fn for_each_used_range<F: FnMut(usize, usize)>(&self, f: &mut F) {
192 let total = self.total_items;
193 let final_start = (0..total).fold(None::<usize>, |run_start, idx| {
194 let used = self.is_used(idx);
195 match (used, run_start) {
196 (true, None) => Some(idx),
197 (true, Some(_)) => run_start,
198 (false, Some(s)) => {
199 f(s, idx);
200 None
201 }
202 (false, None) => None,
203 }
204 });
205 if let Some(s) = final_start {
206 f(s, total);
207 }
208 }
209}
210
211#[cfg(kani)]
212mod kani_proofs {
213 use super::*;
214
215 fn make_bitmap(total: usize) -> BitmapAllocator<[u64; 1]> {
216 let mut bm = BitmapAllocator::<[u64; 1]>::new();
217 bm.init(total);
218 bm.mark_range_free(0..total);
219 bm
220 }
221
222 #[kani::proof]
223 #[kani::unwind(12)]
224 fn invariant_holds_after_op_sequence() {
225 let total: usize = kani::any();
226 kani::assume(total > 0 && total <= 8);
227 let mut bm = make_bitmap(total);
228
229 let ops: [u8; 8] = kani::any();
230 let mut last_allocated: Option<usize> = None;
231
232 ops.iter().for_each(|&op| {
233 match op % 3 {
234 0 => {
235 last_allocated = bm.allocate();
236 }
237 1 => {
238 last_allocated.take().into_iter().for_each(|idx| {
239 let _ = bm.deallocate(idx);
240 });
241 }
242 _ => {}
243 }
244 assert!(bm.used_items() + bm.free_items() == bm.total_items());
245 });
246 }
247
248 #[kani::proof]
249 #[kani::unwind(12)]
250 fn allocate_returns_in_range() {
251 let total: usize = kani::any();
252 kani::assume(total > 0 && total <= 8);
253 let mut bm = make_bitmap(total);
254
255 (0..8).for_each(|_| {
256 bm.allocate().into_iter().for_each(|idx| {
257 assert!(idx < total);
258 });
259 });
260 }
261
262 #[kani::proof]
263 #[kani::unwind(12)]
264 fn allocate_exhaustion() {
265 let total: usize = kani::any();
266 kani::assume(total > 0 && total <= 8);
267 let mut bm = make_bitmap(total);
268
269 (0..total).for_each(|_| {
270 assert!(bm.allocate().is_some());
271 });
272 assert!(bm.allocate().is_none());
273 assert!(bm.free_items() == 0);
274 }
275
276 #[kani::proof]
277 #[kani::unwind(12)]
278 fn double_free_detected() {
279 let total: usize = kani::any();
280 kani::assume(total > 0 && total <= 8);
281 let mut bm = make_bitmap(total);
282
283 let idx = bm.allocate().unwrap();
284 assert!(bm.deallocate(idx).is_ok());
285 assert!(bm.deallocate(idx).is_err());
286 }
287
288 #[kani::proof]
289 #[kani::unwind(12)]
290 fn mark_free_idempotent() {
291 let total: usize = kani::any();
292 kani::assume(total > 0 && total <= 8);
293 let mut bm = make_bitmap(total);
294
295 let idx: usize = kani::any();
296 kani::assume(idx < total);
297
298 let _ = bm.allocate();
299 bm.mark_free(idx);
300 let free_after_first = bm.free_items();
301 bm.mark_free(idx);
302 let free_after_second = bm.free_items();
303
304 assert!(free_after_first == free_after_second);
305 assert!(bm.used_items() + bm.free_items() == bm.total_items());
306 }
307
308 #[kani::proof]
309 #[kani::unwind(12)]
310 fn contiguous_allocation_returns_free_run() {
311 let total: usize = kani::any();
312 kani::assume(total >= 4 && total <= 8);
313
314 let count: usize = kani::any();
315 kani::assume(count >= 2 && count <= 4);
316
317 let mut bm = make_bitmap(total);
318
319 bm.allocate_contiguous(count).into_iter().for_each(|base| {
320 assert!(base + count <= total);
321 (base..base + count).for_each(|i| {
322 assert!(bm.is_used(i));
323 });
324 });
325
326 assert!(bm.used_items() + bm.free_items() == bm.total_items());
327 }
328
329 #[kani::proof]
330 #[kani::unwind(12)]
331 fn deallocate_out_of_range_returns_error() {
332 let total: usize = kani::any();
333 kani::assume(total > 0 && total <= 8);
334 let mut bm = make_bitmap(total);
335
336 let idx: usize = kani::any();
337 kani::assume(idx >= total);
338 assert!(bm.deallocate(idx).is_err());
339 }
340}