kvstructs/
key.rs

1use crate::key_mut::KeyMut;
2use crate::raw_key_pointer::RawKeyPointer;
3use crate::{u64_big_endian, TIMESTAMP_SIZE};
4use alloc::borrow::Cow;
5use alloc::boxed::Box;
6use alloc::string::String;
7use alloc::vec::Vec;
8use bytes::{Buf, Bytes, BytesMut};
9use core::cmp::Ordering;
10use core::hash::{Hash, Hasher};
11use core::ops::{Deref, DerefMut};
12use core::slice::from_raw_parts;
13#[cfg(feature = "std")]
14use std::time::{SystemTime, UNIX_EPOCH};
15
16/// A general Key for key-value storage, the underlying is u8 slice.
17#[derive(Debug, Clone)]
18#[repr(transparent)]
19pub struct Key {
20    data: Bytes,
21}
22
23impl Default for Key {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl Key {
30    /// Returns a empty key
31    #[inline]
32    pub const fn new() -> Self {
33        Self { data: Bytes::new() }
34    }
35
36    /// Returns a Key with data and timestamp.
37    #[inline]
38    pub fn from_with_timestamp(data: Vec<u8>, ts: u64) -> Self {
39        Self::from(data).with_timestamp(ts)
40    }
41
42    /// Returns a Key with data and system time as timestamp.
43    #[cfg(feature = "std")]
44    #[inline]
45    pub fn from_with_system_time(data: Vec<u8>, st: SystemTime) -> Self {
46        Self::from(data).with_system_time(st)
47    }
48
49    /// Returns a Key with data and the current time as timestamp
50    #[cfg(feature = "std")]
51    #[inline]
52    pub fn from_with_now(data: Vec<u8>) -> Self {
53        Self::from(data).with_now()
54    }
55
56    /// Returns a Key by copying the slice data.
57    #[inline]
58    pub fn copy_from_slice(data: &[u8]) -> Self {
59        Bytes::copy_from_slice(data).into()
60    }
61
62    /// Generates a new key by appending timestamp to key.
63    #[inline]
64    pub fn with_timestamp(self, ts: u64) -> Self {
65        let len = self.data.len() + TIMESTAMP_SIZE;
66        let ts = Bytes::from(Box::from((u64::MAX - ts).to_be_bytes()));
67        self.data.chain(ts).copy_to_bytes(len).into()
68    }
69
70    /// Generates a new key by appending the given UNIX system time to key.
71    #[inline]
72    #[cfg(feature = "std")]
73    pub fn with_system_time(self, st: SystemTime) -> Self {
74        let len = self.data.len() + TIMESTAMP_SIZE;
75        let ts = Bytes::from(Box::from(
76            st.duration_since(UNIX_EPOCH)
77                .unwrap()
78                .as_secs()
79                .to_be_bytes(),
80        ));
81        self.data.chain(ts).copy_to_bytes(len).into()
82    }
83
84    /// Generates a new key by appending the current UNIX system time to key.
85    #[inline]
86    #[cfg(feature = "std")]
87    pub fn with_now(self) -> Self {
88        let len = self.data.len() + TIMESTAMP_SIZE;
89        let ts = Bytes::from(Box::from(
90            SystemTime::now()
91                .duration_since(UNIX_EPOCH)
92                .unwrap()
93                .as_secs()
94                .to_be_bytes(),
95        ));
96        self.data.chain(ts).copy_to_bytes(len).into()
97    }
98
99    /// Returns a new Key without timestamp.
100    #[inline]
101    pub fn parse_new_key(&self) -> Self {
102        let sz = self.len();
103        match sz.checked_sub(TIMESTAMP_SIZE) {
104            None => Self {
105                data: self.data.clone(),
106            },
107            Some(sz) => Self {
108                data: self.data.slice(..sz),
109            },
110        }
111    }
112
113    /// Returns the number of bytes contained in this Key.
114    #[inline]
115    pub fn len(&self) -> usize {
116        self.data.len()
117    }
118
119    /// Returns true if the Key has a length of 0.
120    #[inline]
121    pub fn is_empty(&self) -> bool {
122        self.data.is_empty()
123    }
124
125    /// Returns the underlying bytes
126    #[inline]
127    pub fn as_slice(&self) -> &[u8] {
128        self.data.as_ref()
129    }
130
131    /// Remove the timestamp(if exists) from the key
132    pub fn truncate_timestamp(&mut self) {
133        if let Some(sz) = self.data.len().checked_sub(TIMESTAMP_SIZE) {
134            self.data.truncate(sz)
135        }
136    }
137}
138
139impl From<Key> for Bytes {
140    #[inline]
141    fn from(key: Key) -> Self {
142        key.data
143    }
144}
145
146impl<'a, K: KeyExt> PartialEq<K> for KeyRef<'a> {
147    fn eq(&self, other: &K) -> bool {
148        same_key_in(self.as_bytes(), other.as_bytes())
149    }
150}
151
152impl<K: KeyExt> PartialEq<K> for Key {
153    fn eq(&self, other: &K) -> bool {
154        same_key_in(self.as_bytes(), other.as_bytes())
155    }
156}
157
158impl<K: KeyExt> PartialOrd<K> for Key {
159    fn partial_cmp(&self, other: &K) -> Option<Ordering> {
160        Some(compare_key_in(self.as_bytes(), other.as_bytes()))
161    }
162}
163
164impl<'a, K: KeyExt> PartialOrd<K> for KeyRef<'a> {
165    fn partial_cmp(&self, other: &K) -> Option<Ordering> {
166        Some(compare_key_in(self.as_bytes(), other.as_bytes()))
167    }
168}
169
170impl Eq for Key {}
171
172impl Hash for Key {
173    fn hash<H: Hasher>(&self, state: &mut H) {
174        self.data.hash(state)
175    }
176}
177
178impl Ord for Key {
179    /// Checks the key without timestamp and checks the timestamp if keyNoTs
180    /// is same.
181    /// a<timestamp> would be sorted higher than aa<timestamp> if we use bytes.compare
182    /// All keys should have timestamp.
183    fn cmp(&self, other: &Self) -> Ordering {
184        compare_key_in(self.data.as_ref(), other.data.as_ref())
185    }
186}
187
188#[inline(always)]
189pub(crate) fn compare_key_in(me: &[u8], other: &[u8]) -> Ordering {
190    let sb = me.len().saturating_sub(TIMESTAMP_SIZE);
191    let ob = other.len().saturating_sub(TIMESTAMP_SIZE);
192    let (s_key_part, s_ts_part) = me.split_at(sb);
193    let (o_key_part, o_ts_part) = other.split_at(ob);
194
195    match s_key_part.cmp(o_key_part) {
196        Ordering::Less => Ordering::Less,
197        Ordering::Equal => s_ts_part.cmp(o_ts_part),
198        Ordering::Greater => Ordering::Greater,
199    }
200}
201
202/// Checks the key without timestamp and checks the timestamp if keyNoTs
203/// is same.
204/// a<timestamp> would be sorted higher than aa<timestamp> if we use bytes.compare
205/// All keys should have timestamp.
206#[inline(always)]
207pub fn compare_key(a: impl KeyExt, b: impl KeyExt) -> Ordering {
208    let me = a.as_bytes();
209    let other = b.as_bytes();
210    compare_key_in(me, other)
211}
212
213#[inline(always)]
214pub(crate) fn same_key_in(me: &[u8], other: &[u8]) -> bool {
215    let sl = me.len();
216    let ol = other.len();
217    if sl != ol {
218        false
219    } else {
220        let s = match sl.checked_sub(TIMESTAMP_SIZE) {
221            None => me,
222            Some(sz) => me[..sz].as_ref(),
223        };
224        let o = match ol.checked_sub(TIMESTAMP_SIZE) {
225            None => me,
226            Some(sz) => other[..sz].as_ref(),
227        };
228        s.eq(o)
229    }
230}
231
232/// Checks for key equality ignoring the version timestamp.
233#[inline(always)]
234pub fn same_key(a: impl KeyExt, b: impl KeyExt) -> bool {
235    let me = a.as_bytes();
236    let other = b.as_bytes();
237    same_key_in(me, other)
238}
239
240impl<const N: usize> From<[u8; N]> for Key {
241    fn from(data: [u8; N]) -> Self {
242        Self {
243            data: Bytes::from(data.to_vec()),
244        }
245    }
246}
247
248macro_rules! impl_from_for_key {
249    ($($ty: ty), +$(,)?) => {
250        $(
251        impl From<$ty> for Key {
252            fn from(val: $ty) -> Self {
253                Self {
254                    data: Bytes::from(val),
255                }
256            }
257        }
258        )*
259    };
260}
261
262impl_from_for_key! {
263    String,
264    &'static str,
265    Vec<u8>,
266    Box<[u8]>,
267}
268
269impl Deref for Key {
270    type Target = Bytes;
271
272    fn deref(&self) -> &Self::Target {
273        &self.data
274    }
275}
276
277impl DerefMut for Key {
278    fn deref_mut(&mut self) -> &mut Self::Target {
279        &mut self.data
280    }
281}
282
283impl From<Bytes> for Key {
284    fn from(data: Bytes) -> Self {
285        Self { data }
286    }
287}
288
289impl From<BytesMut> for Key {
290    fn from(data: BytesMut) -> Self {
291        Self {
292            data: data.freeze(),
293        }
294    }
295}
296
297impl From<&[u8]> for Key {
298    fn from(data: &[u8]) -> Self {
299        Key::copy_from_slice(data)
300    }
301}
302
303impl AsRef<[u8]> for Key {
304    fn as_ref(&self) -> &[u8] {
305        self.data.as_ref()
306    }
307}
308
309/// KeyRef can only contains a underlying u8 slice of Key
310#[derive(Debug, Copy, Clone)]
311#[repr(transparent)]
312pub struct KeyRef<'a> {
313    data: &'a [u8],
314}
315
316impl<'a> Eq for KeyRef<'a> {}
317
318impl<'a> Ord for KeyRef<'a> {
319    /// Checks the key without timestamp and checks the timestamp if keyNoTs
320    /// is same.
321    /// a<timestamp> would be sorted higher than aa<timestamp> if we use bytes.compare
322    /// All keys should have timestamp.
323    fn cmp(&self, other: &Self) -> Ordering {
324        compare_key(self, other)
325    }
326}
327
328impl<'a> From<&'a [u8]> for KeyRef<'a> {
329    fn from(data: &'a [u8]) -> Self {
330        Self { data }
331    }
332}
333
334impl<'a> KeyRef<'a> {
335    /// Returns a KeyRef from byte slice
336    #[inline]
337    pub const fn new(data: &'a [u8]) -> Self {
338        Self { data }
339    }
340
341    /// Returns a KeyRef from [`RawKeyPointer`]
342    ///
343    /// # Safety
344    /// The inner raw pointer of [`RawKeyPointer`] must be valid.
345    ///
346    /// [`RawKeyPointer`]: struct.RawKeyPointer.html
347    #[inline]
348    pub unsafe fn from_raw_key_pointer(rp: RawKeyPointer) -> Self {
349        Self {
350            data: from_raw_parts(rp.as_ptr(), rp.len()),
351        }
352    }
353
354    /// Returns a KeyRef from raw pointer and length
355    ///
356    /// # Safety
357    /// The raw pointer must be valid.
358    #[inline]
359    pub unsafe fn from_raw_pointer(ptr: *const u8, len: usize) -> Self {
360        Self {
361            data: from_raw_parts(ptr, len),
362        }
363    }
364
365    /// Copy KeyRef to a new Key.
366    #[inline]
367    pub fn to_key(&self) -> Key {
368        Key::copy_from_slice(self.data)
369    }
370
371    /// Returns the number of bytes contained in this Key.
372    #[inline]
373    pub fn len(&self) -> usize {
374        self.data.len()
375    }
376
377    /// Returns true if the Key has a length of 0.
378    #[inline]
379    pub fn is_empty(&self) -> bool {
380        self.data.is_empty()
381    }
382
383    /// Returns the underlying bytes
384    #[inline]
385    pub fn as_slice(&self) -> &[u8] {
386        self.data
387    }
388}
389
390impl KeyExt for &'_ KeyRef<'_> {
391    #[inline]
392    fn as_bytes(&self) -> &[u8] {
393        self.data
394    }
395}
396
397impl KeyExt for &'_ mut KeyRef<'_> {
398    #[inline]
399    fn as_bytes(&self) -> &[u8] {
400        self.data
401    }
402}
403
404impl KeyExt for KeyRef<'_> {
405    #[inline]
406    fn as_bytes(&self) -> &[u8] {
407        self.data
408    }
409}
410
411impl Hash for KeyRef<'_> {
412    fn hash<H: Hasher>(&self, state: &mut H) {
413        self.data.hash(state)
414    }
415}
416
417/// Extensions for Key
418pub trait KeyExt {
419    /// Returns raw pointer of the underlying byte slice
420    #[inline]
421    fn as_ptr(&self) -> *const u8 {
422        self.as_bytes().as_ptr()
423    }
424
425    /// Returns a KeyRef.
426    #[inline]
427    fn as_key_ref(&self) -> KeyRef {
428        KeyRef {
429            data: self.as_bytes(),
430        }
431    }
432
433    /// Returns the underlying slice of key (with timestamp data).
434    fn as_bytes(&self) -> &[u8];
435
436    /// Parses the actual key from the key bytes.
437    #[inline]
438    fn parse_key(&self) -> &[u8] {
439        let data = self.as_bytes();
440        let sz = data.len();
441        match sz.checked_sub(TIMESTAMP_SIZE) {
442            None => data,
443            Some(sz) => data[..sz].as_ref(),
444        }
445    }
446
447    /// Parses the timestamp from the key bytes.
448    ///
449    /// # Panics
450    /// If the length of key less than 8.
451    #[inline]
452    fn parse_timestamp(&self) -> u64 {
453        let data = self.as_bytes();
454        let data_len = data.len();
455        if data_len <= TIMESTAMP_SIZE {
456            0
457        } else {
458            u64::MAX - u64_big_endian(&data[data_len - TIMESTAMP_SIZE..])
459        }
460    }
461
462    /// Checks for key equality ignoring the version timestamp.
463    #[inline]
464    fn same_key(&self, other: impl KeyExt) -> bool {
465        let me = self.as_bytes();
466        let other = other.as_bytes();
467        same_key_in(me, other)
468    }
469
470    /// Checks the key without timestamp and checks the timestamp if keyNoTs
471    /// is same.
472    /// a<timestamp> would be sorted higher than aa<timestamp> if we use bytes.compare
473    /// All keys should have timestamp.
474    #[inline]
475    fn compare_key(&self, other: impl KeyExt) -> Ordering {
476        let me = self.as_bytes();
477        let other = other.as_bytes();
478        compare_key_in(me, other)
479    }
480
481    impl_psfix_suites!(KeyExt::parse_key, u8, "u8");
482}
483
484macro_rules! impl_partial_eq_ord {
485    ($($ty:ty), +$(,)?) => {
486        $(
487        impl PartialEq<Key> for $ty {
488            fn eq(&self, other: &Key) -> bool {
489                other.same_key(self)
490            }
491        }
492
493        impl<'a> PartialEq<KeyRef<'a>> for $ty {
494            fn eq(&self, other: &KeyRef<'a>) -> bool {
495                other.same_key(self)
496            }
497        }
498
499        // impl<'a> PartialEq<$ty> for KeyRef<'a> {
500        //     fn eq(&self, other: &$ty) -> bool {
501        //         self.same_key(other)
502        //     }
503        // }
504
505        impl PartialOrd<Key> for $ty {
506            fn partial_cmp(&self, other: &Key) -> Option<Ordering> {
507                Some(compare_key(other, self))
508            }
509        }
510
511        impl<'a> PartialOrd<KeyRef<'a>> for $ty {
512            fn partial_cmp(&self, other: &KeyRef<'a>) -> Option<Ordering> {
513                Some(compare_key(other, self))
514            }
515        }
516        )*
517    };
518}
519
520macro_rules! impl_key_ext {
521    ($($ty:tt::$conv:tt), +$(,)?) => {
522        $(
523        impl KeyExt for $ty {
524            #[inline]
525            fn as_bytes(&self) -> &[u8] {
526                $ty::$conv(self)
527            }
528        }
529
530        impl<'a> KeyExt for &'a $ty {
531            #[inline]
532            fn as_bytes(&self) -> &[u8] {
533                $ty::$conv(self)
534            }
535        }
536
537        impl<'a> KeyExt for &'a mut $ty {
538            #[inline]
539            fn as_bytes(&self) -> &[u8] {
540                $ty::$conv(self)
541            }
542        }
543        )*
544    };
545}
546
547type VecBytes = Vec<u8>;
548type U8Bytes = [u8];
549type BoxBytes = Box<[u8]>;
550
551impl_partial_eq_ord! {
552    Bytes,
553    BytesMut,
554    BoxBytes,
555    KeyMut,
556    U8Bytes,
557    VecBytes,
558    str,
559    String,
560}
561
562impl_key_ext! {
563    Bytes::as_ref,
564    BytesMut::as_ref,
565    BoxBytes::as_ref,
566    Key::as_ref,
567    U8Bytes::as_ref,
568    VecBytes::as_slice,
569    str::as_bytes,
570    String::as_bytes,
571}
572
573impl<const N: usize> PartialEq<Key> for [u8; N] {
574    fn eq(&self, other: &Key) -> bool {
575        other.same_key(self)
576    }
577}
578
579// impl<const N: usize> PartialEq<[u8; N]> for Key {
580//     fn eq(&self, other: &[u8; N]) -> bool {
581//         self.same_key(other)
582//     }
583// }
584
585impl<const N: usize> KeyExt for [u8; N] {
586    #[inline]
587    fn as_bytes(&self) -> &[u8] {
588        self
589    }
590}
591
592impl<'a, const N: usize> KeyExt for &'a [u8; N] {
593    #[inline]
594    fn as_bytes(&self) -> &[u8] {
595        self.as_slice()
596    }
597}
598
599impl<'a, const N: usize> KeyExt for &'a mut [u8; N] {
600    #[inline]
601    fn as_bytes(&self) -> &[u8] {
602        self.as_slice()
603    }
604}
605
606#[cfg(test)]
607mod test {
608    use super::*;
609    use alloc::vec;
610
611    #[test]
612    fn key_integration_test() {
613        // test key_with_ts
614        let key = vec![0, 1, 2, 3, 4, 5, 6, 7];
615        let nk = Key::from(key.clone()).with_timestamp(10);
616        assert_eq!(
617            vec![0, 1, 2, 3, 4, 5, 6, 7, 255, 255, 255, 255, 255, 255, 255, 245],
618            nk
619        );
620
621        // test parse_ts
622        assert_eq!(nk.parse_timestamp(), 10);
623
624        // test parse_key
625        let nk2 = Key::from(key).with_timestamp(1000);
626        assert_eq!(nk.parse_key(), nk2.parse_key());
627
628        // test cmp
629        assert!(nk.cmp(&nk2).is_gt());
630
631        // test same key
632        assert_eq!(nk, nk2);
633    }
634}
kvstructs/key.rs

kvstructs/
key.rs
