1use arrow_buffer::i256;
5use reifydb_value::{
6 Result,
7 error::{Error, TypeError},
8 value::{
9 Value,
10 blob::Blob,
11 constraint::{precision::Precision, scale::Scale},
12 date::Date,
13 datetime::DateTime,
14 decimal::Decimal,
15 dictionary::DictionaryEntryId,
16 duration::Duration,
17 identity::IdentityId,
18 row_number::RowNumber,
19 time::Time,
20 uuid::{Uuid4, Uuid7},
21 value_type::ValueType,
22 },
23};
24
25use super::{
26 ByteSink, CONTAINER_END, encode_bool, encode_bytes, encode_f32, encode_f64, encode_fixed, encode_i8,
27 encode_i16, encode_i32, encode_i64, encode_i128, encode_i256, encode_u8, encode_u16, encode_u32, encode_u64,
28 encode_u128, encode_u128_varint,
29};
30use crate::{
31 key::{buf::KeyBuf, encoded::EncodedKey},
32 tag::{TypeTag, ValueKind},
33 unscaled::{NARROW, decimal_unscaled, width},
34};
35
36pub(super) fn keycode_type_descending(ty: &ValueType) -> bool {
37 matches!(
38 ty,
39 ValueType::Boolean
40 | ValueType::Float4 | ValueType::Float8
41 | ValueType::Int1 | ValueType::Int2
42 | ValueType::Int4 | ValueType::Int8
43 | ValueType::Int16 | ValueType::Uint1
44 | ValueType::Uint2 | ValueType::Uint4
45 | ValueType::Uint8 | ValueType::Uint16
46 | ValueType::Date | ValueType::DateTime
47 | ValueType::Time | ValueType::Duration
48 | ValueType::Utf8 | ValueType::Blob
49 | ValueType::Uuid4 | ValueType::Uuid7
50 | ValueType::IdentityId
51 | ValueType::Decimal { .. }
52 | ValueType::DictionaryId
53 | ValueType::List(_) | ValueType::Record(_)
54 | ValueType::Tuple(_)
55 )
56}
57
58pub struct KeySerializer {
59 buffer: KeyBuf,
60}
61
62impl KeySerializer {
63 pub fn new() -> Self {
64 Self {
65 buffer: KeyBuf::new(),
66 }
67 }
68
69 pub fn with_capacity(capacity: usize) -> Self {
70 Self {
71 buffer: KeyBuf::with_capacity(capacity),
72 }
73 }
74
75 pub fn extend_bool(&mut self, value: bool) -> &mut Self {
76 self.buffer.push(encode_bool(value));
77 self
78 }
79
80 pub fn extend_f32(&mut self, value: f32) -> &mut Self {
81 self.buffer.extend_from_slice(&encode_f32(value));
82 self
83 }
84
85 pub fn extend_f64(&mut self, value: f64) -> &mut Self {
86 self.buffer.extend_from_slice(&encode_f64(value));
87 self
88 }
89
90 pub fn extend_i8<T: Into<i8>>(&mut self, value: T) -> &mut Self {
91 self.buffer.extend_from_slice(&encode_i8(value.into()));
92 self
93 }
94
95 pub fn extend_i16<T: Into<i16>>(&mut self, value: T) -> &mut Self {
96 self.buffer.extend_from_slice(&encode_i16(value.into()));
97 self
98 }
99
100 pub fn extend_i32<T: Into<i32>>(&mut self, value: T) -> &mut Self {
101 self.buffer.extend_from_slice(&encode_i32(value.into()));
102 self
103 }
104
105 pub fn extend_i64<T: Into<i64>>(&mut self, value: T) -> &mut Self {
106 self.buffer.extend_from_slice(&encode_i64(value.into()));
107 self
108 }
109
110 pub fn extend_i128<T: Into<i128>>(&mut self, value: T) -> &mut Self {
111 self.buffer.extend_from_slice(&encode_i128(value.into()));
112 self
113 }
114
115 pub fn extend_u8<T: Into<u8>>(&mut self, value: T) -> &mut Self {
116 self.buffer.push(encode_u8(value.into()));
117 self
118 }
119
120 pub fn extend_u16<T: Into<u16>>(&mut self, value: T) -> &mut Self {
121 self.buffer.extend_from_slice(&encode_u16(value.into()));
122 self
123 }
124
125 pub fn extend_u32<T: Into<u32>>(&mut self, value: T) -> &mut Self {
126 self.buffer.extend_from_slice(&encode_u32(value.into()));
127 self
128 }
129
130 pub fn extend_u64<T: Into<u64>>(&mut self, value: T) -> &mut Self {
131 self.buffer.extend_from_slice(&encode_u64(value.into()));
132 self
133 }
134
135 pub fn extend_u128<T: Into<u128>>(&mut self, value: T) -> &mut Self {
136 self.buffer.extend_from_slice(&encode_u128(value.into()));
137 self
138 }
139
140 pub fn extend_fixed<const N: usize>(&mut self, value: [u8; N]) -> &mut Self {
141 self.buffer.extend_from_slice(&encode_fixed(value));
142 self
143 }
144
145 pub fn extend_u128_varint<T: Into<u128>>(&mut self, value: T) -> &mut Self {
146 encode_u128_varint(value.into(), &mut self.buffer);
147 self
148 }
149
150 pub fn extend_bytes<T: AsRef<[u8]>>(&mut self, bytes: T) -> &mut Self {
151 encode_bytes(bytes.as_ref(), &mut self.buffer);
152 self
153 }
154
155 pub fn extend_str<T: AsRef<str>>(&mut self, s: T) -> &mut Self {
156 self.extend_bytes(s.as_ref().as_bytes())
157 }
158
159 pub fn finish(self) -> EncodedKey {
160 self.buffer.finish()
161 }
162
163 pub fn to_encoded_key(self) -> EncodedKey {
164 self.buffer.finish()
165 }
166
167 pub fn extend_raw(&mut self, bytes: &[u8]) -> &mut Self {
168 self.buffer.extend_from_slice(bytes);
169 self
170 }
171
172 pub fn extend_kind(&mut self, kind: ValueKind) -> &mut Self {
173 self.buffer.push(kind.byte());
174 self
175 }
176
177 pub fn len(&self) -> usize {
178 self.buffer.len()
179 }
180
181 pub(crate) fn complement_from(&mut self, start: usize) {
182 for byte in &mut self.buffer.as_mut_slice()[start..] {
183 *byte = !*byte;
184 }
185 }
186
187 pub fn is_empty(&self) -> bool {
188 self.buffer.is_empty()
189 }
190
191 pub fn extend_date(&mut self, date: &Date) -> &mut Self {
192 self.extend_i32(date.to_days_since_epoch())
193 }
194
195 pub fn extend_datetime(&mut self, datetime: &DateTime) -> &mut Self {
196 self.extend_i64(datetime.to_nanos())
197 }
198
199 pub fn extend_time(&mut self, time: &Time) -> &mut Self {
200 self.extend_u64(time.to_nanos_since_midnight())
201 }
202
203 pub fn extend_duration(&mut self, duration: &Duration) -> &mut Self {
204 self.extend_i32(duration.get_months()).extend_i32(duration.get_days()).extend_i64(duration.get_nanos())
205 }
206
207 pub fn extend_row_number(&mut self, row_number: &RowNumber) -> &mut Self {
208 self.extend_u64(row_number.0)
209 }
210
211 pub fn extend_identity_id(&mut self, id: &IdentityId) -> &mut Self {
212 self.extend_bytes(id.as_bytes())
213 }
214
215 pub fn extend_uuid4(&mut self, uuid: &Uuid4) -> &mut Self {
216 self.extend_bytes(uuid.as_bytes())
217 }
218
219 pub fn extend_uuid7(&mut self, uuid: &Uuid7) -> &mut Self {
220 self.extend_bytes(uuid.as_bytes())
221 }
222
223 pub fn extend_blob(&mut self, blob: &Blob) -> &mut Self {
224 self.extend_bytes(blob.as_ref() as &[u8])
225 }
226
227 pub fn extend_decimal(&mut self, value: &Decimal, precision: Precision, scale: Scale) -> Result<&mut Self> {
228 let unscaled = decimal_unscaled(value, precision, scale).ok_or_else(|| {
229 key_error(format!(
230 "decimal {value} does not fit precision {} and scale {}",
231 precision.value(),
232 scale.value()
233 ))
234 })?;
235 self.buffer.push(encode_u8(precision.value()));
236 self.buffer.push(encode_u8(scale.value()));
237 Ok(self.extend_unscaled(unscaled, precision))
238 }
239
240 fn extend_unscaled(&mut self, unscaled: i256, precision: Precision) -> &mut Self {
241 if width(precision) == NARROW {
242 let narrow = unscaled.to_i128().expect("a value within precision 38 fits i128");
243 self.buffer.extend_from_slice(&encode_i128(narrow));
244 } else {
245 self.buffer.extend_from_slice(&encode_i256(unscaled));
246 }
247 self
248 }
249
250 pub fn extend_value(&mut self, value: &Value) -> &mut Self {
251 if let Err(error) = self.try_extend_value(value) {
252 panic!("key serialization of a {} value failed: {error}", value.get_type());
253 }
254 self
255 }
256
257 pub fn try_extend_value(&mut self, value: &Value) -> Result<&mut Self> {
258 match value {
259 Value::None {
260 inner,
261 ..
262 } => {
263 self.buffer.push(ValueKind::None.byte());
264 if ValueKind::of_type(inner) == ValueKind::Digest {
265 return Err(Error::from(TypeError::SerdeKeycode {
266 message: format!(
267 "a none of type {inner} cannot be serialized in a key"
268 ),
269 }));
270 }
271 let tag = TypeTag::of_type(inner)
272 .expect("option nesting in a key none inner exceeds the supported depth");
273 self.buffer.push(tag.byte());
274 }
275 Value::Boolean(b) => {
276 self.buffer.push(ValueKind::Boolean.byte());
277 self.extend_bool(*b);
278 }
279 Value::Float4(f) => {
280 self.buffer.push(ValueKind::Float4.byte());
281 self.extend_f32(**f);
282 }
283 Value::Float8(f) => {
284 self.buffer.push(ValueKind::Float8.byte());
285 self.extend_f64(**f);
286 }
287 Value::Int1(i) => {
288 self.buffer.push(ValueKind::Int1.byte());
289 self.extend_i8(*i);
290 }
291 Value::Int2(i) => {
292 self.buffer.push(ValueKind::Int2.byte());
293 self.extend_i16(*i);
294 }
295 Value::Int4(i) => {
296 self.buffer.push(ValueKind::Int4.byte());
297 self.extend_i32(*i);
298 }
299 Value::Int8(i) => {
300 self.buffer.push(ValueKind::Int8.byte());
301 self.extend_i64(*i);
302 }
303 Value::Int16(i) => {
304 self.buffer.push(ValueKind::Int16.byte());
305 self.extend_i128(*i);
306 }
307 Value::Utf8(s) => {
308 self.buffer.push(ValueKind::Utf8.byte());
309 self.extend_str(s);
310 }
311
312 Value::Uint1(u) => {
313 self.buffer.push(ValueKind::Uint1.byte());
314 self.extend_u8(*u);
315 }
316 Value::Uint2(u) => {
317 self.buffer.push(ValueKind::Uint2.byte());
318 self.extend_u16(*u);
319 }
320 Value::Uint4(u) => {
321 self.buffer.push(ValueKind::Uint4.byte());
322 self.extend_u32(*u);
323 }
324 Value::Uint8(u) => {
325 self.buffer.push(ValueKind::Uint8.byte());
326 self.extend_u64(*u);
327 }
328 Value::Uint16(u) => {
329 self.buffer.push(ValueKind::Uint16.byte());
330 self.extend_u128(*u);
331 }
332 Value::Date(d) => {
333 self.buffer.push(ValueKind::Date.byte());
334 self.extend_date(d);
335 }
336 Value::DateTime(dt) => {
337 self.buffer.push(ValueKind::DateTime.byte());
338 self.extend_datetime(dt);
339 }
340 Value::Time(t) => {
341 self.buffer.push(ValueKind::Time.byte());
342 self.extend_time(t);
343 }
344 Value::Duration(i) => {
345 self.buffer.push(ValueKind::Duration.byte());
346 self.extend_duration(i);
347 }
348 Value::IdentityId(id) => {
349 self.buffer.push(ValueKind::IdentityId.byte());
350 self.extend_identity_id(id);
351 }
352 Value::Uuid4(uuid) => {
353 self.buffer.push(ValueKind::Uuid4.byte());
354 self.extend_uuid4(uuid);
355 }
356 Value::Uuid7(uuid) => {
357 self.buffer.push(ValueKind::Uuid7.byte());
358 self.extend_uuid7(uuid);
359 }
360 Value::Blob(b) => {
361 self.buffer.push(ValueKind::Blob.byte());
362 self.extend_blob(b);
363 }
364 Value::Decimal(d) => {
365 self.buffer.push(ValueKind::Decimal.byte());
366 self.extend_decimal(d, Precision::MAX, Scale::new(d.scale()))?;
367 }
368 Value::List(items) => {
369 self.buffer.push(ValueKind::List.byte());
370 for item in items {
371 self.try_extend_value(item)?;
372 }
373 self.buffer.push(CONTAINER_END);
374 }
375 Value::Tuple(items) => {
376 self.buffer.push(ValueKind::Tuple.byte());
377 for item in items {
378 self.try_extend_value(item)?;
379 }
380 self.buffer.push(CONTAINER_END);
381 }
382 Value::Record(fields) => {
383 self.buffer.push(ValueKind::Record.byte());
384 for (name, value) in fields {
385 self.extend_bytes(name.as_bytes());
386 self.try_extend_value(value)?;
387 }
388 self.buffer.push(CONTAINER_END);
389 }
390 Value::Any(_) | Value::Type(_) => {
391 return Err(Error::from(TypeError::SerdeKeycode {
392 message: format!("a {} value cannot be serialized in a key", value.get_type()),
393 }));
394 }
395 Value::Digest(_) => {
396 return Err(Error::from(TypeError::SerdeKeycode {
397 message: "a digest value cannot be serialized in a key".to_string(),
398 }));
399 }
400 Value::Enum(_) => {
401 return Err(Error::from(TypeError::SerdeKeycode {
402 message: "an enum value cannot be serialized in a key".to_string(),
403 }));
404 }
405 Value::DictionaryId(id) => {
406 self.buffer.push(ValueKind::DictionaryId.byte());
407 match id {
408 DictionaryEntryId::U1(v) => {
409 self.buffer.push(0x00);
410 self.extend_u8(*v);
411 }
412 DictionaryEntryId::U2(v) => {
413 self.buffer.push(0x01);
414 self.extend_u16(*v);
415 }
416 DictionaryEntryId::U4(v) => {
417 self.buffer.push(0x02);
418 self.extend_u32(*v);
419 }
420 DictionaryEntryId::U8(v) => {
421 self.buffer.push(0x03);
422 self.extend_u64(*v);
423 }
424 DictionaryEntryId::U16(v) => {
425 self.buffer.push(0x04);
426 self.extend_u128(*v);
427 }
428 }
429 }
430 }
431 Ok(self)
432 }
433}
434
435impl Default for KeySerializer {
436 fn default() -> Self {
437 Self::new()
438 }
439}
440
441impl ByteSink for KeySerializer {
442 fn push(&mut self, byte: u8) {
443 self.buffer.push(byte);
444 }
445
446 fn extend_from_slice(&mut self, slice: &[u8]) {
447 self.buffer.extend_from_slice(slice);
448 }
449}
450
451fn key_error(message: String) -> Error {
452 Error::from(TypeError::SerdeKeycode {
453 message,
454 })
455}