1mod set;
10#[cfg(test)]
11mod tests;
12
13pub use set::{OwnedSet, OwnedSetIntoIter, OwnedSetIter};
14
15use std::{borrow::Borrow, ops::Bound};
16
17use super::{tree, BudgetedMapIter, Link, Node, OwnedNode, MAX_HEIGHT};
18
19pub struct OwnedMap<K, V> {
21 root: Link<K, V>,
22 len: usize,
23}
24
25impl<K, V> Default for OwnedMap<K, V> {
26 fn default() -> Self {
27 Self { root: None, len: 0 }
28 }
29}
30
31impl<K, V> OwnedMap<K, V> {
32 pub fn new() -> Self {
33 Self::default()
34 }
35
36 pub const fn entry_bytes() -> usize {
38 size_of::<Node<K, V>>()
39 }
40
41 pub fn allocated_bytes(&self) -> usize {
42 self.len * Self::entry_bytes()
43 }
44
45 pub fn len(&self) -> usize {
46 self.len
47 }
48
49 pub fn is_empty(&self) -> bool {
50 self.len == 0
51 }
52
53 pub fn iter(&self) -> BudgetedMapIter<'_, K, V> {
54 BudgetedMapIter::new(&self.root, self.len)
55 }
56
57 pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> + std::iter::FusedIterator {
58 self.iter().map(|(key, _)| key)
59 }
60
61 pub fn values(&self) -> impl ExactSizeIterator<Item = &V> + std::iter::FusedIterator {
62 self.iter().map(|(_, value)| value)
63 }
64}
65
66impl<K: Ord, V> OwnedMap<K, V> {
67 pub fn get<Q: Ord + ?Sized>(&self, key: &Q) -> Option<&V>
68 where
69 K: Borrow<Q>,
70 {
71 tree::get(&self.root, key).map(|node| &node.value)
72 }
73
74 pub fn get_mut<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<&mut V>
75 where
76 K: Borrow<Q>,
77 {
78 tree::get_mut(&mut self.root, key).map(|node| &mut node.value)
79 }
80
81 pub fn contains_key<Q: Ord + ?Sized>(&self, key: &Q) -> bool
82 where
83 K: Borrow<Q>,
84 {
85 self.get(key).is_some()
86 }
87
88 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
90 if let Some(previous) = self.get_mut(&key) {
91 return Some(std::mem::replace(previous, value));
92 }
93 let entry = OwnedNode {
94 value: Box::new(Node {
95 key,
96 value,
97 left: None,
98 right: None,
99 height: 1,
100 }),
101 memory: None,
102 };
103 let previous = tree::insert(&mut self.root, entry);
104 self.len += 1;
105 previous
106 }
107
108 pub fn append(&mut self, other: Self) {
110 if self.is_empty() {
111 *self = other;
112 return;
113 }
114 let mut entries = other.into_iter();
115 while let Some(node) = entries.next_node() {
116 self.len += usize::from(tree::insert(&mut self.root, node).is_none());
117 }
118 }
119
120 pub fn remove<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<V>
121 where
122 K: Borrow<Q>,
123 {
124 self.remove_entry(key).map(|(_, value)| value)
125 }
126
127 pub fn remove_entry<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<(K, V)>
128 where
129 K: Borrow<Q>,
130 {
131 let removed = tree::remove(&mut self.root, key);
132 self.len -= usize::from(removed.is_some());
133 removed
134 }
135
136 pub fn first_from<Q: Ord + ?Sized>(&self, start: Bound<&Q>) -> Option<(&K, &V)>
138 where
139 K: Borrow<Q>,
140 {
141 let mut link = &self.root;
142 let mut candidate = None;
143 while let Some(node) = link {
144 let included = match start {
145 Bound::Unbounded => true,
146 Bound::Included(key) => node.key.borrow() >= key,
147 Bound::Excluded(key) => node.key.borrow() > key,
148 };
149 if included {
150 candidate = Some((&node.key, &node.value.value));
151 link = &node.left;
152 } else {
153 link = &node.right;
154 }
155 }
156 candidate
157 }
158}
159
160impl<K: Ord + Borrow<Q>, V, Q: Ord + ?Sized> std::ops::Index<&Q> for OwnedMap<K, V> {
161 type Output = V;
162
163 fn index(&self, key: &Q) -> &V {
164 self.get(key).expect("missing owned map key")
165 }
166}
167
168impl<K: Ord + Clone, V: Clone> Clone for OwnedMap<K, V> {
169 fn clone(&self) -> Self {
170 self.iter()
171 .map(|(key, value)| (key.clone(), value.clone()))
172 .collect()
173 }
174}
175
176impl<K: std::fmt::Debug, V: std::fmt::Debug> std::fmt::Debug for OwnedMap<K, V> {
177 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178 f.debug_map().entries(self.iter()).finish()
179 }
180}
181
182impl<K: PartialEq, V: PartialEq> PartialEq for OwnedMap<K, V> {
183 fn eq(&self, other: &Self) -> bool {
184 self.len == other.len && self.iter().eq(other.iter())
185 }
186}
187
188impl<K: Eq, V: Eq> Eq for OwnedMap<K, V> {}
189
190impl<K: Ord, V> FromIterator<(K, V)> for OwnedMap<K, V> {
191 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
192 let mut map = Self::new();
193 for (key, value) in iter {
194 map.insert(key, value);
195 }
196 map
197 }
198}
199
200impl<'a, K, V> IntoIterator for &'a OwnedMap<K, V> {
201 type Item = (&'a K, &'a V);
202 type IntoIter = BudgetedMapIter<'a, K, V>;
203
204 fn into_iter(self) -> Self::IntoIter {
205 self.iter()
206 }
207}
208
209pub struct OwnedMapIntoIter<K, V> {
211 stack: [Option<OwnedNode<K, V>>; MAX_HEIGHT],
212 depth: usize,
213 remaining: usize,
214}
215
216impl<K, V> OwnedMapIntoIter<K, V> {
217 fn push_left(&mut self, mut link: Link<K, V>) {
218 while let Some(mut node) = link {
219 link = node.value.left.take();
220 self.stack[self.depth] = Some(node);
221 self.depth += 1;
222 }
223 }
224
225 fn next_node(&mut self) -> Option<OwnedNode<K, V>> {
226 self.depth = self.depth.checked_sub(1)?;
227 let mut node = self.stack[self.depth].take().expect("owned traversal node");
228 self.push_left(node.value.right.take());
229 node.value.height = 1;
230 self.remaining -= 1;
231 Some(node)
232 }
233}
234
235impl<K, V> Iterator for OwnedMapIntoIter<K, V> {
236 type Item = (K, V);
237
238 fn next(&mut self) -> Option<Self::Item> {
239 self.next_node().map(tree::into_entry)
240 }
241
242 fn size_hint(&self) -> (usize, Option<usize>) {
243 (self.remaining, Some(self.remaining))
244 }
245}
246
247impl<K, V> ExactSizeIterator for OwnedMapIntoIter<K, V> {}
248impl<K, V> std::iter::FusedIterator for OwnedMapIntoIter<K, V> {}
249
250impl<K, V> IntoIterator for OwnedMap<K, V> {
251 type Item = (K, V);
252 type IntoIter = OwnedMapIntoIter<K, V>;
253
254 fn into_iter(self) -> Self::IntoIter {
255 let mut iter = OwnedMapIntoIter {
256 stack: std::array::from_fn(|_| None),
257 depth: 0,
258 remaining: self.len,
259 };
260 iter.push_left(self.root);
261 iter
262 }
263}