Skip to content

Commit 1b087c9

Browse files
ekinimoekinimo
committed
iterators revamp
1 parent d152cbb commit 1b087c9

File tree

2 files changed

+344
-56
lines changed

2 files changed

+344
-56
lines changed

Diff for: ‎src/iter.rs

+270-27
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,104 @@
1-
use crate::TrieMap;
1+
use crate::{node::test_bit, TrieMap};
22

33
/// An iterator over the key-value pairs of a `TrieMap`.
44
///
55
/// This struct is created by the [`iter`] method on [`TrieMap`].
66
///
77
/// [`iter`]: TrieMap::iter
88
pub struct Iter<'a, T> {
9-
pub(crate) pairs: Vec<(Vec<u8>, &'a T)>,
10-
pub(crate) position: usize,
9+
pub(crate) trie: &'a TrieMap<T>,
10+
11+
pub(crate) stack: Vec<IterState<'a>>,
12+
13+
pub(crate) current_path: Vec<u8>,
14+
15+
pub(crate) remaining: usize,
16+
}
17+
18+
/// State for a single node in the traversal stack
19+
pub(crate) struct IterState<'a> {
20+
pub(crate) node: &'a crate::node::TrieNode,
21+
22+
pub(crate) byte_index: u16,
23+
24+
pub(crate) value_emitted: bool,
1125
}
1226

1327
impl<'a, T> Iterator for Iter<'a, T> {
1428
type Item = (Vec<u8>, &'a T);
1529

1630
fn next(&mut self) -> Option<Self::Item> {
17-
if self.position < self.pairs.len() {
18-
let result = (
19-
self.pairs[self.position].0.clone(),
20-
self.pairs[self.position].1,
21-
);
22-
self.position += 1;
23-
Some(result)
24-
} else {
25-
None
31+
// If we've yielded all items, we're done
32+
if self.remaining == 0 {
33+
return None;
2634
}
35+
36+
while !self.stack.is_empty() {
37+
// Check the top of the stack
38+
let (node, byte_index, value_emitted) = {
39+
let state = self.stack.last_mut().unwrap();
40+
(state.node, state.byte_index, state.value_emitted)
41+
};
42+
43+
if !value_emitted && node.data_idx.is_some() {
44+
let data_idx = node.data_idx.unwrap();
45+
if let Some(value) = self.trie.data[data_idx].as_ref() {
46+
self.stack.last_mut().unwrap().value_emitted = true;
47+
self.remaining -= 1;
48+
return Some((self.current_path.clone(), value));
49+
} else {
50+
self.stack.last_mut().unwrap().value_emitted = true;
51+
}
52+
}
53+
54+
let mut found_next = false;
55+
let mut current_byte = byte_index;
56+
57+
while current_byte <= 255 {
58+
let byte = current_byte as u8;
59+
60+
self.stack.last_mut().unwrap().byte_index = current_byte + 1;
61+
62+
if test_bit(&node.is_present, byte) {
63+
let idx = crate::node::popcount(&node.is_present, byte) as usize;
64+
if idx < node.children.len() {
65+
// Get the child node
66+
let child_node = &node.children[idx];
67+
68+
self.current_path.push(byte);
69+
self.stack.push(IterState {
70+
node: child_node,
71+
byte_index: 0,
72+
value_emitted: false,
73+
});
74+
found_next = true;
75+
break;
76+
}
77+
}
78+
79+
current_byte += 1;
80+
}
81+
82+
if found_next {
83+
continue;
84+
}
85+
86+
self.stack.pop();
87+
if !self.current_path.is_empty() {
88+
self.current_path.pop();
89+
}
90+
}
91+
92+
None
2793
}
2894

2995
fn size_hint(&self) -> (usize, Option<usize>) {
30-
let remaining = self.pairs.len() - self.position;
31-
(remaining, Some(remaining))
96+
(self.remaining, Some(self.remaining))
3297
}
3398
}
3499

100+
impl<'a, T> ExactSizeIterator for Iter<'a, T> {}
101+
35102
/// An iterator over the keys of a `TrieMap`.
36103
///
37104
/// This struct is created by the [`keys`] method on [`TrieMap`].
@@ -115,32 +182,84 @@ impl<T> Drop for DrainIter<'_, T> {
115182

116183
/// An iterator over entries with keys that start with a specific prefix.
117184
pub struct PrefixIter<'a, T> {
118-
pub(crate) pairs: Vec<(Vec<u8>, &'a T)>,
119-
pub(crate) position: usize,
185+
pub(crate) trie: &'a TrieMap<T>,
186+
pub(crate) stack: Vec<IterState<'a>>,
187+
pub(crate) current_path: Vec<u8>,
188+
pub(crate) remaining: usize,
189+
pub(crate) prefix: Vec<u8>,
120190
}
121191

122192
impl<'a, T> Iterator for PrefixIter<'a, T> {
123193
type Item = (Vec<u8>, &'a T);
124194

125195
fn next(&mut self) -> Option<Self::Item> {
126-
if self.position < self.pairs.len() {
127-
let result = (
128-
self.pairs[self.position].0.clone(),
129-
self.pairs[self.position].1,
130-
);
131-
self.position += 1;
132-
Some(result)
133-
} else {
134-
None
196+
if self.remaining == 0 {
197+
return None;
198+
}
199+
200+
while !self.stack.is_empty() {
201+
let (node, byte_index, value_emitted) = {
202+
let state = self.stack.last_mut().unwrap();
203+
(state.node, state.byte_index, state.value_emitted)
204+
};
205+
206+
if !value_emitted && node.data_idx.is_some() {
207+
let data_idx = node.data_idx.unwrap();
208+
if let Some(value) = self.trie.data[data_idx].as_ref() {
209+
self.stack.last_mut().unwrap().value_emitted = true;
210+
self.remaining -= 1;
211+
return Some((self.current_path.clone(), value));
212+
} else {
213+
self.stack.last_mut().unwrap().value_emitted = true;
214+
}
215+
}
216+
217+
let mut found_next = false;
218+
let mut current_byte = byte_index;
219+
220+
while current_byte <= 255 {
221+
let byte = current_byte as u8;
222+
223+
self.stack.last_mut().unwrap().byte_index = current_byte + 1;
224+
225+
if test_bit(&node.is_present, byte) {
226+
let idx = crate::node::popcount(&node.is_present, byte) as usize;
227+
if idx < node.children.len() {
228+
let child_node = &node.children[idx];
229+
230+
self.current_path.push(byte);
231+
self.stack.push(IterState {
232+
node: child_node,
233+
byte_index: 0,
234+
value_emitted: false,
235+
});
236+
found_next = true;
237+
break;
238+
}
239+
}
240+
241+
current_byte += 1;
242+
}
243+
244+
if found_next {
245+
continue;
246+
}
247+
248+
self.stack.pop();
249+
if !self.current_path.is_empty() {
250+
self.current_path.pop();
251+
}
135252
}
253+
254+
None
136255
}
137256

138257
fn size_hint(&self) -> (usize, Option<usize>) {
139-
let remaining = self.pairs.len() - self.position;
140-
(remaining, Some(remaining))
258+
(self.remaining, Some(self.remaining))
141259
}
142260
}
143261

262+
impl<'a, T> ExactSizeIterator for PrefixIter<'a, T> {}
144263
/// Iterator for keys that start with a specific prefix.
145264
pub struct PrefixKeys<'a, T> {
146265
pub(crate) inner: PrefixIter<'a, T>,
@@ -174,3 +293,127 @@ impl<'a, T> Iterator for PrefixValues<'a, T> {
174293
self.inner.size_hint()
175294
}
176295
}
296+
297+
/// A consuming iterator over the key-value pairs of a `TrieMap`.
298+
pub struct IntoIter<T> {
299+
data: Vec<Option<T>>,
300+
stack: Vec<TraversalState>,
301+
current_path: Vec<u8>,
302+
remaining: usize,
303+
}
304+
305+
/// State for traversing the trie in IntoIter
306+
struct TraversalState {
307+
node: crate::node::TrieNode,
308+
byte_index: u16,
309+
value_emitted: bool,
310+
}
311+
312+
impl<T> Iterator for IntoIter<T> {
313+
type Item = (Vec<u8>, T);
314+
315+
fn next(&mut self) -> Option<Self::Item> {
316+
if self.remaining == 0 {
317+
return None;
318+
}
319+
320+
loop {
321+
if self.stack.is_empty() {
322+
return None;
323+
}
324+
325+
let has_value = {
326+
let state = self.stack.last_mut().unwrap();
327+
if !state.value_emitted && state.node.data_idx.is_some() {
328+
let idx = state.node.data_idx.unwrap();
329+
if idx < self.data.len() && self.data[idx].is_some() {
330+
state.value_emitted = true;
331+
true
332+
} else {
333+
state.value_emitted = true;
334+
false
335+
}
336+
} else {
337+
false
338+
}
339+
};
340+
341+
if has_value {
342+
let idx = self.stack.last().unwrap().node.data_idx.unwrap();
343+
let value = self.data[idx].take().unwrap();
344+
self.remaining -= 1;
345+
return Some((self.current_path.clone(), value));
346+
}
347+
348+
let (found_next, byte) = {
349+
let state = self.stack.last_mut().unwrap();
350+
let mut found = false;
351+
let mut next_byte = 0;
352+
353+
while state.byte_index <= 255 {
354+
let b = state.byte_index as u8;
355+
state.byte_index += 1;
356+
357+
if crate::node::test_bit(&state.node.is_present, b) {
358+
let idx = crate::node::popcount(&state.node.is_present, b) as usize;
359+
if idx < state.node.children.len() {
360+
found = true;
361+
next_byte = b;
362+
break;
363+
}
364+
}
365+
}
366+
367+
(found, next_byte)
368+
};
369+
370+
if found_next {
371+
let child = {
372+
let state = self.stack.last().unwrap();
373+
let idx = crate::node::popcount(&state.node.is_present, byte) as usize;
374+
state.node.children[idx].clone()
375+
};
376+
377+
self.current_path.push(byte);
378+
self.stack.push(TraversalState {
379+
node: child,
380+
byte_index: 0,
381+
value_emitted: false,
382+
});
383+
} else {
384+
self.stack.pop();
385+
if !self.current_path.is_empty() {
386+
self.current_path.pop();
387+
}
388+
}
389+
}
390+
}
391+
392+
fn size_hint(&self) -> (usize, Option<usize>) {
393+
(self.remaining, Some(self.remaining))
394+
}
395+
}
396+
397+
impl<T> ExactSizeIterator for IntoIter<T> {}
398+
399+
impl<T> IntoIterator for TrieMap<T> {
400+
type Item = (Vec<u8>, T);
401+
type IntoIter = IntoIter<T>;
402+
403+
fn into_iter(self) -> Self::IntoIter {
404+
let TrieMap {
405+
root, data, size, ..
406+
} = self;
407+
408+
IntoIter {
409+
data,
410+
stack: vec![TraversalState {
411+
node: root,
412+
byte_index: 0,
413+
value_emitted: false,
414+
}],
415+
current_path: Vec::new(),
416+
remaining: size,
417+
}
418+
}
419+
}

0 commit comments

Comments
 (0)