Skip to main content

tikv_client/kv/
key.rs

1// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2
3use std::fmt;
4use std::ops::Bound;
5
6#[allow(unused_imports)]
7#[cfg(test)]
8use proptest::arbitrary::any_with;
9#[allow(unused_imports)]
10#[cfg(test)]
11use proptest::collection::size_range;
12#[cfg(test)]
13use proptest_derive::Arbitrary;
14
15use super::HexRepr;
16use crate::kv::codec::BytesEncoder;
17use crate::kv::codec::{self};
18use crate::proto::kvrpcpb;
19use crate::proto::kvrpcpb::KvPair;
20
21const _PROPTEST_KEY_MAX: usize = 1024 * 2; // 2 KB
22
23/// The key part of a key/value pair.
24///
25/// In TiKV, keys are an ordered sequence of bytes. This has an advantage over choosing `String` as
26/// valid `UTF-8` is not required. This means that the user is permitted to store any data they wish,
27/// as long as it can be represented by bytes. (Which is to say, pretty much anything!)
28///
29/// This type wraps around an owned value, so it should be treated it like `String` or `Vec<u8>`.
30///
31/// # Examples
32/// ```rust
33/// use tikv_client::Key;
34///
35/// let static_str: &'static str = "TiKV";
36/// let from_static_str = Key::from(static_str.to_owned());
37///
38/// let string: String = String::from(static_str);
39/// let from_string = Key::from(string);
40/// assert_eq!(from_static_str, from_string);
41///
42/// let vec: Vec<u8> = static_str.as_bytes().to_vec();
43/// let from_vec = Key::from(vec);
44/// assert_eq!(from_static_str, from_vec);
45///
46/// let bytes = static_str.as_bytes().to_vec();
47/// let from_bytes = Key::from(bytes);
48/// assert_eq!(from_static_str, from_bytes);
49/// ```
50///
51/// While `.into()` is usually sufficient for obtaining the buffer itself, sometimes type inference
52/// isn't able to determine the correct type. Notably in the `assert_eq!()` and `==` cases. In
53/// these cases using the fully-qualified-syntax is useful:
54///
55/// # Examples
56/// ```rust
57/// use tikv_client::Key;
58///
59/// let buf = "TiKV".as_bytes().to_owned();
60/// let key = Key::from(buf.clone());
61/// assert_eq!(Into::<Vec<u8>>::into(key), buf);
62/// ```
63///
64/// Many functions which accept a `Key` accept an `Into<Key>`, which means all of the above types
65/// can be passed directly to those functions.
66#[derive(Default, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
67#[cfg_attr(test, derive(Arbitrary))]
68#[repr(transparent)]
69pub struct Key(
70    #[cfg_attr(
71        test,
72        proptest(strategy = "any_with::<Vec<u8>>((size_range(_PROPTEST_KEY_MAX), ()))")
73    )]
74    pub(crate) Vec<u8>,
75);
76
77impl AsRef<Key> for kvrpcpb::Mutation {
78    fn as_ref(&self) -> &Key {
79        self.key.as_ref()
80    }
81}
82
83pub struct KvPairTTL(pub KvPair, pub u64);
84
85impl AsRef<Key> for KvPairTTL {
86    fn as_ref(&self) -> &Key {
87        self.0.key.as_ref()
88    }
89}
90
91impl From<KvPairTTL> for (KvPair, u64) {
92    fn from(value: KvPairTTL) -> Self {
93        (value.0, value.1)
94    }
95}
96
97impl Key {
98    /// The empty key.
99    pub const EMPTY: Self = Key(Vec::new());
100
101    /// Return whether the key is empty.
102    #[inline]
103    pub fn is_empty(&self) -> bool {
104        self.0.is_empty()
105    }
106
107    /// Return whether the last byte of key is 0.
108    #[inline]
109    pub(super) fn zero_terminated(&self) -> bool {
110        self.0.last().map(|i| *i == 0).unwrap_or(false)
111    }
112
113    /// Push a zero to the end of the key.
114    ///
115    /// Extending a zero makes the new key the smallest key that is greater than than the original one.
116    #[inline]
117    pub(crate) fn next_key(mut self) -> Self {
118        self.0.push(0);
119        self
120    }
121
122    /// Convert the key to a lower bound. The key is treated as inclusive.
123    #[inline]
124    pub(super) fn into_lower_bound(mut self) -> Bound<Key> {
125        if self.zero_terminated() {
126            self.0.pop().unwrap();
127            Bound::Excluded(self)
128        } else {
129            Bound::Included(self)
130        }
131    }
132
133    /// Convert the key to an upper bound. The key is treated as exclusive.
134    #[inline]
135    pub(super) fn into_upper_bound(mut self) -> Bound<Key> {
136        if self.zero_terminated() {
137            self.0.pop().unwrap();
138            Bound::Included(self)
139        } else {
140            Bound::Excluded(self)
141        }
142    }
143
144    /// Return the MVCC-encoded representation of the key.
145    #[inline]
146    #[must_use]
147    pub fn to_encoded(&self) -> Key {
148        let len = codec::max_encoded_bytes_size(self.0.len());
149        let mut encoded = Vec::with_capacity(len);
150        encoded.encode_bytes(&self.0, false).unwrap();
151        Key(encoded)
152    }
153
154    pub fn len(&self) -> usize {
155        self.0.len()
156    }
157}
158
159impl From<Vec<u8>> for Key {
160    fn from(v: Vec<u8>) -> Self {
161        Key(v)
162    }
163}
164
165impl From<String> for Key {
166    fn from(v: String) -> Key {
167        Key(v.into_bytes())
168    }
169}
170
171impl From<Key> for Vec<u8> {
172    fn from(key: Key) -> Self {
173        key.0
174    }
175}
176
177impl<'a> From<&'a Key> for &'a [u8] {
178    fn from(key: &'a Key) -> Self {
179        &key.0
180    }
181}
182
183impl<'a> From<&'a Vec<u8>> for &'a Key {
184    fn from(key: &'a Vec<u8>) -> Self {
185        unsafe { &*(key as *const Vec<u8> as *const Key) }
186    }
187}
188impl AsRef<Key> for Key {
189    fn as_ref(&self) -> &Key {
190        self
191    }
192}
193
194impl AsRef<Key> for Vec<u8> {
195    fn as_ref(&self) -> &Key {
196        unsafe { &*(self as *const Vec<u8> as *const Key) }
197    }
198}
199
200impl fmt::Debug for Key {
201    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
202        write!(f, "Key({})", HexRepr(&self.0))
203    }
204}