redb_extras/buckets/
iterator.rs

1//! Bucket range iterator implementation.
2//!
3//! Provides efficient iteration over bucket ranges for specific base keys.
4
5use crate::buckets::key::{BucketedKey, KeyBuilder};
6use crate::buckets::BucketError;
7use redb::{MultimapRange, MultimapValue, ReadOnlyMultimapTable, ReadOnlyTable};
8
9/// Iterator over a range of buckets for a specific base key.
10///
11/// BucketRangeIterator enables efficient traversal of all values for a given
12/// base key across a specified range of sequence values.
13pub struct BucketRangeIterator<V>
14where
15    V: redb::Value + 'static,
16    for<'b> V: From<V::SelfType<'b>>,
17{
18    range: redb::Range<'static, BucketedKey<u64>, V>,
19    base_key: u64,
20    start_bucket: u64,
21    end_bucket: u64,
22    done: bool,
23}
24
25impl<V> BucketRangeIterator<V>
26where
27    V: redb::Value + 'static,
28    for<'b> V: From<V::SelfType<'b>>,
29{
30    /// Create a new bucket range iterator.
31    pub fn new(
32        table: &ReadOnlyTable<BucketedKey<u64>, V>,
33        key_builder: &KeyBuilder,
34        base_key: u64,
35        start_sequence: u64,
36        end_sequence: u64,
37    ) -> Result<Self, BucketError> {
38        if start_sequence > end_sequence {
39            return Err(BucketError::InvalidRange {
40                start: start_sequence,
41                end: end_sequence,
42            });
43        }
44
45        let bucket_size = key_builder.bucket_size();
46        let start_bucket = start_sequence / bucket_size;
47        let end_bucket = end_sequence / bucket_size;
48        let start_key = BucketedKey::new(base_key, start_bucket);
49        let range = table.range(start_key..).map_err(|err| {
50            BucketError::IterationError(format!("Failed to create range iterator: {}", err))
51        })?;
52
53        Ok(Self {
54            range,
55            base_key,
56            start_bucket,
57            end_bucket,
58            done: false,
59        })
60    }
61
62    /// Get the bucket range.
63    pub fn bucket_range(&self) -> (u64, u64) {
64        (self.start_bucket, self.end_bucket)
65    }
66}
67
68impl<V> Iterator for BucketRangeIterator<V>
69where
70    V: redb::Value + 'static,
71    for<'b> V: From<V::SelfType<'b>>,
72{
73    type Item = Result<V, BucketError>;
74
75    fn next(&mut self) -> Option<Self::Item> {
76        if self.done {
77            return None;
78        }
79
80        loop {
81            match self.range.next() {
82                Some(Ok((key_guard, value_guard))) => {
83                    let key = key_guard.value();
84                    if key.bucket > self.end_bucket {
85                        self.done = true;
86                        return None;
87                    }
88                    if key.base_key == self.base_key {
89                        return Some(Ok(V::from(value_guard.value())));
90                    }
91                }
92                Some(Err(err)) => {
93                    self.done = true;
94                    return Some(Err(BucketError::IterationError(format!(
95                        "Database error during iteration: {}",
96                        err
97                    ))));
98                }
99                None => {
100                    self.done = true;
101                    return None;
102                }
103            }
104        }
105    }
106}
107
108/// Iterator over a range of buckets for a specific base key in multimap tables.
109///
110/// This iterator flattens the multimap values, yielding each value in order
111/// across the requested bucket range.
112///
113/// ```
114/// use redb::{Database, MultimapTableDefinition, ReadableDatabase};
115/// use redb_extras::buckets::{BucketMultimapIterExt, BucketedKey, KeyBuilder};
116///
117/// const TABLE: MultimapTableDefinition<'static, BucketedKey<u64>, u64> =
118///     MultimapTableDefinition::new("bucketed_values");
119///
120/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
121/// let db = Database::create("example.redb")?;
122/// let key_builder = KeyBuilder::new(100)?;
123///
124/// let write_txn = db.begin_write()?;
125/// {
126///     let mut table = write_txn.open_multimap_table(TABLE)?;
127///     table.insert(key_builder.bucketed_key(42u64, 10), 1u64)?;
128///     table.insert(key_builder.bucketed_key(42u64, 110), 2u64)?;
129/// }
130/// write_txn.commit()?;
131///
132/// let read_txn = db.begin_read()?;
133/// let table = read_txn.open_multimap_table(TABLE)?;
134/// let values: Vec<u64> = table
135///     .bucket_range(&key_builder, 42u64, 0, 199)?
136///     .collect::<Result<_, _>>()?;
137/// assert_eq!(values, vec![1u64, 2u64]);
138/// # Ok(())
139/// # }
140/// ```
141pub struct BucketRangeMultimapIterator<V>
142where
143    V: redb::Key + 'static,
144    for<'b> V: From<V::SelfType<'b>>,
145{
146    range: MultimapRange<'static, BucketedKey<u64>, V>,
147    base_key: u64,
148    start_bucket: u64,
149    end_bucket: u64,
150    done: bool,
151    current_values: Option<MultimapValue<'static, V>>,
152}
153
154impl<V> BucketRangeMultimapIterator<V>
155where
156    V: redb::Key + 'static,
157    for<'b> V: From<V::SelfType<'b>>,
158{
159    /// Create a new bucket range iterator for a multimap table.
160    pub fn new(
161        table: &ReadOnlyMultimapTable<BucketedKey<u64>, V>,
162        key_builder: &KeyBuilder,
163        base_key: u64,
164        start_sequence: u64,
165        end_sequence: u64,
166    ) -> Result<Self, BucketError> {
167        if start_sequence > end_sequence {
168            return Err(BucketError::InvalidRange {
169                start: start_sequence,
170                end: end_sequence,
171            });
172        }
173
174        let bucket_size = key_builder.bucket_size();
175        let start_bucket = start_sequence / bucket_size;
176        let end_bucket = end_sequence / bucket_size;
177        let start_key = BucketedKey::new(base_key, start_bucket);
178        let range = table.range(start_key..).map_err(|err| {
179            BucketError::IterationError(format!("Failed to create range iterator: {}", err))
180        })?;
181
182        Ok(Self {
183            range,
184            base_key,
185            start_bucket,
186            end_bucket,
187            done: false,
188            current_values: None,
189        })
190    }
191
192    /// Get the bucket range.
193    pub fn bucket_range(&self) -> (u64, u64) {
194        (self.start_bucket, self.end_bucket)
195    }
196}
197
198impl<V> Iterator for BucketRangeMultimapIterator<V>
199where
200    V: redb::Key + 'static,
201    for<'b> V: From<V::SelfType<'b>>,
202{
203    type Item = Result<V, BucketError>;
204
205    fn next(&mut self) -> Option<Self::Item> {
206        if self.done {
207            return None;
208        }
209
210        loop {
211            if let Some(values) = self.current_values.as_mut() {
212                match values.next() {
213                    Some(Ok(value_guard)) => {
214                        return Some(Ok(V::from(value_guard.value())));
215                    }
216                    Some(Err(err)) => {
217                        self.done = true;
218                        return Some(Err(BucketError::IterationError(format!(
219                            "Database error during iteration: {}",
220                            err
221                        ))));
222                    }
223                    None => {
224                        self.current_values = None;
225                    }
226                }
227            }
228
229            match self.range.next() {
230                Some(Ok((key_guard, values))) => {
231                    let key = key_guard.value();
232                    if key.bucket > self.end_bucket {
233                        self.done = true;
234                        return None;
235                    }
236                    if key.base_key == self.base_key {
237                        self.current_values = Some(values);
238                    }
239                }
240                Some(Err(err)) => {
241                    self.done = true;
242                    return Some(Err(BucketError::IterationError(format!(
243                        "Database error during iteration: {}",
244                        err
245                    ))));
246                }
247                None => {
248                    self.done = true;
249                    return None;
250                }
251            }
252        }
253    }
254}
255
256/// Extension trait for convenient bucket iteration on read-only tables.
257pub trait BucketIterExt<V>
258where
259    V: redb::Value + 'static,
260    for<'b> V: From<V::SelfType<'b>>,
261{
262    fn bucket_range(
263        &self,
264        key_builder: &KeyBuilder,
265        base_key: u64,
266        start_sequence: u64,
267        end_sequence: u64,
268    ) -> Result<BucketRangeIterator<V>, BucketError>;
269}
270
271impl<V> BucketIterExt<V> for ReadOnlyTable<BucketedKey<u64>, V>
272where
273    V: redb::Value + 'static,
274    for<'b> V: From<V::SelfType<'b>>,
275{
276    fn bucket_range(
277        &self,
278        key_builder: &KeyBuilder,
279        base_key: u64,
280        start_sequence: u64,
281        end_sequence: u64,
282    ) -> Result<BucketRangeIterator<V>, BucketError> {
283        BucketRangeIterator::new(self, key_builder, base_key, start_sequence, end_sequence)
284    }
285}
286
287/// Extension trait for convenient bucket iteration on read-only multimap tables.
288///
289/// Returns a flattened iterator over values for the base key within the
290/// requested bucket range.
291pub trait BucketMultimapIterExt<V>
292where
293    V: redb::Key + 'static,
294    for<'b> V: From<V::SelfType<'b>>,
295{
296    fn bucket_range(
297        &self,
298        key_builder: &KeyBuilder,
299        base_key: u64,
300        start_sequence: u64,
301        end_sequence: u64,
302    ) -> Result<BucketRangeMultimapIterator<V>, BucketError>;
303}
304
305impl<V> BucketMultimapIterExt<V> for ReadOnlyMultimapTable<BucketedKey<u64>, V>
306where
307    V: redb::Key + 'static,
308    for<'b> V: From<V::SelfType<'b>>,
309{
310    fn bucket_range(
311        &self,
312        key_builder: &KeyBuilder,
313        base_key: u64,
314        start_sequence: u64,
315        end_sequence: u64,
316    ) -> Result<BucketRangeMultimapIterator<V>, BucketError> {
317        BucketRangeMultimapIterator::new(self, key_builder, base_key, start_sequence, end_sequence)
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use redb::{Database, MultimapTableDefinition, ReadableDatabase, TableDefinition};
325    use tempfile::NamedTempFile;
326
327    const TEST_TABLE: TableDefinition<'static, BucketedKey<u64>, String> =
328        TableDefinition::new("test_table");
329    const TEST_MULTIMAP: MultimapTableDefinition<'static, BucketedKey<u64>, u64> =
330        MultimapTableDefinition::new("test_multimap");
331
332    #[test]
333    fn test_basic_functionality() -> Result<(), Box<dyn std::error::Error>> {
334        let temp_file = NamedTempFile::new()?;
335        let db = Database::create(temp_file.path())?;
336        let key_builder = KeyBuilder::new(100)?;
337
338        // Insert test data
339        {
340            let write_txn = db.begin_write()?;
341            {
342                let mut table = write_txn.open_table(TEST_TABLE)?;
343
344                // Insert values for user123 in different buckets
345                table.insert(key_builder.bucketed_key(123u64, 50), "value_50".to_string())?;
346                table.insert(
347                    key_builder.bucketed_key(123u64, 150),
348                    "value_150".to_string(),
349                )?;
350                table.insert(
351                    key_builder.bucketed_key(123u64, 250),
352                    "value_250".to_string(),
353                )?;
354
355                // Insert values for user456 in same buckets (should not appear in iteration)
356                table.insert(key_builder.bucketed_key(456u64, 50), "other_50".to_string())?;
357                table.insert(
358                    key_builder.bucketed_key(456u64, 150),
359                    "other_150".to_string(),
360                )?;
361            }
362            write_txn.commit()?;
363        }
364
365        // Test that we can create bucket range iterators
366        {
367            let read_txn = db.begin_read()?;
368            let table = read_txn.open_table(TEST_TABLE)?;
369            let iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 0, 199)?;
370            assert_eq!(iter.bucket_range(), (0, 1));
371
372            // Test invalid range
373            let invalid_iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 200, 100);
374            assert!(invalid_iter.is_err());
375        }
376
377        // Test value iteration and base key filtering
378        {
379            let read_txn = db.begin_read()?;
380            let table = read_txn.open_table(TEST_TABLE)?;
381            let iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 0, 299)?;
382            let values: Vec<String> = iter.collect::<Result<_, _>>()?;
383            assert_eq!(
384                values,
385                vec![
386                    "value_50".to_string(),
387                    "value_150".to_string(),
388                    "value_250".to_string()
389                ]
390            );
391
392            let iter = table.bucket_range(&key_builder, 456u64, 0, 299)?;
393            let values: Vec<String> = iter.collect::<Result<_, _>>()?;
394            assert_eq!(
395                values,
396                vec!["other_50".to_string(), "other_150".to_string()]
397            );
398        }
399
400        Ok(())
401    }
402
403    #[test]
404    fn test_multimap_functionality() -> Result<(), Box<dyn std::error::Error>> {
405        let temp_file = NamedTempFile::new()?;
406        let db = Database::create(temp_file.path())?;
407        let key_builder = KeyBuilder::new(100)?;
408
409        {
410            let write_txn = db.begin_write()?;
411            {
412                let mut table = write_txn.open_multimap_table(TEST_MULTIMAP)?;
413
414                table.insert(key_builder.bucketed_key(123u64, 50), 10u64)?;
415                table.insert(key_builder.bucketed_key(123u64, 50), 20u64)?;
416                table.insert(key_builder.bucketed_key(123u64, 150), 30u64)?;
417                table.insert(key_builder.bucketed_key(123u64, 150), 40u64)?;
418
419                table.insert(key_builder.bucketed_key(456u64, 50), 99u64)?;
420                table.insert(key_builder.bucketed_key(456u64, 50), 100u64)?;
421            }
422            write_txn.commit()?;
423        }
424
425        {
426            let read_txn = db.begin_read()?;
427            let table = read_txn.open_multimap_table(TEST_MULTIMAP)?;
428            let iter = BucketRangeMultimapIterator::new(&table, &key_builder, 123u64, 0, 199)?;
429            assert_eq!(iter.bucket_range(), (0, 1));
430
431            let values: Vec<u64> = iter.collect::<Result<_, _>>()?;
432            assert_eq!(values, vec![10u64, 20u64, 30u64, 40u64]);
433
434            let iter = table.bucket_range(&key_builder, 456u64, 0, 99)?;
435            let values: Vec<u64> = iter.collect::<Result<_, _>>()?;
436            assert_eq!(values, vec![99u64, 100u64]);
437        }
438
439        Ok(())
440    }
441}
redb_extras/buckets/iterator.rs

redb_extras/buckets/
iterator.rs
