miniconf/
iter.rs

1use core::marker::PhantomData;
2
3use crate::{IntoKeys, KeyLookup, Keys, Metadata, Node, Transcode, Traversal, TreeKey};
4
5/// Counting wrapper for iterators with known exact size
6#[derive(Clone, Debug, PartialEq, Eq)]
7pub struct ExactSize<T> {
8    iter: T,
9    count: usize,
10}
11
12impl<T: Iterator> Iterator for ExactSize<T> {
13    type Item = T::Item;
14
15    #[inline]
16    fn next(&mut self) -> Option<Self::Item> {
17        if let Some(v) = self.iter.next() {
18            self.count -= 1; // checks for overflow in debug
19            Some(v)
20        } else {
21            debug_assert!(self.count == 0);
22            None
23        }
24    }
25
26    #[inline]
27    fn size_hint(&self) -> (usize, Option<usize>) {
28        (self.count, Some(self.count))
29    }
30}
31
32impl<T> ExactSize<T> {
33    /// Return a reference to the inner iterator
34    #[inline]
35    pub fn inner(&self) -> &T {
36        &self.iter
37    }
38}
39
40// Even though general TreeKey iterations may well be longer than usize::MAX
41// we are sure that the aren't in this case since self.count <= usize::MAX
42impl<T: Iterator> ExactSizeIterator for ExactSize<T> {}
43
44// `Iterator` is sufficient to fuse
45impl<T: Iterator> core::iter::FusedIterator for ExactSize<T> {}
46
47// https://github.com/rust-lang/rust/issues/37572
48// unsafe impl<T: Iterator> core::iter::TrustedLen for ExactSize<T> {}
49
50/// A Keys wrapper that can always finalize()
51pub(crate) struct Consume<T>(pub(crate) T);
52impl<T: Keys> Keys for Consume<T> {
53    #[inline]
54    fn next(&mut self, lookup: &KeyLookup) -> Result<usize, Traversal> {
55        self.0.next(lookup)
56    }
57
58    #[inline]
59    fn finalize(&mut self) -> Result<(), Traversal> {
60        Ok(())
61    }
62}
63
64impl<T: Keys> IntoKeys for Consume<T> {
65    type IntoKeys = Self;
66
67    #[inline]
68    fn into_keys(self) -> Self::IntoKeys {
69        self
70    }
71}
72
73/// Node iterator
74///
75/// A managed indices state for iteration of nodes `N` in a `TreeKey`.
76///
77/// `D` is the depth limit. Internal nodes will be returned on iteration where
78/// the depth limit is exceeded.
79///
80/// The `Err(usize)` variant of the `Iterator::Item` indicates that `N` does
81/// not have sufficient capacity and failed to encode the key at the given depth.
82#[derive(Clone, Debug, PartialEq, Eq)]
83pub struct NodeIter<M: ?Sized, N, const D: usize> {
84    // We can't use Packed as state since we need to be able to modify the
85    // indices directly. Packed erases knowledge of the bit widths of the individual
86    // indices.
87    state: [usize; D],
88    root: usize,
89    depth: usize,
90    _n: PhantomData<N>,
91    _m: PhantomData<M>,
92}
93
94impl<M: ?Sized, N, const D: usize> Default for NodeIter<M, N, D> {
95    fn default() -> Self {
96        Self {
97            state: [0; D],
98            root: 0,
99            // Marker to prevent initial index increment in `next()`
100            depth: D + 1,
101            _n: PhantomData,
102            _m: PhantomData,
103        }
104    }
105}
106
107impl<M: TreeKey + ?Sized, N, const D: usize> NodeIter<M, N, D> {
108    /// Limit and start iteration to at and below the provided root key.
109    ///
110    /// This requires moving `self` to ensure `FusedIterator`.
111    pub fn root<K: IntoKeys>(mut self, root: K) -> Result<Self, Traversal> {
112        let node = self.state.transcode::<M, _>(root)?;
113        self.root = node.depth();
114        self.depth = D + 1;
115        Ok(self)
116    }
117
118    /// Wrap the iterator in an exact size counting iterator that is
119    /// `FusedIterator` and `ExactSizeIterator`.
120    ///
121    /// Note(panic): Panics, if the iterator had `next()` called or
122    /// if the iteration depth has been limited or if the iteration root
123    /// is not the tree root.
124    pub fn exact_size(self) -> ExactSize<Self> {
125        assert_eq!(self.depth, D + 1, "NodeIter partially consumed");
126        assert_eq!(self.root, 0, "NodeIter on sub-tree");
127        debug_assert_eq!(&self.state, &[0; D]); // ensured by depth = D + 1 marker and contract
128        let meta: Metadata = M::traverse_all();
129        assert!(
130            D >= meta.max_depth,
131            "depth D = {D} must be at least {}",
132            meta.max_depth
133        );
134        ExactSize {
135            iter: self,
136            count: meta.count.get(),
137        }
138    }
139
140    /// Return the current iteration depth
141    pub fn current_depth(&self) -> usize {
142        self.depth
143    }
144
145    /// Return the root depth
146    pub fn root_depth(&self) -> usize {
147        self.root
148    }
149}
150
151impl<M, N, const D: usize> Iterator for NodeIter<M, N, D>
152where
153    M: TreeKey + ?Sized,
154    N: Transcode + Default,
155{
156    type Item = Result<(N, Node), usize>;
157
158    fn next(&mut self) -> Option<Self::Item> {
159        loop {
160            debug_assert!(self.depth >= self.root);
161            debug_assert!(self.depth <= D + 1);
162            if self.depth == self.root {
163                // Iteration done
164                return None;
165            }
166            if self.depth <= D {
167                // Not initial state: increment
168                self.state[self.depth - 1] += 1;
169            }
170            return match M::transcode(Consume(self.state.iter().into_keys())) {
171                Err(Traversal::NotFound(depth)) => {
172                    // Reset index at current depth, then retry with incremented index at depth - 1 or terminate
173                    // Key lookup was performed and failed: depth is always >= 1
174                    self.state[depth - 1] = 0;
175                    self.depth = (depth - 1).max(self.root);
176                    continue;
177                }
178                Ok((path, node)) => {
179                    // Leaf or internal node found, save depth for increment at next iteration
180                    self.depth = node.depth();
181                    Some(Ok((path, node)))
182                }
183                Err(Traversal::TooShort(depth)) => {
184                    // Target type can not hold keys
185                    Some(Err(depth))
186                }
187                // TooLong: impossible due to Consume
188                // Absent, Finalization, Invalid, Access: not returned by transcode (traverse_by_key())
189                _ => unreachable!(),
190            };
191        }
192    }
193}
194
195// Contract: Do not allow manipulation of `depth` other than through iteration.
196impl<M: TreeKey + ?Sized, N: Transcode + Default, const D: usize> core::iter::FusedIterator
197    for NodeIter<M, N, D>
198{
199}
miniconf/iter.rs

miniconf/
iter.rs
