1use yo_common::num::{parse_i64, push_i64};
37
38pub const EMBSTR_MAX: usize = 44;
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum Encoding {
47 Int,
49 Embstr,
51 Raw,
53}
54
55impl Encoding {
56 #[inline]
58 pub const fn name(self) -> &'static str {
59 match self {
60 Encoding::Int => "int",
61 Encoding::Embstr => "embstr",
62 Encoding::Raw => "raw",
63 }
64 }
65
66 #[inline]
73 pub fn of(bytes: &[u8]) -> Encoding {
74 if parse_i64(bytes).is_some() {
75 Encoding::Int
76 } else if bytes.len() <= EMBSTR_MAX {
77 Encoding::Embstr
78 } else {
79 Encoding::Raw
80 }
81 }
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96pub enum Kind {
97 String = 0,
99 Hash = 1,
101 Set = 2,
103 Zset = 3,
105 List = 4,
107 Stream = 5,
109 Array = 6,
111}
112
113impl Kind {
114 #[inline]
116 pub const fn name(self) -> &'static str {
117 match self {
118 Kind::String => "string",
119 Kind::Hash => "hash",
120 Kind::Set => "set",
121 Kind::Zset => "zset",
122 Kind::List => "list",
123 Kind::Stream => "stream",
124 Kind::Array => "array",
125 }
126 }
127
128 #[inline]
135 const fn from_bits(bits: u8) -> Kind {
136 match bits {
137 KIND_HASH => Kind::Hash,
138 KIND_SET => Kind::Set,
139 KIND_ZSET => Kind::Zset,
140 KIND_LIST => Kind::List,
141 KIND_STREAM => Kind::Stream,
142 KIND_ARRAY => Kind::Array,
143 _ => Kind::String,
144 }
145 }
146}
147
148const ENC_MASK: u8 = 0b0000_0011;
150const ENC_INT: u8 = 0;
151const ENC_EMBSTR: u8 = 1;
152const ENC_RAW: u8 = 2;
153const HAS_EXPIRY: u8 = 0b0000_0100;
155const KIND_MASK: u8 = 0b0011_1000;
160const KIND_SHIFT: u32 = 3;
161const KIND_HASH: u8 = 1;
162const KIND_SET: u8 = 2;
163const KIND_ZSET: u8 = 3;
164const KIND_LIST: u8 = 4;
165const KIND_STREAM: u8 = 5;
166const KIND_ARRAY: u8 = 6;
167
168const INT_LEN: usize = 8;
170
171#[derive(Debug, Clone, Copy, PartialEq, Eq)]
173pub struct Meta(u8);
174
175impl Meta {
176 #[inline]
178 pub const fn new(kind: Kind, enc: Encoding, has_expiry: bool) -> Meta {
179 let bits = match enc {
180 Encoding::Int => ENC_INT,
181 Encoding::Embstr => ENC_EMBSTR,
182 Encoding::Raw => ENC_RAW,
183 };
184 Meta(bits | ((kind as u8) << KIND_SHIFT) | if has_expiry { HAS_EXPIRY } else { 0 })
185 }
186
187 #[inline]
189 pub const fn string(enc: Encoding, has_expiry: bool) -> Meta {
190 Meta::new(Kind::String, enc, has_expiry)
191 }
192
193 #[inline]
203 pub const fn slot(kind: Kind, has_expiry: bool) -> Meta {
204 Meta::new(kind, Encoding::Int, has_expiry)
205 }
206
207 #[inline]
213 pub const fn from_byte(b: u8) -> Meta {
214 Meta(b)
215 }
216
217 #[inline]
219 pub const fn byte(self) -> u8 {
220 self.0
221 }
222
223 #[inline]
225 pub const fn encoding(self) -> Encoding {
226 match self.0 & ENC_MASK {
227 ENC_INT => Encoding::Int,
228 ENC_EMBSTR => Encoding::Embstr,
229 _ => Encoding::Raw,
230 }
231 }
232
233 #[inline]
235 pub const fn kind(self) -> Kind {
236 Kind::from_bits((self.0 & KIND_MASK) >> KIND_SHIFT)
237 }
238
239 #[inline]
241 pub const fn has_expiry(self) -> bool {
242 self.0 & HAS_EXPIRY != 0
243 }
244
245 #[inline]
247 pub const fn payload_at(self) -> usize {
248 if self.has_expiry() { 1 + 8 } else { 1 }
249 }
250}
251
252#[derive(Debug, Clone, Copy, PartialEq, Eq)]
260pub enum Str<'a> {
261 Int(i64),
263 Bytes(&'a [u8]),
265}
266
267impl Str<'_> {
268 #[inline]
270 pub fn len(&self) -> usize {
271 match self {
272 Str::Int(n) => yo_common::num::i64_len(*n),
273 Str::Bytes(b) => b.len(),
274 }
275 }
276
277 #[inline]
279 pub fn is_empty(&self) -> bool {
280 match self {
281 Str::Int(_) => false,
283 Str::Bytes(b) => b.is_empty(),
284 }
285 }
286
287 #[inline]
289 pub fn write_to(&self, out: &mut Vec<u8>) {
290 match self {
291 Str::Int(n) => push_i64(out, *n),
292 Str::Bytes(b) => out.extend_from_slice(b),
293 }
294 }
295
296 pub fn to_vec(&self) -> Vec<u8> {
298 let mut v = Vec::with_capacity(self.len());
299 self.write_to(&mut v);
300 v
301 }
302
303 #[inline]
310 pub fn as_int(&self) -> Option<i64> {
311 match self {
312 Str::Int(n) => Some(*n),
313 Str::Bytes(b) => parse_i64(b),
314 }
315 }
316
317 #[must_use]
323 pub fn digest(&self) -> u64 {
324 match self {
325 Str::Bytes(b) => yo_common::xxh3::hash64(b),
326 Str::Int(_) => yo_common::xxh3::hash64(&self.to_vec()),
330 }
331 }
332
333 #[inline]
339 pub(crate) fn eq_bytes(&self, want: &[u8]) -> bool {
340 match self {
341 Str::Bytes(b) => *b == want,
342 Str::Int(n) => parse_i64(want) == Some(*n),
343 }
344 }
345}
346
347#[inline]
349pub fn record_len(enc: Encoding, payload: usize, has_expiry: bool) -> usize {
350 let head = if has_expiry { 1 + 8 } else { 1 };
351 head + if enc == Encoding::Int {
352 INT_LEN
353 } else {
354 payload
355 }
356}
357
358#[inline]
364pub fn write_record(out: &mut [u8], enc: Encoding, bytes: &[u8], expire_at: Option<u64>) {
365 out[0] = Meta::string(enc, expire_at.is_some()).byte();
366 let mut at = 1;
367 if let Some(ms) = expire_at {
368 out[at..at + 8].copy_from_slice(&ms.to_le_bytes());
369 at += 8;
370 }
371 match enc {
372 Encoding::Int => {
373 let n =
374 parse_i64(bytes).expect("int encoding was chosen for bytes that are not an int");
375 out[at..at + INT_LEN].copy_from_slice(&n.to_le_bytes());
376 }
377 _ => out[at..].copy_from_slice(bytes),
378 }
379}
380
381#[inline]
383pub fn write_int_record(out: &mut [u8], n: i64, expire_at: Option<u64>) {
384 out[0] = Meta::string(Encoding::Int, expire_at.is_some()).byte();
385 let mut at = 1;
386 if let Some(ms) = expire_at {
387 out[at..at + 8].copy_from_slice(&ms.to_le_bytes());
388 at += 8;
389 }
390 out[at..at + INT_LEN].copy_from_slice(&n.to_le_bytes());
391}
392
393const SLOT_LEN: usize = 4;
395
396#[inline]
398pub fn slot_record_len(has_expiry: bool) -> usize {
399 (if has_expiry { 1 + 8 } else { 1 }) + SLOT_LEN
400}
401
402#[inline]
406pub fn write_slot_record(out: &mut [u8], kind: Kind, slot: u32, expire_at: Option<u64>) {
407 out[0] = Meta::slot(kind, expire_at.is_some()).byte();
408 let mut at = 1;
409 if let Some(ms) = expire_at {
410 out[at..at + 8].copy_from_slice(&ms.to_le_bytes());
411 at += 8;
412 }
413 out[at..at + SLOT_LEN].copy_from_slice(&slot.to_le_bytes());
414}
415
416#[inline]
423pub fn slot(rec: &[u8]) -> u32 {
424 let at = Meta::from_byte(rec[0]).payload_at();
425 let mut b = [0u8; SLOT_LEN];
426 b.copy_from_slice(&rec[at..at + SLOT_LEN]);
427 u32::from_le_bytes(b)
428}
429
430#[inline]
432pub fn kind(rec: &[u8]) -> Kind {
433 Meta::from_byte(rec[0]).kind()
434}
435
436#[inline]
438pub fn expire_at(rec: &[u8]) -> Option<u64> {
439 let m = Meta::from_byte(rec[0]);
440 if !m.has_expiry() {
441 return None;
442 }
443 let mut b = [0u8; 8];
444 b.copy_from_slice(&rec[1..9]);
445 Some(u64::from_le_bytes(b))
446}
447
448#[inline]
453pub fn is_expired(rec: &[u8], now_ms: u64) -> bool {
454 match expire_at(rec) {
455 Some(at) => at <= now_ms,
456 None => false,
457 }
458}
459
460#[inline]
462pub fn read(rec: &[u8]) -> Str<'_> {
463 let m = Meta::from_byte(rec[0]);
464 let at = m.payload_at();
465 match m.encoding() {
466 Encoding::Int => {
467 let mut b = [0u8; INT_LEN];
468 b.copy_from_slice(&rec[at..at + INT_LEN]);
469 Str::Int(i64::from_le_bytes(b))
470 }
471 _ => Str::Bytes(&rec[at..]),
472 }
473}
474
475#[inline]
480pub fn read_int_in_place(rec: &[u8]) -> Option<(i64, usize)> {
481 let m = Meta::from_byte(rec[0]);
482 if m.encoding() != Encoding::Int {
483 return None;
484 }
485 let at = m.payload_at();
486 let mut b = [0u8; INT_LEN];
487 b.copy_from_slice(&rec[at..at + INT_LEN]);
488 Some((i64::from_le_bytes(b), at))
489}
490
491#[inline]
493pub fn write_int_in_place(rec: &mut [u8], at: usize, n: i64) {
494 rec[at..at + INT_LEN].copy_from_slice(&n.to_le_bytes());
495}
496
497#[cfg(test)]
498mod tests {
499 use super::*;
500
501 #[test]
502 fn encoding_follows_redis_boundaries() {
503 assert_eq!(Encoding::of(b"0"), Encoding::Int);
504 assert_eq!(Encoding::of(b"-1"), Encoding::Int);
505 assert_eq!(Encoding::of(b"9223372036854775807"), Encoding::Int);
506 assert_eq!(Encoding::of(b"9223372036854775808"), Encoding::Embstr);
508 assert_eq!(Encoding::of(b"007"), Encoding::Embstr);
510 assert_eq!(Encoding::of(b"+1"), Encoding::Embstr);
511 assert_eq!(Encoding::of(b"-0"), Encoding::Embstr);
512 assert_eq!(Encoding::of(b""), Encoding::Embstr);
513 assert_eq!(Encoding::of(&[b'x'; EMBSTR_MAX]), Encoding::Embstr);
514 assert_eq!(Encoding::of(&[b'x'; EMBSTR_MAX + 1]), Encoding::Raw);
515 }
516
517 const KINDS: [Kind; 7] = [
518 Kind::String,
519 Kind::Hash,
520 Kind::Set,
521 Kind::Zset,
522 Kind::List,
523 Kind::Stream,
524 Kind::Array,
525 ];
526
527 #[test]
528 fn the_meta_byte_survives_a_round_trip() {
529 for kind in KINDS {
530 for enc in [Encoding::Int, Encoding::Embstr, Encoding::Raw] {
531 for expiry in [false, true] {
532 let m = Meta::new(kind, enc, expiry);
533 let back = Meta::from_byte(m.byte());
534 assert_eq!(back.kind(), kind);
535 assert_eq!(back.encoding(), enc);
536 assert_eq!(back.has_expiry(), expiry);
537 assert_eq!(back.payload_at(), if expiry { 9 } else { 1 });
538 }
539 }
540 }
541 }
542
543 #[test]
544 fn the_three_fields_of_the_meta_byte_do_not_reach_into_each_other() {
545 let mut seen = std::collections::HashSet::new();
550 for kind in KINDS {
551 for enc in [Encoding::Int, Encoding::Embstr, Encoding::Raw] {
552 for expiry in [false, true] {
553 assert!(
554 seen.insert(Meta::new(kind, enc, expiry).byte()),
555 "{kind:?} {enc:?} {expiry} collides with something else"
556 );
557 }
558 }
559 }
560 assert_eq!(seen.len(), KINDS.len() * 6);
561 }
562
563 #[test]
564 fn a_record_written_before_the_tag_existed_is_a_string() {
565 assert_eq!(Kind::String as u8, 0);
569 assert_eq!(Meta::from_byte(0b0000_0101).kind(), Kind::String);
570 assert_eq!(Meta::from_byte(0b0000_0101).encoding(), Encoding::Embstr);
571 assert!(Meta::from_byte(0b0000_0101).has_expiry());
572 }
573
574 #[test]
575 fn the_tag_is_the_number_the_file_format_uses() {
576 use yo_format::catalog::ValueType;
577 assert_eq!(Kind::String as u8, ValueType::String as u8);
581 assert_eq!(Kind::Hash as u8, ValueType::Hash as u8);
582 assert_eq!(Kind::Set as u8, ValueType::Set as u8);
583 assert_eq!(Kind::Zset as u8, ValueType::Zset as u8);
584 assert_eq!(Kind::List as u8, ValueType::List as u8);
585 assert_eq!(Kind::Stream as u8, ValueType::Stream as u8);
586 assert_eq!(Kind::Array as u8, ValueType::Array as u8);
587 for k in KINDS {
589 let v = ValueType::from_u8(k as u8).expect("the catalog knows this one");
590 assert_eq!(k.name(), v.redis_name(), "{k:?}");
591 }
592 }
593
594 #[test]
595 fn a_string_record_is_tagged_as_one() {
596 for text in [&b"42"[..], b"hello", &[b'z'; 100]] {
597 for expire in [None, Some(9_000u64)] {
598 assert_eq!(kind(&record(text, expire)), Kind::String);
599 }
600 }
601 let mut v = vec![0u8; record_len(Encoding::Int, 0, false)];
602 write_int_record(&mut v, 7, None);
603 assert_eq!(kind(&v), Kind::String);
604 }
605
606 fn record(bytes: &[u8], expire: Option<u64>) -> Vec<u8> {
607 let enc = Encoding::of(bytes);
608 let mut v = vec![0u8; record_len(enc, bytes.len(), expire.is_some())];
609 write_record(&mut v, enc, bytes, expire);
610 v
611 }
612
613 #[test]
614 fn a_record_gives_back_what_went_into_it() {
615 for text in [
616 &b""[..],
617 b"x",
618 b"0",
619 b"-1",
620 b"42",
621 b"007",
622 b"-0",
623 b"hello world",
624 &[b'z'; 100],
625 ] {
626 for expire in [None, Some(1_234_567_890_123u64)] {
627 let r = record(text, expire);
628 assert_eq!(read(&r).to_vec(), text, "{text:?} at {expire:?}");
629 assert_eq!(read(&r).len(), text.len(), "{text:?} length");
630 assert_eq!(expire_at(&r), expire, "{text:?} deadline");
631 }
632 }
633 }
634
635 #[test]
636 fn an_integer_costs_the_same_however_many_digits_it_has() {
637 let small = record(b"1", None);
638 let large = record(b"-9223372036854775808", None);
639 assert_eq!(small.len(), large.len());
640 assert_eq!(read(&large), Str::Int(i64::MIN));
641 assert_eq!(read(&large).to_vec(), b"-9223372036854775808");
642 }
643
644 #[test]
645 fn an_integer_is_incremented_where_it_lies() {
646 let mut r = record(b"41", Some(99));
647 let (n, at) = read_int_in_place(&r).expect("int encoded");
648 assert_eq!(n, 41);
649 write_int_in_place(&mut r, at, n + 1);
650 assert_eq!(read(&r), Str::Int(42));
651 assert_eq!(expire_at(&r), Some(99));
653 }
654
655 #[test]
656 fn a_string_is_not_read_as_an_integer_in_place() {
657 let r = record(b"hello", None);
658 assert!(read_int_in_place(&r).is_none());
659 }
660
661 #[test]
662 fn a_deadline_that_is_now_has_passed() {
663 let r = record(b"v", Some(100));
664 assert!(!is_expired(&r, 99));
665 assert!(is_expired(&r, 100));
666 assert!(is_expired(&r, 101));
667 let forever = record(b"v", None);
668 assert!(!is_expired(&forever, u64::MAX));
669 }
670
671 #[test]
672 fn a_value_that_is_text_can_still_be_a_number() {
673 assert_eq!(Str::Bytes(b"10").as_int(), Some(10));
675 assert_eq!(Str::Bytes(b"10x").as_int(), None);
676 assert_eq!(Str::Int(-5).as_int(), Some(-5));
677 }
678
679 #[test]
680 fn an_unknown_encoding_reads_as_raw_bytes() {
681 let m = Meta::from_byte(0b11);
685 assert_eq!(m.encoding(), Encoding::Raw);
686 }
687
688 #[test]
689 fn an_unknown_type_tag_reads_as_a_string() {
690 assert_eq!(Meta::from_byte(6 << 3).kind(), Kind::Array);
693 assert_eq!(Meta::from_byte(7 << 3).kind(), Kind::String);
694 }
695
696 #[test]
697 fn the_top_two_bits_are_still_free() {
698 assert_eq!(Meta::from_byte(0b1100_0000).kind(), Kind::String);
701 assert_eq!(Meta::from_byte(0b1101_0000).kind(), Kind::Set);
702 }
703}