traaains/utils/alloc_map/
stack.rs

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

traaains/utils/alloc_map/
stack.rs
