reifydb_core/value/index/
range.rs1use std::{collections::Bound, iter};
5
6use reifydb_codec::key::encoded::{EncodedKey, EncodedKeyRange};
7
8use crate::value::index::encoded::EncodedIndexKey;
9
10#[derive(Clone, Debug)]
11pub struct EncodedIndexKeyRange {
12 pub start: Bound<EncodedIndexKey>,
13 pub end: Bound<EncodedIndexKey>,
14}
15
16impl EncodedIndexKeyRange {
17 pub fn new(start: Bound<EncodedIndexKey>, end: Bound<EncodedIndexKey>) -> Self {
18 Self {
19 start,
20 end,
21 }
22 }
23
24 pub fn start_end(start: Option<EncodedIndexKey>, end: Option<EncodedIndexKey>) -> Self {
25 let start = match start {
26 Some(s) => Bound::Included(s),
27 None => Bound::Unbounded,
28 };
29
30 let end = match end {
31 Some(e) => Bound::Excluded(e),
32 None => Bound::Unbounded,
33 };
34
35 Self {
36 start,
37 end,
38 }
39 }
40
41 pub fn start_end_inclusive(start: Option<EncodedIndexKey>, end: Option<EncodedIndexKey>) -> Self {
42 let start = match start {
43 Some(s) => Bound::Included(s),
44 None => Bound::Unbounded,
45 };
46
47 let end = match end {
48 Some(e) => Bound::Included(e),
49 None => Bound::Unbounded,
50 };
51
52 Self {
53 start,
54 end,
55 }
56 }
57
58 pub fn prefix(prefix: &[u8]) -> Self {
59 let start = Bound::Included(EncodedIndexKey::from_bytes(prefix));
60 let end = match prefix.iter().rposition(|&b| b != 0xff) {
61 Some(i) => Bound::Excluded(EncodedIndexKey::from_bytes(
62 &prefix.iter().take(i).copied().chain(iter::once(prefix[i] + 1)).collect::<Vec<_>>(),
63 )),
64 None => Bound::Unbounded,
65 };
66 Self {
67 start,
68 end,
69 }
70 }
71
72 pub fn all() -> Self {
73 Self {
74 start: Bound::Unbounded,
75 end: Bound::Unbounded,
76 }
77 }
78
79 pub fn to_encoded_key_range(&self) -> EncodedKeyRange {
80 let start = match &self.start {
81 Bound::Included(key) => Bound::Included(EncodedKey::new(key.as_slice())),
82 Bound::Excluded(key) => Bound::Excluded(EncodedKey::new(key.as_slice())),
83 Bound::Unbounded => Bound::Unbounded,
84 };
85
86 let end = match &self.end {
87 Bound::Included(key) => Bound::Included(EncodedKey::new(key.as_slice())),
88 Bound::Excluded(key) => Bound::Excluded(EncodedKey::new(key.as_slice())),
89 Bound::Unbounded => Bound::Unbounded,
90 };
91
92 EncodedKeyRange::new(start, end)
93 }
94
95 pub fn from_prefix(key: &EncodedIndexKey) -> Self {
96 Self::prefix(key.as_slice())
97 }
98}
99
100impl From<EncodedIndexKeyRange> for EncodedKeyRange {
101 fn from(range: EncodedIndexKeyRange) -> Self {
102 range.to_encoded_key_range()
103 }
104}
105
106#[cfg(test)]
107pub mod tests {
108 use reifydb_value::value::value_type::ValueType;
109
110 use super::*;
111 use crate::{sort::SortDirection, value::index::shape::IndexShape};
112
113 #[test]
114 fn test_start_end() {
115 let layout = IndexShape::new(&[ValueType::Uint8], &[SortDirection::Asc]).unwrap();
116
117 let mut key1 = layout.allocate_key();
118 layout.set_u64(&mut key1, 0, 100u64);
119
120 let mut key2 = layout.allocate_key();
121 layout.set_u64(&mut key2, 0, 200u64);
122
123 let range = EncodedIndexKeyRange::start_end(Some(key1.clone()), Some(key2.clone()));
124
125 match &range.start {
126 Bound::Included(k) => {
127 assert_eq!(k.as_slice(), key1.as_slice())
128 }
129 _ => panic!("Expected Included start bound"),
130 }
131
132 match &range.end {
133 Bound::Excluded(k) => {
134 assert_eq!(k.as_slice(), key2.as_slice())
135 }
136 _ => panic!("Expected Excluded end bound"),
137 }
138 }
139
140 #[test]
141 fn test_start_end_inclusive() {
142 let layout = IndexShape::new(&[ValueType::Uint8], &[SortDirection::Asc]).unwrap();
143
144 let mut key1 = layout.allocate_key();
145 layout.set_u64(&mut key1, 0, 100u64);
146
147 let mut key2 = layout.allocate_key();
148 layout.set_u64(&mut key2, 0, 200u64);
149
150 let range = EncodedIndexKeyRange::start_end_inclusive(Some(key1.clone()), Some(key2.clone()));
151
152 match &range.start {
153 Bound::Included(k) => {
154 assert_eq!(k.as_slice(), key1.as_slice())
155 }
156 _ => panic!("Expected Included start bound"),
157 }
158
159 match &range.end {
160 Bound::Included(k) => {
161 assert_eq!(k.as_slice(), key2.as_slice())
162 }
163 _ => panic!("Expected Included end bound"),
164 }
165 }
166
167 #[test]
168 fn test_unbounded() {
169 let range = EncodedIndexKeyRange::start_end(None, None);
170 assert!(matches!(range.start, Bound::Unbounded));
171 assert!(matches!(range.end, Bound::Unbounded));
172 }
173
174 #[test]
175 fn test_prefix() {
176 let prefix = &[0x12, 0x34];
177 let range = EncodedIndexKeyRange::prefix(prefix);
178
179 match &range.start {
180 Bound::Included(k) => assert_eq!(k.as_slice(), prefix),
181 _ => panic!("Expected Included start bound"),
182 }
183
184 match &range.end {
185 Bound::Excluded(k) => {
186 assert_eq!(k.as_slice(), &[0x12, 0x35])
187 }
188 _ => panic!("Expected Excluded end bound"),
189 }
190 }
191
192 #[test]
193 fn test_prefix_with_ff() {
194 let prefix = &[0x12, 0xff];
195 let range = EncodedIndexKeyRange::prefix(prefix);
196
197 match &range.start {
198 Bound::Included(k) => assert_eq!(k.as_slice(), prefix),
199 _ => panic!("Expected Included start bound"),
200 }
201
202 match &range.end {
203 Bound::Excluded(k) => assert_eq!(k.as_slice(), &[0x13]),
204 _ => panic!("Expected Excluded end bound"),
205 }
206 }
207
208 #[test]
209 fn test_prefix_all_ff() {
210 let prefix = &[0xff, 0xff];
211 let range = EncodedIndexKeyRange::prefix(prefix);
212
213 match &range.start {
214 Bound::Included(k) => assert_eq!(k.as_slice(), prefix),
215 _ => panic!("Expected Included start bound"),
216 }
217
218 assert!(matches!(range.end, Bound::Unbounded));
219 }
220
221 #[test]
222 fn test_to_encoded_key_range() {
223 let layout = IndexShape::new(&[ValueType::Uint8], &[SortDirection::Asc]).unwrap();
224
225 let mut key = layout.allocate_key();
226 layout.set_u64(&mut key, 0, 100u64);
227
228 let index_range = EncodedIndexKeyRange::start_end(Some(key.clone()), None);
229 let key_range = index_range.to_encoded_key_range();
230
231 match &key_range.start {
232 Bound::Included(k) => {
233 assert_eq!(k.as_slice(), key.as_slice())
234 }
235 _ => panic!("Expected Included start bound"),
236 }
237
238 assert!(matches!(key_range.end, Bound::Unbounded));
239 }
240
241 #[test]
242 fn test_all() {
243 let range = EncodedIndexKeyRange::all();
244 assert!(matches!(range.start, Bound::Unbounded));
245 assert!(matches!(range.end, Bound::Unbounded));
246 }
247}