embed_collections/btree/
iter.rs

1use super::BTreeMap;
2use super::inter::*;
3use super::leaf::*;
4use super::node::*;
5use core::marker::PhantomData;
6use core::mem::needs_drop;
7
8pub(super) struct IterForward<K, V> {
9    /// Current leaf node for forward iteration
10    pub front_leaf: LeafNode<K, V>,
11    /// Current index within the leaf for forward iteration
12    pub idx: u32,
13}
14
15impl<K, V> IterForward<K, V> {
16    #[inline]
17    pub fn next(&mut self) -> Option<(&LeafNode<K, V>, u32)> {
18        if self.idx >= self.front_leaf.key_count() {
19            if let Some(next) = self.front_leaf.get_right_node() {
20                debug_assert!(next.key_count() > 0);
21                self.front_leaf = next;
22                self.idx = 0;
23            } else {
24                return None;
25            }
26        }
27        let current_idx = self.idx;
28        self.idx = current_idx + 1;
29        Some((&self.front_leaf, current_idx))
30    }
31
32    #[inline]
33    pub unsafe fn next_pair(&mut self) -> Option<(*mut K, *mut V)> {
34        let (leaf, idx) = self.next()?;
35        leaf.get_raw_pair(idx)
36    }
37}
38
39pub(super) struct IterBackward<K, V> {
40    /// Current leaf node for backward iteration
41    pub back_leaf: LeafNode<K, V>,
42    /// back_idx - 1 is the next index within the back leaf, initial to key_count
43    /// When back_idx == 0, should go to previous leaf
44    pub back_idx: u32,
45}
46
47impl<K, V> IterBackward<K, V> {
48    #[inline]
49    pub fn prev(&mut self) -> Option<(&LeafNode<K, V>, u32)> {
50        if self.back_idx == 0 {
51            if let Some(prev) = self.back_leaf.get_left_node() {
52                self.back_idx = prev.key_count();
53                self.back_leaf = prev;
54                debug_assert!(self.back_idx > 0);
55            } else {
56                return None;
57            }
58        }
59        self.back_idx -= 1;
60        Some((&self.back_leaf, self.back_idx))
61    }
62
63    #[inline]
64    pub unsafe fn prev_pair(&mut self) -> Option<(*mut K, *mut V)> {
65        let (leaf, idx) = self.prev()?;
66        leaf.get_raw_pair(idx)
67    }
68}
69
70/// Internal base iterator structure that handles the common logic
71/// for iterating over leaf nodes in both forward and backward directions.
72struct IterBase<K, V> {
73    /// Remaining elements to iterate, it serve as a guard in case forward and backward rendezvous
74    remaining: usize,
75    forward: IterForward<K, V>,
76    backward: IterBackward<K, V>,
77}
78
79impl<K, V> IterBase<K, V> {
80    /// Create a new IterBase
81    /// leaves: (front_leaf, back_leaf) - both Some or both None
82    /// remaining: total remaining elements
83    #[inline]
84    fn new(front_leaf: LeafNode<K, V>, back_leaf: LeafNode<K, V>, remaining: usize) -> Self {
85        Self {
86            forward: IterForward { front_leaf, idx: 0 },
87            backward: IterBackward { back_idx: back_leaf.key_count(), back_leaf },
88            remaining,
89        }
90    }
91
92    /// Advance the forward iterator and return the current leaf and index
93    /// Returns None if we've moved past the end
94    #[inline]
95    fn advance_forward(&mut self) -> Option<(&LeafNode<K, V>, u32)> {
96        if self.remaining == 0 {
97            return None;
98        }
99        self.remaining -= 1;
100        self.forward.next()
101    }
102
103    /// Advance the backward iterator and return the current leaf and index
104    /// Returns None if we've moved past the beginning
105    #[inline]
106    fn advance_backward(&mut self) -> Option<(&LeafNode<K, V>, u32)> {
107        if self.remaining == 0 {
108            return None;
109        }
110        self.remaining -= 1;
111        self.backward.prev()
112    }
113
114    /// Get remaining count
115    #[inline]
116    fn remaining(&self) -> usize {
117        self.remaining
118    }
119}
120
121/// An iterator over the entries of a BTreeMap
122pub struct Iter<'a, K: 'a, V: 'a> {
123    base: Option<IterBase<K, V>>,
124    _marker: PhantomData<&'a ()>,
125}
126
127impl<'a, K: 'a, V: 'a> Iter<'a, K, V> {
128    #[inline]
129    pub(super) fn new(leaves: Option<(LeafNode<K, V>, LeafNode<K, V>)>, remaining: usize) -> Self {
130        if let Some((front, back)) = leaves {
131            Self { base: Some(IterBase::new(front, back, remaining)), _marker: Default::default() }
132        } else {
133            Self { base: None, _marker: Default::default() }
134        }
135    }
136}
137
138impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
139    type Item = (&'a K, &'a V);
140
141    #[inline]
142    fn next(&mut self) -> Option<Self::Item> {
143        let base = self.base.as_mut()?;
144        let (leaf, idx) = base.advance_forward()?;
145        unsafe {
146            let key = (*leaf.key_ptr(idx)).assume_init_ref();
147            let value = (*leaf.value_ptr(idx)).assume_init_ref();
148            Some((key, value))
149        }
150    }
151
152    #[inline]
153    fn size_hint(&self) -> (usize, Option<usize>) {
154        let len = self.len();
155        (len, Some(len))
156    }
157}
158
159impl<'a, K: 'a, V: 'a> ExactSizeIterator for Iter<'a, K, V> {
160    #[inline]
161    fn len(&self) -> usize {
162        self.base.as_ref().map(|x| x.remaining()).unwrap_or(0)
163    }
164}
165
166impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> {
167    #[inline]
168    fn next_back(&mut self) -> Option<Self::Item> {
169        let base = self.base.as_mut()?;
170        let (leaf, idx) = base.advance_backward()?;
171        unsafe {
172            let key = (*leaf.key_ptr(idx)).assume_init_ref();
173            let value = (*leaf.value_ptr(idx)).assume_init_ref();
174            Some((key, value))
175        }
176    }
177}
178
179/// A mutable iterator over the entries of a BTreeMap
180pub struct IterMut<'a, K: 'a, V: 'a> {
181    base: Option<IterBase<K, V>>,
182    _marker: PhantomData<&'a mut ()>,
183}
184
185impl<'a, K: 'a, V: 'a> IterMut<'a, K, V> {
186    #[inline]
187    pub(super) fn new(leaves: Option<(LeafNode<K, V>, LeafNode<K, V>)>, remaining: usize) -> Self {
188        if let Some((front, back)) = leaves {
189            Self { base: Some(IterBase::new(front, back, remaining)), _marker: Default::default() }
190        } else {
191            Self { base: None, _marker: Default::default() }
192        }
193    }
194}
195
196impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
197    type Item = (&'a K, &'a mut V);
198
199    #[inline]
200    fn next(&mut self) -> Option<Self::Item> {
201        let base = self.base.as_mut()?;
202        let (leaf, idx) = base.advance_forward()?;
203        unsafe {
204            let key = (*leaf.key_ptr(idx)).assume_init_ref();
205            let value = (*leaf.value_ptr(idx)).assume_init_mut();
206            Some((key, value))
207        }
208    }
209
210    #[inline]
211    fn size_hint(&self) -> (usize, Option<usize>) {
212        let len = self.len();
213        (len, Some(len))
214    }
215}
216
217impl<'a, K: 'a, V: 'a> ExactSizeIterator for IterMut<'a, K, V> {
218    #[inline]
219    fn len(&self) -> usize {
220        self.base.as_ref().map(|x| x.remaining()).unwrap_or(0)
221    }
222}
223
224impl<'a, K: 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> {
225    #[inline]
226    fn next_back(&mut self) -> Option<Self::Item> {
227        let base = self.base.as_mut()?;
228        let (leaf, idx) = base.advance_backward()?;
229        unsafe {
230            let key = (*leaf.key_ptr(idx)).assume_init_ref();
231            let value = (*leaf.value_ptr(idx)).assume_init_mut();
232            Some((key, value))
233        }
234    }
235}
236
237/// An iterator over the keys of a BTreeMap
238pub struct Keys<'a, K: 'a, V: 'a> {
239    inner: Iter<'a, K, V>,
240}
241
242impl<'a, K: 'a, V: 'a> Keys<'a, K, V> {
243    #[inline]
244    pub(super) fn new(inner: Iter<'a, K, V>) -> Self {
245        Self { inner }
246    }
247}
248
249impl<'a, K: 'a, V: 'a> Iterator for Keys<'a, K, V> {
250    type Item = &'a K;
251
252    #[inline]
253    fn next(&mut self) -> Option<Self::Item> {
254        self.inner.next().map(|(k, _)| k)
255    }
256
257    #[inline]
258    fn size_hint(&self) -> (usize, Option<usize>) {
259        self.inner.size_hint()
260    }
261}
262
263impl<'a, K: 'a, V: 'a> ExactSizeIterator for Keys<'a, K, V> {
264    #[inline]
265    fn len(&self) -> usize {
266        self.inner.len()
267    }
268}
269
270impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Keys<'a, K, V> {
271    #[inline]
272    fn next_back(&mut self) -> Option<Self::Item> {
273        self.inner.next_back().map(|(k, _)| k)
274    }
275}
276
277/// An iterator over the values of a BTreeMap
278pub struct Values<'a, K: 'a, V: 'a> {
279    inner: Iter<'a, K, V>,
280}
281
282impl<'a, K: 'a, V: 'a> Values<'a, K, V> {
283    #[inline]
284    pub(super) fn new(inner: Iter<'a, K, V>) -> Self {
285        Self { inner }
286    }
287}
288
289impl<'a, K: 'a, V: 'a> Iterator for Values<'a, K, V> {
290    type Item = &'a V;
291
292    #[inline]
293    fn next(&mut self) -> Option<Self::Item> {
294        self.inner.next().map(|(_, v)| v)
295    }
296
297    #[inline]
298    fn size_hint(&self) -> (usize, Option<usize>) {
299        self.inner.size_hint()
300    }
301}
302
303impl<'a, K: 'a, V: 'a> ExactSizeIterator for Values<'a, K, V> {
304    #[inline]
305    fn len(&self) -> usize {
306        self.inner.len()
307    }
308}
309
310impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Values<'a, K, V> {
311    #[inline]
312    fn next_back(&mut self) -> Option<Self::Item> {
313        self.inner.next_back().map(|(_, v)| v)
314    }
315}
316
317/// A mutable iterator over the values of a BTreeMap
318pub struct ValuesMut<'a, K: 'a, V: 'a> {
319    inner: IterMut<'a, K, V>,
320}
321
322impl<'a, K: 'a, V: 'a> ValuesMut<'a, K, V> {
323    #[inline]
324    pub(super) fn new(inner: IterMut<'a, K, V>) -> Self {
325        Self { inner }
326    }
327}
328
329impl<'a, K: 'a, V: 'a> Iterator for ValuesMut<'a, K, V> {
330    type Item = &'a mut V;
331
332    #[inline]
333    fn next(&mut self) -> Option<Self::Item> {
334        self.inner.next().map(|(_, v)| v)
335    }
336
337    #[inline]
338    fn size_hint(&self) -> (usize, Option<usize>) {
339        self.inner.size_hint()
340    }
341}
342
343impl<'a, K: 'a, V: 'a> ExactSizeIterator for ValuesMut<'a, K, V> {
344    #[inline]
345    fn len(&self) -> usize {
346        self.inner.len()
347    }
348}
349
350impl<'a, K: 'a, V: 'a> DoubleEndedIterator for ValuesMut<'a, K, V> {
351    #[inline]
352    fn next_back(&mut self) -> Option<Self::Item> {
353        self.inner.next_back().map(|(_, v)| v)
354    }
355}
356
357// Range and RangeMut implementations with DoubleEndedIterator support
358
359/// Internal base structure for range iteration
360/// Manages forward and backward iteration over a range of entries
361/// When the range is exhausted, the entire struct should be set to None
362pub(super) struct RangeBase<'a, K: 'a, V: 'a> {
363    front_leaf: LeafNode<K, V>,
364    back_leaf: LeafNode<K, V>,
365    front_idx: u32,
366    // back_idx -1 is the next pos, initial to be key_count
367    back_idx: u32,
368    _marker: PhantomData<&'a ()>,
369}
370
371impl<'a, K: 'a, V: 'a> RangeBase<'a, K, V> {
372    #[inline]
373    pub fn new(
374        front_leaf: LeafNode<K, V>, front_idx: u32, back_leaf: LeafNode<K, V>, back_idx: u32,
375    ) -> Self {
376        Self { front_leaf, front_idx, back_leaf, back_idx, _marker: PhantomData }
377    }
378
379    /// Advance forward and return the current leaf and index
380    /// Returns None when range is exhausted, caller should set RangeBase to None
381    #[inline]
382    fn advance_forward(&mut self) -> Option<(&LeafNode<K, V>, u32)> {
383        loop {
384            let idx = self.front_idx;
385            if self.front_leaf == self.back_leaf {
386                if idx < self.back_idx {
387                    self.front_idx = idx + 1;
388                    return Some((&self.front_leaf, idx));
389                } else {
390                    return None;
391                }
392            } else {
393                if idx < self.front_leaf.key_count() {
394                    self.front_idx = idx + 1;
395                    return Some((&self.front_leaf, idx));
396                } else {
397                    self.front_leaf = self.front_leaf.get_right_node().unwrap();
398                    self.front_idx = 0;
399                    // Because we always need to compare front_leaf and back_leaf after cursor
400                    // moves, cannot use IterForward & IterBackward here.
401                }
402            }
403        }
404    }
405
406    /// Advance backward and return the current leaf and index
407    /// Returns None when range is exhausted, caller should set RangeBase to None
408    #[inline]
409    fn advance_backward(&mut self) -> Option<(&LeafNode<K, V>, u32)> {
410        loop {
411            if self.back_leaf == self.front_leaf {
412                if self.back_idx == 0 {
413                    return None;
414                }
415                self.back_idx -= 1;
416                if self.back_idx >= self.front_idx {
417                    return Some((&self.back_leaf, self.back_idx));
418                } else {
419                    return None;
420                }
421            } else {
422                if self.back_idx == 0 {
423                    self.back_leaf = self.back_leaf.get_left_node().unwrap();
424                    self.back_idx = self.back_leaf.key_count();
425                    // Because we always need to compare front_leaf and back_leaf after cursor
426                    // moves, cannot use IterForward & IterBackward here.
427                } else {
428                    self.back_idx -= 1;
429                    return Some((&self.back_leaf, self.back_idx));
430                }
431            }
432        }
433    }
434}
435
436/// An iterator over a sub-range of entries in a BTreeMap
437/// RangeBase is wrapped in Option - when exhausted, it becomes None
438pub struct Range<'a, K: 'a, V: 'a> {
439    base: Option<RangeBase<'a, K, V>>,
440}
441
442impl<'a, K: 'a, V: 'a> Range<'a, K, V> {
443    #[inline]
444    pub(super) fn new(base: Option<RangeBase<'a, K, V>>) -> Self {
445        Self { base }
446    }
447}
448
449impl<'a, K: 'a, V: 'a> Iterator for Range<'a, K, V> {
450    type Item = (&'a K, &'a V);
451
452    #[inline]
453    fn next(&mut self) -> Option<Self::Item> {
454        let base = self.base.as_mut()?;
455        if let Some((leaf, idx)) = base.advance_forward() {
456            unsafe {
457                let key = (*leaf.key_ptr(idx)).assume_init_ref();
458                let value = (*leaf.value_ptr(idx)).assume_init_ref();
459                Some((key, value))
460            }
461        } else {
462            self.base = None;
463            None
464        }
465    }
466}
467
468impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Range<'a, K, V> {
469    #[inline]
470    fn next_back(&mut self) -> Option<Self::Item> {
471        let base = self.base.as_mut()?;
472        if let Some((leaf, idx)) = base.advance_backward() {
473            unsafe {
474                let key = (*leaf.key_ptr(idx)).assume_init_ref();
475                let value = (*leaf.value_ptr(idx)).assume_init_ref();
476                Some((key, value))
477            }
478        } else {
479            self.base = None;
480            None
481        }
482    }
483}
484
485/// A mutable iterator over a sub-range of entries in a BTreeMap
486/// RangeBase is wrapped in Option - when exhausted, it becomes None
487pub struct RangeMut<'a, K: 'a, V: 'a> {
488    base: Option<RangeBase<'a, K, V>>,
489}
490
491impl<'a, K: 'a, V: 'a> RangeMut<'a, K, V> {
492    #[inline]
493    pub(super) fn new(base: Option<RangeBase<'a, K, V>>) -> Self {
494        Self { base }
495    }
496}
497
498impl<'a, K: 'a, V: 'a> Iterator for RangeMut<'a, K, V> {
499    type Item = (&'a K, &'a mut V);
500
501    #[inline]
502    fn next(&mut self) -> Option<Self::Item> {
503        let base = self.base.as_mut()?;
504        if let Some((leaf, idx)) = base.advance_forward() {
505            unsafe {
506                let key = (*leaf.key_ptr(idx)).assume_init_ref();
507                let value = (*leaf.value_ptr(idx)).assume_init_mut();
508                Some((key, value))
509            }
510        } else {
511            self.base = None;
512            None
513        }
514    }
515}
516
517impl<'a, K: 'a, V: 'a> DoubleEndedIterator for RangeMut<'a, K, V> {
518    #[inline]
519    fn next_back(&mut self) -> Option<Self::Item> {
520        let base = self.base.as_mut()?;
521        if let Some((leaf, idx)) = base.advance_backward() {
522            unsafe {
523                let key = (*leaf.key_ptr(idx)).assume_init_ref();
524                let value = (*leaf.value_ptr(idx)).assume_init_mut();
525                Some((key, value))
526            }
527        } else {
528            self.base = None;
529            None
530        }
531    }
532}
533
534/// An owning iterator over the entries of a BTreeMap
535/// Uses PathCache to manage tree traversal and safe deallocation
536///
537/// NOTE: In order to keep the logic simple, we does not implement `DoubleEndedIterator` here
538pub struct IntoIter<K: Ord + Clone + Sized, V: Sized> {
539    /// Path cache for deallocating internal nodes after iteration
540    cache: PathCache<K, V>,
541    /// Current leaf being iterated
542    leaf: Option<LeafNode<K, V>>,
543    /// Current index within the leaf
544    idx: u32,
545    /// Remaining elements to iterate
546    remaining: usize,
547    /// If true, iterate in reverse order
548    is_forward: bool,
549}
550
551impl<K: Ord + Clone + Sized, V: Sized> IntoIter<K, V> {
552    #[inline]
553    pub(super) fn new(mut tree: BTreeMap<K, V>, is_forward: bool) -> Self {
554        if let Some(root) = tree.root.take() {
555            let mut cache = tree.get_cache().take();
556            let leaf = if is_forward {
557                root.find_first_leaf(Some(&mut cache))
558            } else {
559                root.find_last_leaf(Some(&mut cache))
560            };
561            Self {
562                cache,
563                idx: if is_forward { 0 } else { leaf.key_count() },
564                leaf: Some(leaf),
565                remaining: tree.len,
566                is_forward,
567            }
568        } else {
569            Self { cache: PathCache::new(), leaf: None, idx: 0, remaining: 0, is_forward }
570        }
571    }
572
573    #[inline(always)]
574    fn advance_forward(&mut self) -> LeafNode<K, V> {
575        let _leaf = self.leaf.take().unwrap();
576        _leaf.dealloc::<false>();
577        let (parent, idx) = self.cache.move_right_and_pop_l1(InterNode::dealloc::<true>).unwrap();
578        self.cache.push(parent.clone(), idx);
579        let new_leaf = parent.get_child(idx).into_leaf();
580        self.idx = 0;
581        new_leaf
582    }
583
584    #[inline(always)]
585    fn advance_backward(&mut self) -> LeafNode<K, V> {
586        let _leaf = self.leaf.take().unwrap();
587        _leaf.dealloc::<false>();
588        let (parent, idx) = self.cache.move_left_and_pop_l1(InterNode::dealloc::<true>).unwrap();
589        self.cache.push(parent.clone(), idx);
590        let new_leaf = parent.get_child(idx).into_leaf();
591        self.idx = new_leaf.key_count();
592        new_leaf
593    }
594}
595
596impl<K: Ord + Clone + Sized, V: Sized> Iterator for IntoIter<K, V> {
597    type Item = (K, V);
598
599    #[inline]
600    fn next(&mut self) -> Option<Self::Item> {
601        if self.remaining == 0 {
602            return None;
603        }
604        self.remaining -= 1;
605
606        // Check if we need to advance to next/prev leaf
607        if self.is_forward {
608            if let Some(ref leaf) = self.leaf {
609                if self.idx >= leaf.key_count() {
610                    let new_leaf = self.advance_forward();
611                    self.leaf = Some(new_leaf);
612                }
613            } else {
614                return None;
615            }
616            let idx = self.idx;
617            self.idx += 1;
618            if let Some(ref leaf) = self.leaf {
619                unsafe {
620                    let key = (*leaf.key_ptr(idx)).assume_init_read();
621                    let value = (*leaf.value_ptr(idx)).assume_init_read();
622                    return Some((key, value));
623                }
624            }
625        } else {
626            if self.idx == 0 {
627                if let Some(ref leaf) = self.leaf {
628                    leaf.get_left_node()?;
629                }
630                let new_leaf = self.advance_backward();
631                self.leaf = Some(new_leaf);
632            }
633            if let Some(ref leaf) = self.leaf {
634                self.idx -= 1;
635                let idx = self.idx;
636                unsafe {
637                    let key = (*leaf.key_ptr(idx)).assume_init_read();
638                    let value = (*leaf.value_ptr(idx)).assume_init_read();
639                    return Some((key, value));
640                }
641            }
642        }
643        None
644    }
645
646    #[inline]
647    fn size_hint(&self) -> (usize, Option<usize>) {
648        (self.remaining, Some(self.remaining))
649    }
650}
651
652impl<K: Ord + Clone + Sized, V: Sized> ExactSizeIterator for IntoIter<K, V> {
653    #[inline]
654    fn len(&self) -> usize {
655        self.remaining
656    }
657}
658
659impl<K: Ord + Clone + Sized, V: Sized> Drop for IntoIter<K, V> {
660    fn drop(&mut self) {
661        // NOTE: if the original tree has root, then self.leaf always exists after iteration done
662        if let Some(leaf) = self.leaf.take() {
663            if needs_drop::<K>() {
664                if self.is_forward {
665                    for idx in self.idx..leaf.key_count() {
666                        unsafe { (*leaf.key_ptr(idx)).assume_init_drop() };
667                    }
668                } else {
669                    for idx in 0..self.idx {
670                        unsafe { (*leaf.key_ptr(idx)).assume_init_drop() };
671                    }
672                }
673            }
674            if needs_drop::<V>() {
675                if self.is_forward {
676                    for idx in self.idx..leaf.key_count() {
677                        unsafe { (*leaf.value_ptr(idx)).assume_init_drop() };
678                    }
679                } else {
680                    for idx in 0..self.idx {
681                        unsafe { (*leaf.value_ptr(idx)).assume_init_drop() };
682                    }
683                }
684            }
685            leaf.dealloc::<false>();
686            // We should free the rest internal nodes even after leaf iteration done
687            if self.is_forward {
688                while let Some((parent, idx)) =
689                    self.cache.move_right_and_pop_l1(InterNode::dealloc::<true>)
690                {
691                    self.cache.push(parent.clone(), idx);
692                    if let Node::Leaf(leaf) = parent.get_child(idx) {
693                        leaf.dealloc::<true>();
694                    } else {
695                        unreachable!();
696                    }
697                }
698            } else {
699                while let Some((parent, idx)) =
700                    self.cache.move_left_and_pop_l1(InterNode::dealloc::<true>)
701                {
702                    self.cache.push(parent.clone(), idx);
703                    if let Node::Leaf(leaf) = parent.get_child(idx) {
704                        leaf.dealloc::<true>();
705                    } else {
706                        unreachable!();
707                    }
708                }
709            }
710        }
711    }
712}
embed_collections/btree/iter.rs

embed_collections/btree/
iter.rs
