traaains/utils/alloc_map/
heap.rs

1use crate::utils::interval_tree::{
2    IntervalBitmapVec, interval_to_internal, leaves_to_depth, left_child, right_child,
3};
4
5/// Allocation map, backed by an interval tree on top of a bitmap.
6/// `N` must be a power of two. Steps are taken to treat incorrect input (rounded to nearest power).
7#[derive(Clone, Default)]
8pub struct AllocMapVec(IntervalBitmapVec);
9
10impl AllocMapVec {
11    /// Returns interval index of first unallocated item, if exists
12    fn find_first_free(&self) -> Option<usize> {
13        // if the root node is true, then there are no free spaces left
14        if self.0.get_at_tree(0) {
15            return None;
16        }
17
18        // find the smallest unallocated space
19        let mut curr_index = 0;
20        for _ in 0..leaves_to_depth(self.0.interval_size()) {
21            let left_child = left_child(curr_index);
22            curr_index = if !self.0.get_at_tree(left_child) {
23                left_child
24            } else {
25                let right_child = right_child(curr_index);
26                debug_assert!(
27                    !self.0.get_at_tree(right_child),
28                    "both children allocated while parent is not"
29                );
30                right_child
31            };
32        }
33
34        Some(curr_index - interval_to_internal(self.0.interval_size()))
35    }
36
37    /// Allocates the smallest available space
38    pub fn allocate(&mut self) -> Option<usize> {
39        self.find_first_free().inspect(|index| {
40            self.0.set_at(*index, true);
41        })
42    }
43
44    /// Deallocates the target index
45    pub fn deallocate(&mut self, index: usize) {
46        self.0.set_at(index, false);
47    }
48
49    pub fn is_allocated(&mut self, index: usize) -> bool {
50        self.0.get_at(index)
51    }
52
53    pub fn has_free(&self) -> bool {
54        self.find_first_free().is_some()
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use crate::utils::alloc_map::AllocMap;
61
62    #[test]
63    fn alloc_map() {
64        let mut map = AllocMap::<2>::default();
65        assert_eq!(map.allocate(), Some(0));
66        assert_eq!(map.allocate(), Some(1));
67        assert_eq!(map.allocate(), None);
68        map.deallocate(0);
69        assert_eq!(map.allocate(), Some(0));
70    }
71}
traaains/utils/alloc_map/heap.rs

traaains/utils/alloc_map/
heap.rs
