bitmap_allocator/
lib.rs

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

bitmap_allocator/
lib.rs
