trie_rs/iter/
prefix_iter.rs

1use crate::map::Trie;
2use crate::try_collect::{TryCollect, TryFromIterator};
3use louds_rs::LoudsNodeNum;
4use std::marker::PhantomData;
5
6#[derive(Debug, Clone)]
7/// Iterates through all the common prefixes of a given query.
8pub 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}