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.
13///
14/// Implements `DoubleEndedIterator` for reverse iteration.
15pub struct BucketRangeIterator<V>
16where
17    V: redb::Value + 'static,
18    for<'b> V: From<V::SelfType<'b>>,
19{
20    range: redb::Range<'static, BucketedKey<u64>, V>,
21    base_key: u64,
22    start_bucket: u64,
23    end_bucket: u64,
24    finished: bool,
25}
26
27impl<V> BucketRangeIterator<V>
28where
29    V: redb::Value + 'static,
30    for<'b> V: From<V::SelfType<'b>>,
31{
32    /// Create a new bucket range iterator.
33    pub fn new(
34        table: &ReadOnlyTable<BucketedKey<u64>, V>,
35        key_builder: &KeyBuilder,
36        base_key: u64,
37        start_sequence: u64,
38        end_sequence: u64,
39    ) -> Result<Self, BucketError> {
40        if start_sequence > end_sequence {
41            return Err(BucketError::InvalidRange {
42                start: start_sequence,
43                end: end_sequence,
44            });
45        }
46
47        let bucket_size = key_builder.bucket_size();
48        let start_bucket = start_sequence / bucket_size;
49        let end_bucket = end_sequence / bucket_size;
50        let start_key = BucketedKey::new(base_key, start_bucket);
51        let end_key = BucketedKey::new(base_key, end_bucket);
52        let range = table.range(start_key..=end_key).map_err(|err| {
53            BucketError::IterationError(format!("Failed to create range iterator: {}", err))
54        })?;
55
56        Ok(Self {
57            range,
58            base_key,
59            start_bucket,
60            end_bucket,
61            finished: false,
62        })
63    }
64
65    /// Get the bucket range.
66    pub fn bucket_range(&self) -> (u64, u64) {
67        (self.start_bucket, self.end_bucket)
68    }
69}
70
71impl<V> Iterator for BucketRangeIterator<V>
72where
73    V: redb::Value + 'static,
74    for<'b> V: From<V::SelfType<'b>>,
75{
76    type Item = Result<V, BucketError>;
77
78    fn next(&mut self) -> Option<Self::Item> {
79        if self.finished {
80            return None;
81        }
82
83        loop {
84            match self.range.next() {
85                Some(Ok((key_guard, value_guard))) => {
86                    let key = key_guard.value();
87                    if key.base_key == self.base_key {
88                        return Some(Ok(V::from(value_guard.value())));
89                    }
90                }
91                Some(Err(err)) => {
92                    self.finished = true;
93                    return Some(Err(BucketError::IterationError(format!(
94                        "Database error during iteration: {}",
95                        err
96                    ))));
97                }
98                None => {
99                    self.finished = true;
100                    return None;
101                }
102            }
103        }
104    }
105}
106
107impl<V> DoubleEndedIterator for BucketRangeIterator<V>
108where
109    V: redb::Value + 'static,
110    for<'b> V: From<V::SelfType<'b>>,
111{
112    fn next_back(&mut self) -> Option<Self::Item> {
113        if self.finished {
114            return None;
115        }
116
117        loop {
118            match self.range.next_back() {
119                Some(Ok((key_guard, value_guard))) => {
120                    let key = key_guard.value();
121                    if key.base_key == self.base_key {
122                        return Some(Ok(V::from(value_guard.value())));
123                    }
124                }
125                Some(Err(err)) => {
126                    self.finished = true;
127                    return Some(Err(BucketError::IterationError(format!(
128                        "Database error during iteration: {}",
129                        err
130                    ))));
131                }
132                None => {
133                    self.finished = true;
134                    return None;
135                }
136            }
137        }
138    }
139}
140
141/// Iterator over a range of buckets for a specific base key in multimap tables.
142///
143/// This iterator flattens the multimap values, yielding each value in order
144/// across the requested bucket range.
145///
146/// Implements `DoubleEndedIterator` to iterate buckets and values in reverse.
147///
148/// ```
149/// use redb::{Database, MultimapTableDefinition, ReadableDatabase};
150/// use redb_extras::buckets::{BucketMultimapIterExt, BucketedKey, KeyBuilder};
151///
152/// const TABLE: MultimapTableDefinition<'static, BucketedKey<u64>, u64> =
153///     MultimapTableDefinition::new("bucketed_values");
154///
155/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
156/// let db = Database::create("example.redb")?;
157/// let key_builder = KeyBuilder::new(100)?;
158///
159/// let write_txn = db.begin_write()?;
160/// {
161///     let mut table = write_txn.open_multimap_table(TABLE)?;
162///     table.insert(key_builder.bucketed_key(42u64, 10), 1u64)?;
163///     table.insert(key_builder.bucketed_key(42u64, 110), 2u64)?;
164/// }
165/// write_txn.commit()?;
166///
167/// let read_txn = db.begin_read()?;
168/// let table = read_txn.open_multimap_table(TABLE)?;
169/// let values: Vec<u64> = table
170///     .bucket_range(&key_builder, 42u64, 0, 199)?
171///     .collect::<Result<_, _>>()?;
172/// assert_eq!(values, vec![1u64, 2u64]);
173/// # Ok(())
174/// # }
175/// ```
176pub struct BucketRangeMultimapIterator<V>
177where
178    V: redb::Key + 'static,
179    for<'b> V: From<V::SelfType<'b>>,
180{
181    range: MultimapRange<'static, BucketedKey<u64>, V>,
182    base_key: u64,
183    start_bucket: u64,
184    end_bucket: u64,
185    finished: bool,
186    front_values: Option<MultimapValue<'static, V>>,
187    back_values: Option<MultimapValue<'static, V>>,
188}
189
190impl<V> BucketRangeMultimapIterator<V>
191where
192    V: redb::Key + 'static,
193    for<'b> V: From<V::SelfType<'b>>,
194{
195    /// Create a new bucket range iterator for a multimap table.
196    pub fn new(
197        table: &ReadOnlyMultimapTable<BucketedKey<u64>, V>,
198        key_builder: &KeyBuilder,
199        base_key: u64,
200        start_sequence: u64,
201        end_sequence: u64,
202    ) -> Result<Self, BucketError> {
203        if start_sequence > end_sequence {
204            return Err(BucketError::InvalidRange {
205                start: start_sequence,
206                end: end_sequence,
207            });
208        }
209
210        let bucket_size = key_builder.bucket_size();
211        let start_bucket = start_sequence / bucket_size;
212        let end_bucket = end_sequence / bucket_size;
213        let start_key = BucketedKey::new(base_key, start_bucket);
214        let end_key = BucketedKey::new(base_key, end_bucket);
215        let range = table.range(start_key..=end_key).map_err(|err| {
216            BucketError::IterationError(format!("Failed to create range iterator: {}", err))
217        })?;
218
219        Ok(Self {
220            range,
221            base_key,
222            start_bucket,
223            end_bucket,
224            finished: false,
225            front_values: None,
226            back_values: None,
227        })
228    }
229
230    /// Get the bucket range.
231    pub fn bucket_range(&self) -> (u64, u64) {
232        (self.start_bucket, self.end_bucket)
233    }
234}
235
236impl<V> Iterator for BucketRangeMultimapIterator<V>
237where
238    V: redb::Key + 'static,
239    for<'b> V: From<V::SelfType<'b>>,
240{
241    type Item = Result<V, BucketError>;
242
243    fn next(&mut self) -> Option<Self::Item> {
244        if self.finished {
245            return None;
246        }
247
248        loop {
249            if let Some(values) = self.front_values.as_mut() {
250                match values.next() {
251                    Some(Ok(value_guard)) => {
252                        return Some(Ok(V::from(value_guard.value())));
253                    }
254                    Some(Err(err)) => {
255                        self.finished = true;
256                        return Some(Err(BucketError::IterationError(format!(
257                            "Database error during iteration: {}",
258                            err
259                        ))));
260                    }
261                    None => {
262                        self.front_values = None;
263                    }
264                }
265            }
266
267            match self.range.next() {
268                Some(Ok((key_guard, values))) => {
269                    let key = key_guard.value();
270                    if key.base_key == self.base_key {
271                        self.front_values = Some(values);
272                    }
273                }
274                Some(Err(err)) => {
275                    self.finished = true;
276                    return Some(Err(BucketError::IterationError(format!(
277                        "Database error during iteration: {}",
278                        err
279                    ))));
280                }
281                None => {
282                    if let Some(values) = self.back_values.as_mut() {
283                        match values.next() {
284                            Some(Ok(value_guard)) => {
285                                return Some(Ok(V::from(value_guard.value())));
286                            }
287                            Some(Err(err)) => {
288                                self.finished = true;
289                                return Some(Err(BucketError::IterationError(format!(
290                                    "Database error during iteration: {}",
291                                    err
292                                ))));
293                            }
294                            None => {
295                                self.back_values = None;
296                            }
297                        }
298                    }
299
300                    if self.front_values.is_none() && self.back_values.is_none() {
301                        self.finished = true;
302                        return None;
303                    }
304                }
305            }
306        }
307    }
308}
309
310impl<V> DoubleEndedIterator for BucketRangeMultimapIterator<V>
311where
312    V: redb::Key + 'static,
313    for<'b> V: From<V::SelfType<'b>>,
314{
315    fn next_back(&mut self) -> Option<Self::Item> {
316        if self.finished {
317            return None;
318        }
319
320        loop {
321            if let Some(values) = self.back_values.as_mut() {
322                match values.next_back() {
323                    Some(Ok(value_guard)) => {
324                        return Some(Ok(V::from(value_guard.value())));
325                    }
326                    Some(Err(err)) => {
327                        self.finished = true;
328                        return Some(Err(BucketError::IterationError(format!(
329                            "Database error during iteration: {}",
330                            err
331                        ))));
332                    }
333                    None => {
334                        self.back_values = None;
335                    }
336                }
337            }
338
339            match self.range.next_back() {
340                Some(Ok((key_guard, values))) => {
341                    let key = key_guard.value();
342                    if key.base_key == self.base_key {
343                        self.back_values = Some(values);
344                    }
345                }
346                Some(Err(err)) => {
347                    self.finished = true;
348                    return Some(Err(BucketError::IterationError(format!(
349                        "Database error during iteration: {}",
350                        err
351                    ))));
352                }
353                None => {
354                    if let Some(values) = self.front_values.as_mut() {
355                        match values.next_back() {
356                            Some(Ok(value_guard)) => {
357                                return Some(Ok(V::from(value_guard.value())));
358                            }
359                            Some(Err(err)) => {
360                                self.finished = true;
361                                return Some(Err(BucketError::IterationError(format!(
362                                    "Database error during iteration: {}",
363                                    err
364                                ))));
365                            }
366                            None => {
367                                self.front_values = None;
368                            }
369                        }
370                    }
371
372                    if self.front_values.is_none() && self.back_values.is_none() {
373                        self.finished = true;
374                        return None;
375                    }
376                }
377            }
378        }
379    }
380}
381
382/// Extension trait for convenient bucket iteration on read-only tables.
383pub trait BucketIterExt<V>
384where
385    V: redb::Value + 'static,
386    for<'b> V: From<V::SelfType<'b>>,
387{
388    fn bucket_range(
389        &self,
390        key_builder: &KeyBuilder,
391        base_key: u64,
392        start_sequence: u64,
393        end_sequence: u64,
394    ) -> Result<BucketRangeIterator<V>, BucketError>;
395}
396
397impl<V> BucketIterExt<V> for ReadOnlyTable<BucketedKey<u64>, V>
398where
399    V: redb::Value + 'static,
400    for<'b> V: From<V::SelfType<'b>>,
401{
402    fn bucket_range(
403        &self,
404        key_builder: &KeyBuilder,
405        base_key: u64,
406        start_sequence: u64,
407        end_sequence: u64,
408    ) -> Result<BucketRangeIterator<V>, BucketError> {
409        BucketRangeIterator::new(self, key_builder, base_key, start_sequence, end_sequence)
410    }
411}
412
413/// Extension trait for convenient bucket iteration on read-only multimap tables.
414///
415/// Returns a flattened iterator over values for the base key within the
416/// requested bucket range.
417pub trait BucketMultimapIterExt<V>
418where
419    V: redb::Key + 'static,
420    for<'b> V: From<V::SelfType<'b>>,
421{
422    fn bucket_range(
423        &self,
424        key_builder: &KeyBuilder,
425        base_key: u64,
426        start_sequence: u64,
427        end_sequence: u64,
428    ) -> Result<BucketRangeMultimapIterator<V>, BucketError>;
429}
430
431impl<V> BucketMultimapIterExt<V> for ReadOnlyMultimapTable<BucketedKey<u64>, V>
432where
433    V: redb::Key + 'static,
434    for<'b> V: From<V::SelfType<'b>>,
435{
436    fn bucket_range(
437        &self,
438        key_builder: &KeyBuilder,
439        base_key: u64,
440        start_sequence: u64,
441        end_sequence: u64,
442    ) -> Result<BucketRangeMultimapIterator<V>, BucketError> {
443        BucketRangeMultimapIterator::new(self, key_builder, base_key, start_sequence, end_sequence)
444    }
445}
446
447#[cfg(test)]
448mod tests {
449    use super::*;
450    use redb::{Database, MultimapTableDefinition, ReadableDatabase, TableDefinition};
451    use tempfile::NamedTempFile;
452
453    const TEST_TABLE: TableDefinition<'static, BucketedKey<u64>, String> =
454        TableDefinition::new("test_table");
455    const TEST_MULTIMAP: MultimapTableDefinition<'static, BucketedKey<u64>, u64> =
456        MultimapTableDefinition::new("test_multimap");
457
458    #[test]
459    fn test_basic_functionality() -> Result<(), Box<dyn std::error::Error>> {
460        let temp_file = NamedTempFile::new()?;
461        let db = Database::create(temp_file.path())?;
462        let key_builder = KeyBuilder::new(100)?;
463
464        // Insert test data
465        {
466            let write_txn = db.begin_write()?;
467            {
468                let mut table = write_txn.open_table(TEST_TABLE)?;
469
470                // Insert values for user123 in different buckets
471                table.insert(key_builder.bucketed_key(123u64, 50), "value_50".to_string())?;
472                table.insert(
473                    key_builder.bucketed_key(123u64, 150),
474                    "value_150".to_string(),
475                )?;
476                table.insert(
477                    key_builder.bucketed_key(123u64, 250),
478                    "value_250".to_string(),
479                )?;
480
481                // Insert values for user456 in same buckets (should not appear in iteration)
482                table.insert(key_builder.bucketed_key(456u64, 50), "other_50".to_string())?;
483                table.insert(
484                    key_builder.bucketed_key(456u64, 150),
485                    "other_150".to_string(),
486                )?;
487            }
488            write_txn.commit()?;
489        }
490
491        // Test that we can create bucket range iterators
492        {
493            let read_txn = db.begin_read()?;
494            let table = read_txn.open_table(TEST_TABLE)?;
495            let iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 0, 199)?;
496            assert_eq!(iter.bucket_range(), (0, 1));
497
498            // Test invalid range
499            let invalid_iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 200, 100);
500            assert!(invalid_iter.is_err());
501        }
502
503        // Test value iteration and base key filtering
504        {
505            let read_txn = db.begin_read()?;
506            let table = read_txn.open_table(TEST_TABLE)?;
507            let iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 0, 299)?;
508            let values: Vec<String> = iter.collect::<Result<_, _>>()?;
509            assert_eq!(
510                values,
511                vec![
512                    "value_50".to_string(),
513                    "value_150".to_string(),
514                    "value_250".to_string()
515                ]
516            );
517
518            let iter = BucketRangeIterator::new(&table, &key_builder, 123u64, 0, 299)?;
519            let values: Vec<String> = iter.rev().collect::<Result<_, _>>()?;
520            assert_eq!(
521                values,
522                vec![
523                    "value_250".to_string(),
524                    "value_150".to_string(),
525                    "value_50".to_string()
526                ]
527            );
528
529            let iter = table.bucket_range(&key_builder, 456u64, 0, 299)?;
530            let values: Vec<String> = iter.collect::<Result<_, _>>()?;
531            assert_eq!(
532                values,
533                vec!["other_50".to_string(), "other_150".to_string()]
534            );
535        }
536
537        Ok(())
538    }
539
540    #[test]
541    fn test_multimap_functionality() -> Result<(), Box<dyn std::error::Error>> {
542        let temp_file = NamedTempFile::new()?;
543        let db = Database::create(temp_file.path())?;
544        let key_builder = KeyBuilder::new(100)?;
545
546        {
547            let write_txn = db.begin_write()?;
548            {
549                let mut table = write_txn.open_multimap_table(TEST_MULTIMAP)?;
550
551                table.insert(key_builder.bucketed_key(123u64, 50), 10u64)?;
552                table.insert(key_builder.bucketed_key(123u64, 50), 20u64)?;
553                table.insert(key_builder.bucketed_key(123u64, 150), 30u64)?;
554                table.insert(key_builder.bucketed_key(123u64, 150), 40u64)?;
555
556                table.insert(key_builder.bucketed_key(456u64, 50), 99u64)?;
557                table.insert(key_builder.bucketed_key(456u64, 50), 100u64)?;
558            }
559            write_txn.commit()?;
560        }
561
562        {
563            let read_txn = db.begin_read()?;
564            let table = read_txn.open_multimap_table(TEST_MULTIMAP)?;
565            let iter = BucketRangeMultimapIterator::new(&table, &key_builder, 123u64, 0, 199)?;
566            assert_eq!(iter.bucket_range(), (0, 1));
567
568            let values: Vec<u64> = iter.collect::<Result<_, _>>()?;
569            assert_eq!(values, vec![10u64, 20u64, 30u64, 40u64]);
570
571            let iter = BucketRangeMultimapIterator::new(&table, &key_builder, 123u64, 0, 199)?;
572            let values: Vec<u64> = iter.rev().collect::<Result<_, _>>()?;
573            assert_eq!(values, vec![40u64, 30u64, 20u64, 10u64]);
574
575            let iter = table.bucket_range(&key_builder, 456u64, 0, 99)?;
576            let values: Vec<u64> = iter.collect::<Result<_, _>>()?;
577            assert_eq!(values, vec![99u64, 100u64]);
578        }
579
580        Ok(())
581    }
582}
redb_extras/buckets/iterator.rs

redb_extras/buckets/
iterator.rs
