Skip to main content

traaains/utils/interval_tree/
stack.rs

1use bitmap::bitmap::Bitmap;
2
3use crate::utils::interval_tree::{
4    interval_to_bitfields, interval_to_internal, left_child, parent_for, right_child,
5};
6
7/// A statically allocated interval tree backed by a bitmap spanning an interval of size `N`.
8/// `N` must be a power of two. Steps are taken to treat incorrect input (rounded to nearest power).
9#[derive(Clone, Default)]
10pub struct IntervalBitmap<const N: usize>(Bitmap<usize, { interval_to_bitfields(N) }>)
11where
12    [(); interval_to_bitfields(N)]:;
13
14impl<const N: usize> IntervalBitmap<N>
15where
16    Bitmap<usize, { interval_to_bitfields(N) }>:,
17{
18    /// Sets the item at position `index` from the underlying interval to `value`.
19    /// - `index` - index in the underlying interval
20    /// - `value` - the specified value
21    pub fn set_at(&mut self, index: usize, value: bool) {
22        assert!(index < N);
23
24        let mut curr_index = index + interval_to_internal(N);
25
26        // set leaf
27        self.0.set(curr_index, value);
28
29        while let Some(parent) = parent_for(curr_index) {
30            let value = self.0.get(left_child(parent)) && self.0.get(right_child(parent));
31            self.0.set(parent, value);
32            curr_index = parent;
33        }
34    }
35
36    /// Gets the item at position `index` from the underlying interval
37    pub fn get_at(&self, index: usize) -> bool {
38        self.0.get(index + interval_to_internal(N))
39    }
40
41    /// Gets the item at position `index` from the tree
42    pub fn get_at_tree(&self, index: usize) -> bool {
43        self.0.get(index)
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use crate::utils::interval_tree::IntervalBitmap;
50
51    #[test]
52    fn test_interval_set_get() {
53        let mut interval = IntervalBitmap::<128>::default();
54
55        for i in 0..127 {
56            interval.set_at(i, true);
57        }
58
59        assert!(!interval.get_at_tree(0));
60    }
61}