miniconf/
iter.rs

1use core::marker::PhantomData;
2
3use crate::{DescendError, IntoKeys, KeyError, Keys, Schema, Short, Track, Transcode};
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
32// Even though general TreeSchema iterations may well be longer than usize::MAX
33// we are sure that the aren't in this case since self.count <= usize::MAX
34impl<T: Iterator> ExactSizeIterator for ExactSize<T> {}
35
36// `Iterator` is sufficient to fuse
37impl<T: Iterator> core::iter::FusedIterator for ExactSize<T> {}
38
39// https://github.com/rust-lang/rust/issues/37572
40// unsafe impl<T: Iterator> core::iter::TrustedLen for ExactSize<T> {}
41
42/// Node iterator
43///
44/// A managed indices state for iteration of nodes `N` in a `TreeSchema`.
45///
46/// `D` is the depth limit. Internal nodes will be returned on iteration where
47/// the depth limit is exceeded.
48///
49/// The `Err(usize)` variant of the `Iterator::Item` indicates that `N` does
50/// not have sufficient capacity and failed to encode the key at the given depth.
51#[derive(Clone, Debug, PartialEq)]
52pub struct NodeIter<N, const D: usize> {
53    // We can't use Packed as state since we need to be able to modify the
54    // indices directly. Packed erases knowledge of the bit widths of the individual
55    // indices.
56    schema: &'static Schema,
57    state: [usize; D],
58    root: usize,
59    depth: usize,
60    _n: PhantomData<N>,
61}
62
63impl<N, const D: usize> NodeIter<N, D> {
64    /// Create a new iterator.
65    ///
66    /// # Panic
67    /// If the root depth exceeds the state length.
68    #[inline]
69    pub const fn with(schema: &'static Schema, state: [usize; D], root: usize) -> Self {
70        assert!(root <= D);
71        Self {
72            schema,
73            state,
74            root,
75            // Marker to prevent initial index increment in `next()`
76            depth: D + 1,
77            _n: PhantomData,
78        }
79    }
80
81    /// Create a new iterator with default root and initial state.
82    #[inline]
83    pub const fn new(schema: &'static Schema) -> Self {
84        Self::with(schema, [0; D], 0)
85    }
86
87    /// Limit and start iteration to at and below the provided root key.
88    ///
89    /// This requires moving `self` to ensure `FusedIterator`.
90    pub fn with_root(
91        schema: &'static Schema,
92        root: impl IntoKeys,
93    ) -> Result<Self, DescendError<()>> {
94        let mut state = [0; D];
95        let mut root = root.into_keys().track();
96        let mut tr = Short::new(state.as_mut());
97        tr.transcode(schema, &mut root)?;
98        Ok(Self::with(schema, state, root.depth()))
99    }
100
101    /// Wrap the iterator in an exact size counting iterator that is
102    /// `FusedIterator` and `ExactSizeIterator`.
103    ///
104    /// Note(panic): Panics, if the iterator had `next()` called or
105    /// if the iteration depth has been limited or if the iteration root
106    /// is not the tree root.
107    ///
108    #[inline]
109    pub const fn exact_size(schema: &'static Schema) -> ExactSize<Self> {
110        let shape = schema.shape();
111        if D < shape.max_depth {
112            panic!("insufficient depth for exact size iteration");
113        }
114        ExactSize {
115            iter: Self::new(schema),
116            count: shape.count.get(),
117        }
118    }
119
120    /// Return the underlying schema
121    #[inline]
122    pub const fn schema(&self) -> &'static Schema {
123        self.schema
124    }
125
126    /// Return the current state
127    #[inline]
128    pub fn state(&self) -> Option<&[usize]> {
129        self.state.get(..self.depth)
130    }
131
132    /// Return the root depth
133    #[inline]
134    pub const fn root(&self) -> usize {
135        self.root
136    }
137}
138
139impl<N: Transcode + Default, const D: usize> Iterator for NodeIter<N, D> {
140    type Item = Result<N, N::Error>;
141
142    fn next(&mut self) -> Option<Self::Item> {
143        loop {
144            debug_assert!(self.depth >= self.root);
145            debug_assert!(self.depth <= D + 1);
146            if self.depth == self.root {
147                // Iteration done
148                return None;
149            }
150            if self.depth <= D {
151                // Not initial state: increment
152                self.state[self.depth - 1] += 1;
153            }
154            let mut item = Track::new(N::default());
155            let ret = item.transcode(self.schema, &self.state[..]);
156            // Track<N> counts is the number of successful Keys::next()
157            let (item, depth) = item.into_inner();
158            return match ret {
159                Err(DescendError::Key(KeyError::NotFound)) => {
160                    // Reset index at NotFound depth, then retry with incremented earlier index or terminate
161                    self.state[depth] = 0;
162                    self.depth = depth.max(self.root);
163                    continue;
164                }
165                Err(DescendError::Key(KeyError::TooLong)) | Ok(()) => {
166                    // Leaf node found, save depth for increment at next iteration
167                    self.depth = depth;
168                    Some(Ok(item))
169                }
170                Err(DescendError::Key(KeyError::TooShort)) => {
171                    // Use Short<N> to suppress this branch and also get internal short nodes
172                    self.depth = depth;
173                    continue;
174                }
175                Err(DescendError::Inner(e)) => {
176                    // Target type can not hold keys
177                    Some(Err(e))
178                }
179            };
180        }
181    }
182}
183
184// Contract: Do not allow manipulation of `depth` other than through iteration.
185impl<N: Transcode + Default, const D: usize> core::iter::FusedIterator for NodeIter<N, D> {}