immutable_chunkmap/
chunk.rs

Help
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 = mem::replace(
697            Arc::make_mut(&mut self.0),
698            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 = Arc::make_mut(&mut self.0);
720        (&inner.keys).into_iter().zip(&mut inner.vals)
721    }
722}
immutable_chunkmap/chunk.rs

immutable_chunkmap/
chunk.rs
