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> {}