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