mini_lsm/iterators/
two_merge_iterator.rs

1use anyhow::Result;
2
3use super::StorageIterator;
4
5/// Merges two iterators of different types into one. If the two iterators have the same key, only
6/// produce the key once and prefer the entry from A.
7pub struct TwoMergeIterator<A: StorageIterator, B: StorageIterator> {
8    a: A,
9    b: B,
10    choose_a: bool,
11}
12
13impl<
14        A: 'static + StorageIterator,
15        B: 'static + for<'a> StorageIterator<KeyType<'a> = A::KeyType<'a>>,
16    > TwoMergeIterator<A, B>
17{
18    fn choose_a(a: &A, b: &B) -> bool {
19        if !a.is_valid() {
20            return false;
21        }
22        if !b.is_valid() {
23            return true;
24        }
25        a.key() < b.key()
26    }
27
28    fn skip_b(&mut self) -> Result<()> {
29        if self.a.is_valid() && self.b.is_valid() && self.b.key() == self.a.key() {
30            self.b.next()?;
31        }
32        Ok(())
33    }
34
35    pub fn create(a: A, b: B) -> Result<Self> {
36        let mut iter = Self {
37            choose_a: false,
38            a,
39            b,
40        };
41        iter.skip_b()?;
42        iter.choose_a = Self::choose_a(&iter.a, &iter.b);
43        Ok(iter)
44    }
45}
46
47impl<
48        A: 'static + StorageIterator,
49        B: 'static + for<'a> StorageIterator<KeyType<'a> = A::KeyType<'a>>,
50    > StorageIterator for TwoMergeIterator<A, B>
51{
52    type KeyType<'a> = A::KeyType<'a>;
53
54    fn key(&self) -> Self::KeyType<'_> {
55        if self.choose_a {
56            self.a.key()
57        } else {
58            self.b.key()
59        }
60    }
61
62    fn value(&self) -> &[u8] {
63        if self.choose_a {
64            self.a.value()
65        } else {
66            self.b.value()
67        }
68    }
69
70    fn is_valid(&self) -> bool {
71        if self.choose_a {
72            self.a.is_valid()
73        } else {
74            self.b.is_valid()
75        }
76    }
77
78    fn next(&mut self) -> Result<()> {
79        if self.choose_a {
80            self.a.next()?;
81        } else {
82            self.b.next()?;
83        }
84        self.skip_b()?;
85        self.choose_a = Self::choose_a(&self.a, &self.b);
86        Ok(())
87    }
88
89    fn num_active_iterators(&self) -> usize {
90        self.a.num_active_iterators() + self.b.num_active_iterators()
91    }
92}
mini_lsm/iterators/two_merge_iterator.rs

mini_lsm/iterators/
two_merge_iterator.rs
