1use std::{
5 borrow::Borrow,
6 cmp::Ordering,
7 collections::{
8 Bound,
9 Bound::{Excluded, Included, Unbounded},
10 },
11 fmt,
12 hash::{Hash, Hasher},
13 iter, mem,
14 ops::{Deref, RangeBounds},
15 sync::Arc,
16};
17
18use reifydb_value::{
19 util::hash::Hash128,
20 value::{
21 Value,
22 blob::Blob,
23 date::Date,
24 datetime::DateTime,
25 decimal::Decimal,
26 duration::Duration,
27 identity::IdentityId,
28 int::Int,
29 row_number::RowNumber,
30 time::Time,
31 uint::Uint,
32 uuid::{Uuid4, Uuid7},
33 },
34};
35use serde::{
36 de::{Deserialize, Deserializer},
37 ser::{Serialize, Serializer},
38};
39
40use crate::key::serializer::KeySerializer;
41
42fn decode_binary(s: &str) -> Vec<u8> {
43 let mut buf = [0; 4];
44 let mut bytes = Vec::new();
45 for c in s.chars() {
46 match c as u32 {
47 b @ 0x80..=0xff => bytes.push(b as u8),
48 _ => bytes.extend(c.encode_utf8(&mut buf).as_bytes()),
49 }
50 }
51 bytes
52}
53
54pub const INLINE_CAP: usize = 39;
55
56#[derive(Clone)]
57pub enum EncodedKey {
58 Inline {
59 len: u8,
60 buf: [u8; INLINE_CAP],
61 },
62 Shared(Arc<[u8]>),
63}
64
65const _: () = assert!(mem::size_of::<EncodedKey>() == INLINE_CAP + 1 + mem::size_of::<usize>());
66
67impl EncodedKey {
68 pub fn new(key: impl AsRef<[u8]>) -> Self {
69 let key = key.as_ref();
70 if key.len() <= INLINE_CAP {
71 let mut buf = [0u8; INLINE_CAP];
72 buf[..key.len()].copy_from_slice(key);
73 EncodedKey::Inline {
74 len: key.len() as u8,
75 buf,
76 }
77 } else {
78 EncodedKey::Shared(Arc::from(key))
79 }
80 }
81
82 pub fn builder() -> EncodedKeyBuilder {
83 EncodedKeyBuilder::new()
84 }
85
86 pub fn as_bytes(&self) -> &[u8] {
87 self.as_slice()
88 }
89
90 pub fn as_slice(&self) -> &[u8] {
91 match self {
92 EncodedKey::Inline {
93 len,
94 buf,
95 } => &buf[..*len as usize],
96 EncodedKey::Shared(bytes) => bytes,
97 }
98 }
99
100 pub fn to_vec(&self) -> Vec<u8> {
101 self.as_slice().to_vec()
102 }
103
104 pub fn heap_bytes(&self) -> usize {
105 match self {
106 EncodedKey::Inline {
107 ..
108 } => 0,
109 EncodedKey::Shared(bytes) => bytes.len() + 2 * mem::size_of::<usize>(),
110 }
111 }
112}
113
114impl Deref for EncodedKey {
115 type Target = [u8];
116
117 fn deref(&self) -> &[u8] {
118 self.as_slice()
119 }
120}
121
122impl AsRef<[u8]> for EncodedKey {
123 fn as_ref(&self) -> &[u8] {
124 self.as_slice()
125 }
126}
127
128impl Borrow<[u8]> for EncodedKey {
129 fn borrow(&self) -> &[u8] {
130 self.as_slice()
131 }
132}
133
134impl PartialEq for EncodedKey {
135 fn eq(&self, other: &Self) -> bool {
136 self.as_slice() == other.as_slice()
137 }
138}
139
140impl Eq for EncodedKey {}
141
142impl PartialOrd for EncodedKey {
143 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
144 Some(self.cmp(other))
145 }
146}
147
148impl Ord for EncodedKey {
149 fn cmp(&self, other: &Self) -> Ordering {
150 self.as_slice().cmp(other.as_slice())
151 }
152}
153
154impl Hash for EncodedKey {
155 fn hash<H: Hasher>(&self, state: &mut H) {
156 self.as_slice().hash(state);
157 }
158}
159
160impl PartialEq<Vec<u8>> for EncodedKey {
161 fn eq(&self, other: &Vec<u8>) -> bool {
162 self.as_slice() == other.as_slice()
163 }
164}
165
166impl PartialEq<[u8]> for EncodedKey {
167 fn eq(&self, other: &[u8]) -> bool {
168 self.as_slice() == other
169 }
170}
171
172impl Serialize for EncodedKey {
173 fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
174 self.as_slice().serialize(serializer)
175 }
176}
177
178impl<'de> Deserialize<'de> for EncodedKey {
179 fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
180 let vec = Vec::<u8>::deserialize(deserializer)?;
181 Ok(EncodedKey::new(vec))
182 }
183}
184
185impl fmt::Debug for EncodedKey {
186 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187 write!(f, "EncodedKey({:02x?})", self.as_slice())
188 }
189}
190
191pub struct EncodedKeyBuilder {
192 serializer: KeySerializer,
193}
194
195impl EncodedKeyBuilder {
196 pub fn new() -> Self {
197 Self {
198 serializer: KeySerializer::new(),
199 }
200 }
201
202 pub fn with_capacity(capacity: usize) -> Self {
203 Self {
204 serializer: KeySerializer::with_capacity(capacity),
205 }
206 }
207
208 pub fn build(self) -> EncodedKey {
209 self.serializer.to_encoded_key()
210 }
211
212 pub fn bool(mut self, value: bool) -> Self {
213 self.serializer.extend_bool(value);
214 self
215 }
216
217 pub fn f32(mut self, value: f32) -> Self {
218 self.serializer.extend_f32(value);
219 self
220 }
221
222 pub fn f64(mut self, value: f64) -> Self {
223 self.serializer.extend_f64(value);
224 self
225 }
226
227 pub fn i8<T: Into<i8>>(mut self, value: T) -> Self {
228 self.serializer.extend_i8(value);
229 self
230 }
231
232 pub fn i16<T: Into<i16>>(mut self, value: T) -> Self {
233 self.serializer.extend_i16(value);
234 self
235 }
236
237 pub fn i32<T: Into<i32>>(mut self, value: T) -> Self {
238 self.serializer.extend_i32(value);
239 self
240 }
241
242 pub fn i64<T: Into<i64>>(mut self, value: T) -> Self {
243 self.serializer.extend_i64(value);
244 self
245 }
246
247 pub fn i128<T: Into<i128>>(mut self, value: T) -> Self {
248 self.serializer.extend_i128(value);
249 self
250 }
251
252 pub fn u8<T: Into<u8>>(mut self, value: T) -> Self {
253 self.serializer.extend_u8(value);
254 self
255 }
256
257 pub fn u16<T: Into<u16>>(mut self, value: T) -> Self {
258 self.serializer.extend_u16(value);
259 self
260 }
261
262 pub fn u32<T: Into<u32>>(mut self, value: T) -> Self {
263 self.serializer.extend_u32(value);
264 self
265 }
266
267 pub fn u64<T: Into<u64>>(mut self, value: T) -> Self {
268 self.serializer.extend_u64(value);
269 self
270 }
271
272 pub fn u128<T: Into<u128>>(mut self, value: T) -> Self {
273 self.serializer.extend_u128(value);
274 self
275 }
276
277 pub fn bytes<T: AsRef<[u8]>>(mut self, bytes: T) -> Self {
278 self.serializer.extend_bytes(bytes);
279 self
280 }
281
282 pub fn str<T: AsRef<str>>(mut self, s: T) -> Self {
283 self.serializer.extend_str(s);
284 self
285 }
286
287 pub fn serialize<T: Serialize>(mut self, value: &T) -> Self {
288 self.serializer.extend_serialize(value);
289 self
290 }
291
292 pub fn raw(mut self, bytes: &[u8]) -> Self {
293 self.serializer.extend_raw(bytes);
294 self
295 }
296
297 pub fn len(&self) -> usize {
298 self.serializer.len()
299 }
300
301 pub fn is_empty(&self) -> bool {
302 self.serializer.is_empty()
303 }
304
305 pub fn date(mut self, date: &Date) -> Self {
306 self.serializer.extend_date(date);
307 self
308 }
309
310 pub fn datetime(mut self, datetime: &DateTime) -> Self {
311 self.serializer.extend_datetime(datetime);
312 self
313 }
314
315 pub fn time(mut self, time: &Time) -> Self {
316 self.serializer.extend_time(time);
317 self
318 }
319
320 pub fn duration(mut self, duration: &Duration) -> Self {
321 self.serializer.extend_duration(duration);
322 self
323 }
324
325 pub fn row_number(mut self, row_number: &RowNumber) -> Self {
326 self.serializer.extend_row_number(row_number);
327 self
328 }
329
330 pub fn identity_id(mut self, id: &IdentityId) -> Self {
331 self.serializer.extend_identity_id(id);
332 self
333 }
334
335 pub fn uuid4(mut self, uuid: &Uuid4) -> Self {
336 self.serializer.extend_uuid4(uuid);
337 self
338 }
339
340 pub fn uuid7(mut self, uuid: &Uuid7) -> Self {
341 self.serializer.extend_uuid7(uuid);
342 self
343 }
344
345 pub fn blob(mut self, blob: &Blob) -> Self {
346 self.serializer.extend_blob(blob);
347 self
348 }
349
350 pub fn int(mut self, int: &Int) -> Self {
351 self.serializer.extend_int(int);
352 self
353 }
354
355 pub fn uint(mut self, uint: &Uint) -> Self {
356 self.serializer.extend_uint(uint);
357 self
358 }
359
360 pub fn decimal(mut self, decimal: &Decimal) -> Self {
361 self.serializer.extend_decimal(decimal);
362 self
363 }
364
365 pub fn value(mut self, value: &Value) -> Self {
366 self.serializer.extend_value(value);
367 self
368 }
369}
370
371impl Default for EncodedKeyBuilder {
372 fn default() -> Self {
373 Self::new()
374 }
375}
376
377pub trait IntoEncodedKey {
378 fn into_encoded_key(self) -> EncodedKey;
379}
380
381impl IntoEncodedKey for EncodedKey {
382 fn into_encoded_key(self) -> EncodedKey {
383 self
384 }
385}
386
387impl IntoEncodedKey for &str {
388 fn into_encoded_key(self) -> EncodedKey {
389 let mut serializer = KeySerializer::new();
390 serializer.extend_str(self);
391 serializer.to_encoded_key()
392 }
393}
394
395impl IntoEncodedKey for String {
396 fn into_encoded_key(self) -> EncodedKey {
397 let mut serializer = KeySerializer::new();
398 serializer.extend_str(&self);
399 serializer.to_encoded_key()
400 }
401}
402
403impl IntoEncodedKey for Vec<u8> {
404 fn into_encoded_key(self) -> EncodedKey {
405 let mut serializer = KeySerializer::new();
406 serializer.extend_bytes(&self);
407 serializer.to_encoded_key()
408 }
409}
410
411impl IntoEncodedKey for &[u8] {
412 fn into_encoded_key(self) -> EncodedKey {
413 let mut serializer = KeySerializer::new();
414 serializer.extend_bytes(self);
415 serializer.to_encoded_key()
416 }
417}
418
419impl IntoEncodedKey for u64 {
420 fn into_encoded_key(self) -> EncodedKey {
421 let mut serializer = KeySerializer::with_capacity(8);
422 serializer.extend_u64(self);
423 serializer.to_encoded_key()
424 }
425}
426
427impl IntoEncodedKey for i64 {
428 fn into_encoded_key(self) -> EncodedKey {
429 let mut serializer = KeySerializer::with_capacity(8);
430 serializer.extend_i64(self);
431 serializer.to_encoded_key()
432 }
433}
434
435impl IntoEncodedKey for u32 {
436 fn into_encoded_key(self) -> EncodedKey {
437 let mut serializer = KeySerializer::with_capacity(4);
438 serializer.extend_u32(self);
439 serializer.to_encoded_key()
440 }
441}
442
443impl IntoEncodedKey for i32 {
444 fn into_encoded_key(self) -> EncodedKey {
445 let mut serializer = KeySerializer::with_capacity(4);
446 serializer.extend_i32(self);
447 serializer.to_encoded_key()
448 }
449}
450
451impl IntoEncodedKey for u16 {
452 fn into_encoded_key(self) -> EncodedKey {
453 let mut serializer = KeySerializer::with_capacity(2);
454 serializer.extend_u16(self);
455 serializer.to_encoded_key()
456 }
457}
458
459impl IntoEncodedKey for i16 {
460 fn into_encoded_key(self) -> EncodedKey {
461 let mut serializer = KeySerializer::with_capacity(2);
462 serializer.extend_i16(self);
463 serializer.to_encoded_key()
464 }
465}
466
467impl IntoEncodedKey for u8 {
468 fn into_encoded_key(self) -> EncodedKey {
469 let mut serializer = KeySerializer::with_capacity(1);
470 serializer.extend_u8(self);
471 serializer.to_encoded_key()
472 }
473}
474
475impl IntoEncodedKey for i8 {
476 fn into_encoded_key(self) -> EncodedKey {
477 let mut serializer = KeySerializer::with_capacity(1);
478 serializer.extend_i8(self);
479 serializer.to_encoded_key()
480 }
481}
482
483impl IntoEncodedKey for &u64 {
484 fn into_encoded_key(self) -> EncodedKey {
485 (*self).into_encoded_key()
486 }
487}
488
489impl IntoEncodedKey for &i64 {
490 fn into_encoded_key(self) -> EncodedKey {
491 (*self).into_encoded_key()
492 }
493}
494
495impl IntoEncodedKey for &u32 {
496 fn into_encoded_key(self) -> EncodedKey {
497 (*self).into_encoded_key()
498 }
499}
500
501impl IntoEncodedKey for &i32 {
502 fn into_encoded_key(self) -> EncodedKey {
503 (*self).into_encoded_key()
504 }
505}
506
507impl IntoEncodedKey for &u16 {
508 fn into_encoded_key(self) -> EncodedKey {
509 (*self).into_encoded_key()
510 }
511}
512
513impl IntoEncodedKey for &i16 {
514 fn into_encoded_key(self) -> EncodedKey {
515 (*self).into_encoded_key()
516 }
517}
518
519impl IntoEncodedKey for &u8 {
520 fn into_encoded_key(self) -> EncodedKey {
521 (*self).into_encoded_key()
522 }
523}
524
525impl IntoEncodedKey for &i8 {
526 fn into_encoded_key(self) -> EncodedKey {
527 (*self).into_encoded_key()
528 }
529}
530
531impl IntoEncodedKey for &Hash128 {
532 fn into_encoded_key(self) -> EncodedKey {
533 let mut serializer = KeySerializer::with_capacity(16);
534 serializer.extend_u128(self.0);
535 serializer.to_encoded_key()
536 }
537}
538
539impl IntoEncodedKey for RowNumber {
540 fn into_encoded_key(self) -> EncodedKey {
541 self.0.into_encoded_key()
542 }
543}
544
545impl IntoEncodedKey for &RowNumber {
546 fn into_encoded_key(self) -> EncodedKey {
547 self.0.into_encoded_key()
548 }
549}
550
551impl IntoEncodedKey for Value {
552 fn into_encoded_key(self) -> EncodedKey {
553 let mut serializer = KeySerializer::new();
554 serializer.extend_value(&self);
555 serializer.to_encoded_key()
556 }
557}
558
559impl IntoEncodedKey for &Value {
560 fn into_encoded_key(self) -> EncodedKey {
561 let mut serializer = KeySerializer::new();
562 serializer.extend_value(self);
563 serializer.to_encoded_key()
564 }
565}
566
567impl IntoEncodedKey for Vec<Value> {
568 fn into_encoded_key(self) -> EncodedKey {
569 let mut serializer = KeySerializer::new();
570 for value in self.iter() {
571 serializer.extend_value(value);
572 }
573 serializer.to_encoded_key()
574 }
575}
576
577impl IntoEncodedKey for &[Value] {
578 fn into_encoded_key(self) -> EncodedKey {
579 let mut serializer = KeySerializer::new();
580 for value in self.iter() {
581 serializer.extend_value(value);
582 }
583 serializer.to_encoded_key()
584 }
585}
586
587impl IntoEncodedKey for (String, String) {
588 fn into_encoded_key(self) -> EncodedKey {
589 let mut serializer = KeySerializer::new();
590 serializer.extend_str(&self.0);
591 serializer.extend_str(&self.1);
592 serializer.to_encoded_key()
593 }
594}
595
596impl IntoEncodedKey for &(String, String) {
597 fn into_encoded_key(self) -> EncodedKey {
598 let mut serializer = KeySerializer::new();
599 serializer.extend_str(&self.0);
600 serializer.extend_str(&self.1);
601 serializer.to_encoded_key()
602 }
603}
604
605impl IntoEncodedKey for (&str, &str) {
606 fn into_encoded_key(self) -> EncodedKey {
607 let mut serializer = KeySerializer::new();
608 serializer.extend_str(self.0);
609 serializer.extend_str(self.1);
610 serializer.to_encoded_key()
611 }
612}
613
614impl IntoEncodedKey for (Arc<str>, Arc<str>) {
615 fn into_encoded_key(self) -> EncodedKey {
616 let mut serializer = KeySerializer::new();
617 serializer.extend_str(&self.0);
618 serializer.extend_str(&self.1);
619 serializer.to_encoded_key()
620 }
621}
622
623impl IntoEncodedKey for &(Arc<str>, Arc<str>) {
624 fn into_encoded_key(self) -> EncodedKey {
625 let mut serializer = KeySerializer::new();
626 serializer.extend_str(&self.0);
627 serializer.extend_str(&self.1);
628 serializer.to_encoded_key()
629 }
630}
631
632impl IntoEncodedKey for (String, String, String) {
633 fn into_encoded_key(self) -> EncodedKey {
634 let mut serializer = KeySerializer::new();
635 serializer.extend_str(&self.0);
636 serializer.extend_str(&self.1);
637 serializer.extend_str(&self.2);
638 serializer.to_encoded_key()
639 }
640}
641
642impl IntoEncodedKey for &(String, String, String) {
643 fn into_encoded_key(self) -> EncodedKey {
644 let mut serializer = KeySerializer::new();
645 serializer.extend_str(&self.0);
646 serializer.extend_str(&self.1);
647 serializer.extend_str(&self.2);
648 serializer.to_encoded_key()
649 }
650}
651
652impl IntoEncodedKey for &String {
653 fn into_encoded_key(self) -> EncodedKey {
654 let mut serializer = KeySerializer::new();
655 serializer.extend_str(self);
656 serializer.to_encoded_key()
657 }
658}
659
660#[derive(Clone, Debug)]
661pub struct EncodedKeyRange {
662 pub start: Bound<EncodedKey>,
663 pub end: Bound<EncodedKey>,
664}
665
666impl EncodedKeyRange {
667 pub fn new(start: Bound<EncodedKey>, end: Bound<EncodedKey>) -> Self {
668 Self {
669 start,
670 end,
671 }
672 }
673
674 pub fn prefix(prefix: &[u8]) -> Self {
675 let start = Bound::Included(EncodedKey::new(prefix));
676 let end = match prefix.iter().rposition(|&b| b != 0xff) {
677 Some(i) => Bound::Excluded(EncodedKey::new(
678 prefix.iter().take(i).copied().chain(iter::once(prefix[i] + 1)).collect::<Vec<_>>(),
679 )),
680 None => Bound::Unbounded,
681 };
682 Self {
683 start,
684 end,
685 }
686 }
687
688 pub fn with_prefix(&self, prefix: EncodedKey) -> Self {
689 let start = match self.start_bound() {
690 Included(key) => {
691 let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
692 prefixed.extend_from_slice(prefix.as_ref());
693 prefixed.extend_from_slice(key.as_ref());
694 Included(EncodedKey::new(prefixed))
695 }
696 Excluded(key) => {
697 let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
698 prefixed.extend_from_slice(prefix.as_ref());
699 prefixed.extend_from_slice(key.as_ref());
700 Excluded(EncodedKey::new(prefixed))
701 }
702 Unbounded => Included(prefix.clone()),
703 };
704
705 let end = match self.end_bound() {
706 Included(key) => {
707 let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
708 prefixed.extend_from_slice(prefix.as_ref());
709 prefixed.extend_from_slice(key.as_ref());
710 Included(EncodedKey::new(prefixed))
711 }
712 Excluded(key) => {
713 let mut prefixed = Vec::with_capacity(prefix.len() + key.len());
714 prefixed.extend_from_slice(prefix.as_ref());
715 prefixed.extend_from_slice(key.as_ref());
716 Excluded(EncodedKey::new(prefixed))
717 }
718 Unbounded => match prefix.as_ref().iter().rposition(|&b| b != 0xff) {
719 Some(i) => {
720 let mut next_prefix = prefix.as_ref()[..=i].to_vec();
721 next_prefix[i] += 1;
722 Excluded(EncodedKey::new(next_prefix))
723 }
724 None => Unbounded,
725 },
726 };
727
728 EncodedKeyRange::new(start, end)
729 }
730
731 pub fn start_end(start: Option<EncodedKey>, end: Option<EncodedKey>) -> Self {
732 let start = match start {
733 Some(s) => Bound::Included(s),
734 None => Bound::Unbounded,
735 };
736
737 let end = match end {
738 Some(e) => Bound::Included(e),
739 None => Bound::Unbounded,
740 };
741
742 Self {
743 start,
744 end,
745 }
746 }
747
748 pub fn all() -> Self {
749 Self {
750 start: Bound::Unbounded,
751 end: Bound::Unbounded,
752 }
753 }
754
755 pub fn parse(str: &str) -> Self {
756 let (mut start, mut end) = (Bound::<EncodedKey>::Unbounded, Bound::<EncodedKey>::Unbounded);
757
758 if let Some(dot_pos) = str.find("..") {
759 let start_part = &str[..dot_pos];
760 let end_part = &str[dot_pos + 2..];
761
762 if !start_part.is_empty() {
763 start = Bound::Included(EncodedKey::new(decode_binary(start_part)));
764 }
765
766 if let Some(end_str) = end_part.strip_prefix('=') {
767 if !end_str.is_empty() {
768 end = Bound::Included(EncodedKey::new(decode_binary(end_str)));
769 }
770 } else if !end_part.is_empty() {
771 end = Bound::Excluded(EncodedKey::new(decode_binary(end_part)));
772 }
773
774 Self {
775 start,
776 end,
777 }
778 } else {
779 Self {
780 start: Bound::Included(EncodedKey::new([0xff])),
781 end: Bound::Excluded(EncodedKey::new([0xff])),
782 }
783 }
784 }
785}
786
787impl RangeBounds<EncodedKey> for EncodedKeyRange {
788 fn start_bound(&self) -> Bound<&EncodedKey> {
789 self.start.as_ref()
790 }
791
792 fn end_bound(&self) -> Bound<&EncodedKey> {
793 self.end.as_ref()
794 }
795}