1use core::{
2 borrow::Borrow,
3 marker::PhantomData,
4 ops::{Bound, Deref, DerefMut, RangeBounds},
5};
6
7use crate::{
8 alloc_proxy::proxy::Allocator, layout::AugmentedRBTreeLayout, node::internal_details::NodeRef,
9 policy::internal_details::TreePolicy,
10};
11
12#[derive(Debug)]
17pub struct NodeGuard<'a, K, V, S, P>
18where
19 P: TreePolicy<K = K, V = V, S = S>,
20{
21 node: NodeRef<K, V, S>,
22 marker: PhantomData<(&'a mut (K, V, S), &'a P)>,
23}
24
25impl<K, V, S, P> NodeGuard<'_, K, V, S, P>
26where
27 P: TreePolicy<K = K, V = V, S = S>,
28{
29 pub(crate) fn new(node: NodeRef<K, V, S>) -> Self {
30 Self {
31 node,
32 marker: PhantomData,
33 }
34 }
35 #[must_use]
37 pub fn key(&self) -> &K {
38 unsafe { self.node.key() }
39 }
40
41 #[must_use]
43 pub fn stats(&self) -> &S {
44 unsafe { self.node.stats() }
45 }
46
47 #[must_use]
49 pub fn value_mut(&mut self) -> &mut V {
50 unsafe { self.node.value_mut() }
51 }
52 #[must_use]
54 pub fn value(&self) -> &V {
55 unsafe { self.node.value() }
56 }
57}
58
59impl<K, V, S, P> AsMut<V> for NodeGuard<'_, K, V, S, P>
60where
61 P: TreePolicy<K = K, V = V, S = S>,
62{
63 fn as_mut(&mut self) -> &mut V {
64 unsafe { self.node.value_mut() }
65 }
66}
67
68impl<K, V, S, P> Deref for NodeGuard<'_, K, V, S, P>
69where
70 P: TreePolicy<K = K, V = V, S = S>,
71{
72 type Target = V;
73
74 #[inline]
75 fn deref(&self) -> &Self::Target {
76 unsafe { self.node.value() }
77 }
78}
79
80impl<K, V, S, P> DerefMut for NodeGuard<'_, K, V, S, P>
81where
82 P: TreePolicy<K = K, V = V, S = S>,
83{
84 #[inline]
85 fn deref_mut(&mut self) -> &mut Self::Target {
86 unsafe { self.node.value_mut() }
87 }
88}
89
90impl<K, V, S, P> Drop for NodeGuard<'_, K, V, S, P>
91where
92 P: TreePolicy<K = K, V = V, S = S>,
93{
94 fn drop(&mut self) {
95 P::augment(self.node);
96 P::augment_upstream(self.node);
97 }
98}
99
100#[derive(Debug)]
105pub struct ValueGuard<'a, K, V, S, P>
106where
107 P: TreePolicy<K = K, V = V, S = S>,
108{
109 node: NodeRef<K, V, S>,
110 marker: PhantomData<(&'a mut (K, V, S), &'a P)>,
111}
112
113impl<K, V, S, P> ValueGuard<'_, K, V, S, P>
114where
115 P: TreePolicy<K = K, V = V, S = S>,
116{
117 pub(crate) fn new(node: NodeRef<K, V, S>) -> Self {
118 Self {
119 node,
120 marker: PhantomData,
121 }
122 }
123}
124
125impl<K, V, S, P> Deref for ValueGuard<'_, K, V, S, P>
126where
127 P: TreePolicy<K = K, V = V, S = S>,
128{
129 type Target = V;
130
131 #[inline]
132 fn deref(&self) -> &Self::Target {
133 unsafe { self.node.value() }
134 }
135}
136
137impl<K, V, S, P> DerefMut for ValueGuard<'_, K, V, S, P>
138where
139 P: TreePolicy<K = K, V = V, S = S>,
140{
141 #[inline]
142 fn deref_mut(&mut self) -> &mut Self::Target {
143 unsafe { self.node.value_mut() }
144 }
145}
146
147impl<K, V, S, P> Drop for ValueGuard<'_, K, V, S, P>
148where
149 P: TreePolicy<K = K, V = V, S = S>,
150{
151 fn drop(&mut self) {
152 P::augment(self.node);
153 P::augment_upstream(self.node);
154 }
155}
156
157#[derive(Debug)]
161pub struct IterMut<'a, K, V, S, P>
162where
163 P: TreePolicy<K = K, V = V, S = S>,
164{
165 next: Option<NodeRef<K, V, S>>,
166 back: Option<NodeRef<K, V, S>>,
167 len: usize,
168 _marker: PhantomData<(&'a mut (K, V, S), &'a P)>,
169}
170
171impl<K, V, S, P> IterMut<'_, K, V, S, P>
172where
173 P: TreePolicy<K = K, V = V, S = S>,
174{
175 pub(crate) fn new(root: Option<NodeRef<K, V, S>>, len: usize) -> Self {
176 let next = root.map(NodeRef::leftmost);
177 let back = root.map(NodeRef::rightmost);
178 Self {
179 next,
180 back,
181 len,
182 _marker: PhantomData,
183 }
184 }
185}
186
187impl<'a, K, V, S, P> Iterator for IterMut<'a, K, V, S, P>
188where
189 P: TreePolicy<K = K, V = V, S = S>,
190{
191 type Item = NodeGuard<'a, K, V, S, P>;
192
193 fn next(&mut self) -> Option<Self::Item> {
194 if self.len == 0 {
195 return None;
196 }
197
198 let node = self.next?;
199
200 if self.next == self.back {
202 self.next = None;
203 self.back = None;
204 } else {
205 self.next = node.next_node();
206 }
207
208 self.len -= 1;
209
210 let guard = NodeGuard::new(node);
211 Some(guard)
212 }
213
214 fn size_hint(&self) -> (usize, Option<usize>) {
215 (self.len, Some(self.len))
216 }
217}
218
219impl<'a, K, V, S: 'a, P> DoubleEndedIterator for IterMut<'a, K, V, S, P>
220where
221 P: TreePolicy<K = K, V = V, S = S>,
222{
223 fn next_back(&mut self) -> Option<Self::Item> {
224 if self.len == 0 {
225 return None;
226 }
227
228 let node = self.back?;
229
230 if self.next == self.back {
232 self.next = None;
233 self.back = None;
234 } else {
235 self.back = node.prev_node();
236 }
237
238 self.len -= 1;
239
240 let guard = NodeGuard::new(node);
241 Some(guard)
242 }
243}
244
245impl<'a, K, V, S: 'a, P> ExactSizeIterator for IterMut<'a, K, V, S, P>
246where
247 P: TreePolicy<K = K, V = V, S = S>,
248{
249 fn len(&self) -> usize {
250 self.len
251 }
252}
253
254impl<'a, K, V, S: 'a, P> core::iter::FusedIterator for IterMut<'a, K, V, S, P> where
255 P: TreePolicy<K = K, V = V, S = S>
256{
257}
258
259#[derive(Debug)]
263pub struct ValuesMut<'a, K, V, S, P>
264where
265 P: TreePolicy<K = K, V = V, S = S>,
266{
267 next: Option<NodeRef<K, V, S>>,
268 back: Option<NodeRef<K, V, S>>,
269 len: usize,
270 _marker: PhantomData<(&'a mut (K, V, S), &'a P)>,
271}
272
273impl<K, V, S, P> ValuesMut<'_, K, V, S, P>
274where
275 P: TreePolicy<K = K, V = V, S = S>,
276{
277 pub(crate) fn new(root: Option<NodeRef<K, V, S>>, len: usize) -> Self {
278 let next = root.map(NodeRef::leftmost);
279 let back = root.map(NodeRef::rightmost);
280 Self {
281 next,
282 back,
283 len,
284 _marker: PhantomData,
285 }
286 }
287}
288
289impl<'a, K, V, S, P> Iterator for ValuesMut<'a, K, V, S, P>
290where
291 P: TreePolicy<K = K, V = V, S = S>,
292{
293 type Item = ValueGuard<'a, K, V, S, P>;
294
295 fn next(&mut self) -> Option<Self::Item> {
296 if self.len == 0 {
297 return None;
298 }
299
300 let node = self.next?;
301
302 if self.next == self.back {
304 self.next = None;
305 self.back = None;
306 } else {
307 self.next = node.next_node();
308 }
309
310 self.len -= 1;
311
312 let guard = ValueGuard::new(node);
313 Some(guard)
314 }
315
316 fn size_hint(&self) -> (usize, Option<usize>) {
317 (self.len, Some(self.len))
318 }
319}
320
321impl<'a, K, V, S: 'a, P> DoubleEndedIterator for ValuesMut<'a, K, V, S, P>
322where
323 P: TreePolicy<K = K, V = V, S = S>,
324{
325 fn next_back(&mut self) -> Option<Self::Item> {
326 if self.len == 0 {
327 return None;
328 }
329
330 let node = self.back?;
331
332 if self.next == self.back {
334 self.next = None;
335 self.back = None;
336 } else {
337 self.back = node.prev_node();
338 }
339
340 self.len -= 1;
341
342 let guard = ValueGuard::new(node);
343 Some(guard)
344 }
345}
346
347impl<'a, K, V, S: 'a, P> ExactSizeIterator for ValuesMut<'a, K, V, S, P>
348where
349 P: TreePolicy<K = K, V = V, S = S>,
350{
351 fn len(&self) -> usize {
352 self.len
353 }
354}
355
356impl<'a, K, V, S: 'a, P> core::iter::FusedIterator for ValuesMut<'a, K, V, S, P> where
357 P: TreePolicy<K = K, V = V, S = S>
358{
359}
360
361#[derive(Debug)]
365pub struct Iter<'a, K, V, S> {
366 next: Option<NodeRef<K, V, S>>,
367 back: Option<NodeRef<K, V, S>>,
368 len: usize,
369 _marker: PhantomData<&'a (K, V, S)>,
370}
371
372impl<K, V, S> Iter<'_, K, V, S> {
373 pub(crate) fn new(root: Option<NodeRef<K, V, S>>, len: usize) -> Self {
374 let next = root.map(NodeRef::leftmost);
375 let back = root.map(NodeRef::rightmost);
376 Self {
377 next,
378 back,
379 len,
380 _marker: PhantomData,
381 }
382 }
383}
384
385impl<'a, K, V, S> Iterator for Iter<'a, K, V, S> {
386 type Item = (&'a K, &'a V, &'a S);
387
388 fn next(&mut self) -> Option<Self::Item> {
389 if self.len == 0 {
390 return None;
391 }
392
393 let node = self.next?;
394
395 if self.next == self.back {
397 self.next = None;
398 self.back = None;
399 } else {
400 self.next = node.next_node();
401 }
402
403 self.len -= 1;
404
405 unsafe { Some((node.key(), node.value(), node.stats())) }
406 }
407
408 fn size_hint(&self) -> (usize, Option<usize>) {
409 (self.len, Some(self.len))
410 }
411}
412
413impl<'a, K, V, S: 'a> DoubleEndedIterator for Iter<'a, K, V, S> {
414 fn next_back(&mut self) -> Option<Self::Item> {
415 if self.len == 0 {
416 return None;
417 }
418
419 let node = self.back?;
420
421 if self.next == self.back {
423 self.next = None;
424 self.back = None;
425 } else {
426 self.back = node.prev_node();
427 }
428
429 self.len -= 1;
430
431 unsafe { Some((node.key(), node.value(), node.stats())) }
432 }
433}
434
435impl<'a, K, V, S: 'a> ExactSizeIterator for Iter<'a, K, V, S> {
436 fn len(&self) -> usize {
437 self.len
438 }
439}
440
441impl<'a, K, V, S: 'a> core::iter::FusedIterator for Iter<'a, K, V, S> {}
442
443#[derive(Debug)]
451pub struct IntoIter<K, V, S, A, P>
452where
453 P: TreePolicy<K = K, V = V, S = S>,
454 A: Allocator,
455{
456 next: Option<NodeRef<K, V, S>>,
457 back: Option<NodeRef<K, V, S>>,
458 layout: AugmentedRBTreeLayout<K, V, S, A, P>,
459 len: usize,
460}
461
462impl<K, V, S, A: Allocator, P> IntoIter<K, V, S, A, P>
463where
464 P: TreePolicy<K = K, V = V, S = S>,
465 A: Allocator,
466{
467 pub(crate) fn new(layout: AugmentedRBTreeLayout<K, V, S, A, P>) -> Self {
468 let next = layout.root.map(NodeRef::leftmost);
469 let back = layout.root.map(NodeRef::rightmost);
470 let len = layout.len;
471 Self {
472 next,
473 back,
474 layout,
475 len,
476 }
477 }
478}
479
480impl<K, V, S, A, P> Drop for IntoIter<K, V, S, A, P>
481where
482 P: TreePolicy<K = K, V = V, S = S>,
483 A: Allocator,
484{
485 fn drop(&mut self) {
486 self.layout.clear();
487 }
488}
489
490impl<K, V, S, A: Allocator, P> Iterator for IntoIter<K, V, S, A, P>
491where
492 P: TreePolicy<K = K, V = V, S = S>,
493{
494 type Item = (K, V);
495
496 fn next(&mut self) -> Option<Self::Item> {
497 if self.len == 0 {
498 return None;
499 }
500
501 let current = self
502 .next
503 .take()
504 .expect("This node must exists because len > 0");
505
506 let next_node = current.next_node();
507
508 self.next = next_node;
509
510 let (key, value) = self.layout.delete_node_no_fixup(current);
511
512 self.len -= 1;
513
514 Some((key, value))
515 }
516}
517
518impl<K, V, S, A: Allocator, P> DoubleEndedIterator for IntoIter<K, V, S, A, P>
519where
520 P: TreePolicy<K = K, V = V, S = S>,
521{
522 fn next_back(&mut self) -> Option<Self::Item> {
523 if self.len == 0 {
524 return None;
525 }
526
527 let current = self
528 .back
529 .take()
530 .expect("This node must exists because len > 0");
531
532 let back_node = current.prev_node();
533
534 self.back = back_node;
535
536 let (key, value) = self.layout.delete_node_no_fixup(current);
537
538 self.len -= 1;
539
540 Some((key, value))
541 }
542}
543
544impl<K, V, S, A: Allocator, P: TreePolicy<K = K, V = V, S = S>> ExactSizeIterator
545 for IntoIter<K, V, S, A, P>
546{
547 fn len(&self) -> usize {
548 self.len
549 }
550}
551
552impl<K, V, S, A: Allocator, P: TreePolicy<K = K, V = V, S = S>> core::iter::FusedIterator
553 for IntoIter<K, V, S, A, P>
554{
555}
556
557#[derive(Debug)]
561pub struct Keys<'a, K, V, S> {
562 inner: Iter<'a, K, V, S>,
563}
564
565impl<'a, K, V, S> Keys<'a, K, V, S> {
566 pub(crate) fn new(inner: Iter<'a, K, V, S>) -> Self {
567 Self { inner }
568 }
569}
570
571impl<'a, K, V, S: 'a> Iterator for Keys<'a, K, V, S> {
572 type Item = &'a K;
573
574 fn next(&mut self) -> Option<Self::Item> {
575 self.inner.next().map(|(k, _, _)| k)
576 }
577
578 fn size_hint(&self) -> (usize, Option<usize>) {
579 self.inner.size_hint()
580 }
581}
582
583impl<'a, K, V, S: 'a> DoubleEndedIterator for Keys<'a, K, V, S> {
584 fn next_back(&mut self) -> Option<Self::Item> {
585 self.inner.next_back().map(|(k, _, _)| k)
586 }
587}
588
589impl<'a, K, V, S: 'a> ExactSizeIterator for Keys<'a, K, V, S> {
590 fn len(&self) -> usize {
591 self.inner.len()
592 }
593}
594
595impl<'a, K, V, S: 'a> core::iter::FusedIterator for Keys<'a, K, V, S> {}
596
597#[derive(Debug)]
601pub struct Values<'a, K, V, S> {
602 inner: Iter<'a, K, V, S>,
603}
604
605impl<'a, K, V, S> Values<'a, K, V, S> {
606 pub(crate) fn new(inner: Iter<'a, K, V, S>) -> Self {
607 Self { inner }
608 }
609}
610
611impl<'a, K, V, S: 'a> Iterator for Values<'a, K, V, S> {
612 type Item = &'a V;
613
614 fn next(&mut self) -> Option<Self::Item> {
615 self.inner.next().map(|(_, v, _)| v)
616 }
617
618 fn size_hint(&self) -> (usize, Option<usize>) {
619 self.inner.size_hint()
620 }
621}
622
623impl<'a, K, V, S: 'a> DoubleEndedIterator for Values<'a, K, V, S> {
624 fn next_back(&mut self) -> Option<Self::Item> {
625 self.inner.next_back().map(|(_, v, _)| v)
626 }
627}
628
629impl<'a, K, V, S: 'a> ExactSizeIterator for Values<'a, K, V, S> {
630 fn len(&self) -> usize {
631 self.inner.len()
632 }
633}
634
635impl<'a, K, V, S: 'a> core::iter::FusedIterator for Values<'a, K, V, S> {}
636
637#[derive(Debug)]
641pub struct Stats<'a, K, V, S> {
642 inner: Iter<'a, K, V, S>,
643}
644
645impl<'a, K, V, S> Stats<'a, K, V, S> {
646 pub(crate) fn new(inner: Iter<'a, K, V, S>) -> Self {
647 Self { inner }
648 }
649}
650
651impl<'a, K, V, S: 'a> Iterator for Stats<'a, K, V, S> {
652 type Item = &'a S;
653
654 fn next(&mut self) -> Option<Self::Item> {
655 self.inner.next().map(|(_, _, s)| s)
656 }
657
658 fn size_hint(&self) -> (usize, Option<usize>) {
659 self.inner.size_hint()
660 }
661}
662
663impl<'a, K, V, S: 'a> DoubleEndedIterator for Stats<'a, K, V, S> {
664 fn next_back(&mut self) -> Option<Self::Item> {
665 self.inner.next_back().map(|(_, _, s)| s)
666 }
667}
668
669impl<'a, K, V, S: 'a> ExactSizeIterator for Stats<'a, K, V, S> {
670 fn len(&self) -> usize {
671 self.inner.len()
672 }
673}
674
675impl<'a, K, V, S: 'a> core::iter::FusedIterator for Stats<'a, K, V, S> {}
676
677#[derive(Debug)]
685pub struct Range<'a, K, V, S> {
686 front: Option<NodeRef<K, V, S>>,
687 back: Option<NodeRef<K, V, S>>,
688 exhausted: bool,
689 _marker: PhantomData<&'a (K, V, S)>,
690}
691
692impl<'a, K, V, S> Range<'a, K, V, S>
693where
694 K: Ord,
695{
696 pub(crate) fn new<Q, R>(layout: &'a dyn RangeBoundsLimits<K, V, S, Q>, range: R) -> Self
697 where
698 K: Borrow<Q> + Ord,
699 Q: Ord + ?Sized + 'a,
700 R: RangeBounds<Q>,
701 {
702 let (front, back, exhausted) = range_bounds_to_nodes(layout, &range);
703 Self {
704 front,
705 back,
706 exhausted,
707 _marker: PhantomData,
708 }
709 }
710}
711
712impl<'a, K, V, S: 'a> Iterator for Range<'a, K, V, S> {
713 type Item = (&'a K, &'a V, &'a S);
714
715 fn next(&mut self) -> Option<Self::Item> {
716 if self.exhausted {
717 return None;
718 }
719 let node = self.front?;
720 if self.front == self.back {
721 self.exhausted = true;
722 } else {
723 self.front = node.next_node();
724 }
725 unsafe { Some((node.key(), node.value(), node.stats())) }
726 }
727}
728
729impl<'a, K, V, S: 'a> DoubleEndedIterator for Range<'a, K, V, S> {
730 fn next_back(&mut self) -> Option<Self::Item> {
731 if self.exhausted {
732 return None;
733 }
734 let node = self.back?;
735 if self.front == self.back {
736 self.exhausted = true;
737 } else {
738 self.back = node.prev_node();
739 }
740 unsafe { Some((node.key(), node.value(), node.stats())) }
741 }
742}
743
744impl<'a, K, V, S: 'a> core::iter::FusedIterator for Range<'a, K, V, S> {}
745
746#[derive(Debug)]
750pub struct RangeMut<'a, K, V, S, P>
751where
752 P: TreePolicy<K = K, V = V, S = S>,
753{
754 front: Option<NodeRef<K, V, S>>,
755 back: Option<NodeRef<K, V, S>>,
756 exhausted: bool,
757 _marker: PhantomData<(&'a mut (K, V, S), &'a P)>,
758}
759
760impl<'a, K, V, S, P> RangeMut<'a, K, V, S, P>
761where
762 K: Ord,
763 P: TreePolicy<K = K, V = V, S = S>,
764{
765 pub(crate) fn new<Q, R>(layout: &'a dyn RangeBoundsLimits<K, V, S, Q>, range: R) -> Self
766 where
767 K: Borrow<Q> + Ord,
768 Q: Ord + ?Sized,
769 R: RangeBounds<Q>,
770 {
771 let (front, back, exhausted) = range_bounds_to_nodes(layout, &range);
772
773 Self {
774 front,
775 back,
776 exhausted,
777 _marker: PhantomData,
778 }
779 }
780}
781
782impl<'a, K, V, S, P> Iterator for RangeMut<'a, K, V, S, P>
783where
784 P: TreePolicy<K = K, V = V, S = S>,
785{
786 type Item = NodeGuard<'a, K, V, S, P>;
787
788 fn next(&mut self) -> Option<Self::Item> {
789 if self.exhausted {
790 return None;
791 }
792 let node = self.front?;
793 if self.front == self.back {
794 self.exhausted = true;
795 } else {
796 self.front = node.next_node();
797 }
798 let guard = NodeGuard::new(node);
799 Some(guard)
800 }
801}
802
803impl<'a, K, V, S: 'a, P> DoubleEndedIterator for RangeMut<'a, K, V, S, P>
804where
805 P: TreePolicy<K = K, V = V, S = S>,
806{
807 fn next_back(&mut self) -> Option<Self::Item> {
808 if self.exhausted {
809 return None;
810 }
811 let node = self.back?;
812 if self.front == self.back {
813 self.exhausted = true;
814 } else {
815 self.back = node.prev_node();
816 }
817 let guard = NodeGuard::new(node);
818 Some(guard)
819 }
820}
821
822impl<'a, K, V, S: 'a, P> core::iter::FusedIterator for RangeMut<'a, K, V, S, P> where
823 P: TreePolicy<K = K, V = V, S = S>
824{
825}
826
827type RangeBoundsResult<K, V, S> = (Option<NodeRef<K, V, S>>, Option<NodeRef<K, V, S>>, bool);
829
830pub(crate) trait RangeBoundsLimits<K, V, S, Q: ?Sized> {
831 fn lower_bound(&self, key: &Q) -> Option<NodeRef<K, V, S>>;
832 fn lower_bound_excluded(&self, key: &Q) -> Option<NodeRef<K, V, S>>;
833 fn floor(&self, key: &Q) -> Option<NodeRef<K, V, S>>;
834 fn floor_excluded(&self, key: &Q) -> Option<NodeRef<K, V, S>>;
835 #[allow(unused)]
836 fn leftmost(&self) -> Option<NodeRef<K, V, S>>;
837 fn rightmost(&self) -> Option<NodeRef<K, V, S>>;
838}
839
840fn range_bounds_to_nodes<K, V, S, R, Q>(
841 layout: &dyn RangeBoundsLimits<K, V, S, Q>,
842 range: &R,
843) -> RangeBoundsResult<K, V, S>
844where
845 R: RangeBounds<Q>,
846 K: Borrow<Q> + Ord,
847 Q: Ord + ?Sized,
848{
849 let front = match range.start_bound() {
850 Bound::Included(k) => layout.lower_bound(k),
851 Bound::Excluded(k) => layout.lower_bound_excluded(k),
852 Bound::Unbounded => layout.leftmost(),
853 };
854
855 let back = match range.end_bound() {
856 Bound::Included(k) => layout.floor(k),
857 Bound::Excluded(k) => layout.floor_excluded(k),
858 Bound::Unbounded => layout.rightmost(),
859 };
860
861 let exhausted = match (front, back) {
863 (Some(f), Some(b)) => unsafe { f.key() > b.key() },
864 (Some(_), None) | (None, _) => true,
865 };
866
867 (front, back, exhausted)
868}