kermit_ds/ds/column_trie/
column_trie_iter.rs

1use {
2    super::implementation::ColumnTrie,
3    crate::relation::Relation,
4    kermit_derive::IntoTrieIter,
5    kermit_iters::{LinearIterator, TrieIterable, TrieIterator, TrieIteratorWrapper},
6};
7
8#[derive(IntoTrieIter)]
9pub struct ColumnTrieIter<'a> {
10    /// Index of current layer.
11    layer_number: usize,
12    /// Index of interval start at current layer
13    interval_i: usize,
14    /// Relative index based on interval
15    rel_data_i: usize,
16    /// Relative data at current layer
17    rel_data: Option<&'a [usize]>,
18    /// The trie being iterated.
19    trie: &'a ColumnTrie,
20}
21
22impl<'a> ColumnTrieIter<'a> {
23    pub fn new(trie: &'a ColumnTrie) -> Self {
24        ColumnTrieIter {
25            interval_i: 0,
26            rel_data: None,
27            rel_data_i: 0,
28            layer_number: 0,
29            trie,
30        }
31    }
32}
33
34impl LinearIterator for ColumnTrieIter<'_> {
35    fn key(&self) -> Option<usize> {
36        if let Some(data) = self.rel_data {
37            data.get(self.rel_data_i).copied()
38        } else {
39            None
40        }
41    }
42
43    fn next(&mut self) -> Option<usize> {
44        if let Some(data) = self.rel_data {
45            if self.rel_data_i > data.len() - 1 {
46                return None;
47            }
48            self.rel_data_i += 1;
49            data.get(self.rel_data_i).copied()
50        } else {
51            None
52        }
53    }
54
55    fn seek(&mut self, seek_key: usize) -> bool {
56        while let Some(key) = self.next() {
57            if key >= seek_key {
58                break;
59            }
60        }
61        self.at_end()
62    }
63
64    fn at_end(&self) -> bool {
65        if let Some(data) = self.rel_data {
66            self.rel_data_i >= data.len()
67        } else {
68            true
69        }
70    }
71}
72
73impl TrieIterator for ColumnTrieIter<'_> {
74    fn open(&mut self) -> bool {
75        if self.layer_number == self.trie.header().arity() {
76            // If at leaf, return false
77            false
78        } else if self.layer_number == 0 {
79            let next_layer = self.trie.layer(0);
80            if next_layer.data.is_empty() {
81                // If the first layer is empty, we are in an empty trie and return false
82                return false;
83            }
84            // If at root, initialize the first layer
85            self.layer_number = 1;
86            self.rel_data_i = 0;
87            self.interval_i = 0;
88            self.rel_data = Some(&self.trie.layer(0).data);
89            true
90        } else {
91            // If not at root or leaf, compute next interval index
92            let curr_layer = self.trie.layer(self.layer_number - 1);
93            let prev_start_index = curr_layer.interval[self.interval_i];
94            self.interval_i = prev_start_index + self.rel_data_i;
95            // increment layer number
96            self.layer_number += 1;
97            // get new layer
98            let next_layer = self.trie.layer(self.layer_number - 1);
99            // get next start and end indices
100            let next_start_index = next_layer.interval[self.interval_i];
101            let next_end_index = if self.interval_i + 1 < next_layer.interval.len() {
102                next_layer.interval[self.interval_i + 1]
103            } else {
104                next_layer.data.len()
105            };
106            // set new relative data
107            self.rel_data = Some(&next_layer.data[next_start_index..next_end_index]);
108            self.rel_data_i = 0;
109            true
110        }
111    }
112
113    fn up(&mut self) -> bool {
114        if self.layer_number == 0 {
115            // If already at root, cannot go up
116            false
117        } else if self.layer_number == 1 {
118            // If moving to root, reset all indices
119            self.layer_number = 0;
120            self.interval_i = 0;
121            self.rel_data_i = 0;
122            self.rel_data = None;
123            true
124        } else {
125            // If moving up, decrement layer index
126            self.layer_number -= 1;
127            let layer = self.trie.layer(self.layer_number - 1);
128            // Our global data index is interval_i, so we must find the start index
129
130            // Data index of parent is equivalent to current interval index
131            let data_index = self.interval_i;
132            // We need to find the interval index of the data in the previous layer
133            // The start indexes are ordered, so we need to find the point at which the
134            // data index is less than the current start index. Then the previous index
135            // is our new interval index
136            for (i, start_index) in layer.interval.iter().enumerate() {
137                if data_index < *start_index {
138                    break;
139                } else {
140                    self.interval_i = i;
141                }
142            }
143
144            // Our new start index is at the new interval index
145            let start_index = layer.interval[self.interval_i];
146            // The end index is either the next start index, or the length of the data
147            let end_index = if self.interval_i + 1 < layer.interval.len() {
148                layer.interval[self.interval_i + 1]
149            } else {
150                layer.data.len()
151            };
152
153            // Set new relative data
154            self.rel_data = Some(&layer.data[start_index..end_index]);
155            self.rel_data_i = data_index - start_index;
156
157            true
158        }
159    }
160}
161
162/// Implementation of the `TrieIterable` trait for `TreeTrie`.
163impl TrieIterable for ColumnTrie {
164    fn trie_iter(&self) -> impl TrieIterator + IntoIterator<Item = Vec<usize>> {
165        ColumnTrieIter::new(self)
166    }
167}
kermit_ds/ds/column_trie/column_trie_iter.rs

kermit_ds/ds/column_trie/
column_trie_iter.rs
