risingwave_hummock_sdk/
prost_key_range.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp;
16
17use risingwave_pb::hummock::KeyRange;
18
19use crate::key_range::KeyRangeCommon;
20use crate::{KeyComparator, impl_key_range_common, key_range_cmp};
21
22impl_key_range_common!(KeyRange);
23
24pub trait KeyRangeExt {
25    fn inf() -> Self;
26    fn new(left: Vec<u8>, right: Vec<u8>) -> Self;
27    fn compare(&self, other: &Self) -> cmp::Ordering;
28    fn start_bound_inf(&self) -> bool;
29    fn end_bound_inf(&self) -> bool;
30}
31
32impl KeyRangeExt for KeyRange {
33    fn inf() -> Self {
34        Self {
35            left: vec![],
36            right: vec![],
37            right_exclusive: false,
38        }
39    }
40
41    fn new(left: Vec<u8>, right: Vec<u8>) -> Self {
42        Self {
43            left,
44            right,
45            right_exclusive: false,
46        }
47    }
48
49    fn compare(&self, other: &Self) -> cmp::Ordering {
50        key_range_cmp!(self, other)
51    }
52
53    #[inline]
54    fn start_bound_inf(&self) -> bool {
55        self.left.is_empty()
56    }
57
58    #[inline]
59    fn end_bound_inf(&self) -> bool {
60        self.right.is_empty()
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use std::cmp;
67
68    use risingwave_pb::hummock::KeyRange;
69
70    use crate::key::key_with_epoch;
71    use crate::key_range::KeyRangeCommon;
72    use crate::prost_key_range::KeyRangeExt;
73
74    #[test]
75    fn test_cmp() {
76        let a1 = key_with_epoch(Vec::from("a"), 1);
77        let a2 = key_with_epoch(Vec::from("a"), 2);
78        let b1 = key_with_epoch(Vec::from("b"), 1);
79
80        assert_eq!(a1.cmp(&a1), cmp::Ordering::Equal);
81        assert_eq!(a2.cmp(&a2), cmp::Ordering::Equal);
82        assert_eq!(b1.cmp(&b1), cmp::Ordering::Equal);
83        assert_eq!(a1.cmp(&a2), cmp::Ordering::Less);
84        assert_eq!(a2.cmp(&a1), cmp::Ordering::Greater);
85        assert_eq!(a1.cmp(&b1), cmp::Ordering::Less);
86        assert_eq!(b1.cmp(&a1), cmp::Ordering::Greater);
87        assert_eq!(a2.cmp(&b1), cmp::Ordering::Less);
88        assert_eq!(b1.cmp(&a2), cmp::Ordering::Greater);
89
90        let kr1 = KeyRange::new(a2.clone(), a1.clone());
91        let kr2 = KeyRange::new(a2, b1.clone());
92        let kr3 = KeyRange::new(a1, b1);
93        assert_eq!(kr1.compare(&kr1), cmp::Ordering::Equal);
94        assert_eq!(kr2.compare(&kr2), cmp::Ordering::Equal);
95        assert_eq!(kr3.compare(&kr3), cmp::Ordering::Equal);
96        assert_eq!(kr1.compare(&kr2), cmp::Ordering::Less);
97        assert_eq!(kr1.compare(&kr3), cmp::Ordering::Less);
98        assert_eq!(kr2.compare(&kr3), cmp::Ordering::Less);
99        assert_eq!(kr2.compare(&kr1), cmp::Ordering::Greater);
100        assert_eq!(kr3.compare(&kr1), cmp::Ordering::Greater);
101        assert_eq!(kr3.compare(&kr2), cmp::Ordering::Greater);
102
103        let kr_inf = KeyRange::inf();
104        assert_eq!(kr_inf.compare(&kr_inf), cmp::Ordering::Equal);
105        assert_eq!(kr1.compare(&kr_inf), cmp::Ordering::Greater);
106    }
107
108    #[test]
109    fn test_full_key_overlap() {
110        let a1 = key_with_epoch(Vec::from("a"), 1);
111        let a2 = key_with_epoch(Vec::from("a"), 2);
112        let b1 = key_with_epoch(Vec::from("b"), 1);
113        let b2 = key_with_epoch(Vec::from("b"), 2);
114
115        let kr1 = KeyRange::new(a2.clone(), a1.clone());
116        let kr2 = KeyRange::new(a2, b1.clone());
117        let kr3 = KeyRange::new(a1.clone(), b2.clone());
118        let kr4 = KeyRange::new(a1, b1.clone());
119        let kr5 = KeyRange::new(b2, b1);
120
121        assert!(kr1.full_key_overlap(&kr2));
122        assert!(kr1.full_key_overlap(&kr3));
123        assert!(kr1.full_key_overlap(&kr4));
124        assert!(kr2.full_key_overlap(&kr3));
125        assert!(kr2.full_key_overlap(&kr4));
126        assert!(kr3.full_key_overlap(&kr4));
127        assert!(!kr1.full_key_overlap(&kr5));
128        assert!(kr2.full_key_overlap(&kr5));
129        assert!(kr3.full_key_overlap(&kr5));
130        assert!(kr4.full_key_overlap(&kr5));
131
132        let kr_inf = KeyRange::inf();
133        assert!(kr_inf.full_key_overlap(&kr_inf));
134        assert!(kr1.full_key_overlap(&kr_inf));
135    }
136
137    #[test]
138    fn full_key_extend() {
139        let a1 = key_with_epoch(Vec::from("a"), 1);
140        let a2 = key_with_epoch(Vec::from("a"), 2);
141        let b1 = key_with_epoch(Vec::from("b"), 1);
142        let b2 = key_with_epoch(Vec::from("b"), 2);
143
144        let kr1 = KeyRange::new(a2.clone(), a1.clone());
145        let kr2 = KeyRange::new(a2.clone(), b2);
146        let kr3 = KeyRange::new(a1, b1.clone());
147
148        let mut kr = kr1;
149        kr.full_key_extend(&kr2);
150        assert_eq!(kr.compare(&kr2), cmp::Ordering::Equal);
151        kr.full_key_extend(&kr3);
152        assert_eq!(kr.compare(&KeyRange::new(a2, b1)), cmp::Ordering::Equal);
153
154        let kr_inf = KeyRange::inf();
155        kr.full_key_extend(&kr_inf);
156        assert_eq!(kr.compare(&kr_inf), cmp::Ordering::Equal);
157
158        kr.full_key_extend(&kr_inf);
159        assert_eq!(kr.compare(&kr_inf), cmp::Ordering::Equal);
160    }
161}
risingwave_hummock_sdk/prost_key_range.rs

risingwave_hummock_sdk/
prost_key_range.rs
