Skip to main content

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;