1use alloc::collections::{
4 btree_map::{Iter, IterMut, Keys, Values, ValuesMut},
5 BTreeMap,
6};
7use core::{marker::PhantomData, usize};
8use owner_monad::{Owner, OwnerMut};
9
10#[derive(Debug, PartialEq, Eq)]
13pub struct Map<K, V>(BTreeMap<K, V>);
14
15impl<K: Ord, V> Map<K, V> {
16 pub fn new() -> Self {
18 Self::default()
19 }
20
21 pub fn len(&self) -> usize {
23 self.0.len()
24 }
25
26 pub fn is_empty(&self) -> bool {
28 self.0.is_empty()
29 }
30
31 pub fn iter(&self) -> Iter<K, V> {
33 self.0.iter()
34 }
35
36 pub fn iter_mut(&mut self) -> IterMut<K, V> {
38 self.0.iter_mut()
39 }
40
41 pub fn keys(&self) -> Keys<K, V> {
43 self.0.keys()
44 }
45
46 pub fn values(&self) -> Values<K, V> {
48 self.0.values()
49 }
50
51 pub fn values_mut(&mut self) -> ValuesMut<K, V> {
53 self.0.values_mut()
54 }
55}
56
57impl<K: Ord, V> Default for Map<K, V> {
58 fn default() -> Self {
59 Self(Default::default())
60 }
61}
62
63impl<K: Ord, V> IntoIterator for Map<K, V> {
64 type Item = (K, V);
65
66 type IntoIter = <BTreeMap<K, V> as IntoIterator>::IntoIter;
67
68 fn into_iter(self) -> Self::IntoIter {
69 self.0.into_iter()
70 }
71}
72
73impl<'a, K: Ord, V> IntoIterator for &'a Map<K, V> {
74 type Item = (&'a K, &'a V);
75
76 type IntoIter = <&'a BTreeMap<K, V> as IntoIterator>::IntoIter;
77
78 fn into_iter(self) -> Self::IntoIter {
79 #[allow(clippy::into_iter_on_ref)]
80 (&self.0).into_iter()
81 }
82}
83
84impl<'a, K: Ord, V> IntoIterator for &'a mut Map<K, V> {
85 type Item = (&'a K, &'a mut V);
86
87 type IntoIter = <&'a mut BTreeMap<K, V> as IntoIterator>::IntoIter;
88
89 fn into_iter(self) -> Self::IntoIter {
90 #[allow(clippy::into_iter_on_ref)]
91 (&mut self.0).into_iter()
92 }
93}
94
95pub struct MapHandle<K: Ord, V, O: OwnerMut<Map<K, V>>> {
97 key: K,
98 owner: O,
99 _phantom: PhantomData<V>,
100}
101
102impl<K: Ord, V, O: OwnerMut<Map<K, V>>> Owner<O> for MapHandle<K, V, O> {
103 fn with<'a, U>(&'a self, f: impl FnOnce(&O) -> U) -> Option<U>
104 where
105 O: 'a,
106 {
107 Some(f(&self.owner))
108 }
109}
110
111impl<K: Ord, V, O: OwnerMut<Map<K, V>>> Drop for MapHandle<K, V, O> {
112 fn drop(&mut self) {
113 let key = &self.key;
114 self.owner.with(|map| map.0.remove(key));
115 }
116}
117
118pub fn insert<K: Clone + Ord, V, O: OwnerMut<Map<K, V>>>(
122 mut owner: O,
123 key: K,
124 value: V,
125) -> Option<MapHandle<K, V, O>> {
126 let r = {
127 let key = key.clone();
128 owner.with(move |map| map.0.insert(key, value))?
129 };
130 if r.is_none() {
131 Some(MapHandle {
132 key,
133 owner,
134 _phantom: PhantomData,
135 })
136 } else {
137 None
138 }
139}