bitmap_allocator/
lib.rs

1#![no_std]
2#![doc = include_str!("../README.md")]
3
4use core::ops::Range;
5
6use bit_field::BitField;
7
8/// Allocator of a bitmap, able to allocate / free bits.
9pub trait BitAlloc: Default {
10    /// The bitmap has a total of CAP bits, numbered from 0 to CAP-1 inclusively.
11    const CAP: usize;
12
13    /// The default value. Workaround for `const fn new() -> Self`.
14    const DEFAULT: Self;
15
16    /// Allocate a free bit.
17    fn alloc(&mut self) -> Option<usize>;
18
19    /// Allocate a free block with a given size, and return the first bit position.
20    ///
21    /// If `base` is not `None`, the allocator will try to allocate the block at the given base.
22    fn alloc_contiguous(
23        &mut self,
24        base: Option<usize>,
25        size: usize,
26        align_log2: usize,
27    ) -> Option<usize>;
28
29    /// Find a index not less than a given key, where the bit is free.
30    fn next(&self, key: usize) -> Option<usize>;
31
32    /// Free an allocated bit.
33    ///
34    /// Returns true if successful, false if the bit is already free.
35    fn dealloc(&mut self, key: usize) -> bool;
36
37    /// Free a contiguous block of bits.
38    ///
39    /// Returns true if successful, false if the bits in the range are already free.
40    fn dealloc_contiguous(&mut self, base: usize, size: usize) -> bool;
41
42    /// Mark bits in the range as unallocated (available)
43    fn insert(&mut self, range: Range<usize>);
44
45    /// Reverse of insert
46    fn remove(&mut self, range: Range<usize>);
47
48    /// Whether there are free bits remaining
49    #[deprecated = "use `!self.is_empty()` instead"]
50    fn any(&self) -> bool;
51
52    /// Returns true if no bits is available.
53    fn is_empty(&self) -> bool;
54
55    /// Whether a specific bit is free
56    fn test(&self, key: usize) -> bool;
57}
58
59/// A bitmap of 256 bits
60///
61/// ## Example
62///
63/// ```rust
64/// use bitmap_allocator::{BitAlloc, BitAlloc256};
65///
66/// let mut ba = BitAlloc256::default();
67/// ba.insert(0..16);
68/// for i in 0..16 {
69///     assert!(ba.test(i));
70/// }
71/// ba.remove(2..8);
72/// assert_eq!(ba.alloc(), Some(0));
73/// assert_eq!(ba.alloc(), Some(1));
74/// assert_eq!(ba.alloc(), Some(8));
75/// ba.dealloc(0);
76/// ba.dealloc(1);
77/// ba.dealloc(8);
78///
79/// assert!(!ba.is_empty());
80/// ```
81pub type BitAlloc256 = BitAllocCascade16<BitAlloc16>;
82/// A bitmap of 4K bits
83pub type BitAlloc4K = BitAllocCascade16<BitAlloc256>;
84/// A bitmap of 64K bits
85pub type BitAlloc64K = BitAllocCascade16<BitAlloc4K>;
86/// A bitmap of 1M bits
87pub type BitAlloc1M = BitAllocCascade16<BitAlloc64K>;
88/// A bitmap of 16M bits
89pub type BitAlloc16M = BitAllocCascade16<BitAlloc1M>;
90/// A bitmap of 256M bits
91pub type BitAlloc256M = BitAllocCascade16<BitAlloc16M>;
92
93/// Implement the bit allocator by segment tree algorithm.
94#[derive(Default)]
95pub struct BitAllocCascade16<T: BitAlloc> {
96    /// for each bit, 1 indicates available, 0 indicates inavailable
97    bitset: u16,
98    sub: [T; 16],
99}
100
101impl<T: BitAlloc> BitAlloc for BitAllocCascade16<T> {
102    const CAP: usize = T::CAP * 16;
103
104    const DEFAULT: Self = BitAllocCascade16 {
105        bitset: 0,
106        sub: [T::DEFAULT; 16],
107    };
108
109    fn alloc(&mut self) -> Option<usize> {
110        if !self.is_empty() {
111            let i = self.bitset.trailing_zeros() as usize;
112            let res = self.sub[i].alloc().unwrap() + i * T::CAP;
113            self.bitset.set_bit(i, !self.sub[i].is_empty());
114            Some(res)
115        } else {
116            None
117        }
118    }
119
120    fn alloc_contiguous(
121        &mut self,
122        base: Option<usize>,
123        size: usize,
124        align_log2: usize,
125    ) -> Option<usize> {
126        match base {
127            Some(base) => check_contiguous(self, base, Self::CAP, size, align_log2).then(|| {
128                self.remove(base..base + size);
129                base
130            }),
131            None => find_contiguous(self, Self::CAP, size, align_log2).inspect(|&base| {
132                self.remove(base..base + size);
133            }),
134        }
135    }
136
137    fn dealloc(&mut self, key: usize) -> bool {
138        let i = key / T::CAP;
139        self.bitset.set_bit(i, true);
140        self.sub[i].dealloc(key % T::CAP)
141    }
142
143    fn dealloc_contiguous(&mut self, base: usize, size: usize) -> bool {
144        let mut success = true;
145        let Range { start, end } = base..base + size;
146
147        // Check if the range is valid.
148        if end > Self::CAP {
149            return false;
150        }
151
152        for i in start / T::CAP..=(end - 1) / T::CAP {
153            let begin = if start / T::CAP == i {
154                start % T::CAP
155            } else {
156                0
157            };
158            let end = if end / T::CAP == i {
159                end % T::CAP
160            } else {
161                T::CAP
162            };
163            success = success && self.sub[i].dealloc_contiguous(begin, end - begin);
164            self.bitset.set_bit(i, !self.sub[i].is_empty());
165        }
166        success
167    }
168
169    fn insert(&mut self, range: Range<usize>) {
170        self.for_range(range, |sub: &mut T, range| sub.insert(range));
171    }
172    fn remove(&mut self, range: Range<usize>) {
173        self.for_range(range, |sub: &mut T, range| sub.remove(range));
174    }
175    fn any(&self) -> bool {
176        !self.is_empty()
177    }
178    fn is_empty(&self) -> bool {
179        self.bitset == 0
180    }
181    fn test(&self, key: usize) -> bool {
182        self.sub[key / T::CAP].test(key % T::CAP)
183    }
184    fn next(&self, key: usize) -> Option<usize> {
185        let idx = key / T::CAP;
186        (idx..16).find_map(|i| {
187            if self.bitset.get_bit(i) {
188                let key = if i == idx { key - T::CAP * idx } else { 0 };
189                self.sub[i].next(key).map(|x| x + T::CAP * i)
190            } else {
191                None
192            }
193        })
194    }
195}
196
197impl<T: BitAlloc> BitAllocCascade16<T> {
198    fn for_range(&mut self, range: Range<usize>, f: impl Fn(&mut T, Range<usize>)) {
199        let Range { start, end } = range;
200        assert!(start <= end);
201        assert!(end <= Self::CAP);
202        for i in start / T::CAP..=(end - 1) / T::CAP {
203            let begin = if start / T::CAP == i {
204                start % T::CAP
205            } else {
206                0
207            };
208            let end = if end / T::CAP == i {
209                end % T::CAP
210            } else {
211                T::CAP
212            };
213            f(&mut self.sub[i], begin..end);
214            self.bitset.set_bit(i, !self.sub[i].is_empty());
215        }
216    }
217}
218
219/// A bitmap consisting of only 16 bits.
220/// BitAlloc16 acts as the leaf (except the leaf bits of course) nodes in the segment trees.
221///
222/// ## Example
223///
224/// ```rust
225/// use bitmap_allocator::{BitAlloc, BitAlloc16};
226///
227/// let mut ba = BitAlloc16::default();
228/// assert_eq!(BitAlloc16::CAP, 16);
229/// ba.insert(0..16);
230/// for i in 0..16 {
231///     assert!(ba.test(i));
232/// }
233/// ba.remove(2..8);
234/// assert_eq!(ba.alloc(), Some(0));
235/// assert_eq!(ba.alloc(), Some(1));
236/// assert_eq!(ba.alloc(), Some(8));
237/// ba.dealloc(0);
238/// ba.dealloc(1);
239/// ba.dealloc(8);
240///
241/// assert!(!ba.is_empty());
242/// ```
243#[derive(Default)]
244pub struct BitAlloc16(u16);
245
246impl BitAlloc for BitAlloc16 {
247    const CAP: usize = u16::BITS as usize;
248
249    const DEFAULT: Self = Self(0);
250
251    fn alloc(&mut self) -> Option<usize> {
252        let i = self.0.trailing_zeros() as usize;
253        if i < Self::CAP {
254            self.0.set_bit(i, false);
255            Some(i)
256        } else {
257            None
258        }
259    }
260    fn alloc_contiguous(
261        &mut self,
262        base: Option<usize>,
263        size: usize,
264        align_log2: usize,
265    ) -> Option<usize> {
266        match base {
267            Some(base) => check_contiguous(self, base, Self::CAP, size, align_log2).then(|| {
268                self.remove(base..base + size);
269                base
270            }),
271            None => find_contiguous(self, Self::CAP, size, align_log2).inspect(|&base| {
272                self.remove(base..base + size);
273            }),
274        }
275    }
276
277    fn dealloc(&mut self, key: usize) -> bool {
278        let success = !self.test(key);
279        self.0.set_bit(key, true);
280        success
281    }
282
283    fn dealloc_contiguous(&mut self, base: usize, size: usize) -> bool {
284        if self.0.get_bits(base..base + size) == 0 {
285            self.insert(base..base + size);
286            return true;
287        }
288        false
289    }
290
291    fn insert(&mut self, range: Range<usize>) {
292        self.0.set_bits(range.clone(), 0xffff.get_bits(range));
293    }
294    fn remove(&mut self, range: Range<usize>) {
295        self.0.set_bits(range, 0);
296    }
297    fn any(&self) -> bool {
298        !self.is_empty()
299    }
300    fn is_empty(&self) -> bool {
301        self.0 == 0
302    }
303    fn test(&self, key: usize) -> bool {
304        self.0.get_bit(key)
305    }
306    fn next(&self, key: usize) -> Option<usize> {
307        (key..Self::CAP).find(|&i| self.0.get_bit(i))
308    }
309}
310
311fn find_contiguous(
312    ba: &impl BitAlloc,
313    capacity: usize,
314    size: usize,
315    align_log2: usize,
316) -> Option<usize> {
317    if align_log2 >= 64 || capacity < (1 << align_log2) || ba.is_empty() {
318        return None;
319    }
320
321    let mut base = 0;
322    // First, we need to make sure that base is aligned.
323    let start = ba.next(base)?;
324    base = align_up_log2(start, align_log2);
325
326    let mut offset = base;
327
328    while offset < capacity {
329        let next = ba.next(offset)?;
330        if next != offset {
331            // it can be guarenteed that no bit in (offset..next) is free
332            // move to next aligned position after next-1
333            assert!(next > offset);
334            base = (((next - 1) >> align_log2) + 1) << align_log2;
335            assert_ne!(offset, next);
336            offset = base;
337            continue;
338        }
339        offset += 1;
340        if offset - base == size {
341            return Some(base);
342        }
343    }
344    None
345}
346
347fn check_contiguous(
348    ba: &impl BitAlloc,
349    base: usize,
350    capacity: usize,
351    size: usize,
352    align_log2: usize,
353) -> bool {
354    if align_log2 >= 64 || capacity < (1 << align_log2) || ba.is_empty() {
355        return false;
356    }
357
358    // First, we need to make sure that base is aligned.
359    if !is_aligned_log2(base, align_log2) {
360        return false;
361    }
362
363    let mut offset = base;
364    while offset < capacity {
365        if let Some(next) = ba.next(offset) {
366            if next != offset {
367                return false;
368            }
369            offset += 1;
370            if offset - base == size {
371                return true;
372            }
373        } else {
374            return false;
375        }
376    }
377    false
378}
379
380fn align_up_log2(base: usize, align_log2: usize) -> usize {
381    (base + ((1 << align_log2) - 1)) & !((1 << align_log2) - 1)
382}
383
384fn is_aligned_log2(base: usize, align_log2: usize) -> bool {
385    (base & ((1 << align_log2) - 1)) == 0
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391
392    #[test]
393    fn bitalloc16() {
394        let mut ba = BitAlloc16::default();
395        assert_eq!(BitAlloc16::CAP, 16);
396        ba.insert(0..16);
397        for i in 0..16 {
398            assert!(ba.test(i));
399        }
400        ba.remove(2..8);
401        assert_eq!(ba.alloc(), Some(0));
402        assert_eq!(ba.alloc(), Some(1));
403        assert_eq!(ba.alloc(), Some(8));
404        ba.dealloc(0);
405        ba.dealloc(1);
406        ba.dealloc(8);
407
408        assert!(!ba.is_empty());
409        for _ in 0..10 {
410            assert!(ba.alloc().is_some());
411        }
412        assert!(ba.is_empty());
413        assert!(ba.alloc().is_none());
414
415        for key in 0..16 {
416            assert!(ba.dealloc(key));
417        }
418
419        assert!(!ba.dealloc(10));
420        assert!(!ba.dealloc(0));
421
422        assert_eq!(ba.alloc(), Some(0));
423        assert_eq!(ba.test(0), false);
424        assert_eq!(ba.alloc_contiguous(None, 2, 0), Some(1));
425        assert_eq!(ba.test(1), false);
426        assert_eq!(ba.test(2), false);
427
428        // Test alloc alignment.
429        assert_eq!(ba.alloc_contiguous(None, 2, 1), Some(4));
430        // Bit 3 is free due to alignment.
431        assert_eq!(ba.test(3), true);
432        assert_eq!(ba.test(4), false);
433        assert_eq!(ba.test(5), false);
434        assert_eq!(ba.next(5), Some(6));
435
436        // Test alloc alignment.
437        assert_eq!(ba.alloc_contiguous(None, 3, 3), Some(8));
438        assert_eq!(ba.next(8), Some(11));
439
440        assert_eq!(ba.alloc_contiguous(Some(2), 2, 1), None);
441        assert_eq!(ba.alloc_contiguous(Some(6), 2, 1), Some(6));
442        assert_eq!(ba.alloc_contiguous(Some(6), 2, 1), None);
443
444        assert!(ba.dealloc_contiguous(8, 3));
445
446        assert_eq!(ba.alloc_contiguous(Some(8), 3, 2), Some(8));
447        assert_eq!(ba.next(8), Some(11));
448        assert_eq!(ba.alloc_contiguous(Some(11), 1, 0), Some(11));
449        assert_eq!(ba.next(11), Some(12));
450
451        assert_eq!(ba.alloc_contiguous(Some(12), 3, 2), Some(12));
452
453        assert_eq!(ba.next(12), Some(15));
454
455        assert_eq!(ba.alloc(), Some(3));
456        assert_eq!(ba.alloc(), Some(15));
457
458        assert!(ba.is_empty());
459        assert!(ba.alloc().is_none());
460
461        assert!(ba.dealloc_contiguous(6, 2));
462        assert!(ba.dealloc_contiguous(8, 3));
463        assert!(ba.dealloc_contiguous(11, 1));
464        assert!(ba.dealloc_contiguous(12, 3));
465    }
466
467    #[test]
468    fn bitalloc4k() {
469        let mut ba = BitAlloc4K::default();
470        assert_eq!(BitAlloc4K::CAP, 4096);
471        ba.insert(0..4096);
472        for i in 0..4096 {
473            assert!(ba.test(i));
474        }
475        ba.remove(2..4094);
476        for i in 0..4096 {
477            assert_eq!(ba.test(i), !(2..4094).contains(&i));
478        }
479        assert_eq!(ba.alloc(), Some(0));
480        assert_eq!(ba.alloc(), Some(1));
481        assert_eq!(ba.alloc(), Some(4094));
482        ba.dealloc(0);
483        ba.dealloc(1);
484        ba.dealloc(4094);
485
486        assert!(!ba.is_empty());
487        for _ in 0..4 {
488            assert!(ba.alloc().is_some());
489        }
490        assert!(ba.is_empty());
491        assert!(ba.alloc().is_none());
492    }
493
494    #[test]
495    fn bitalloc_contiguous() {
496        let mut ba0 = BitAlloc16::default();
497        ba0.insert(0..BitAlloc16::CAP);
498        ba0.remove(3..6);
499        assert_eq!(ba0.next(0), Some(0));
500        assert_eq!(ba0.alloc_contiguous(None, 1, 1), Some(0));
501        assert_eq!(find_contiguous(&ba0, BitAlloc4K::CAP, 2, 0), Some(1));
502
503        let mut ba = BitAlloc4K::default();
504        assert_eq!(BitAlloc4K::CAP, 4096);
505        ba.insert(0..BitAlloc4K::CAP);
506        ba.remove(3..6);
507        assert_eq!(ba.next(0), Some(0));
508        assert_eq!(ba.alloc_contiguous(None, 1, 1), Some(0));
509        assert_eq!(ba.next(0), Some(1));
510        assert_eq!(ba.next(1), Some(1));
511        assert_eq!(ba.next(2), Some(2));
512        assert_eq!(find_contiguous(&ba, BitAlloc4K::CAP, 2, 0), Some(1));
513        assert_eq!(ba.alloc_contiguous(None, 2, 0), Some(1));
514        assert_eq!(ba.alloc_contiguous(None, 2, 3), Some(8));
515        ba.remove(0..4096 - 64);
516        assert_eq!(ba.alloc_contiguous(None, 128, 7), None);
517        assert_eq!(ba.alloc_contiguous(None, 7, 3), Some(4096 - 64));
518        ba.insert(321..323);
519        assert_eq!(ba.alloc_contiguous(None, 2, 1), Some(4096 - 64 + 8));
520        assert_eq!(ba.alloc_contiguous(None, 2, 0), Some(321));
521        assert_eq!(ba.alloc_contiguous(None, 64, 6), None);
522        assert_eq!(ba.alloc_contiguous(None, 32, 4), Some(4096 - 48));
523        for i in 0..4096 - 64 + 7 {
524            assert!(ba.dealloc(i));
525        }
526        for i in 4096 - 64 + 8..4096 - 64 + 10 {
527            assert!(ba.dealloc(i));
528        }
529        for i in 4096 - 48..4096 - 16 {
530            assert!(ba.dealloc(i));
531        }
532    }
533}
bitmap_allocator/lib.rs

bitmap_allocator/
lib.rs
