1use std::{error::Error, fmt};
7
8const ROW_KEY_TAG: u8 = 0x01;
9const TABLE_PREFIX_LEN: usize = 1 + std::mem::size_of::<u32>();
10const ROW_KEY_LEN: usize = TABLE_PREFIX_LEN + std::mem::size_of::<u64>();
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum RowKeyRangeError {
15 InvalidKeyEncoding,
17 InvalidBounds,
19 CrossTableBounds,
21}
22
23impl fmt::Display for RowKeyRangeError {
24 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25 match self {
26 Self::InvalidKeyEncoding => f.write_str("invalid canonical primary row-key encoding"),
27 Self::InvalidBounds => {
28 f.write_str("row-key range bounds are not a non-empty half-open interval")
29 }
30 Self::CrossTableBounds => f.write_str("row-key range bounds refer to different tables"),
31 }
32 }
33}
34
35impl Error for RowKeyRangeError {}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
39pub struct CanonicalRowKey {
40 table_id: u32,
41 row_id: u64,
42}
43
44impl CanonicalRowKey {
45 pub const fn new(table_id: u32, row_id: u64) -> Self {
47 Self { table_id, row_id }
48 }
49
50 pub fn decode(encoded: &[u8]) -> Result<Self, RowKeyRangeError> {
52 if encoded.len() != ROW_KEY_LEN || encoded[0] != ROW_KEY_TAG {
53 return Err(RowKeyRangeError::InvalidKeyEncoding);
54 }
55 let table_id = u32::from_be_bytes(
56 encoded[1..TABLE_PREFIX_LEN]
57 .try_into()
58 .map_err(|_| RowKeyRangeError::InvalidKeyEncoding)?,
59 );
60 let row_id = u64::from_be_bytes(
61 encoded[TABLE_PREFIX_LEN..]
62 .try_into()
63 .map_err(|_| RowKeyRangeError::InvalidKeyEncoding)?,
64 );
65 Ok(Self::new(table_id, row_id))
66 }
67
68 pub fn encode(self) -> Vec<u8> {
70 let mut encoded = Vec::with_capacity(ROW_KEY_LEN);
71 encoded.extend_from_slice(&Self::table_prefix(self.table_id));
72 encoded.extend_from_slice(&self.row_id.to_be_bytes());
73 encoded
74 }
75
76 pub const fn table_id(self) -> u32 {
78 self.table_id
79 }
80
81 pub const fn row_id(self) -> u64 {
83 self.row_id
84 }
85
86 pub fn table_prefix(table_id: u32) -> Vec<u8> {
88 let mut prefix = Vec::with_capacity(TABLE_PREFIX_LEN);
89 prefix.push(ROW_KEY_TAG);
90 prefix.extend_from_slice(&table_id.to_be_bytes());
91 prefix
92 }
93
94 pub fn table_prefix_end(table_id: u32) -> Vec<u8> {
97 prefix_successor(&Self::table_prefix(table_id))
98 .expect("a primary row-key prefix beginning with 0x01 always has a successor")
99 }
100}
101
102#[derive(Debug, Clone, PartialEq, Eq)]
104pub struct EncodedRowKeyRange {
105 pub lower_inclusive: Vec<u8>,
107 pub upper_exclusive: Vec<u8>,
109}
110
111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub struct RowKeyRange {
114 table_id: u32,
115 lower_inclusive: Option<u64>,
116 upper_exclusive: Option<u64>,
117}
118
119impl RowKeyRange {
120 pub const fn full_table(table_id: u32) -> Self {
122 Self {
123 table_id,
124 lower_inclusive: None,
125 upper_exclusive: None,
126 }
127 }
128
129 pub fn new(
131 table_id: u32,
132 lower_inclusive: Option<u64>,
133 upper_exclusive: Option<u64>,
134 ) -> Result<Self, RowKeyRangeError> {
135 if let (Some(lower), Some(upper)) = (lower_inclusive, upper_exclusive) {
136 if lower >= upper {
137 return Err(RowKeyRangeError::InvalidBounds);
138 }
139 }
140 Ok(Self {
141 table_id,
142 lower_inclusive,
143 upper_exclusive,
144 })
145 }
146
147 pub fn from_keys(
149 lower_inclusive: Option<CanonicalRowKey>,
150 upper_exclusive: Option<CanonicalRowKey>,
151 table_id: u32,
152 ) -> Result<Self, RowKeyRangeError> {
153 for key in [lower_inclusive, upper_exclusive].into_iter().flatten() {
154 if key.table_id() != table_id {
155 return Err(RowKeyRangeError::CrossTableBounds);
156 }
157 }
158 Self::new(
159 table_id,
160 lower_inclusive.map(CanonicalRowKey::row_id),
161 upper_exclusive.map(CanonicalRowKey::row_id),
162 )
163 }
164
165 pub const fn table_id(self) -> u32 {
167 self.table_id
168 }
169
170 pub const fn lower_inclusive(self) -> Option<u64> {
172 self.lower_inclusive
173 }
174
175 pub const fn upper_exclusive(self) -> Option<u64> {
177 self.upper_exclusive
178 }
179
180 pub fn encoded_bounds(self) -> EncodedRowKeyRange {
182 EncodedRowKeyRange {
183 lower_inclusive: CanonicalRowKey::new(self.table_id, self.lower_inclusive.unwrap_or(0))
184 .encode(),
185 upper_exclusive: self.upper_exclusive.map_or_else(
186 || CanonicalRowKey::table_prefix_end(self.table_id),
187 |row_id| CanonicalRowKey::new(self.table_id, row_id).encode(),
188 ),
189 }
190 }
191
192 pub fn contains(self, key: CanonicalRowKey) -> bool {
194 key.table_id() == self.table_id
195 && self
196 .lower_inclusive
197 .is_none_or(|lower| key.row_id() >= lower)
198 && self
199 .upper_exclusive
200 .is_none_or(|upper| key.row_id() < upper)
201 }
202}
203
204fn prefix_successor(prefix: &[u8]) -> Option<Vec<u8>> {
205 let mut successor = prefix.to_vec();
206 for index in (0..successor.len()).rev() {
207 if successor[index] != u8::MAX {
208 successor[index] += 1;
209 successor.truncate(index + 1);
210 return Some(successor);
211 }
212 }
213 None
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219
220 #[test]
221 fn canonical_key_roundtrip_and_lexicographic_order_are_stable() {
222 let first = CanonicalRowKey::new(7, 1);
223 let later = CanonicalRowKey::new(7, u64::MAX);
224
225 assert_eq!(CanonicalRowKey::decode(&first.encode()).unwrap(), first);
226 assert!(first.encode() < later.encode());
227 }
228
229 #[test]
230 fn full_table_range_uses_prefix_successor_for_the_upper_bound() {
231 let range = RowKeyRange::full_table(u32::MAX);
232 let encoded = range.encoded_bounds();
233
234 assert_eq!(
235 encoded.lower_inclusive,
236 CanonicalRowKey::new(u32::MAX, 0).encode()
237 );
238 assert_eq!(encoded.upper_exclusive, vec![0x02]);
239 assert!(range.contains(CanonicalRowKey::new(u32::MAX, u64::MAX)));
240 }
241
242 #[test]
243 fn half_open_bounds_include_lower_and_exclude_upper() {
244 let range = RowKeyRange::new(3, Some(10), Some(20)).unwrap();
245
246 assert!(!range.contains(CanonicalRowKey::new(3, 9)));
247 assert!(range.contains(CanonicalRowKey::new(3, 10)));
248 assert!(range.contains(CanonicalRowKey::new(3, 19)));
249 assert!(!range.contains(CanonicalRowKey::new(3, 20)));
250 assert!(!range.contains(CanonicalRowKey::new(4, 10)));
251 }
252
253 #[test]
254 fn invalid_or_cross_table_bounds_are_rejected() {
255 assert_eq!(
256 CanonicalRowKey::decode(&[0x02, 0, 0, 0, 1]),
257 Err(RowKeyRangeError::InvalidKeyEncoding)
258 );
259 assert_eq!(
260 RowKeyRange::new(3, Some(20), Some(10)),
261 Err(RowKeyRangeError::InvalidBounds)
262 );
263 assert_eq!(
264 RowKeyRange::from_keys(Some(CanonicalRowKey::new(4, 1)), None, 3),
265 Err(RowKeyRangeError::CrossTableBounds)
266 );
267 }
268}