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, R>
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, R>
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, R>
393where
394 K: Copy + Ord,
395 R: RangeBounds<K>,
396{
397 inner: cranelift_bforest::MapRange<'a, K, Id, R, ()>,
398 values: &'a Slab<V>,
399}
400
401impl<'a, K, V, R> Iterator for BTreeMapRange<'a, K, V, R>
402where
403 K: Copy + Ord,
404 R: RangeBounds<K>,
405{
406 type Item = (K, &'a V);
407
408 fn next(&mut self) -> Option<Self::Item> {
409 let (key, id) = self.inner.next()?;
410 Some((key, &self.values[id]))
411 }
412}
413
414pub struct BTreeMapRangeMut<'a, K, V, R>
416where
417 K: Copy + Ord,
418 R: RangeBounds<K>,
419{
420 inner: cranelift_bforest::MapRange<'a, K, Id, R, ()>,
421 values: &'a mut Slab<V>,
422}
423
424impl<'a, K, V, R> Iterator for BTreeMapRangeMut<'a, K, V, R>
425where
426 K: Copy + Ord,
427 R: RangeBounds<K>,
428{
429 type Item = (K, &'a mut V);
430
431 fn next(&mut self) -> Option<Self::Item> {
432 let (key, id) = self.inner.next()?;
433 let val = &mut self.values[id];
434
435 let val = unsafe { NonNull::from(val).as_mut() };
439
440 Some((key, val))
441 }
442}
443
444#[allow(missing_docs, reason = "self explanatory")]
446pub enum Entry<'a, K, V>
447where
448 K: Copy + Ord,
449{
450 Occupied(OccupiedEntry<'a, K, V>),
451 Vacant(VacantEntry<'a, K, V>),
452}
453
454impl<'a, K, V> Entry<'a, K, V>
455where
456 K: Copy + Ord,
457{
458 pub fn or_insert(self, default: V) -> Result<&'a mut V, OutOfMemory> {
461 self.or_insert_with(|| default)
462 }
463
464 pub fn or_insert_with<F>(self, default: F) -> Result<&'a mut V, OutOfMemory>
467 where
468 F: FnOnce() -> V,
469 {
470 self.or_insert_with_key(|_| default())
471 }
472
473 pub fn or_insert_with_key<F>(self, default: F) -> Result<&'a mut V, OutOfMemory>
476 where
477 F: FnOnce(K) -> V,
478 {
479 match self {
480 Entry::Occupied(e) => {
481 let id = e.cursor.value().unwrap();
482 Ok(&mut e.values[id])
483 }
484 Entry::Vacant(mut e) => {
485 let id = e.values.alloc(default(e.key))?;
486 match e.cursor.try_insert(e.key, id) {
487 Ok(old) => {
488 debug_assert!(old.is_none());
489 Ok(&mut e.values[id])
490 }
491 Err(oom) => {
492 e.values.dealloc(id);
493 Err(oom)
494 }
495 }
496 }
497 }
498 }
499
500 pub fn key(&self) -> K {
502 match self {
503 Entry::Occupied(e) => e.cursor.key().unwrap(),
504 Entry::Vacant(e) => e.key,
505 }
506 }
507
508 pub fn and_modify<F>(self, f: F) -> Self
510 where
511 F: FnOnce(&mut V),
512 {
513 match self {
514 Entry::Occupied(mut e) => {
515 f(e.get_mut());
516 Entry::Occupied(e)
517 }
518 e @ Entry::Vacant(_) => e,
519 }
520 }
521
522 pub fn insert_entry(self, value: V) -> Result<OccupiedEntry<'a, K, V>, OutOfMemory> {
525 match self {
526 Entry::Occupied(e) => Ok(e),
527 Entry::Vacant(e) => e.insert_entry(value),
528 }
529 }
530
531 pub fn or_default(self) -> Result<&'a mut V, OutOfMemory>
534 where
535 V: Default,
536 {
537 self.or_insert_with(Default::default)
538 }
539}
540
541pub struct OccupiedEntry<'a, K, V>
543where
544 K: Copy + Ord,
545{
546 cursor: cranelift_bforest::MapCursorMut<'a, K, Id, ()>,
547 values: &'a mut Slab<V>,
548}
549
550impl<'a, K, V> OccupiedEntry<'a, K, V>
551where
552 K: Copy + Ord,
553{
554 pub fn key(&self) -> K {
556 self.cursor.key().unwrap()
557 }
558
559 pub fn remove_entry(mut self) -> (K, V) {
561 let key = self.key();
562 let id = self.cursor.remove().unwrap();
563 let value = self.values.dealloc(id);
564 (key, value)
565 }
566
567 pub fn get(&self) -> &V {
569 let id = self.cursor.value().unwrap();
570 &self.values[id]
571 }
572
573 pub fn get_mut(&mut self) -> &mut V {
575 let id = self.cursor.value().unwrap();
576 &mut self.values[id]
577 }
578
579 pub fn into_mut(self) -> &'a mut V {
581 let id = self.cursor.value().unwrap();
582 &mut self.values[id]
583 }
584
585 pub fn insert(self, value: V) -> V {
587 let id = self.cursor.value().unwrap();
588 mem::replace(&mut self.values[id], value)
589 }
590
591 pub fn remove(mut self) -> V {
593 let id = self.cursor.remove().unwrap();
594 self.values.dealloc(id)
595 }
596}
597
598pub struct VacantEntry<'a, K, V>
600where
601 K: Copy + Ord,
602{
603 key: K,
604 cursor: cranelift_bforest::MapCursorMut<'a, K, Id, ()>,
605 values: &'a mut Slab<V>,
606}
607
608impl<'a, K, V> VacantEntry<'a, K, V>
609where
610 K: Copy + Ord,
611{
612 pub fn key(&self) -> K {
614 self.key
615 }
616
617 pub fn into_key(self) -> K {
619 self.key
620 }
621
622 pub fn insert(self, value: V) -> Result<&'a mut V, OutOfMemory> {
625 Ok(self.insert_entry(value)?.into_mut())
626 }
627
628 pub fn insert_entry(mut self, value: V) -> Result<OccupiedEntry<'a, K, V>, OutOfMemory> {
631 let id = self.values.alloc(value)?;
632 match self.cursor.try_insert(self.key, id) {
633 Ok(old) => {
634 debug_assert!(old.is_none());
635 Ok(OccupiedEntry {
636 cursor: self.cursor,
637 values: self.values,
638 })
639 }
640 Err(oom) => {
641 self.values.dealloc(id);
642 Err(oom)
643 }
644 }
645 }
646}
647
648#[cfg(test)]
649mod tests {
650 use super::*;
651 use crate::error::Result;
652 use alloc::{vec, vec::Vec};
653
654 type M = TryBTreeMap<usize, f32>;
655
656 #[test]
657 fn new() -> Result<()> {
658 let m = M::new();
659 assert!(m.is_empty());
660 Ok(())
661 }
662
663 #[test]
664 fn len() -> Result<()> {
665 let mut m = M::new();
666 for i in 0..10 {
667 assert_eq!(m.len(), i);
668 m.insert(i, i as f32)?;
669 }
670 assert_eq!(m.len(), 10);
671 for i in 0..10 {
672 assert_eq!(m.len(), 10 - i);
673 m.remove(i);
674 }
675 assert_eq!(m.len(), 0);
676 Ok(())
677 }
678
679 #[test]
680 fn is_empty() -> Result<()> {
681 let mut m = M::new();
682 assert!(m.is_empty());
683 m.insert(42, 42.0)?;
684 assert!(!m.is_empty());
685 m.remove(42);
686 assert!(m.is_empty());
687 Ok(())
688 }
689
690 #[test]
691 fn contains_key() -> Result<()> {
692 let mut m = M::new();
693 assert!(!m.contains_key(36));
694 assert!(!m.contains_key(42));
695 m.insert(42, 42.0)?;
696 assert!(!m.contains_key(36));
697 assert!(m.contains_key(42));
698 m.remove(42);
699 assert!(!m.contains_key(36));
700 assert!(!m.contains_key(42));
701 Ok(())
702 }
703
704 #[test]
705 fn get() -> Result<()> {
706 let mut m = M::new();
707 assert!(m.get(36).is_none());
708 assert!(m.get(42).is_none());
709 m.insert(42, 42.0)?;
710 assert!(m.get(36).is_none());
711 assert_eq!(m.get(42), Some(&42.0));
712 m.remove(42);
713 assert!(m.get(36).is_none());
714 assert!(m.get(42).is_none());
715 Ok(())
716 }
717
718 #[test]
719 fn get_mut() -> Result<()> {
720 let mut m = M::new();
721 assert!(m.get_mut(36).is_none());
722 assert!(m.get_mut(42).is_none());
723 m.insert(42, 42.0)?;
724 assert!(m.get_mut(36).is_none());
725 assert_eq!(m.get_mut(42).copied(), Some(42.0));
726 *m.get_mut(42).unwrap() = 99.0;
727 assert_eq!(m.get_mut(42).copied(), Some(99.0));
728 m.remove(42);
729 assert!(m.get_mut(36).is_none());
730 assert!(m.get_mut(42).is_none());
731 Ok(())
732 }
733
734 #[test]
735 fn insert() -> Result<()> {
736 let mut m = M::new();
737 let old = m.insert(11, 0.0)?;
738 assert!(old.is_none());
739 let old = m.insert(11, 1.0)?;
740 assert_eq!(old, Some(0.0));
741 Ok(())
742 }
743
744 #[test]
745 fn remove() -> Result<()> {
746 let mut m = M::new();
747 let old = m.remove(10);
748 assert!(old.is_none());
749 m.insert(10, 123.0)?;
750 let old = m.remove(10);
751 assert_eq!(old, Some(123.0));
752 Ok(())
753 }
754
755 #[test]
756 fn clear() -> Result<()> {
757 let mut m = M::new();
758 for i in 0..10 {
759 m.insert(i, i as f32)?;
760 }
761 m.clear();
762 assert!(m.is_empty());
763 for i in 0..10 {
764 assert!(m.get(i).is_none());
765 }
766 Ok(())
767 }
768
769 #[test]
770 fn iter() -> Result<()> {
771 let mut m = M::new();
772 for i in 0..5 {
773 m.insert(i, i as f32)?;
774 }
775 assert_eq!(
776 m.iter().collect::<Vec<_>>(),
777 vec![(0, &0.0), (1, &1.0), (2, &2.0), (3, &3.0), (4, &4.0)],
778 );
779 Ok(())
780 }
781
782 #[test]
783 fn iter_mut() -> Result<()> {
784 let mut m = M::new();
785 for i in 0..5 {
786 m.insert(i, i as f32)?;
787 }
788 assert_eq!(
789 m.iter_mut()
790 .map(|(k, v)| (k, mem::replace(v, 0.0)))
791 .collect::<Vec<_>>(),
792 vec![(0, 0.0), (1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0)],
793 );
794 for i in 0..5 {
795 assert_eq!(m.get(i), Some(&0.0));
796 }
797 Ok(())
798 }
799
800 #[test]
801 fn keys() -> Result<()> {
802 let mut m = M::new();
803 for i in 0..5 {
804 m.insert(i, i as f32)?;
805 }
806 assert_eq!(m.keys().collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
807 Ok(())
808 }
809
810 #[test]
811 fn values() -> Result<()> {
812 let mut m = M::new();
813 for i in 0..5 {
814 m.insert(i, i as f32)?;
815 }
816 assert_eq!(
817 m.values().collect::<Vec<_>>(),
818 vec![&0.0, &1.0, &2.0, &3.0, &4.0],
819 );
820 Ok(())
821 }
822
823 #[test]
824 fn values_mut() -> Result<()> {
825 let mut m = M::new();
826 for i in 0..5 {
827 m.insert(i, i as f32)?;
828 }
829 assert_eq!(
830 m.values_mut()
831 .map(|v| mem::replace(v, 0.0))
832 .collect::<Vec<_>>(),
833 vec![0.0, 1.0, 2.0, 3.0, 4.0],
834 );
835 assert_eq!(
836 m.values().collect::<Vec<_>>(),
837 vec![&0.0, &0.0, &0.0, &0.0, &0.0],
838 );
839 Ok(())
840 }
841
842 #[test]
843 fn range() -> Result<()> {
844 let mut m = M::new();
845 for i in 0..5 {
846 m.insert(i, i as f32)?;
847 }
848 assert_eq!(m.range(..2).collect::<Vec<_>>(), vec![(0, &0.0), (1, &1.0)]);
849 assert_eq!(m.range(3..).collect::<Vec<_>>(), vec![(3, &3.0), (4, &4.0)]);
850 assert_eq!(
851 m.range(1..3).collect::<Vec<_>>(),
852 vec![(1, &1.0), (2, &2.0)],
853 );
854 assert_eq!(
855 m.range(2..=3).collect::<Vec<_>>(),
856 vec![(2, &2.0), (3, &3.0)],
857 );
858 assert_eq!(m.range(5..).collect::<Vec<_>>(), vec![]);
859 assert_eq!(m.range(..0).collect::<Vec<_>>(), vec![]);
860 Ok(())
861 }
862
863 #[test]
864 fn range_mut() -> Result<()> {
865 let mut m = M::new();
866 for i in 0..5 {
867 m.insert(i, i as f32)?;
868 }
869 assert_eq!(
870 m.range_mut(1..3)
871 .map(|(k, v)| (k, mem::replace(v, 99.0)))
872 .collect::<Vec<_>>(),
873 vec![(1, 1.0), (2, 2.0)],
874 );
875 assert_eq!(
876 m.values().copied().collect::<Vec<_>>(),
877 vec![0.0, 99.0, 99.0, 3.0, 4.0]
878 );
879 Ok(())
880 }
881
882 #[test]
883 fn entry() -> Result<()> {
884 let mut m = M::new();
885
886 let v = m.entry(0).or_insert(99.0)?;
887 assert_eq!(*v, 99.0);
888
889 let v = m.entry(0).or_insert(0.0)?;
890 assert_eq!(*v, 99.0);
891
892 let e = match m.entry(1) {
893 Entry::Occupied(_) => unreachable!(),
894 Entry::Vacant(e) => e.insert_entry(42.0)?,
895 };
896 assert_eq!(e.key(), 1);
897 assert_eq!(e.get(), &42.0);
898 let old = e.insert(43.0);
899 assert_eq!(old, 42.0);
900 assert_eq!(m.get(1), Some(&43.0));
901
902 let e = match m.entry(1) {
903 Entry::Occupied(e) => e,
904 Entry::Vacant(_) => unreachable!(),
905 };
906 assert_eq!(e.key(), 1);
907 assert_eq!(e.get(), &43.0);
908 let (k, v) = e.remove_entry();
909 assert_eq!(k, 1);
910 assert_eq!(v, 43.0);
911 assert!(m.get(1).is_none());
912
913 let e = match m.entry(2) {
914 Entry::Occupied(_) => unreachable!(),
915 Entry::Vacant(e) => e,
916 };
917 assert_eq!(e.key(), 2);
918 let v = e.insert(99.0)?;
919 assert_eq!(*v, 99.0);
920
921 Ok(())
922 }
923}