subms_merge_iterator/features/reverse.rs
1//! Descending k-way merge - the backward half of a cursor. Sources must
2//! be sorted DESCENDING; output is the global descending union.
3//!
4//! A freshly built `ReverseMergeIterator` sits on the largest value
5//! across every source, which is the position RocksDB calls
6//! `SeekToLast`. `seek_for_prev(target)` moves it to the largest value
7//! `<= target`, and `set_lower_bound(lo)` stops the walk at `lo`
8//! inclusive, matching RocksDB's `iterate_lower_bound`.
9//!
10//! What this is NOT is a bidirectional cursor. RocksDB's `Prev()` can
11//! reverse mid-scan because each child is a seekable file cursor; the
12//! sources here are one-shot iterators that only move forward through
13//! their own order, so a direction flip would have to re-read them.
14//! Pick the direction when you open the merge.
15
16use std::collections::BinaryHeap;
17
18pub struct ReverseMergeIterator<T: Ord, I: Iterator<Item = T>> {
19 streams: Vec<I>,
20 /// Max-heap on (value, stream index) - the plain `BinaryHeap` order,
21 /// which is what a descending merge wants.
22 heap: BinaryHeap<(T, usize)>,
23 lower_bound: Option<T>,
24}
25
26impl<T: Ord, I: Iterator<Item = T>> ReverseMergeIterator<T, I> {
27 pub fn new<S: IntoIterator<Item = I>>(streams: S) -> Self {
28 let mut streams: Vec<I> = streams.into_iter().collect();
29 let mut heap = BinaryHeap::with_capacity(streams.len());
30 for (i, s) in streams.iter_mut().enumerate() {
31 if let Some(v) = s.next() {
32 heap.push((v, i));
33 }
34 }
35 Self {
36 streams,
37 heap,
38 lower_bound: None,
39 }
40 }
41
42 /// Retreat past every entry strictly greater than `target`. After this
43 /// call the next `next()` yields the largest value `<= target`, or the
44 /// iterator is exhausted.
45 pub fn seek_for_prev(&mut self, target: &T) {
46 let mut new_heads: Vec<(T, usize)> = Vec::new();
47 while let Some((value, _)) = self.heap.peek() {
48 if value > target {
49 let (_value, idx) = self.heap.pop().unwrap();
50 let mut found: Option<T> = None;
51 for v in self.streams[idx].by_ref() {
52 if &v <= target {
53 found = Some(v);
54 break;
55 }
56 }
57 if let Some(v) = found {
58 new_heads.push((v, idx));
59 }
60 } else {
61 break;
62 }
63 }
64 for head in new_heads {
65 self.heap.push(head);
66 }
67 }
68
69 /// Stop the descending scan at `bound`. The bound is inclusive: a value
70 /// equal to it is the last one yielded.
71 pub fn set_lower_bound(&mut self, bound: T) {
72 self.lower_bound = Some(bound);
73 }
74
75 /// Drop the lower bound and let the scan run to the end of every source.
76 pub fn clear_lower_bound(&mut self) {
77 self.lower_bound = None;
78 }
79
80 /// The value the next `next()` will yield, without consuming it.
81 pub fn peek(&self) -> Option<&T> {
82 let head = self.heap.peek().map(|(value, _)| value)?;
83 match &self.lower_bound {
84 Some(lo) if head < lo => None,
85 _ => Some(head),
86 }
87 }
88
89 /// Streams still holding a head in the heap, ignoring the lower bound.
90 pub fn live_streams(&self) -> usize {
91 self.heap.len()
92 }
93
94 /// Streams the merge was constructed over, live or not.
95 pub fn num_streams(&self) -> usize {
96 self.streams.len()
97 }
98}
99
100impl<T: Ord, I: Iterator<Item = T>> Iterator for ReverseMergeIterator<T, I> {
101 type Item = T;
102 fn next(&mut self) -> Option<T> {
103 self.peek()?;
104 let (value, idx) = self.heap.pop()?;
105 if let Some(next_value) = self.streams[idx].next() {
106 self.heap.push((next_value, idx));
107 }
108 Some(value)
109 }
110}
111
112#[cfg(test)]
113#[path = "reverse_tests.rs"]
114mod tests;