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