1use std::collections::Bound;
5
6use reifydb_codec::key::{
7 deserializer::KeyDeserializer,
8 encoded::{EncodedKey, EncodedKeyRange},
9 serializer::KeySerializer,
10};
11
12use super::{EncodableKey, EncodableKeyRange, KeyKind};
13use crate::{
14 interface::catalog::{id::IndexId, object::ObjectId},
15 key::catalog::{KeyDeserializerCatalogExt, KeySerializerCatalogExt},
16 value::index::{encoded::EncodedIndexKey, range::EncodedIndexKeyRange},
17};
18
19#[derive(Debug, Clone, PartialEq)]
20pub struct IndexEntryKey {
21 pub object: ObjectId,
22 pub index: IndexId,
23 pub key: EncodedIndexKey,
24}
25
26impl IndexEntryKey {
27 pub fn new(object: impl Into<ObjectId>, index: IndexId, key: EncodedIndexKey) -> Self {
28 Self {
29 object: object.into(),
30 index,
31 key,
32 }
33 }
34
35 pub fn encoded(object: impl Into<ObjectId>, index: IndexId, key: EncodedIndexKey) -> EncodedKey {
36 Self::new(object, index, key).encode()
37 }
38}
39
40#[derive(Debug, Clone, PartialEq)]
41pub struct IndexEntryKeyRange {
42 pub object: ObjectId,
43 pub index: IndexId,
44}
45
46impl IndexEntryKeyRange {
47 fn decode_key(key: &EncodedKey) -> Option<Self> {
48 let mut de = KeyDeserializer::from_bytes(key.as_slice());
49
50 let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
51 if kind != Self::KIND {
52 return None;
53 }
54
55 let object = de.read_object_id().ok()?;
56 let index = de.read_index_id().ok()?;
57
58 Some(IndexEntryKeyRange {
59 object,
60 index,
61 })
62 }
63}
64
65impl EncodableKeyRange for IndexEntryKeyRange {
66 const KIND: KeyKind = KeyKind::IndexEntry;
67
68 fn start(&self) -> Option<EncodedKey> {
69 let mut serializer = KeySerializer::with_capacity(19);
70 serializer.extend_u8(Self::KIND as u8).extend_object_id(self.object).extend_index_id(self.index);
71 Some(serializer.to_encoded_key())
72 }
73
74 fn end(&self) -> Option<EncodedKey> {
75 let mut serializer = KeySerializer::with_capacity(19);
76 serializer.extend_u8(Self::KIND as u8).extend_object_id(self.object).extend_index_id(self.index.prev());
77 Some(serializer.to_encoded_key())
78 }
79
80 fn decode(range: &EncodedKeyRange) -> (Option<Self>, Option<Self>)
81 where
82 Self: Sized,
83 {
84 let start_key = match &range.start {
85 Bound::Included(key) | Bound::Excluded(key) => Self::decode_key(key),
86 Bound::Unbounded => None,
87 };
88
89 let end_key = match &range.end {
90 Bound::Included(key) | Bound::Excluded(key) => Self::decode_key(key),
91 Bound::Unbounded => None,
92 };
93
94 (start_key, end_key)
95 }
96}
97
98impl EncodableKey for IndexEntryKey {
99 const KIND: KeyKind = KeyKind::IndexEntry;
100
101 fn encode(&self) -> EncodedKey {
102 let mut serializer = KeySerializer::with_capacity(20 + self.key.len());
103 serializer
104 .extend_u8(Self::KIND as u8)
105 .extend_object_id(self.object)
106 .extend_index_id(self.index)
107 .extend_raw(self.key.as_slice());
108 serializer.to_encoded_key()
109 }
110
111 fn decode(key: &EncodedKey) -> Option<Self> {
112 let mut de = KeyDeserializer::from_bytes(key.as_slice());
113
114 let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
115 if kind != Self::KIND {
116 return None;
117 }
118
119 let object = de.read_object_id().ok()?;
120 let index = de.read_index_id().ok()?;
121
122 let remaining = de.remaining();
123 if remaining > 0 {
124 let remaining_bytes = de.read_raw(remaining).ok()?;
125 let index_key = EncodedIndexKey::new(remaining_bytes);
126 Some(Self {
127 object,
128 index,
129 key: index_key,
130 })
131 } else {
132 None
133 }
134 }
135}
136
137impl IndexEntryKey {
138 pub fn index_range(object: impl Into<ObjectId>, index: IndexId) -> EncodedKeyRange {
139 let range = IndexEntryKeyRange {
140 object: object.into(),
141 index,
142 };
143 EncodedKeyRange::new(Bound::Included(range.start().unwrap()), Bound::Excluded(range.end().unwrap()))
144 }
145
146 pub fn object_range(object: impl Into<ObjectId>) -> EncodedKeyRange {
147 let object = object.into();
148 let mut start_serializer = KeySerializer::with_capacity(10);
149 start_serializer.extend_u8(KeyKind::IndexEntry as u8).extend_object_id(object);
150
151 let next_object = object.next();
152 let mut end_serializer = KeySerializer::with_capacity(10);
153 end_serializer.extend_u8(KeyKind::IndexEntry as u8).extend_object_id(next_object);
154
155 EncodedKeyRange {
156 start: Bound::Included(start_serializer.to_encoded_key()),
157 end: Bound::Excluded(end_serializer.to_encoded_key()),
158 }
159 }
160
161 pub fn key_prefix_range(object: impl Into<ObjectId>, index: IndexId, key_prefix: &[u8]) -> EncodedKeyRange {
162 let object = object.into();
163 let mut serializer = KeySerializer::with_capacity(20 + key_prefix.len());
164 serializer
165 .extend_u8(KeyKind::IndexEntry as u8)
166 .extend_object_id(object)
167 .extend_index_id(index)
168 .extend_raw(key_prefix);
169 EncodedKeyRange::prefix(serializer.to_encoded_key().as_slice())
170 }
171
172 pub fn key_range(
173 object: impl Into<ObjectId>,
174 index: IndexId,
175 index_range: EncodedIndexKeyRange,
176 ) -> EncodedKeyRange {
177 let object = object.into();
178
179 let mut prefix_serializer = KeySerializer::with_capacity(19);
180 prefix_serializer.extend_u8(KeyKind::IndexEntry as u8).extend_object_id(object).extend_index_id(index);
181 let prefix = prefix_serializer.to_encoded_key().to_vec();
182
183 let start = match index_range.start {
184 Bound::Included(key) => {
185 let mut bytes = prefix.clone();
186 bytes.extend_from_slice(key.as_slice());
187 Bound::Included(EncodedKey::new(bytes))
188 }
189 Bound::Excluded(key) => {
190 let mut bytes = prefix.clone();
191 bytes.extend_from_slice(key.as_slice());
192 Bound::Excluded(EncodedKey::new(bytes))
193 }
194 Bound::Unbounded => Bound::Included(EncodedKey::new(prefix.clone())),
195 };
196
197 let end = match index_range.end {
198 Bound::Included(key) => {
199 let mut bytes = prefix.clone();
200 bytes.extend_from_slice(key.as_slice());
201 Bound::Included(EncodedKey::new(bytes))
202 }
203 Bound::Excluded(key) => {
204 let mut bytes = prefix.clone();
205 bytes.extend_from_slice(key.as_slice());
206 Bound::Excluded(EncodedKey::new(bytes))
207 }
208 Bound::Unbounded => {
209 let mut serializer = KeySerializer::with_capacity(19);
210 serializer
211 .extend_u8(KeyKind::IndexEntry as u8)
212 .extend_object_id(object)
213 .extend_index_id(index.prev());
214 Bound::Excluded(serializer.to_encoded_key())
215 }
216 };
217
218 EncodedKeyRange {
219 start,
220 end,
221 }
222 }
223}
224
225#[cfg(test)]
226pub mod tests {
227 use reifydb_value::value::value_type::ValueType;
228
229 use super::*;
230 use crate::{sort::SortDirection, value::index::shape::IndexShape};
231
232 #[test]
233 fn test_encode_decode() {
234 let layout = IndexShape::new(
235 &[ValueType::Uint8, ValueType::Uint8],
236 &[SortDirection::Asc, SortDirection::Asc],
237 )
238 .unwrap();
239
240 let mut index_key = layout.allocate_key();
241 layout.set_u64(&mut index_key, 0, 100u64);
242 layout.set_row_number(&mut index_key, 1, 1u64);
243
244 let entry = IndexEntryKey {
245 object: ObjectId::table(42),
246 index: IndexId::primary(7),
247 key: index_key.clone(),
248 };
249
250 let encoded = entry.encode();
251 let decoded = IndexEntryKey::decode(&encoded).unwrap();
252
253 assert_eq!(decoded.object, ObjectId::table(42));
254 assert_eq!(decoded.index, IndexId::primary(7));
255 assert_eq!(decoded.key.as_slice(), index_key.as_slice());
256 }
257
258 #[test]
259 fn test_ordering() {
260 let layout = IndexShape::new(&[ValueType::Uint8], &[SortDirection::Asc]).unwrap();
261
262 let mut key1 = layout.allocate_key();
263 layout.set_u64(&mut key1, 0, 100u64);
264
265 let mut key2 = layout.allocate_key();
266 layout.set_u64(&mut key2, 0, 200u64);
267
268 let entry1 = IndexEntryKey {
269 object: ObjectId::table(1),
270 index: IndexId::primary(1),
271 key: key1,
272 };
273
274 let entry2 = IndexEntryKey {
275 object: ObjectId::table(1),
276 index: IndexId::primary(1),
277 key: key2,
278 };
279
280 let encoded1 = entry1.encode();
281 let encoded2 = entry2.encode();
282
283 assert!(encoded1.as_slice() < encoded2.as_slice());
284 }
285
286 #[test]
287 fn test_index_range() {
288 let range = IndexEntryKey::index_range(ObjectId::table(10), IndexId::primary(5));
289
290 let layout = IndexShape::new(&[ValueType::Uint8], &[SortDirection::Asc]).unwrap();
291
292 let mut key = layout.allocate_key();
293 layout.set_u64(&mut key, 0, 50u64);
294
295 let entry = IndexEntryKey {
296 object: ObjectId::table(10),
297 index: IndexId::primary(5),
298 key,
299 };
300
301 let encoded = entry.encode();
302
303 if let (Bound::Included(start), Bound::Excluded(end)) = (&range.start, &range.end) {
304 assert!(encoded.as_slice() >= start.as_slice());
305 assert!(encoded.as_slice() < end.as_slice());
306 } else {
307 panic!("Expected Included/Excluded bounds");
308 }
309
310 let entry2 = IndexEntryKey {
311 object: ObjectId::table(10),
312 index: IndexId::primary(6),
313 key: layout.allocate_key(),
314 };
315
316 let encoded2 = entry2.encode();
317
318 if let (Bound::Included(start), Bound::Excluded(end)) = (&range.start, &range.end) {
319 assert!(encoded2.as_slice() < start.as_slice() || encoded2.as_slice() >= end.as_slice());
320 }
321 }
322
323 #[test]
324 fn test_key_prefix_range() {
325 let layout = IndexShape::new(
326 &[ValueType::Uint8, ValueType::Uint8],
327 &[SortDirection::Asc, SortDirection::Asc],
328 )
329 .unwrap();
330
331 let mut key = layout.allocate_key();
332 layout.set_u64(&mut key, 0, 100u64);
333 layout.set_row_number(&mut key, 1, 0u64);
334
335 let prefix = &key.as_slice()[..layout.fields[1].offset];
336 let range = IndexEntryKey::key_prefix_range(ObjectId::table(1), IndexId::primary(1), prefix);
337
338 layout.set_row_number(&mut key, 1, 999u64);
339 let entry = IndexEntryKey {
340 object: ObjectId::table(1),
341 index: IndexId::primary(1),
342 key: key.clone(),
343 };
344
345 let encoded = entry.encode();
346
347 if let (Bound::Included(start), Bound::Excluded(end)) = (&range.start, &range.end) {
348 assert!(encoded.as_slice() >= start.as_slice());
349 assert!(encoded.as_slice() < end.as_slice());
350 }
351
352 let mut key2 = layout.allocate_key();
353 layout.set_u64(&mut key2, 0, 200u64);
354 layout.set_row_number(&mut key2, 1, 1u64);
355
356 let entry2 = IndexEntryKey {
357 object: ObjectId::table(1),
358 index: IndexId::primary(1),
359 key: key2,
360 };
361
362 let encoded2 = entry2.encode();
363
364 if let Bound::Excluded(end) = &range.end {
365 assert!(encoded2.as_slice() >= end.as_slice());
366 }
367 }
368}