miniconf/
iter.rs

1use crate::{DescendError, IntoKeys, KeyError, Keys, Schema, Seeded, Short, Track, Transcode};
2
3/// Counting wrapper for iterators with known exact size
4#[derive(Clone, Debug, PartialEq, Eq)]
5pub struct ExactSize<T> {
6    iter: T,
7    count: usize,
8}
9
10impl<T: Iterator> Iterator for ExactSize<T> {
11    type Item = T::Item;
12
13    fn next(&mut self) -> Option<Self::Item> {
14        if let Some(v) = self.iter.next() {
15            self.count -= 1; // checks for overflow in debug
16            Some(v)
17        } else {
18            debug_assert!(self.count == 0);
19            None
20        }
21    }
22
23    fn size_hint(&self) -> (usize, Option<usize>) {
24        (self.count, Some(self.count))
25    }
26}
27
28// Even though general TreeSchema iterations may well be longer than usize::MAX
29// we are sure that the aren't in this case since self.count <= usize::MAX
30impl<T: Iterator> ExactSizeIterator for ExactSize<T> {}
31
32// `Iterator` is sufficient to fuse
33impl<T: Iterator> core::iter::FusedIterator for ExactSize<T> {}
34
35// https://github.com/rust-lang/rust/issues/37572
36// unsafe impl<T: Iterator> core::iter::TrustedLen for ExactSize<T> {}
37
38/// Node iterator
39///
40/// A managed indices state for iteration of nodes `N` in a `TreeSchema`.
41///
42/// `D` is the depth limit. Internal nodes will be returned on iteration where
43/// the depth limit is exceeded.
44///
45/// The `Err(usize)` variant of the `Iterator::Item` indicates that `N` does
46/// not have sufficient capacity and failed to encode the key at the given depth.
47#[derive(Clone, Debug, PartialEq)]
48pub struct NodeIter<N: Seeded, const D: usize> {
49    // We can't use Packed as state since we need to be able to modify the
50    // indices directly. Packed erases knowledge of the bit widths of the individual
51    // indices.
52    schema: &'static Schema,
53    state: [usize; D],
54    root: usize,
55    depth: usize,
56    seed: N::Seed,
57}
58
59impl<N: Seeded, const D: usize> NodeIter<N, D> {
60    /// Create a new iterator.
61    ///
62    /// # Panic
63    /// If the root depth exceeds the state length.
64    pub const fn new(
65        schema: &'static Schema,
66        state: [usize; D],
67        root: usize,
68        seed: N::Seed,
69    ) -> 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            seed,
78        }
79    }
80
81    /// Limit and start iteration to at and below the provided root key.
82    ///
83    /// This requires moving `self` to ensure `FusedIterator`.
84    pub fn with_root(
85        schema: &'static Schema,
86        root: impl IntoKeys,
87        seed: N::Seed,
88    ) -> Result<Self, DescendError<()>> {
89        let mut state = [0; D];
90        let mut root = root.into_keys().track();
91        let mut tr = Short::new(state.as_mut());
92        tr.transcode(schema, &mut root)?;
93        Ok(Self::new(schema, state, root.depth(), seed))
94    }
95
96    /// Wrap the iterator in an exact size counting iterator that is
97    /// `FusedIterator` and `ExactSizeIterator`.
98    ///
99    /// Note(panic): Panics, if the iterator had `next()` called or
100    /// if the iteration depth has been limited or if the iteration root
101    /// is not the tree root.
102    ///
103    pub const fn exact_size(self) -> ExactSize<Self> {
104        let shape = self.schema.shape();
105        if D < shape.max_depth {
106            panic!("insufficient depth for exact size iteration");
107        }
108        let mut i = 0;
109        while i < D {
110            if self.state[i] != 0 {
111                panic!("exact size requires a fresh root iterator");
112            }
113            i += 1;
114        }
115        if self.root != 0 || self.depth != D + 1 {
116            panic!("exact size requires a fresh root iterator");
117        }
118        ExactSize {
119            iter: self,
120            count: shape.count.get(),
121        }
122    }
123
124    /// Return the underlying schema
125    pub const fn schema(&self) -> &'static Schema {
126        self.schema
127    }
128
129    /// Return the current state
130    pub fn state(&self) -> Option<&[usize]> {
131        self.state.get(..self.depth)
132    }
133
134    /// Return the root depth
135    pub const fn root(&self) -> usize {
136        self.root
137    }
138}
139
140impl<N: Transcode + Seeded, const D: usize> Iterator for NodeIter<N, D> {
141    type Item = Result<N, N::Error>;
142
143    fn next(&mut self) -> Option<Self::Item> {
144        loop {
145            debug_assert!(self.depth >= self.root);
146            debug_assert!(self.depth <= self.state.len() + 1);
147            if self.depth == self.root {
148                return None;
149            }
150            if self.depth <= self.state.len() {
151                self.state[self.depth - 1] += 1;
152            }
153            let mut item = Track::new(N::from_seed(&self.seed));
154            let ret = item.transcode(self.schema, &self.state[..]);
155            let (item, depth) = item.into_inner();
156            match ret {
157                Err(DescendError::Key(KeyError::NotFound)) => {
158                    self.state[depth] = 0;
159                    self.depth = depth.max(self.root);
160                }
161                Err(DescendError::Key(KeyError::TooLong)) | Ok(()) => {
162                    self.depth = depth;
163                    return Some(Ok(item));
164                }
165                Err(DescendError::Key(KeyError::TooShort)) => {
166                    self.depth = depth;
167                }
168                Err(DescendError::Inner(e)) => {
169                    self.depth = depth;
170                    return Some(Err(e));
171                }
172            }
173        }
174    }
175}
176
177// Contract: Do not allow manipulation of `depth` other than through iteration.
178impl<N: Transcode + Seeded, const D: usize> core::iter::FusedIterator for NodeIter<N, D> {}
miniconf/iter.rs

miniconf/
iter.rs
