1use std::{
2 borrow::Borrow,
3 cmp::Ordering,
4 marker::PhantomData,
5 ops::{Bound, RangeBounds},
6};
7
8#[derive(Clone, Copy)]
9enum Dir {
10 Fwd,
11 Bwd,
12}
13
14pub struct Range<'a, R, K, V, Q> {
15 map: &'a ArrayMap<K, V>,
16 range: R,
17 cur_fwd: usize,
18 cur_bwd: usize,
19 phantom: PhantomData<Q>,
20}
21
22impl<'a, R, K, V, Q> Range<'a, R, K, V, Q>
23where
24 R: RangeBounds<Q>,
25 K: Ord + Borrow<Q>,
26 Q: Ord,
27{
28 fn step(&mut self, dir: Dir) -> Option<(&'a K, &'a V)> {
29 if self.cur_fwd <= self.cur_bwd {
30 let idx = match dir {
31 Dir::Fwd => self.cur_fwd,
32 Dir::Bwd => self.cur_bwd,
33 };
34 match self.map.get_kv_at(idx) {
35 Some((k, v)) => {
36 if self.range.contains(k.borrow()) {
37 match dir {
38 Dir::Fwd => {
39 self.cur_fwd += 1;
40 }
41 Dir::Bwd => {
42 self.cur_bwd -= 1;
43 }
44 }
45 Some((k, v))
46 } else {
47 None
48 }
49 }
50 None => None,
51 }
52 } else {
53 None
54 }
55 }
56}
57
58impl<'a, R, K, V, Q> Iterator for Range<'a, R, K, V, Q>
59where
60 R: RangeBounds<Q>,
61 K: Ord + Borrow<Q>,
62 Q: Ord,
63{
64 type Item = (&'a K, &'a V);
65
66 fn next(&mut self) -> Option<Self::Item> {
67 self.step(Dir::Fwd)
68 }
69}
70
71impl<'a, R, K, V, Q> DoubleEndedIterator for Range<'a, R, K, V, Q>
72where
73 R: RangeBounds<Q>,
74 K: Ord + Borrow<Q>,
75 Q: Ord,
76{
77 fn next_back(&mut self) -> Option<Self::Item> {
78 self.step(Dir::Bwd)
79 }
80}
81
82#[derive(Debug)]
83pub struct ArrayMap<K, V> {
84 keys: Vec<K>,
85 vals: Vec<V>,
86}
87
88impl<K, V> ArrayMap<K, V>
89where
90 K: Ord,
91{
92 pub fn new() -> Self {
93 Self { keys: vec![], vals: vec![] }
94 }
95
96 pub fn len(&self) -> usize {
97 self.keys.len()
98 }
99
100 pub fn iter(&self) -> impl DoubleEndedIterator<Item = (&K, &V)> {
101 self.keys.iter().zip(self.vals.iter())
102 }
103
104 pub fn keys(&self) -> impl DoubleEndedIterator<Item = &K> {
105 self.keys.iter()
106 }
107
108 pub fn values(&self) -> impl DoubleEndedIterator<Item = &V> {
109 self.vals.iter()
110 }
111
112 fn pos_for_bound<Q>(&self, bound: Bound<&Q>, dir: Dir) -> usize
113 where
114 Q: Ord,
115 K: Borrow<Q>,
116 {
117 match bound {
118 Bound::Unbounded => match dir {
119 Dir::Fwd => 0,
120 Dir::Bwd => self.keys.len().saturating_sub(1),
121 },
122 Bound::Excluded(k) => {
123 match self.keys.binary_search_by(|key| key.borrow().cmp(k)) {
124 Err(i) => i,
125 Ok(i) => match dir {
126 Dir::Fwd => i + 1,
127 Dir::Bwd => i - 1,
128 },
129 }
130 }
131 Bound::Included(k) => {
132 match self.keys.binary_search_by(|key| key.borrow().cmp(k)) {
133 Ok(i) | Err(i) => i,
134 }
135 }
136 }
137 }
138
139 pub fn range<'a, R, Q>(&'a self, range: R) -> Range<'a, R, K, V, Q>
140 where
141 R: RangeBounds<Q>,
142 Q: Ord,
143 K: Borrow<Q>,
144 {
145 let cur_fwd = self.pos_for_bound(range.start_bound(), Dir::Fwd);
146 let cur_bwd = self.pos_for_bound(range.end_bound(), Dir::Bwd);
147 Range { map: self, range, cur_fwd, cur_bwd, phantom: PhantomData }
148 }
149
150 pub fn reserve(&mut self, i: usize) {
151 self.keys.reserve(i);
152 self.vals.reserve(i);
153 }
154
155 pub fn shrink_to_fit(&mut self) {
156 self.keys.shrink_to_fit();
157 self.vals.shrink_to_fit();
158 }
159
160 pub fn get_at(&self, i: usize) -> Option<&V> {
161 self.vals.get(i)
162 }
163
164 pub fn get_kv_at(&self, i: usize) -> Option<(&K, &V)> {
165 self.keys.get(i).map(|k| (k, &self.vals[i]))
166 }
167
168 pub fn insert(&mut self, k: K, v: V) {
169 match self.keys.last() {
170 Some(last_k) => match k.cmp(last_k) {
171 Ordering::Greater | Ordering::Equal => {
172 self.keys.push(k);
173 self.vals.push(v);
174 }
175 Ordering::Less => match self.keys.binary_search(&k) {
176 Ok(i) => {
177 self.keys[i] = k;
178 self.vals[i] = v;
179 }
180 Err(i) => {
181 self.keys.insert(i, k);
182 self.vals.insert(i, v);
183 }
184 },
185 },
186 None => {
187 self.keys.push(k);
188 self.vals.push(v);
189 }
190 }
191 }
192
193 pub fn remove<Q>(&mut self, k: &Q)
194 where
195 K: Borrow<Q>,
196 Q: Ord,
197 {
198 match self.keys.binary_search_by(|key| key.borrow().cmp(k)) {
199 Ok(i) => {
200 self.keys.remove(i);
201 self.vals.remove(i);
202 }
203 Err(_) => (),
204 }
205 }
206
207 pub fn get<Q>(&self, k: &Q) -> Option<&V>
208 where
209 K: Borrow<Q>,
210 Q: Ord,
211 {
212 match self.keys.binary_search_by(|cur| cur.borrow().cmp(k)) {
213 Ok(i) => Some(&self.vals[i]),
214 Err(_) => None,
215 }
216 }
217
218 pub fn first_key_value(&self) -> Option<(&K, &V)> {
219 self.keys.get(0).map(|k| (k, &self.vals[0]))
220 }
221
222 pub fn last_key_value(&self) -> Option<(&K, &V)> {
223 self.keys.last().map(|k| (k, self.vals.last().unwrap()))
224 }
225}
226
227#[cfg(test)]
228mod test {
229 use super::*;
230 use rand::{rng, seq::SliceRandom, RngExt};
231 use std::collections::BTreeMap;
232
233 #[test]
234 fn test_insert_lookup_remove() {
235 let mut model: BTreeMap<u32, u32> = BTreeMap::new();
236 let mut index: ArrayMap<u32, u32> = ArrayMap::new();
237 let mut rng = rng();
238 for _ in 0..10_000 {
239 let kv = rng.random();
240 model.insert(kv, kv);
241 index.insert(kv, kv);
242 for (k, v) in model.iter() {
243 assert_eq!(index.get(&k), Some(v))
244 }
245 for (k, v) in index.iter() {
246 assert_eq!(model.get(&k), Some(v));
247 }
248 }
249 let mut keys: Vec<u32> = model.keys().copied().collect();
250 keys.shuffle(&mut rng);
251 for k in &keys {
252 model.remove(k);
253 index.remove(k);
254 for (k, v) in model.iter() {
255 assert_eq!(index.get(&k), Some(v));
256 }
257 for (k, v) in index.iter() {
258 assert_eq!(model.get(&k), Some(v));
259 }
260 }
261 }
262
263 #[test]
264 fn test_range() {
265 let mut model: BTreeMap<u32, u32> = BTreeMap::new();
266 let mut index: ArrayMap<u32, u32> = ArrayMap::new();
267 let mut keys = vec![];
268 let mut rng = rng();
269 for _ in 0..10_000 {
270 let kv = rng.random();
271 model.insert(kv, kv);
272 index.insert(kv, kv);
273 keys.push(kv);
274 }
275 keys.sort();
276 while keys.len() > 0 {
277 let mid = keys.len() / 2;
278 let li = rng.random_range(0..=mid);
279 let ui = rng.random_range(mid..keys.len());
280 let mut le = false;
281 let lower = if rng.random_bool(0.1) {
282 Bound::Unbounded
283 } else if rng.random_bool(0.5) {
284 Bound::Included(keys.remove(li))
285 } else {
286 le = true;
287 Bound::Excluded(keys.remove(li))
288 };
289 let upper = if rng.random_bool(0.1) {
290 Bound::Unbounded
291 } else {
292 let c = if le || rng.random_bool(0.5) {
293 Bound::Included
294 } else {
295 Bound::Excluded
296 };
297 match keys.get(ui) {
298 None => Bound::Unbounded,
299 Some(k) => {
300 let r = c(*k);
301 keys.remove(ui);
302 r
303 }
304 }
305 };
306 let mut irange = index.range((lower, upper));
307 let mut mrange = model.range((lower, upper));
308 loop {
309 let (mval, ival) = if rng.random_bool(0.1) {
310 match mrange.next_back() {
311 Some(mval) => (Some(mval), irange.next_back()),
312 None => break,
313 }
314 } else {
315 match mrange.next() {
316 Some(mval) => (Some(mval), irange.next()),
317 None => break,
318 }
319 };
320 assert_eq!(mval, ival)
321 }
322 assert_eq!(irange.next(), None);
323 }
324 }
325}