1use arrayvec::ArrayVec;
2use alloc::{sync::Arc, vec::Vec};
3use core::{
4 borrow::Borrow,
5 cmp::{min, Ord, Ordering},
6 fmt::{self, Debug, Formatter},
7 iter, mem,
8 ops::Deref,
9 slice,
10};
11
12#[derive(PartialEq)]
13pub(crate) enum Loc {
14 InRight,
15 InLeft,
16 NotPresent(usize),
17 Here(usize),
18}
19
20/*
21elts is a sorted array of pairs, increasing the SIZE has several effects;
22
23-- decreases the height of the tree for a given number of elements,
24 decreasing the amount of indirection necessary to get to any given
25 key.
26
27-- decreases the number of objects allocated on the heap each time a
28 key is added or removed
29
30-- increases the size of each allocation
31
32-- icreases the overall amount of memory allocated for each change to
33 the tree
34*/
35pub const DEFAULT_SIZE: usize = 512;
36
37pub(crate) enum UpdateChunk<
38 Q: Ord,
39 K: Ord + Clone + Borrow<Q>,
40 V: Clone,
41 D,
42 const SIZE: usize,
43> {
44 UpdateLeft(Vec<(Q, D)>),
45 UpdateRight(Vec<(Q, D)>),
46 Updated {
47 elts: Chunk<K, V, SIZE>,
48 update_left: Vec<(Q, D)>,
49 update_right: Vec<(Q, D)>,
50 overflow_right: Vec<(K, V)>,
51 },
52 Removed {
53 not_done: Vec<(Q, D)>,
54 update_left: Vec<(Q, D)>,
55 update_right: Vec<(Q, D)>,
56 },
57}
58
59pub(crate) enum Update<Q: Ord, K: Ord + Clone + Borrow<Q>, V: Clone, D, const SIZE: usize>
60{
61 UpdateLeft(Q, D),
62 UpdateRight(Q, D),
63 Updated {
64 elts: Chunk<K, V, SIZE>,
65 overflow: Option<(K, V)>,
66 previous: Option<V>,
67 },
68}
69
70pub(crate) enum MutUpdate<Q: Ord, K: Ord + Clone + Borrow<Q>, V: Clone, D> {
71 UpdateLeft(Q, D),
72 UpdateRight(Q, D),
73 Updated {
74 overflow: Option<(K, V)>,
75 previous: Option<V>,
76 },
77}
78
79#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
80pub(crate) struct ChunkInner<K, V, const SIZE: usize> {
81 keys: ArrayVec<K, SIZE>,
82 vals: ArrayVec<V, SIZE>,
83}
84
85#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
86pub(crate) struct Chunk<K, V, const SIZE: usize>(Arc<ChunkInner<K, V, SIZE>>);
87
88impl<K, V, const SIZE: usize> Deref for Chunk<K, V, SIZE> {
89 type Target = ChunkInner<K, V, SIZE>;
90
91 fn deref(&self) -> &Self::Target {
92 &*self.0
93 }
94}
95
96impl<K, V, const SIZE: usize> Debug for Chunk<K, V, SIZE>
97where
98 K: Debug,
99 V: Debug,
100{
101 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
102 f.debug_map().entries(self.into_iter()).finish()
103 }
104}
105
106impl<K, V, const SIZE: usize> Chunk<K, V, SIZE>
107where
108 K: Ord + Clone,
109 V: Clone,
110{
111 pub(crate) fn singleton(k: K, v: V) -> Self {
112 let mut t = Chunk::empty();
113 let inner = Arc::make_mut(&mut t.0);
114 inner.keys.push(k);
115 inner.vals.push(v);
116 t
117 }
118
119 pub(crate) fn empty() -> Self {
120 Chunk(Arc::new(ChunkInner {
121 keys: ArrayVec::new(),
122 vals: ArrayVec::new(),
123 }))
124 }
125
126 pub(crate) fn create_with<Q, D, F>(chunk: Vec<(Q, D)>, f: &mut F) -> Self
127 where
128 Q: Ord,
129 K: Borrow<Q>,
130 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
131 {
132 let mut t = Chunk::empty();
133 let inner = Arc::make_mut(&mut t.0);
134 for (k, v) in chunk.into_iter().filter_map(|(q, d)| f(q, d, None)) {
135 inner.keys.push(k);
136 inner.vals.push(v);
137 }
138 t
139 }
140
141 pub(crate) fn get_local<Q: ?Sized + Ord>(&self, k: &Q) -> Option<usize>
142 where
143 K: Borrow<Q>,
144 {
145 match self.keys.binary_search_by_key(&k, |k| k.borrow()) {
146 Ok(i) => Some(i),
147 Err(_) => None,
148 }
149 }
150
151 pub(crate) fn get<Q: ?Sized + Ord>(&self, k: &Q) -> Loc
152 where
153 K: Borrow<Q>,
154 {
155 let len = self.len();
156 if len == 0 {
157 Loc::NotPresent(0)
158 } else {
159 let first = k.cmp(&self.keys[0].borrow());
160 let last = k.cmp(&self.keys[len - 1].borrow());
161 match (first, last) {
162 (Ordering::Equal, _) => Loc::Here(0),
163 (_, Ordering::Equal) => Loc::Here(len - 1),
164 (Ordering::Less, _) => Loc::InLeft,
165 (_, Ordering::Greater) => Loc::InRight,
166 (Ordering::Greater, Ordering::Less) => {
167 match self.keys.binary_search_by_key(&k, |k| k.borrow()) {
168 Result::Ok(i) => Loc::Here(i),
169 Result::Err(i) => Loc::NotPresent(i),
170 }
171 }
172 }
173 }
174 }
175
176 // invariant: chunk is sorted and deduped
177 pub(crate) fn update_chunk<Q, D, F>(
178 &self,
179 chunk: Vec<(Q, D)>,
180 leaf: bool,
181 f: &mut F,
182 ) -> UpdateChunk<Q, K, V, D, SIZE>
183 where
184 Q: Ord,
185 K: Borrow<Q>,
186 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
187 {
188 assert!(chunk.len() <= SIZE && chunk.len() > 0 && self.len() > 0);
189 let full = !leaf || self.len() >= SIZE;
190 let in_left = self.get(&chunk[chunk.len() - 1].0) == Loc::InLeft;
191 let in_right = self.get(&chunk[0].0) == Loc::InRight;
192 if full && in_left {
193 UpdateChunk::UpdateLeft(chunk)
194 } else if full && in_right {
195 UpdateChunk::UpdateRight(chunk)
196 } else if leaf && (in_left || in_right) {
197 let iter = chunk.into_iter().filter_map(|(q, d)| f(q, d, None));
198 let mut overflow_right = Vec::new();
199 let mut elts = Chunk::empty();
200 let inner = Arc::make_mut(&mut elts.0);
201 if in_right {
202 inner.clone_from(self);
203 for (k, v) in iter {
204 if inner.keys.len() < SIZE {
205 inner.keys.push(k);
206 inner.vals.push(v);
207 } else {
208 overflow_right.push((k, v));
209 }
210 }
211 } else {
212 for (k, v) in iter {
213 if inner.keys.len() < SIZE {
214 inner.keys.push(k);
215 inner.vals.push(v);
216 } else {
217 overflow_right.push((k, v));
218 }
219 }
220 for (k, v) in self.keys.iter().cloned().zip(self.vals.iter().cloned()) {
221 if inner.keys.len() < SIZE {
222 inner.keys.push(k);
223 inner.vals.push(v);
224 } else {
225 overflow_right.push((k, v));
226 }
227 }
228 }
229 UpdateChunk::Updated {
230 elts,
231 update_left: Vec::new(),
232 update_right: Vec::new(),
233 overflow_right,
234 }
235 } else {
236 #[inline(always)]
237 fn copy_up_to<K, V, const SIZE: usize>(
238 t: &Chunk<K, V, SIZE>,
239 elts: &mut ChunkInner<K, V, SIZE>,
240 overflow_right: &mut Vec<(K, V)>,
241 m: &mut usize,
242 i: usize,
243 ) where
244 K: Ord + Clone,
245 V: Clone,
246 {
247 let n = min(i - *m, SIZE - elts.keys.len());
248 if n > 0 {
249 elts.keys.extend(t.keys[*m..*m + n].iter().cloned());
250 elts.vals.extend(t.vals[*m..*m + n].iter().cloned());
251 *m += n;
252 }
253 if *m < i {
254 overflow_right.extend(
255 t.keys.as_slice()[*m..i]
256 .iter()
257 .cloned()
258 .zip(t.vals.as_slice()[*m..i].iter().cloned()),
259 );
260 *m = i;
261 }
262 }
263 // invariant: don't cross the streams.
264 let mut not_done = Vec::new();
265 let mut update_left = Vec::new();
266 let mut update_right = Vec::new();
267 let mut overflow_right = Vec::new();
268 let mut m = 0;
269 let mut elts = Chunk::empty();
270 let inner = Arc::make_mut(&mut elts.0);
271 let mut chunk = chunk.into_iter();
272 loop {
273 if m == self.len() && inner.keys.len() == 0 {
274 not_done.extend(chunk);
275 break;
276 }
277 match chunk.next() {
278 None => {
279 copy_up_to(self, inner, &mut overflow_right, &mut m, self.len());
280 break;
281 }
282 Some((q, d)) => {
283 match self.get(&q) {
284 Loc::Here(i) => {
285 copy_up_to(self, inner, &mut overflow_right, &mut m, i);
286 let r = f(q, d, Some((&self.keys[i], &self.vals[i])));
287 if let Some((k, v)) = r {
288 if inner.keys.len() < SIZE {
289 inner.keys.push(k);
290 inner.vals.push(v);
291 } else {
292 overflow_right.push((k, v))
293 }
294 }
295 // self[i] was replaced or deleted, skip it
296 m += 1;
297 }
298 Loc::NotPresent(i) => {
299 copy_up_to(self, inner, &mut overflow_right, &mut m, i);
300 if let Some((k, v)) = f(q, d, None) {
301 if inner.keys.len() < SIZE {
302 inner.keys.push(k);
303 inner.vals.push(v);
304 } else {
305 overflow_right.push((k, v));
306 }
307 }
308 }
309 Loc::InLeft => {
310 if leaf && inner.keys.len() < SIZE {
311 if let Some((k, v)) = f(q, d, None) {
312 inner.keys.push(k);
313 inner.vals.push(v);
314 }
315 } else {
316 update_left.push((q, d))
317 }
318 }
319 Loc::InRight => {
320 let len = self.len();
321 copy_up_to(self, inner, &mut overflow_right, &mut m, len);
322 if leaf && inner.keys.len() < SIZE {
323 let iter = iter::once((q, d))
324 .chain(chunk)
325 .filter_map(|(q, d)| f(q, d, None));
326 for (k, v) in iter {
327 if inner.keys.len() < SIZE {
328 inner.keys.push(k);
329 inner.vals.push(v);
330 } else {
331 overflow_right.push((k, v));
332 }
333 }
334 } else {
335 update_right.push((q, d));
336 update_right.extend(chunk);
337 }
338 break;
339 }
340 }
341 }
342 }
343 }
344 if elts.len() > 0 {
345 assert_eq!(not_done.len(), 0);
346 UpdateChunk::Updated {
347 elts,
348 update_left,
349 update_right,
350 overflow_right,
351 }
352 } else {
353 assert_eq!(overflow_right.len(), 0);
354 UpdateChunk::Removed {
355 not_done,
356 update_left,
357 update_right,
358 }
359 }
360 }
361 }
362
363 pub(crate) fn update<Q, D, F>(
364 &self,
365 q: Q,
366 d: D,
367 leaf: bool,
368 f: &mut F,
369 ) -> Update<Q, K, V, D, SIZE>
370 where
371 Q: Ord,
372 K: Borrow<Q>,
373 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
374 {
375 match self.get(&q) {
376 Loc::Here(i) => {
377 let mut elts = Chunk::empty();
378 let inner = Arc::make_mut(&mut elts.0);
379 inner.keys.extend(self.keys[0..i].iter().cloned());
380 inner.vals.extend(self.vals[0..i].iter().cloned());
381 if let Some((k, v)) = f(q, d, Some((&self.keys[i], &self.vals[i]))) {
382 inner.keys.push(k);
383 inner.vals.push(v);
384 }
385 if i + 1 < self.len() {
386 inner
387 .keys
388 .extend(self.keys[i + 1..self.len()].iter().cloned());
389 inner
390 .vals
391 .extend(self.vals[i + 1..self.len()].iter().cloned());
392 }
393 Update::Updated {
394 elts,
395 overflow: None,
396 previous: Some(self.vals[i].clone()),
397 }
398 }
399 Loc::NotPresent(i) => {
400 let mut elts = Chunk::empty();
401 let inner = Arc::make_mut(&mut elts.0);
402 inner.keys.extend(self.keys[0..i].iter().cloned());
403 inner.vals.extend(self.vals[0..i].iter().cloned());
404 let overflow = match f(q, d, None) {
405 None => None,
406 Some((k, v)) => {
407 if inner.keys.len() == SIZE {
408 let (ok, ov) =
409 (inner.keys.pop().unwrap(), inner.vals.pop().unwrap());
410 inner.keys.push(k);
411 inner.vals.push(v);
412 Some((ok, ov))
413 } else {
414 inner.keys.push(k);
415 inner.vals.push(v);
416 let mut iter = self.keys[i..self.len()]
417 .iter()
418 .cloned()
419 .zip(self.vals[i..self.len()].iter().cloned());
420 loop {
421 match iter.next() {
422 None => break None,
423 Some((k, v)) => {
424 if inner.keys.len() < SIZE {
425 inner.keys.push(k);
426 inner.vals.push(v);
427 } else {
428 break Some((k, v));
429 }
430 }
431 }
432 }
433 }
434 }
435 };
436 Update::Updated {
437 elts,
438 overflow,
439 previous: None,
440 }
441 }
442 loc @ Loc::InLeft | loc @ Loc::InRight => {
443 if !leaf || self.len() == SIZE {
444 match loc {
445 Loc::InLeft => Update::UpdateLeft(q, d),
446 Loc::InRight => Update::UpdateRight(q, d),
447 Loc::Here(..) | Loc::NotPresent(..) => unreachable!(),
448 }
449 } else {
450 let mut elts = Chunk::empty();
451 let inner = Arc::make_mut(&mut elts.0);
452 match loc {
453 Loc::InLeft => {
454 if let Some((k, v)) = f(q, d, None) {
455 inner.keys.push(k);
456 inner.vals.push(v);
457 }
458 inner.keys.extend(self.keys[0..self.len()].iter().cloned());
459 inner.vals.extend(self.vals[0..self.len()].iter().cloned());
460 }
461 Loc::InRight => {
462 inner.keys.extend(self.keys[0..self.len()].iter().cloned());
463 inner.vals.extend(self.vals[0..self.len()].iter().cloned());
464 if let Some((k, v)) = f(q, d, None) {
465 inner.keys.push(k);
466 inner.vals.push(v);
467 }
468 }
469 _ => unreachable!("bug"),
470 };
471 Update::Updated {
472 elts,
473 overflow: None,
474 previous: None,
475 }
476 }
477 }
478 }
479 }
480
481 pub(crate) fn update_mut<Q, D, F>(
482 &mut self,
483 q: Q,
484 d: D,
485 leaf: bool,
486 f: &mut F,
487 ) -> MutUpdate<Q, K, V, D>
488 where
489 Q: Ord,
490 K: Borrow<Q>,
491 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
492 {
493 match self.get(&q) {
494 Loc::Here(i) => match f(q, d, Some((&self.keys[i], &self.vals[i]))) {
495 Some((k, v)) => {
496 let inner = Arc::make_mut(&mut self.0);
497 inner.keys[i] = k;
498 MutUpdate::Updated {
499 overflow: None,
500 previous: Some(mem::replace(&mut inner.vals[i], v)),
501 }
502 }
503 None => {
504 let inner = Arc::make_mut(&mut self.0);
505 inner.keys.remove(i);
506 MutUpdate::Updated {
507 overflow: None,
508 previous: Some(inner.vals.remove(i)),
509 }
510 }
511 },
512 Loc::NotPresent(i) => match f(q, d, None) {
513 Some((k, v)) => {
514 let inner = Arc::make_mut(&mut self.0);
515 let overflow = if inner.keys.len() == SIZE {
516 let (ok, ov) =
517 (inner.keys.pop().unwrap(), inner.vals.pop().unwrap());
518 inner.keys.insert(i, k);
519 inner.vals.insert(i, v);
520 Some((ok, ov))
521 } else {
522 inner.keys.insert(i, k);
523 inner.vals.insert(i, v);
524 None
525 };
526 MutUpdate::Updated {
527 overflow,
528 previous: None,
529 }
530 }
531 None => MutUpdate::Updated {
532 overflow: None,
533 previous: None,
534 },
535 },
536 loc @ Loc::InLeft | loc @ Loc::InRight => {
537 if !leaf || self.len() == SIZE {
538 match loc {
539 Loc::InLeft => MutUpdate::UpdateLeft(q, d),
540 Loc::InRight => MutUpdate::UpdateRight(q, d),
541 Loc::Here(..) | Loc::NotPresent(..) => unreachable!(),
542 }
543 } else {
544 let inner = Arc::make_mut(&mut self.0);
545 match loc {
546 Loc::InLeft => {
547 if let Some((k, v)) = f(q, d, None) {
548 inner.keys.insert(0, k);
549 inner.vals.insert(0, v);
550 }
551 }
552 Loc::InRight => {
553 if let Some((k, v)) = f(q, d, None) {
554 inner.keys.push(k);
555 inner.vals.push(v);
556 }
557 }
558 _ => unreachable!("bug"),
559 };
560 MutUpdate::Updated {
561 overflow: None,
562 previous: None,
563 }
564 }
565 }
566 }
567 }
568
569 pub(crate) fn intersect<F>(
570 c0: &Chunk<K, V, SIZE>,
571 c1: &Chunk<K, V, SIZE>,
572 r: &mut Vec<(K, V)>,
573 f: &mut F,
574 ) where
575 F: FnMut(&K, &V, &V) -> Option<V>,
576 {
577 if c0.len() > 0 && c1.len() > 0 {
578 let (c0, c1) = if c0.len() < c1.len() {
579 (c0, c1)
580 } else {
581 (c1, c0)
582 };
583 r.extend(c0.keys.iter().enumerate().filter_map(|(i, k)| {
584 match c1.keys.binary_search(&k) {
585 Err(_) => None,
586 Ok(j) => f(k, &c0.vals[i], &c1.vals[j]).map(|v| (k.clone(), v)),
587 }
588 }))
589 }
590 }
591
592 pub(crate) fn remove_elt_at(&self, i: usize) -> Self {
593 let mut elts = Chunk::empty();
594 let t = Arc::make_mut(&mut elts.0);
595 if self.keys.len() == 0 {
596 panic!("can't remove from an empty chunk")
597 } else if self.len() == 1 {
598 assert_eq!(i, 0);
599 elts
600 } else if i == 0 {
601 t.keys.extend(self.keys[1..self.len()].iter().cloned());
602 t.vals.extend(self.vals[1..self.len()].iter().cloned());
603 elts
604 } else if i == self.keys.len() - 1 {
605 t.keys.extend(self.keys[0..self.len() - 1].iter().cloned());
606 t.vals.extend(self.vals[0..self.len() - 1].iter().cloned());
607 elts
608 } else {
609 t.keys.extend(self.keys[0..i].iter().cloned());
610 t.keys.extend(self.keys[i + 1..self.len()].iter().cloned());
611 t.vals.extend(self.vals[0..i].iter().cloned());
612 t.vals.extend(self.vals[i + 1..self.len()].iter().cloned());
613 elts
614 }
615 }
616
617 pub(crate) fn remove_elt_at_mut(&mut self, i: usize) -> (K, V) {
618 if self.len() == 0 {
619 panic!("can't remove from an empty chunk")
620 } else {
621 let inner = Arc::make_mut(&mut self.0);
622 let k = inner.keys.remove(i);
623 let v = inner.vals.remove(i);
624 (k, v)
625 }
626 }
627
628 pub(crate) fn append<I: IntoIterator<Item = (K, V)>>(&self, other: I) -> Self {
629 let mut elts = self.clone();
630 let inner = Arc::make_mut(&mut elts.0);
631 for (k, v) in other {
632 if inner.keys.len() < SIZE {
633 inner.keys.push(k);
634 inner.vals.push(v);
635 }
636 }
637 elts
638 }
639
640 pub(crate) fn min_max_key(&self) -> Option<(K, K)> {
641 if self.len() == 0 {
642 None
643 } else {
644 Some((self.keys[0].clone(), self.keys[self.len() - 1].clone()))
645 }
646 }
647
648 pub(crate) fn min_elt(&self) -> Option<(&K, &V)> {
649 if self.len() == 0 {
650 None
651 } else {
652 Some((&self.keys[0], &self.vals[0]))
653 }
654 }
655
656 pub(crate) fn max_elt(&self) -> Option<(&K, &V)> {
657 if self.len() == 0 {
658 None
659 } else {
660 let last = self.len() - 1;
661 Some((&self.keys[last], &self.vals[last]))
662 }
663 }
664
665 pub(crate) fn len(&self) -> usize {
666 self.keys.len()
667 }
668
669 pub(crate) fn key(&self, i: usize) -> &K {
670 &self.keys[i]
671 }
672
673 pub(crate) fn val(&self, i: usize) -> &V {
674 &self.vals[i]
675 }
676
677 pub(crate) fn val_mut(&mut self, i: usize) -> &mut V {
678 &mut Arc::make_mut(&mut self.0).vals[i]
679 }
680
681 pub(crate) fn kv(&self, i: usize) -> (&K, &V) {
682 (&self.keys[i], &self.vals[i])
683 }
684
685 pub(crate) fn to_vec(&self) -> Vec<(K, V)> {
686 self.into_iter()
687 .map(|(k, v)| (k.clone(), v.clone()))
688 .collect()
689 }
690}
691
692impl<K: Clone, V: Clone, const SIZE: usize> IntoIterator for Chunk<K, V, SIZE> {
693 type Item = (K, V);
694 type IntoIter = iter::Zip<arrayvec::IntoIter<K, SIZE>, arrayvec::IntoIter<V, SIZE>>;
695 fn into_iter(mut self) -> Self::IntoIter {
696 let inner: ChunkInner = mem::replace(
697 dest:Arc::make_mut(&mut self.0),
698 src:ChunkInner {
699 keys: ArrayVec::new(),
700 vals: ArrayVec::new(),
701 },
702 );
703 inner.keys.into_iter().zip(inner.vals.into_iter())
704 }
705}
706
707impl<'a, K, V, const SIZE: usize> IntoIterator for &'a Chunk<K, V, SIZE> {
708 type Item = (&'a K, &'a V);
709 type IntoIter = iter::Zip<slice::Iter<'a, K>, slice::Iter<'a, V>>;
710 fn into_iter(self) -> Self::IntoIter {
711 (&self.keys).into_iter().zip(&self.vals)
712 }
713}
714
715impl<'a, K: Clone, V: Clone, const SIZE: usize> IntoIterator for &'a mut Chunk<K, V, SIZE> {
716 type Item = (&'a K, &'a mut V);
717 type IntoIter = iter::Zip<slice::Iter<'a, K>, slice::IterMut<'a, V>>;
718 fn into_iter(self) -> Self::IntoIter {
719 let inner: &mut ChunkInner = Arc::make_mut(&mut self.0);
720 (&inner.keys).into_iter().zip(&mut inner.vals)
721 }
722}
723