1use crate::error::OutOfMemory;
4use core::{mem, ops::RangeBounds, ptr::NonNull};
5use wasmtime_core::slab::{Id, Slab};
6
7pub struct TryBTreeMap<K, V>
10where
11 K: Copy,
12{
13 values: Slab<V>,
14 forest: cranelift_bforest::MapForest<K, Id>,
15 map: cranelift_bforest::Map<K, Id>,
16}
17
18impl<K, V> Default for TryBTreeMap<K, V>
19where
20 K: Copy,
21{
22 fn default() -> Self {
23 Self {
24 values: Default::default(),
25 forest: cranelift_bforest::MapForest::new(),
26 map: Default::default(),
27 }
28 }
29}
30
31impl<K, V> TryBTreeMap<K, V>
32where
33 K: Copy,
34{
35 pub fn new() -> Self {
37 Self::default()
38 }
39
40 pub fn len(&self) -> usize {
42 self.values.len()
43 }
44
45 pub fn is_empty(&self) -> bool {
47 self.map.is_empty()
48 }
49
50 fn get_id(&self, key: K) -> Option<Id>
51 where
52 K: Ord,
53 {
54 self.map.get(key, &self.forest, &())
55 }
56
57 pub fn contains_key(&self, key: K) -> bool
59 where
60 K: Ord,
61 {
62 self.get_id(key).is_some()
63 }
64
65 pub fn get(&self, key: K) -> Option<&V>
67 where
68 K: Ord,
69 {
70 let id = self.get_id(key)?;
71 Some(&self.values[id])
72 }
73
74 pub fn get_mut(&mut self, key: K) -> Option<&mut V>
76 where
77 K: Ord,
78 {
79 let id = self.get_id(key)?;
80 Some(&mut self.values[id])
81 }
82
83 pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, OutOfMemory>
86 where
87 K: Ord,
88 {
89 if let Some(id) = self.get_id(key) {
90 return Ok(Some(mem::replace(&mut self.values[id], value)));
91 }
92
93 let id = self.values.alloc(value)?;
94 match self.map.try_insert(key, id, &mut self.forest, &()) {
95 Ok(old) => {
96 debug_assert!(old.is_none());
97 Ok(None)
98 }
99 Err(oom) => {
100 self.values.dealloc(id);
101 Err(oom)
102 }
103 }
104 }
105
106 pub fn remove(&mut self, key: K) -> Option<V>
108 where
109 K: Ord,
110 {
111 let id = self.map.remove(key, &mut self.forest, &())?;
112 Some(self.values.dealloc(id))
113 }
114
115 pub fn clear(&mut self) {
119 self.values.clear();
120 self.forest.clear();
128 self.map = Default::default();
129 }
130
131 pub fn iter(&self) -> BTreeMapIter<'_, K, V> {
133 BTreeMapIter {
134 inner: self.map.iter(&self.forest),
135 values: &self.values,
136 }
137 }
138
139 pub fn iter_mut(&mut self) -> BTreeMapIterMut<'_, K, V> {
141 BTreeMapIterMut {
142 inner: self.map.iter(&self.forest),
143 values: &mut self.values,
144 }
145 }
146
147 pub fn keys(&self) -> BTreeMapKeys<'_, K, V> {
149 BTreeMapKeys { inner: self.iter() }
150 }
151
152 pub fn values(&self) -> BTreeMapValues<'_, K, V> {
154 BTreeMapValues { inner: self.iter() }
155 }
156
157 pub fn values_mut(&mut self) -> BTreeMapValuesMut<'_, K, V> {
159 BTreeMapValuesMut {
160 inner: self.iter_mut(),
161 }
162 }
163
164 pub fn range<R>(&self, range: R) -> BTreeMapRange<'_, K, V>
166 where
167 K: Ord,
168 R: RangeBounds<K>,
169 {
170 BTreeMapRange {
171 inner: self.map.range(range, &self.forest, &()),
172 values: &self.values,
173 }
174 }
175
176 pub fn range_mut<R>(&mut self, range: R) -> BTreeMapRangeMut<'_, K, V>
178 where
179 K: Ord,
180 R: RangeBounds<K>,
181 {
182 BTreeMapRangeMut {
183 inner: self.map.range(range, &self.forest, &()),
184 values: &mut self.values,
185 }
186 }
187
188 pub fn entry(&mut self, key: K) -> Entry<'_, K, V>
190 where
191 K: Ord,
192 {
193 let mut cursor = self.map.cursor_mut(&mut self.forest, &());
194 match cursor.goto(key) {
195 Some(_) => Entry::Occupied(OccupiedEntry {
196 cursor,
197 values: &mut self.values,
198 }),
199 None => Entry::Vacant(VacantEntry {
200 key,
201 cursor,
202 values: &mut self.values,
203 }),
204 }
205 }
206}
207
208impl<'a, K, V> IntoIterator for &'a TryBTreeMap<K, V>
209where
210 K: Copy,
211{
212 type Item = (K, &'a V);
213 type IntoIter = BTreeMapIter<'a, K, V>;
214
215 fn into_iter(self) -> Self::IntoIter {
216 self.iter()
217 }
218}
219
220pub struct BTreeMapIter<'a, K, V>
222where
223 K: Copy,
224{
225 inner: cranelift_bforest::MapIter<'a, K, Id>,
226 values: &'a Slab<V>,
227}
228
229impl<'a, K, V> Iterator for BTreeMapIter<'a, K, V>
230where
231 K: Copy,
232{
233 type Item = (K, &'a V);
234
235 fn next(&mut self) -> Option<Self::Item> {
236 let (key, id) = self.inner.next()?;
237 Some((key, &self.values[id]))
238 }
239
240 fn size_hint(&self) -> (usize, Option<usize>) {
241 self.inner.size_hint()
242 }
243}
244
245impl<'a, K, V> IntoIterator for &'a mut TryBTreeMap<K, V>
246where
247 K: Copy,
248{
249 type Item = (K, &'a mut V);
250 type IntoIter = BTreeMapIterMut<'a, K, V>;
251
252 fn into_iter(self) -> Self::IntoIter {
253 self.iter_mut()
254 }
255}
256
257pub struct BTreeMapIterMut<'a, K, V>
259where
260 K: Copy,
261{
262 inner: cranelift_bforest::MapIter<'a, K, Id>,
263 values: &'a mut Slab<V>,
264}
265
266impl<'a, K, V> Iterator for BTreeMapIterMut<'a, K, V>
267where
268 K: Copy,
269{
270 type Item = (K, &'a mut V);
271
272 fn next(&mut self) -> Option<Self::Item> {
273 let (key, id) = self.inner.next()?;
274 let val = &mut self.values[id];
275
276 let val = unsafe { NonNull::from(val).as_mut() };
280
281 Some((key, val))
282 }
283
284 fn size_hint(&self) -> (usize, Option<usize>) {
285 self.inner.size_hint()
286 }
287}
288
289impl<K, V> IntoIterator for TryBTreeMap<K, V>
290where
291 K: Copy + Ord,
292{
293 type Item = (K, V);
294 type IntoIter = BTreeMapIntoIter<K, V>;
295
296 fn into_iter(self) -> Self::IntoIter {
297 BTreeMapIntoIter {
298 inner: self.map.into_iter(self.forest),
299 values: self.values,
300 }
301 }
302}
303
304pub struct BTreeMapIntoIter<K, V>
306where
307 K: Copy,
308{
309 inner: cranelift_bforest::MapIntoIter<K, Id>,
310 values: Slab<V>,
311}
312
313impl<K, V> Iterator for BTreeMapIntoIter<K, V>
314where
315 K: Copy + Ord,
316{
317 type Item = (K, V);
318
319 fn next(&mut self) -> Option<Self::Item> {
320 let (key, id) = self.inner.next()?;
321 let value = self.values.dealloc(id);
322 Some((key, value))
323 }
324
325 fn size_hint(&self) -> (usize, Option<usize>) {
326 let len = self.values.len();
327 (len, Some(len))
328 }
329}
330
331pub struct BTreeMapKeys<'a, K, V>
333where
334 K: Copy,
335{
336 inner: BTreeMapIter<'a, K, V>,
337}
338
339impl<'a, K, V> Iterator for BTreeMapKeys<'a, K, V>
340where
341 K: Copy,
342{
343 type Item = K;
344
345 fn next(&mut self) -> Option<Self::Item> {
346 let (k, _v) = self.inner.next()?;
347 Some(k)
348 }
349}
350
351pub struct BTreeMapValues<'a, K, V>
353where
354 K: Copy,
355{
356 inner: BTreeMapIter<'a, K, V>,
357}
358
359impl<'a, K, V> Iterator for BTreeMapValues<'a, K, V>
360where
361 K: Copy,
362{
363 type Item = &'a V;
364
365 fn next(&mut self) -> Option<Self::Item> {
366 let (_k, v) = self.inner.next()?;
367 Some(v)
368 }
369}
370
371pub struct BTreeMapValuesMut<'a, K, V>
373where
374 K: Copy,
375{
376 inner: BTreeMapIterMut<'a, K, V>,
377}
378
379impl<'a, K, V> Iterator for BTreeMapValuesMut<'a, K, V>
380where
381 K: Copy,
382{
383 type Item = &'a mut V;
384
385 fn next(&mut self) -> Option<Self::Item> {
386 let (_k, v) = self.inner.next()?;
387 Some(v)
388 }
389}
390
391pub struct BTreeMapRange<'a, K, V>
393where
394 K: Copy + Ord,
395{
396 inner: cranelift_bforest::MapRange<'a, K, Id, ()>,
397 values: &'a Slab<V>,
398}
399
400impl<'a, K, V> Iterator for BTreeMapRange<'a, K, V>
401where
402 K: Copy + Ord,
403{
404 type Item = (K, &'a V);
405
406 fn next(&mut self) -> Option<Self::Item> {
407 let (key, id) = self.inner.next()?;
408 Some((key, &self.values[id]))
409 }
410}
411
412impl<'a, K, V> DoubleEndedIterator for BTreeMapRange<'a, K, V>
413where
414 K: Copy + Ord,
415{
416 fn next_back(&mut self) -> Option<Self::Item> {
417 let (key, id) = self.inner.next_back()?;
418 Some((key, &self.values[id]))
419 }
420}
421
422pub struct BTreeMapRangeMut<'a, K, V>
424where
425 K: Copy + Ord,
426{
427 inner: cranelift_bforest::MapRange<'a, K, Id, ()>,
428 values: &'a mut Slab<V>,
429}
430
431impl<'a, K, V> Iterator for BTreeMapRangeMut<'a, K, V>
432where
433 K: Copy + Ord,
434{
435 type Item = (K, &'a mut V);
436
437 fn next(&mut self) -> Option<Self::Item> {
438 let (key, id) = self.inner.next()?;
439 let val = &mut self.values[id];
440
441 let val = unsafe { NonNull::from(val).as_mut() };
445
446 Some((key, val))
447 }
448}
449
450#[allow(missing_docs, reason = "self explanatory")]
452pub enum Entry<'a, K, V>
453where
454 K: Copy + Ord,
455{
456 Occupied(OccupiedEntry<'a, K, V>),
457 Vacant(VacantEntry<'a, K, V>),
458}
459
460impl<'a, K, V> Entry<'a, K, V>
461where
462 K: Copy + Ord,
463{
464 pub fn or_insert(self, default: V) -> Result<&'a mut V, OutOfMemory> {
467 self.or_insert_with(|| default)
468 }
469
470 pub fn or_insert_with<F>(self, default: F) -> Result<&'a mut V, OutOfMemory>
473 where
474 F: FnOnce() -> V,
475 {
476 self.or_insert_with_key(|_| default())
477 }
478
479 pub fn or_insert_with_key<F>(self, default: F) -> Result<&'a mut V, OutOfMemory>
482 where
483 F: FnOnce(K) -> V,
484 {
485 match self {
486 Entry::Occupied(e) => {
487 let id = e.cursor.value().unwrap();
488 Ok(&mut e.values[id])
489 }
490 Entry::Vacant(mut e) => {
491 let id = e.values.alloc(default(e.key))?;
492 match e.cursor.try_insert(e.key, id) {
493 Ok(old) => {
494 debug_assert!(old.is_none());
495 Ok(&mut e.values[id])
496 }
497 Err(oom) => {
498 e.values.dealloc(id);
499 Err(oom)
500 }
501 }
502 }
503 }
504 }
505
506 pub fn key(&self) -> K {
508 match self {
509 Entry::Occupied(e) => e.cursor.key().unwrap(),
510 Entry::Vacant(e) => e.key,
511 }
512 }
513
514 pub fn and_modify<F>(self, f: F) -> Self
516 where
517 F: FnOnce(&mut V),
518 {
519 match self {
520 Entry::Occupied(mut e) => {
521 f(e.get_mut());
522 Entry::Occupied(e)
523 }
524 e @ Entry::Vacant(_) => e,
525 }
526 }
527
528 pub fn insert_entry(self, value: V) -> Result<OccupiedEntry<'a, K, V>, OutOfMemory> {
531 match self {
532 Entry::Occupied(e) => Ok(e),
533 Entry::Vacant(e) => e.insert_entry(value),
534 }
535 }
536
537 pub fn or_default(self) -> Result<&'a mut V, OutOfMemory>
540 where
541 V: Default,
542 {
543 self.or_insert_with(Default::default)
544 }
545}
546
547pub struct OccupiedEntry<'a, K, V>
549where
550 K: Copy + Ord,
551{
552 cursor: cranelift_bforest::MapCursorMut<'a, K, Id, ()>,
553 values: &'a mut Slab<V>,
554}
555
556impl<'a, K, V> OccupiedEntry<'a, K, V>
557where
558 K: Copy + Ord,
559{
560 pub fn key(&self) -> K {
562 self.cursor.key().unwrap()
563 }
564
565 pub fn remove_entry(mut self) -> (K, V) {
567 let key = self.key();
568 let id = self.cursor.remove().unwrap();
569 let value = self.values.dealloc(id);
570 (key, value)
571 }
572
573 pub fn get(&self) -> &V {
575 let id = self.cursor.value().unwrap();
576 &self.values[id]
577 }
578
579 pub fn get_mut(&mut self) -> &mut V {
581 let id = self.cursor.value().unwrap();
582 &mut self.values[id]
583 }
584
585 pub fn into_mut(self) -> &'a mut V {
587 let id = self.cursor.value().unwrap();
588 &mut self.values[id]
589 }
590
591 pub fn insert(self, value: V) -> V {
593 let id = self.cursor.value().unwrap();
594 mem::replace(&mut self.values[id], value)
595 }
596
597 pub fn remove(mut self) -> V {
599 let id = self.cursor.remove().unwrap();
600 self.values.dealloc(id)
601 }
602}
603
604pub struct VacantEntry<'a, K, V>
606where
607 K: Copy + Ord,
608{
609 key: K,
610 cursor: cranelift_bforest::MapCursorMut<'a, K, Id, ()>,
611 values: &'a mut Slab<V>,
612}
613
614impl<'a, K, V> VacantEntry<'a, K, V>
615where
616 K: Copy + Ord,
617{
618 pub fn key(&self) -> K {
620 self.key
621 }
622
623 pub fn into_key(self) -> K {
625 self.key
626 }
627
628 pub fn insert(self, value: V) -> Result<&'a mut V, OutOfMemory> {
631 Ok(self.insert_entry(value)?.into_mut())
632 }
633
634 pub fn insert_entry(mut self, value: V) -> Result<OccupiedEntry<'a, K, V>, OutOfMemory> {
637 let id = self.values.alloc(value)?;
638 match self.cursor.try_insert(self.key, id) {
639 Ok(old) => {
640 debug_assert!(old.is_none());
641 Ok(OccupiedEntry {
642 cursor: self.cursor,
643 values: self.values,
644 })
645 }
646 Err(oom) => {
647 self.values.dealloc(id);
648 Err(oom)
649 }
650 }
651 }
652}
653
654#[cfg(test)]
655mod tests {
656 use super::*;
657 use crate::error::Result;
658 use alloc::{vec, vec::Vec};
659
660 type M = TryBTreeMap<usize, f32>;
661
662 #[test]
663 fn new() -> Result<()> {
664 let m = M::new();
665 assert!(m.is_empty());
666 Ok(())
667 }
668
669 #[test]
670 fn len() -> Result<()> {
671 let mut m = M::new();
672 for i in 0..10 {
673 assert_eq!(m.len(), i);
674 m.insert(i, i as f32)?;
675 }
676 assert_eq!(m.len(), 10);
677 for i in 0..10 {
678 assert_eq!(m.len(), 10 - i);
679 m.remove(i);
680 }
681 assert_eq!(m.len(), 0);
682 Ok(())
683 }
684
685 #[test]
686 fn is_empty() -> Result<()> {
687 let mut m = M::new();
688 assert!(m.is_empty());
689 m.insert(42, 42.0)?;
690 assert!(!m.is_empty());
691 m.remove(42);
692 assert!(m.is_empty());
693 Ok(())
694 }
695
696 #[test]
697 fn contains_key() -> Result<()> {
698 let mut m = M::new();
699 assert!(!m.contains_key(36));
700 assert!(!m.contains_key(42));
701 m.insert(42, 42.0)?;
702 assert!(!m.contains_key(36));
703 assert!(m.contains_key(42));
704 m.remove(42);
705 assert!(!m.contains_key(36));
706 assert!(!m.contains_key(42));
707 Ok(())
708 }
709
710 #[test]
711 fn get() -> Result<()> {
712 let mut m = M::new();
713 assert!(m.get(36).is_none());
714 assert!(m.get(42).is_none());
715 m.insert(42, 42.0)?;
716 assert!(m.get(36).is_none());
717 assert_eq!(m.get(42), Some(&42.0));
718 m.remove(42);
719 assert!(m.get(36).is_none());
720 assert!(m.get(42).is_none());
721 Ok(())
722 }
723
724 #[test]
725 fn get_mut() -> Result<()> {
726 let mut m = M::new();
727 assert!(m.get_mut(36).is_none());
728 assert!(m.get_mut(42).is_none());
729 m.insert(42, 42.0)?;
730 assert!(m.get_mut(36).is_none());
731 assert_eq!(m.get_mut(42).copied(), Some(42.0));
732 *m.get_mut(42).unwrap() = 99.0;
733 assert_eq!(m.get_mut(42).copied(), Some(99.0));
734 m.remove(42);
735 assert!(m.get_mut(36).is_none());
736 assert!(m.get_mut(42).is_none());
737 Ok(())
738 }
739
740 #[test]
741 fn insert() -> Result<()> {
742 let mut m = M::new();
743 let old = m.insert(11, 0.0)?;
744 assert!(old.is_none());
745 let old = m.insert(11, 1.0)?;
746 assert_eq!(old, Some(0.0));
747 Ok(())
748 }
749
750 #[test]
751 fn remove() -> Result<()> {
752 let mut m = M::new();
753 let old = m.remove(10);
754 assert!(old.is_none());
755 m.insert(10, 123.0)?;
756 let old = m.remove(10);
757 assert_eq!(old, Some(123.0));
758 Ok(())
759 }
760
761 #[test]
762 fn clear() -> Result<()> {
763 let mut m = M::new();
764 for i in 0..10 {
765 m.insert(i, i as f32)?;
766 }
767 m.clear();
768 assert!(m.is_empty());
769 for i in 0..10 {
770 assert!(m.get(i).is_none());
771 }
772 Ok(())
773 }
774
775 #[test]
776 fn iter() -> Result<()> {
777 let mut m = M::new();
778 for i in 0..5 {
779 m.insert(i, i as f32)?;
780 }
781 assert_eq!(
782 m.iter().collect::<Vec<_>>(),
783 vec![(0, &0.0), (1, &1.0), (2, &2.0), (3, &3.0), (4, &4.0)],
784 );
785 Ok(())
786 }
787
788 #[test]
789 fn iter_mut() -> Result<()> {
790 let mut m = M::new();
791 for i in 0..5 {
792 m.insert(i, i as f32)?;
793 }
794 assert_eq!(
795 m.iter_mut()
796 .map(|(k, v)| (k, mem::replace(v, 0.0)))
797 .collect::<Vec<_>>(),
798 vec![(0, 0.0), (1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0)],
799 );
800 for i in 0..5 {
801 assert_eq!(m.get(i), Some(&0.0));
802 }
803 Ok(())
804 }
805
806 #[test]
807 fn keys() -> Result<()> {
808 let mut m = M::new();
809 for i in 0..5 {
810 m.insert(i, i as f32)?;
811 }
812 assert_eq!(m.keys().collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
813 Ok(())
814 }
815
816 #[test]
817 fn values() -> Result<()> {
818 let mut m = M::new();
819 for i in 0..5 {
820 m.insert(i, i as f32)?;
821 }
822 assert_eq!(
823 m.values().collect::<Vec<_>>(),
824 vec![&0.0, &1.0, &2.0, &3.0, &4.0],
825 );
826 Ok(())
827 }
828
829 #[test]
830 fn values_mut() -> Result<()> {
831 let mut m = M::new();
832 for i in 0..5 {
833 m.insert(i, i as f32)?;
834 }
835 assert_eq!(
836 m.values_mut()
837 .map(|v| mem::replace(v, 0.0))
838 .collect::<Vec<_>>(),
839 vec![0.0, 1.0, 2.0, 3.0, 4.0],
840 );
841 assert_eq!(
842 m.values().collect::<Vec<_>>(),
843 vec![&0.0, &0.0, &0.0, &0.0, &0.0],
844 );
845 Ok(())
846 }
847
848 #[test]
849 fn range() -> Result<()> {
850 let mut m = M::new();
851 for i in 0..5 {
852 m.insert(i, i as f32)?;
853 }
854 assert_eq!(m.range(..2).collect::<Vec<_>>(), vec![(0, &0.0), (1, &1.0)]);
855 assert_eq!(m.range(3..).collect::<Vec<_>>(), vec![(3, &3.0), (4, &4.0)]);
856 assert_eq!(
857 m.range(1..3).collect::<Vec<_>>(),
858 vec![(1, &1.0), (2, &2.0)],
859 );
860 assert_eq!(
861 m.range(2..=3).collect::<Vec<_>>(),
862 vec![(2, &2.0), (3, &3.0)],
863 );
864 assert_eq!(m.range(5..).collect::<Vec<_>>(), vec![]);
865 assert_eq!(m.range(..0).collect::<Vec<_>>(), vec![]);
866 Ok(())
867 }
868
869 #[test]
870 fn range_mut() -> Result<()> {
871 let mut m = M::new();
872 for i in 0..5 {
873 m.insert(i, i as f32)?;
874 }
875 assert_eq!(
876 m.range_mut(1..3)
877 .map(|(k, v)| (k, mem::replace(v, 99.0)))
878 .collect::<Vec<_>>(),
879 vec![(1, 1.0), (2, 2.0)],
880 );
881 assert_eq!(
882 m.values().copied().collect::<Vec<_>>(),
883 vec![0.0, 99.0, 99.0, 3.0, 4.0]
884 );
885 Ok(())
886 }
887
888 #[test]
889 fn entry() -> Result<()> {
890 let mut m = M::new();
891
892 let v = m.entry(0).or_insert(99.0)?;
893 assert_eq!(*v, 99.0);
894
895 let v = m.entry(0).or_insert(0.0)?;
896 assert_eq!(*v, 99.0);
897
898 let e = match m.entry(1) {
899 Entry::Occupied(_) => unreachable!(),
900 Entry::Vacant(e) => e.insert_entry(42.0)?,
901 };
902 assert_eq!(e.key(), 1);
903 assert_eq!(e.get(), &42.0);
904 let old = e.insert(43.0);
905 assert_eq!(old, 42.0);
906 assert_eq!(m.get(1), Some(&43.0));
907
908 let e = match m.entry(1) {
909 Entry::Occupied(e) => e,
910 Entry::Vacant(_) => unreachable!(),
911 };
912 assert_eq!(e.key(), 1);
913 assert_eq!(e.get(), &43.0);
914 let (k, v) = e.remove_entry();
915 assert_eq!(k, 1);
916 assert_eq!(v, 43.0);
917 assert!(m.get(1).is_none());
918
919 let e = match m.entry(2) {
920 Entry::Occupied(_) => unreachable!(),
921 Entry::Vacant(e) => e,
922 };
923 assert_eq!(e.key(), 2);
924 let v = e.insert(99.0)?;
925 assert_eq!(*v, 99.0);
926
927 Ok(())
928 }
929}