trie_rs/iter/
prefix_iter.rs1use crate::map::Trie;
2use crate::try_collect::{TryCollect, TryFromIterator};
3use louds_rs::LoudsNodeNum;
4use std::marker::PhantomData;
5
6#[derive(Debug, Clone)]
7pub struct PrefixIter<'a, Label, Value, C, M> {
9 trie: &'a Trie<Label, Value>,
10 query: Vec<Label>,
11 index: usize,
12 node: LoudsNodeNum,
13 buffer: Vec<&'a Label>,
14 consume: Option<&'a Value>,
15 col: PhantomData<(C, M)>,
16}
17
18impl<'a, Label: Ord + Clone, Value, C, M> PrefixIter<'a, Label, Value, C, M> {
19 #[inline]
20 pub(crate) fn new(trie: &'a Trie<Label, Value>, query: impl AsRef<[Label]>) -> Self {
21 Self {
22 trie,
23 query: query.as_ref().to_vec(),
24 index: 0,
25 node: LoudsNodeNum(1),
26 buffer: Vec::new(),
27 consume: None,
28 col: PhantomData,
29 }
30 }
31}
32
33impl<'a, Label: Ord + Clone, Value, C, M> Iterator for PrefixIter<'a, Label, Value, C, M>
34where
35 C: TryFromIterator<Label, M>,
36{
37 type Item = (C, &'a Value);
38 fn next(&mut self) -> Option<Self::Item> {
39 while self.consume.is_none() {
40 if let Some(chr) = self.query.get(self.index) {
41 let children_node_nums: Vec<_> = self.trie.children_node_nums(self.node).collect();
42 let res = self
43 .trie
44 .bin_search_by_children_labels(chr, &children_node_nums[..]);
45 match res {
46 Ok(j) => {
47 let child_node_num = children_node_nums[j];
48 self.buffer.push(self.trie.label(child_node_num));
49 self.consume = self.trie.value(child_node_num);
50 self.node = child_node_num;
51 }
52 Err(_) => break,
53 }
54 } else {
55 return None;
56 }
57 self.index += 1;
58 }
59 if let Some(v) = self.consume.take() {
60 let col = self.buffer.clone();
61 Some((
62 col.into_iter()
63 .cloned()
64 .try_collect()
65 .expect("Could not collect"),
66 v,
67 ))
68 } else {
69 None
70 }
71 }
72}