1#![allow(dead_code)]
2
3use allocated::DropGuard;
4use allocated::FromIteratorIn;
5use core::borrow::Borrow;
6use core::mem::ManuallyDrop;
7use core::ops::Add;
8use core::ops::Mul;
9use core::ptr::NonNull;
10
11extern crate alloc;
12use alloc::vec;
13use alloc::vec::Vec;
14
15use allocator_api2::alloc::Allocator;
16use generic_array::ArrayLength;
17use typenum::Prod;
18use typenum::Sum;
19use typenum::U1;
20use typenum::U2;
21use typenum::U6;
22
23use allocated::AllocResult;
24use allocated::DropGuardResult;
25use allocated::DropIn;
26use allocated::IntoIteratorIn;
27use allocated::RecursiveDropIn;
28
29mod entry;
30mod node;
31mod wrapper;
32
33pub use entry::{Entry, OccupiedEntry, VacantEntry};
34use node::{
35 ChildPtr, IntoIter as NodeIntoIter, Iter as NodeIter, LeafNode, MutNodeRef, Node, NodeEntry,
36 NodeRef, NodeRefT, OccupiedNodeEntry,
37};
38pub use wrapper::CompressedBTreeMap;
39
40pub struct AllocatedBTreeMap<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength = U6>
76where
77 U2: Mul<B>,
78 Prod<U2, B>: ArrayLength,
79 U1: Add<Prod<U2, B>>,
80 Sum<U1, Prod<U2, B>>: ArrayLength,
81{
82 root: Option<ManuallyDrop<Node<K, V, B>>>,
83 n: usize,
84}
85
86impl<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> AllocatedBTreeMap<K, V, B>
87where
88 U2: Mul<B>,
89 Prod<U2, B>: ArrayLength,
90 U1: Add<Prod<U2, B>>,
91 Sum<U1, Prod<U2, B>>: ArrayLength,
92{
93 fn root_node_ref(&self) -> NodeRef<'_, K, V, B> {
94 self.root.as_ref().unwrap().as_ref()
95 }
96
97 fn root_mut_node_ref(&mut self) -> MutNodeRef<'_, K, V, B> {
98 self.root.as_mut().unwrap().as_mut()
99 }
100}
101
102impl<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> AllocatedBTreeMap<K, V, B>
103where
104 U2: Mul<B>,
105 Prod<U2, B>: ArrayLength,
106 U1: Add<Prod<U2, B>>,
107 Sum<U1, Prod<U2, B>>: ArrayLength,
108{
109 pub fn new_in<A: Allocator>(alloc: &A) -> DropGuardResult<Self, &A> {
113 Ok(LeafNode::new_in(alloc).map(|root| AllocatedBTreeMap {
114 root: Some(ManuallyDrop::new(Node::Leaf(root))),
115 n: 0,
116 }))
117 }
118
119 pub fn is_empty(&self) -> bool {
121 self.n == 0
122 }
123
124 pub fn len(&self) -> usize {
126 self.n
127 }
128
129 pub fn contains_key<Q>(&self, key: &Q) -> bool
131 where
132 K: Borrow<Q> + core::cmp::PartialOrd + core::fmt::Debug,
133 Q: core::cmp::PartialOrd + core::fmt::Debug,
134 {
135 let root = self.root_node_ref();
136 match root.ref_entry(key, vec![]) {
137 NodeEntry::Vacant(_) => false,
138 NodeEntry::Occupied(_) => true,
139 }
140 }
141
142 pub unsafe fn insert_in<A: Allocator>(
155 &mut self,
156 alloc: &A,
157 key: K,
158 value: V,
159 ) -> AllocResult<Option<V>> {
160 let entry = unsafe { self.entry_in(alloc, key) };
162 entry.insert(value)
163 }
164
165 pub unsafe fn clear_in<A: Allocator>(&mut self, alloc: &A) -> AllocResult<()> {
175 if let Some(mut r) = self.root.take() {
176 unsafe { r.drop_in(alloc) };
178 }
179 self.n = 0;
180 self.root = Some(LeafNode::new_in(alloc).map(|n| Node::Leaf(n)).into_inner());
181 Ok(())
182 }
183
184 pub fn get<'s, Q>(&'s self, key: &Q) -> Option<&'s V>
186 where
187 K: Borrow<Q> + core::cmp::PartialOrd,
188 Q: core::cmp::PartialOrd + core::fmt::Debug,
189 {
190 let root = self.root_node_ref();
191 match root.ref_entry(key, vec![]) {
192 NodeEntry::Vacant(_) => None,
193 NodeEntry::Occupied(o) => Some(o.into_value()),
194 }
195 }
196
197 pub fn get_key_value<'s, Q>(&'s self, key: &'s Q) -> Option<(&'s K, &'s V)>
199 where
200 K: Borrow<Q> + core::cmp::PartialOrd,
201 Q: core::cmp::PartialOrd + core::fmt::Debug,
202 {
203 let root = self.root_node_ref();
204 match root.ref_entry(key, vec![]) {
205 NodeEntry::Vacant(_) => None,
206 NodeEntry::Occupied(o) => Some(o.into_key_value()),
207 }
208 }
209
210 pub fn get_mut<'s, Q>(&'s mut self, key: &'s Q) -> Option<&'s mut V>
212 where
213 K: Borrow<Q> + core::cmp::PartialOrd,
214 Q: core::cmp::PartialOrd + core::fmt::Debug,
215 {
216 let root = self.root_mut_node_ref();
217 match root.ref_entry(key, vec![]) {
218 NodeEntry::Vacant(_) => None,
219 NodeEntry::Occupied(o) => Some(o.into_mut()),
220 }
221 }
222
223 pub fn iter(&self) -> Iter<'_, K, V, B> {
225 let mut stack = Vec::new();
226 if self.n > 0 {
227 stack.push((ChildPtr::from_node_ref(self.root_node_ref()), 0));
228 }
229 Iter {
230 inner: NodeIter::<K, V, B, NodeRef<K, V, B>>::new(stack),
231 }
232 }
233
234 pub fn keys(&self) -> Keys<'_, K, V, B> {
236 let mut stack = Vec::new();
237 if self.n > 0 {
238 stack.push((ChildPtr::from_node_ref(self.root_node_ref()), 0));
239 }
240 Keys {
241 inner: NodeIter::<K, V, B, NodeRef<K, V, B>>::new(stack),
242 }
243 }
244
245 unsafe fn into_keys_in<A: Allocator>(self, alloc: &A) -> IntoKeys<'_, K, V, B, A> {
249 IntoKeys {
250 inner: NodeIntoIter::new(alloc, ManuallyDrop::into_inner(self.root.unwrap())),
251 }
252 }
253
254 pub fn values(&self) -> Values<'_, K, V, B> {
256 let mut stack = Vec::new();
257 if self.n > 0 {
258 stack.push((ChildPtr::from_node_ref(self.root_node_ref()), 0));
259 }
260 Values {
261 inner: NodeIter::<K, V, B, NodeRef<K, V, B>>::new(stack),
262 }
263 }
264
265 pub fn values_mut(&mut self) -> ValuesMut<'_, K, V, B> {
267 let mut stack = Vec::new();
268 if self.n > 0 {
269 stack.push((
270 ChildPtr::from_mut_node_ref(&mut self.root_mut_node_ref()),
271 0,
272 ));
273 }
274 ValuesMut {
275 inner: NodeIter::<K, V, B, MutNodeRef<K, V, B>>::new(stack),
276 }
277 }
278
279 unsafe fn into_values_in<A: Allocator>(self, alloc: &A) -> IntoValues<'_, K, V, B, A> {
283 IntoValues {
284 inner: NodeIntoIter::new(alloc, ManuallyDrop::into_inner(self.root.unwrap())),
285 }
286 }
287
288 pub unsafe fn entry_in<'a, 's, A: Allocator>(
292 &'s mut self,
293 alloc: &'a A,
294 key: K,
295 ) -> Entry<'a, 's, A, K, V, B> {
296 let map: NonNull<AllocatedBTreeMap<K, V, B>> = NonNull::from_mut(self);
297 let node_ref: MutNodeRef<'s, K, V, B> = self.root_mut_node_ref();
298 let inner_entry: NodeEntry<'s, K, K, V, B, MutNodeRef<'s, K, V, B>> =
300 unsafe { node_ref.entry_in(key, vec![]) };
301 Entry::new(alloc, inner_entry, map)
302 }
303
304 pub unsafe fn first_entry_in<'a, 's, A: Allocator>(
308 &'s mut self,
309 alloc: &'a A,
310 ) -> Option<OccupiedEntry<'a, 's, A, K, V, B>> {
311 if self.n == 0 {
312 return None;
313 }
314
315 let map = NonNull::new(self)?;
316 let root = self.root_mut_node_ref();
317 let inner_entry: OccupiedNodeEntry<'s, K, V, B, MutNodeRef<'s, K, V, B>> =
318 root.first_entry_in(vec![]);
319 unsafe { Some(OccupiedEntry::new(alloc, inner_entry, map)) }
321 }
322
323 pub fn first_key_value<'s>(&'s self) -> Option<(&'s K, &'s V)> {
326 if self.n == 0 {
327 return None;
328 }
329
330 let root = self.root_node_ref();
331 let inner_entry: OccupiedNodeEntry<'s, K, V, B, NodeRef<'s, K, V, B>> =
332 root.first_entry_in(vec![]);
333 Some(inner_entry.into_key_value())
334 }
335
336 pub unsafe fn last_entry_in<'a, 's, A: Allocator>(
340 &'s mut self,
341 alloc: &'a A,
342 ) -> Option<OccupiedEntry<'a, 's, A, K, V, B>> {
343 if self.n == 0 {
344 return None;
345 }
346
347 let map = NonNull::new(self)?;
348 let root = self.root_mut_node_ref();
349 let inner_entry: OccupiedNodeEntry<'s, K, V, B, MutNodeRef<'s, K, V, B>> =
350 root.last_entry_in(vec![]);
351 unsafe { Some(OccupiedEntry::new(alloc, inner_entry, map)) }
353 }
354
355 pub fn last_key_value<'s>(&'s self) -> Option<(&'s K, &'s V)> {
358 if self.n == 0 {
359 return None;
360 }
361
362 let root = self.root_node_ref();
363 let inner_entry: OccupiedNodeEntry<'s, K, V, B, NodeRef<'s, K, V, B>> =
364 root.last_entry_in(vec![]);
365 Some(inner_entry.into_key_value())
366 }
367
368 pub fn first(&self) -> Option<&K> {
371 self.first_key_value().map(|(k, _)| k)
372 }
373
374 pub fn last(&self) -> Option<&K> {
377 self.last_key_value().map(|(k, _)| k)
378 }
379
380 pub unsafe fn entry_ref_in<'a, 's, A: Allocator, Q>(
414 &'s mut self,
415 alloc: &'a A,
416 key: &Q,
417 ) -> Option<OccupiedEntry<'a, 's, A, K, V, B>>
418 where
419 K: Borrow<Q>,
420 Q: PartialOrd + core::fmt::Debug + ?Sized,
421 {
422 let map = NonNull::from_mut(self);
423 let root = self.root_mut_node_ref();
424
425 match root.ref_entry(key, vec![]) {
426 NodeEntry::Vacant(_) => None,
427 NodeEntry::Occupied(inner) => {
428 Some(unsafe { OccupiedEntry::new(alloc, inner, map) })
430 }
431 }
432 }
433
434 pub unsafe fn remove_in<A: Allocator, Q>(&mut self, alloc: &A, key: &Q) -> Option<V>
461 where
462 K: Borrow<Q>,
463 Q: PartialOrd + core::fmt::Debug + ?Sized,
464 {
465 unsafe { self.remove_entry_in(alloc, key).map(|(_, v)| v) }
467 }
468
469 pub unsafe fn remove_entry_in<A: Allocator, Q>(&mut self, alloc: &A, key: &Q) -> Option<(K, V)>
496 where
497 K: Borrow<Q>,
498 Q: PartialOrd + core::fmt::Debug + ?Sized,
499 {
500 unsafe {
502 self.entry_ref_in(alloc, key)
503 .map(|entry| entry.remove_entry())
504 }
505 }
506}
507
508impl<'a, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength, A: Allocator>
528 FromIteratorIn<'a, (K, V), A> for AllocatedBTreeMap<K, V, B>
529where
530 U2: Mul<B>,
531 Prod<U2, B>: ArrayLength,
532 U1: Add<Prod<U2, B>>,
533 Sum<U1, Prod<U2, B>>: ArrayLength,
534{
535 fn from_iter_in<T>(alloc: &'a A, iter: T) -> DropGuardResult<Self, &'a A>
536 where
537 T: IntoIterator<Item = (K, V)>,
538 {
539 let mut btree: DropGuard<Self, &'a A> = Self::new_in(alloc)?;
540
541 for (k, v) in iter {
542 unsafe {
544 btree.insert_in(alloc, k, v)?;
545 }
546 }
547
548 Ok(btree)
549 }
550}
551
552impl<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> DropIn
553 for AllocatedBTreeMap<K, V, B>
554where
555 U2: Mul<B>,
556 Prod<U2, B>: ArrayLength,
557 U1: Add<Prod<U2, B>>,
558 Sum<U1, Prod<U2, B>>: ArrayLength,
559{
560 unsafe fn drop_in<A: Allocator>(&mut self, alloc: &A) {
564 unsafe {
566 if let Some(r) = self.root.as_mut() {
567 r.drop_in(alloc);
568 }
569 }
570 }
571}
572
573impl<K, V, B> RecursiveDropIn for AllocatedBTreeMap<K, V, B>
574where
575 K: core::cmp::PartialOrd + core::fmt::Debug + DropIn,
576 V: DropIn,
577 B: ArrayLength,
578 U2: Mul<B>,
579 Prod<U2, B>: ArrayLength,
580 U1: Add<Prod<U2, B>>,
581 Sum<U1, Prod<U2, B>>: ArrayLength,
582{
583 unsafe fn recursive_drop_in<A: Allocator>(&mut self, alloc: &A) {
587 if let Some(root) = self.root.as_mut() {
588 let root_owned = unsafe { ManuallyDrop::take(root) };
590 for (mut k, mut v) in NodeIntoIter::new(alloc, root_owned) {
592 unsafe { k.drop_in(alloc) };
594 unsafe { v.drop_in(alloc) };
596 }
597 }
598 unsafe { self.drop_in(alloc) };
600 }
601}
602
603impl<'a, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength, A: Allocator + 'a>
604 IntoIteratorIn<'a, A> for AllocatedBTreeMap<K, V, B>
605where
606 U2: Mul<B>,
607 Prod<U2, B>: ArrayLength,
608 U1: Add<Prod<U2, B>>,
609 Sum<U1, Prod<U2, B>>: ArrayLength,
610{
611 type Item = (K, V);
612 type IntoIter = IntoIter<'a, K, V, B, A>;
613
614 unsafe fn into_iter_in(self, alloc: &'a A) -> Self::IntoIter {
615 IntoIter {
616 inner: NodeIntoIter::new(alloc, ManuallyDrop::into_inner(self.root.unwrap())),
617 }
618 }
619}
620
621impl<'s, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> IntoIterator
622 for &'s AllocatedBTreeMap<K, V, B>
623where
624 U2: Mul<B>,
625 Prod<U2, B>: ArrayLength,
626 U1: Add<Prod<U2, B>>,
627 Sum<U1, Prod<U2, B>>: ArrayLength,
628{
629 type IntoIter = Iter<'s, K, V, B>;
630 type Item = (&'s K, &'s V);
631
632 fn into_iter(self) -> Self::IntoIter {
633 self.iter()
634 }
635}
636
637mod iters;
638#[cfg(test)]
639mod tests;
640
641pub use iters::{IntoIter, IntoKeys, IntoValues, Iter, Keys, Values, ValuesMut};