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    if let Some(start) = ba.next(base) {
324        base = align_up_log2(start, align_log2);
325    } else {
326        return None;
327    }
328
329    let mut offset = base;
330
331    while offset < capacity {
332        if let Some(next) = ba.next(offset) {
333            if next != offset {
334                // it can be guarenteed that no bit in (offset..next) is free
335                // move to next aligned position after next-1
336                assert!(next > offset);
337                base = (((next - 1) >> align_log2) + 1) << align_log2;
338                assert_ne!(offset, next);
339                offset = base;
340                continue;
341            }
342        } else {
343            return None;
344        }
345        offset += 1;
346        if offset - base == size {
347            return Some(base);
348        }
349    }
350    None
351}
352
353fn check_contiguous(
354    ba: &impl BitAlloc,
355    base: usize,
356    capacity: usize,
357    size: usize,
358    align_log2: usize,
359) -> bool {
360    if align_log2 >= 64 || capacity < (1 << align_log2) || ba.is_empty() {
361        return false;
362    }
363
364    // First, we need to make sure that base is aligned.
365    if !is_aligned_log2(base, align_log2) {
366        return false;
367    }
368
369    let mut offset = base;
370    while offset < capacity {
371        if let Some(next) = ba.next(offset) {
372            if next != offset {
373                return false;
374            }
375            offset += 1;
376            if offset - base == size {
377                return true;
378            }
379        } else {
380            return false;
381        }
382    }
383    false
384}
385
386fn align_up_log2(base: usize, align_log2: usize) -> usize {
387    (base + ((1 << align_log2) - 1)) & !((1 << align_log2) - 1)
388}
389
390fn is_aligned_log2(base: usize, align_log2: usize) -> bool {
391    (base & ((1 << align_log2) - 1)) == 0
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397
398    #[test]
399    fn bitalloc16() {
400        let mut ba = BitAlloc16::default();
401        assert_eq!(BitAlloc16::CAP, 16);
402        ba.insert(0..16);
403        for i in 0..16 {
404            assert!(ba.test(i));
405        }
406        ba.remove(2..8);
407        assert_eq!(ba.alloc(), Some(0));
408        assert_eq!(ba.alloc(), Some(1));
409        assert_eq!(ba.alloc(), Some(8));
410        ba.dealloc(0);
411        ba.dealloc(1);
412        ba.dealloc(8);
413
414        assert!(!ba.is_empty());
415        for _ in 0..10 {
416            assert!(ba.alloc().is_some());
417        }
418        assert!(ba.is_empty());
419        assert!(ba.alloc().is_none());
420
421        for key in 0..16 {
422            assert!(ba.dealloc(key));
423        }
424
425        assert!(!ba.dealloc(10));
426        assert!(!ba.dealloc(0));
427
428        assert_eq!(ba.alloc(), Some(0));
429        assert_eq!(ba.test(0), false);
430        assert_eq!(ba.alloc_contiguous(None, 2, 0), Some(1));
431        assert_eq!(ba.test(1), false);
432        assert_eq!(ba.test(2), false);
433
434        // Test alloc alignment.
435        assert_eq!(ba.alloc_contiguous(None, 2, 1), Some(4));
436        // Bit 3 is free due to alignment.
437        assert_eq!(ba.test(3), true);
438        assert_eq!(ba.test(4), false);
439        assert_eq!(ba.test(5), false);
440        assert_eq!(ba.next(5), Some(6));
441
442        // Test alloc alignment.
443        assert_eq!(ba.alloc_contiguous(None, 3, 3), Some(8));
444        assert_eq!(ba.next(8), Some(11));
445
446        assert_eq!(ba.alloc_contiguous(Some(2), 2, 1), None);
447        assert_eq!(ba.alloc_contiguous(Some(6), 2, 1), Some(6));
448        assert_eq!(ba.alloc_contiguous(Some(6), 2, 1), None);
449
450        assert!(ba.dealloc_contiguous(8, 3));
451
452        assert_eq!(ba.alloc_contiguous(Some(8), 3, 2), Some(8));
453        assert_eq!(ba.next(8), Some(11));
454        assert_eq!(ba.alloc_contiguous(Some(11), 1, 0), Some(11));
455        assert_eq!(ba.next(11), Some(12));
456
457        assert_eq!(ba.alloc_contiguous(Some(12), 3, 2), Some(12));
458
459        assert_eq!(ba.next(12), Some(15));
460
461        assert_eq!(ba.alloc(), Some(3));
462        assert_eq!(ba.alloc(), Some(15));
463
464        assert!(ba.is_empty());
465        assert!(ba.alloc().is_none());
466
467        assert!(ba.dealloc_contiguous(6, 2));
468        assert!(ba.dealloc_contiguous(8, 3));
469        assert!(ba.dealloc_contiguous(11, 1));
470        assert!(ba.dealloc_contiguous(12, 3));
471    }
472
473    #[test]
474    fn bitalloc4k() {
475        let mut ba = BitAlloc4K::default();
476        assert_eq!(BitAlloc4K::CAP, 4096);
477        ba.insert(0..4096);
478        for i in 0..4096 {
479            assert!(ba.test(i));
480        }
481        ba.remove(2..4094);
482        for i in 0..4096 {
483            assert_eq!(ba.test(i), !(2..4094).contains(&i));
484        }
485        assert_eq!(ba.alloc(), Some(0));
486        assert_eq!(ba.alloc(), Some(1));
487        assert_eq!(ba.alloc(), Some(4094));
488        ba.dealloc(0);
489        ba.dealloc(1);
490        ba.dealloc(4094);
491
492        assert!(!ba.is_empty());
493        for _ in 0..4 {
494            assert!(ba.alloc().is_some());
495        }
496        assert!(ba.is_empty());
497        assert!(ba.alloc().is_none());
498    }
499
500    #[test]
501    fn bitalloc_contiguous() {
502        let mut ba0 = BitAlloc16::default();
503        ba0.insert(0..BitAlloc16::CAP);
504        ba0.remove(3..6);
505        assert_eq!(ba0.next(0), Some(0));
506        assert_eq!(ba0.alloc_contiguous(None, 1, 1), Some(0));
507        assert_eq!(find_contiguous(&ba0, BitAlloc4K::CAP, 2, 0), Some(1));
508
509        let mut ba = BitAlloc4K::default();
510        assert_eq!(BitAlloc4K::CAP, 4096);
511        ba.insert(0..BitAlloc4K::CAP);
512        ba.remove(3..6);
513        assert_eq!(ba.next(0), Some(0));
514        assert_eq!(ba.alloc_contiguous(None, 1, 1), Some(0));
515        assert_eq!(ba.next(0), Some(1));
516        assert_eq!(ba.next(1), Some(1));
517        assert_eq!(ba.next(2), Some(2));
518        assert_eq!(find_contiguous(&ba, BitAlloc4K::CAP, 2, 0), Some(1));
519        assert_eq!(ba.alloc_contiguous(None, 2, 0), Some(1));
520        assert_eq!(ba.alloc_contiguous(None, 2, 3), Some(8));
521        ba.remove(0..4096 - 64);
522        assert_eq!(ba.alloc_contiguous(None, 128, 7), None);
523        assert_eq!(ba.alloc_contiguous(None, 7, 3), Some(4096 - 64));
524        ba.insert(321..323);
525        assert_eq!(ba.alloc_contiguous(None, 2, 1), Some(4096 - 64 + 8));
526        assert_eq!(ba.alloc_contiguous(None, 2, 0), Some(321));
527        assert_eq!(ba.alloc_contiguous(None, 64, 6), None);
528        assert_eq!(ba.alloc_contiguous(None, 32, 4), Some(4096 - 48));
529        for i in 0..4096 - 64 + 7 {
530            assert!(ba.dealloc(i));
531        }
532        for i in 4096 - 64 + 8..4096 - 64 + 10 {
533            assert!(ba.dealloc(i));
534        }
535        for i in 4096 - 48..4096 - 16 {
536            assert!(ba.dealloc(i));
537        }
538    }
539}
bitmap_allocator/lib.rs

bitmap_allocator/
lib.rs
