Skip to main content

brk_types/
range_map.rs

1use std::marker::PhantomData;
2
3/// Direct-mapped cache size. Power of 2 for fast masking.
4/// 1024 entries × ~32 bytes = 32 KB (fits in L1 cache).
5const CACHE_SIZE: usize = 1024;
6const CACHE_MASK: usize = CACHE_SIZE - 1;
7
8/// Cache entry: (range_low, range_high, value, occupied).
9type CacheEntry<I, V> = (I, I, V, bool);
10
11/// Maps ranges of indices to values for efficient reverse lookups.
12///
13/// Stores first_index values in a sorted Vec and uses binary search
14/// to find the value for any index. The value is derived from the position.
15///
16/// Includes a direct-mapped cache for O(1) floor lookups when there's locality.
17pub struct RangeMap<I, V> {
18    first_indexes: Vec<I>,
19    cache: [CacheEntry<I, V>; CACHE_SIZE],
20    _phantom: PhantomData<V>,
21}
22
23impl<I: Default + Copy, V: Default + Copy> Clone for RangeMap<I, V> {
24    fn clone(&self) -> Self {
25        Self {
26            first_indexes: self.first_indexes.clone(),
27            cache: [(I::default(), I::default(), V::default(), false); CACHE_SIZE],
28            _phantom: PhantomData,
29        }
30    }
31}
32
33impl<I: Default + Copy, V: Default + Copy> From<Vec<I>> for RangeMap<I, V> {
34    fn from(first_indexes: Vec<I>) -> Self {
35        Self {
36            first_indexes,
37            cache: [(I::default(), I::default(), V::default(), false); CACHE_SIZE],
38            _phantom: PhantomData,
39        }
40    }
41}
42
43impl<I: Default + Copy, V: Default + Copy> Default for RangeMap<I, V> {
44    fn default() -> Self {
45        Self {
46            first_indexes: Vec::new(),
47            cache: [(I::default(), I::default(), V::default(), false); CACHE_SIZE],
48            _phantom: PhantomData,
49        }
50    }
51}
52
53impl<I: Ord + Copy + Default + Into<usize>, V: From<usize> + Copy + Default> RangeMap<I, V> {
54    /// Number of ranges stored.
55    #[allow(clippy::len_without_is_empty)]
56    pub fn len(&self) -> usize {
57        self.first_indexes.len()
58    }
59
60    /// Truncate to `new_len` ranges and clear the cache.
61    pub fn truncate(&mut self, new_len: usize) {
62        self.first_indexes.truncate(new_len);
63        self.clear_cache();
64    }
65
66    /// Reserve capacity for additional entries.
67    pub fn reserve(&mut self, additional: usize) {
68        self.first_indexes.reserve(additional);
69    }
70
71    /// Push a new first_index. Value is implicitly the current length.
72    /// Must be called in order (first_index must be >= all previous).
73    #[inline]
74    pub fn push(&mut self, first_index: I) {
75        debug_assert!(
76            self.first_indexes
77                .last()
78                .is_none_or(|&last| first_index >= last),
79            "RangeMap: first_index must be monotonically increasing"
80        );
81        self.first_indexes.push(first_index);
82    }
83
84    /// Returns the last pushed first_index, if any.
85    #[inline]
86    pub fn last_key(&self) -> Option<I> {
87        self.first_indexes.last().copied()
88    }
89
90    /// Floor: returns the value (position) of the largest first_index <= given index.
91    #[inline]
92    pub fn get(&mut self, index: I) -> Option<V> {
93        if self.first_indexes.is_empty() {
94            return None;
95        }
96
97        let slot = Self::cache_slot(&index);
98        let entry = &self.cache[slot];
99        if entry.3 && index >= entry.0 && index < entry.1 {
100            return Some(entry.2);
101        }
102
103        let pos = self.first_indexes.partition_point(|&first| first <= index);
104        if pos > 0 {
105            let value = V::from(pos - 1);
106            if pos < self.first_indexes.len() {
107                self.cache[slot] = (
108                    self.first_indexes[pos - 1],
109                    self.first_indexes[pos],
110                    value,
111                    true,
112                );
113            }
114            Some(value)
115        } else {
116            None
117        }
118    }
119
120    /// Ceil: returns the value (position) of the smallest first_index >= given index.
121    #[inline]
122    pub fn ceil(&self, index: I) -> Option<V> {
123        if self.first_indexes.is_empty() {
124            return None;
125        }
126
127        let pos = self.first_indexes.partition_point(|&first| first < index);
128        if pos < self.first_indexes.len() {
129            Some(V::from(pos))
130        } else {
131            None
132        }
133    }
134
135    /// Shared (immutable) floor lookup — binary search only, no cache update.
136    /// Use when you only have `&self` (e.g. read-only clones in the query layer).
137    #[inline]
138    pub fn get_shared(&self, index: I) -> Option<V> {
139        if self.first_indexes.is_empty() {
140            return None;
141        }
142        let pos = self.first_indexes.partition_point(|&first| first <= index);
143        if pos > 0 {
144            Some(V::from(pos - 1))
145        } else {
146            None
147        }
148    }
149
150    #[inline]
151    fn cache_slot(index: &I) -> usize {
152        let v: usize = (*index).into();
153        v & CACHE_MASK
154    }
155
156    fn clear_cache(&mut self) {
157        for entry in self.cache.iter_mut() {
158            entry.3 = false;
159        }
160    }
161}