1use rudb_common::{Error, Result};
46
47pub const ESCAPE: u8 = 255;
50
51pub const MAX_SYMBOLS: usize = 255;
53
54pub const MAX_SYMBOL_LEN: usize = 8;
57
58const GENERATIONS: usize = 5;
60
61const HASH_SLOTS: usize = 1024;
64
65const PROBE: usize = 8;
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
71struct Symbol {
72 value: u64,
73 len: u8,
74}
75
76impl Symbol {
77 fn new(bytes: &[u8]) -> Self {
78 let len = bytes.len().min(MAX_SYMBOL_LEN);
79 let mut value = 0u64;
80 for (index, byte) in bytes[..len].iter().enumerate() {
81 value |= u64::from(*byte) << (8 * index);
82 }
83 Self { value, len: len as u8 }
84 }
85
86 fn single(byte: u8) -> Self {
87 Self { value: u64::from(byte), len: 1 }
88 }
89
90 fn len(self) -> usize {
91 self.len as usize
92 }
93
94 fn mask(self) -> u64 {
95 mask_of(self.len())
96 }
97
98 fn bytes(self) -> Vec<u8> {
99 (0..self.len()).map(|index| (self.value >> (8 * index)) as u8).collect()
100 }
101
102 fn concat(self, other: Self) -> Self {
104 if self.len() >= MAX_SYMBOL_LEN {
105 return self;
106 }
107 let len = (self.len() + other.len()).min(MAX_SYMBOL_LEN);
108 let value = self.value | (other.value << (8 * self.len()));
109 Self { value: value & mask_of(len), len: len as u8 }
110 }
111}
112
113fn mask_of(len: usize) -> u64 {
114 if len >= 8 { u64::MAX } else { (1u64 << (8 * len)) - 1 }
115}
116
117pub struct SymbolTable {
119 symbols: Vec<Symbol>,
121 single: Vec<u8>,
123 pair: Vec<u16>,
125 hash: Vec<Option<(Symbol, u8)>>,
127}
128
129impl std::fmt::Debug for SymbolTable {
130 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131 formatter
133 .debug_struct("SymbolTable")
134 .field("symbols", &self.symbols.len())
135 .field("bytes", &self.serialized_len())
136 .finish()
137 }
138}
139
140impl PartialEq for SymbolTable {
147 fn eq(&self, other: &Self) -> bool {
148 self.symbols == other.symbols
149 }
150}
151
152impl Eq for SymbolTable {}
153
154impl SymbolTable {
155 #[must_use]
162 pub fn footprint(&self) -> usize {
163 size_of::<Self>()
164 + self.symbols.capacity() * size_of::<Symbol>()
165 + self.single.capacity()
166 + self.pair.capacity() * size_of::<u16>()
167 + self.hash.capacity() * size_of::<Option<(Symbol, u8)>>()
168 }
169
170 #[must_use]
173 pub fn empty() -> Self {
174 Self::build(Vec::new())
175 }
176
177 #[must_use]
183 pub fn train(samples: &[&[u8]]) -> Self {
184 let mut table = Self::empty();
185 let mut counts = Counts::new();
186 for _ in 0..GENERATIONS {
187 counts.clear();
188 for sample in samples {
189 table.count(sample, &mut counts);
190 }
191 let next = counts.best(&table);
192 if next.is_empty() {
193 break;
194 }
195 table = Self::build(next);
196 }
197 table
198 }
199
200 #[must_use]
202 pub fn len(&self) -> usize {
203 self.symbols.len()
204 }
205
206 #[must_use]
208 pub fn is_empty(&self) -> bool {
209 self.symbols.is_empty()
210 }
211
212 #[must_use]
216 pub fn serialized_len(&self) -> usize {
217 1 + self.symbols.iter().map(|symbol| 1 + symbol.len()).sum::<usize>()
218 }
219
220 pub fn serialize(&self, out: &mut Vec<u8>) {
222 out.push(self.symbols.len() as u8);
223 for symbol in &self.symbols {
224 out.push(symbol.len);
225 out.extend_from_slice(&symbol.bytes());
226 }
227 }
228
229 pub fn deserialize(bytes: &[u8]) -> Result<(Self, usize)> {
235 let count = *bytes.first().ok_or_else(|| truncated("a symbol table header"))? as usize;
236 let mut at = 1;
237 let mut symbols = Vec::with_capacity(count);
238 for _ in 0..count {
239 let len = *bytes.get(at).ok_or_else(|| truncated("a symbol length"))? as usize;
240 if len == 0 || len > MAX_SYMBOL_LEN {
241 return Err(Error::internal(format!("a symbol of {len} bytes is not a symbol")));
242 }
243 at += 1;
244 let end = at + len;
245 if end > bytes.len() {
246 return Err(truncated("a symbol"));
247 }
248 symbols.push(Symbol::new(&bytes[at..end]));
249 at = end;
250 }
251 Ok((Self::build(symbols), at))
252 }
253
254 pub fn compress(&self, input: &[u8], out: &mut Vec<u8>) {
260 let mut at = 0;
261 while at < input.len() {
262 let (code, len) = self.match_at(input, at);
263 if code == ESCAPE {
264 out.push(ESCAPE);
265 out.push(input[at]);
266 } else {
267 out.push(code);
268 }
269 at += len;
270 }
271 }
272
273 pub fn decompress(&self, input: &[u8], out: &mut Vec<u8>) -> Result<()> {
290 out.reserve(input.len().saturating_mul(2).saturating_add(MAX_SYMBOL_LEN));
294 let mut at = 0;
295 while at < input.len() {
296 let code = input[at];
297 at += 1;
298 if code == ESCAPE {
299 let literal = *input.get(at).ok_or_else(|| truncated("an escaped byte"))?;
300 out.push(literal);
301 at += 1;
302 } else {
303 let symbol = *self
304 .symbols
305 .get(code as usize)
306 .ok_or_else(|| Error::internal(format!("code {code} is not in the table")))?;
307 out.extend_from_slice(&symbol.value.to_le_bytes());
308 out.truncate(out.len() - (MAX_SYMBOL_LEN - symbol.len()));
309 }
310 }
311 Ok(())
312 }
313
314 fn match_at(&self, input: &[u8], at: usize) -> (u8, usize) {
316 let remaining = input.len() - at;
317 let word = load(input, at);
318 if remaining >= 3 {
321 if let Some((symbol, code)) = self.probe(word, remaining) {
322 return (code, symbol.len());
323 }
324 }
325 if remaining >= 2 {
326 let code = self.pair[(word & 0xffff) as usize];
327 if code != u16::MAX {
328 return (code as u8, 2);
329 }
330 }
331 let code = self.single[(word & 0xff) as usize];
332 if code == ESCAPE { (ESCAPE, 1) } else { (code, 1) }
333 }
334
335 fn probe(&self, word: u64, remaining: usize) -> Option<(Symbol, u8)> {
341 let mut slot = hash_of(word);
342 let mut best: Option<(Symbol, u8)> = None;
343 for _ in 0..PROBE {
344 match self.hash[slot] {
345 None => break,
346 Some((symbol, code)) => {
347 if symbol.len() <= remaining
348 && word & symbol.mask() == symbol.value
349 && best.is_none_or(|(found, _)| symbol.len() > found.len())
350 {
351 best = Some((symbol, code));
352 }
353 }
354 }
355 slot = (slot + 1) & (HASH_SLOTS - 1);
356 }
357 best
358 }
359
360 fn count(&self, input: &[u8], counts: &mut Counts) {
363 let mut at = 0;
364 let mut previous: Option<u16> = None;
365 while at < input.len() {
366 let (code, len) = self.match_at(input, at);
367 let id = if code == ESCAPE { 256 + u16::from(input[at]) } else { u16::from(code) };
368 counts.one(id);
369 if let Some(previous) = previous {
370 counts.two(previous, id);
371 }
372 previous = Some(id);
373 at += len;
374 }
375 }
376
377 fn build(symbols: Vec<Symbol>) -> Self {
378 let mut table = Self {
379 symbols,
380 single: vec![ESCAPE; 256],
381 pair: vec![u16::MAX; 65536],
382 hash: vec![None; HASH_SLOTS],
383 };
384 let mut order: Vec<(Symbol, u8)> =
387 table.symbols.iter().enumerate().map(|(code, symbol)| (*symbol, code as u8)).collect();
388 order.sort_by_key(|(symbol, code)| (std::cmp::Reverse(symbol.len()), *code));
389 for (symbol, code) in order {
390 match symbol.len() {
391 1 => {
392 let index = (symbol.value & 0xff) as usize;
393 if table.single[index] == ESCAPE {
394 table.single[index] = code;
395 }
396 }
397 2 => {
398 let index = (symbol.value & 0xffff) as usize;
399 if table.pair[index] == u16::MAX {
400 table.pair[index] = u16::from(code);
401 }
402 }
403 _ => {
404 let mut slot = hash_of(symbol.value);
405 for _ in 0..PROBE {
406 if table.hash[slot].is_none() {
407 table.hash[slot] = Some((symbol, code));
408 break;
409 }
410 slot = (slot + 1) & (HASH_SLOTS - 1);
411 }
412 }
413 }
414 }
415 table
416 }
417}
418
419fn load(input: &[u8], at: usize) -> u64 {
424 if at + 8 <= input.len() {
425 let bytes: [u8; 8] = input[at..at + 8].try_into().expect("eight bytes were checked");
426 u64::from_le_bytes(bytes)
427 } else {
428 let mut word = 0u64;
429 for (index, byte) in input[at..].iter().enumerate() {
430 word |= u64::from(*byte) << (8 * index);
431 }
432 word
433 }
434}
435
436fn hash_of(word: u64) -> usize {
440 let key = word & 0xff_ffff;
441 ((key.wrapping_mul(0x9e37_79b9_7f4a_7c15)) >> (64 - HASH_SLOTS.trailing_zeros())) as usize
442}
443
444struct Counts {
448 single: Vec<u32>,
449 pairs: Vec<u32>,
452 seen: Vec<u32>,
455}
456
457const IDS: usize = 512;
459
460impl Counts {
461 fn new() -> Self {
462 Self { single: vec![0; IDS], pairs: vec![0; IDS * IDS], seen: Vec::new() }
463 }
464
465 fn clear(&mut self) {
466 self.single.fill(0);
467 for slot in self.seen.drain(..) {
468 self.pairs[slot as usize] = 0;
469 }
470 }
471
472 fn one(&mut self, id: u16) {
473 self.single[id as usize] += 1;
474 }
475
476 fn two(&mut self, first: u16, second: u16) {
477 let slot = first as usize * IDS + second as usize;
478 if self.pairs[slot] == 0 {
479 self.seen.push(slot as u32);
480 }
481 self.pairs[slot] += 1;
482 }
483
484 fn best(&self, table: &SymbolTable) -> Vec<Symbol> {
491 let mut gains: Vec<(Symbol, u64)> = Vec::with_capacity(IDS + self.seen.len());
492 for (id, count) in self.single.iter().enumerate() {
493 if *count == 0 {
494 continue;
495 }
496 let symbol = symbol_of(table, id as u16);
497 gains.push((symbol, u64::from(*count) * symbol.len() as u64));
498 }
499 for slot in &self.seen {
500 let slot = *slot as usize;
501 let (first, second) = ((slot / IDS) as u16, (slot % IDS) as u16);
502 let symbol = symbol_of(table, first).concat(symbol_of(table, second));
503 gains.push((symbol, u64::from(self.pairs[slot]) * symbol.len() as u64));
504 }
505 gains.sort_unstable_by_key(|(symbol, _)| *symbol);
508 gains.dedup_by(|next, kept| {
509 let same = next.0 == kept.0;
510 if same {
511 kept.1 += next.1;
512 }
513 same
514 });
515 let order = |left: &(Symbol, u64), right: &(Symbol, u64)| {
519 right.1.cmp(&left.1).then(left.0.cmp(&right.0))
520 };
521 if gains.len() > MAX_SYMBOLS {
522 gains.select_nth_unstable_by(MAX_SYMBOLS - 1, order);
523 gains.truncate(MAX_SYMBOLS);
524 }
525 gains.sort_unstable_by(order);
526 gains.into_iter().map(|(symbol, _)| symbol).collect()
527 }
528}
529
530fn symbol_of(table: &SymbolTable, id: u16) -> Symbol {
531 if (id as usize) < table.symbols.len() {
532 table.symbols[id as usize]
533 } else {
534 Symbol::single((id.saturating_sub(256)) as u8)
535 }
536}
537
538fn truncated(what: &str) -> Error {
539 Error::internal(format!("the input ended in the middle of {what}"))
540}
541
542#[cfg(test)]
543mod tests {
544 use super::*;
545
546 fn urls() -> Vec<Vec<u8>> {
550 let hosts = ["www.example.com", "shop.example.com", "news.other.example.org"];
551 let paths = ["/index.html", "/catalog/item", "/search", "/user/profile/settings"];
552 let mut out = Vec::new();
553 for index in 0..600 {
554 let host = hosts[index % hosts.len()];
555 let path = paths[(index / 3) % paths.len()];
556 out.push(
557 format!("http://{host}{path}?session={}&ref=google&page={}", index * 7, index % 20)
558 .into_bytes(),
559 );
560 }
561 out
562 }
563
564 fn borrow(strings: &[Vec<u8>]) -> Vec<&[u8]> {
565 strings.iter().map(Vec::as_slice).collect()
566 }
567
568 fn round_trip(table: &SymbolTable, strings: &[Vec<u8>]) -> (usize, usize) {
569 let mut raw = 0;
570 let mut compressed = 0;
571 for string in strings {
572 let mut bytes = Vec::new();
573 table.compress(string, &mut bytes);
574 let mut back = Vec::new();
575 table.decompress(&bytes, &mut back).unwrap();
576 assert_eq!(back, *string, "{}", String::from_utf8_lossy(string));
577 raw += string.len();
578 compressed += bytes.len();
579 }
580 (raw, compressed)
581 }
582
583 #[test]
584 fn urls_compress_by_more_than_half_and_come_back_unchanged() {
585 let strings = urls();
588 let table = SymbolTable::train(&borrow(&strings));
589 let (raw, compressed) = round_trip(&table, &strings);
590 let ratio = raw as f64 / compressed as f64;
591 assert!(ratio > 2.5, "{ratio:.2}x, {raw} to {compressed}");
592 assert!(table.len() > 100, "{} symbols", table.len());
593 }
594
595 #[test]
596 fn the_trainer_finds_the_long_repeated_pieces() {
597 let strings = urls();
598 let table = SymbolTable::train(&borrow(&strings));
599 let found: Vec<String> = (0..table.len())
600 .map(|code| String::from_utf8_lossy(&table.symbols[code].bytes()).into_owned())
601 .collect();
602 let long = found.iter().filter(|symbol| symbol.len() >= 6).count();
605 assert!(long > 60, "only {long} symbols of six bytes or more: {found:?}");
606 }
607
608 #[test]
609 fn english_text_round_trips_and_shrinks() {
610 let text: Vec<Vec<u8>> = "the quick brown fox jumps over the lazy dog while the other dog \
611 watches the fox and the dog and the fox go over the hill together"
612 .split(' ')
613 .map(|word| word.as_bytes().to_vec())
614 .collect();
615 let table = SymbolTable::train(&borrow(&text));
616 let (raw, compressed) = round_trip(&table, &text);
617 assert!(compressed < raw, "{raw} to {compressed}");
618 }
619
620 #[test]
621 fn incompressible_bytes_round_trip_and_cost_what_escaping_costs() {
622 let mut state = 0x1234_5678_9abc_def0u64;
626 let strings: Vec<Vec<u8>> = (0..100)
627 .map(|_| {
628 (0..64)
629 .map(|_| {
630 state ^= state << 13;
631 state ^= state >> 7;
632 state ^= state << 17;
633 state as u8
634 })
635 .collect()
636 })
637 .collect();
638 let table = SymbolTable::train(&borrow(&strings));
639 let (raw, compressed) = round_trip(&table, &strings);
640 assert!(compressed < raw * 2, "{raw} to {compressed}");
641 }
642
643 #[test]
644 fn an_empty_table_escapes_everything_and_still_round_trips() {
645 let table = SymbolTable::empty();
646 let strings = vec![b"hello".to_vec(), Vec::new(), b"x".to_vec()];
647 let (raw, compressed) = round_trip(&table, &strings);
648 assert_eq!(compressed, raw * 2);
649 }
650
651 #[test]
652 fn an_empty_string_compresses_to_nothing() {
653 let table = SymbolTable::train(&[b"abcabcabc"]);
654 let mut out = Vec::new();
655 table.compress(b"", &mut out);
656 assert!(out.is_empty());
657 let mut back = Vec::new();
658 table.decompress(&out, &mut back).unwrap();
659 assert!(back.is_empty());
660 }
661
662 #[test]
663 fn a_string_shorter_than_the_symbols_does_not_read_past_its_end() {
664 let table = SymbolTable::train(&[b"ab\0ab\0ab\0ab\0", b"abcdefgh"]);
667 for string in [b"a".to_vec(), b"ab".to_vec(), b"abc".to_vec()] {
668 let mut bytes = Vec::new();
669 table.compress(&string, &mut bytes);
670 let mut back = Vec::new();
671 table.decompress(&bytes, &mut back).unwrap();
672 assert_eq!(back, string);
673 }
674 }
675
676 #[test]
677 fn a_table_survives_being_written_and_read_back() {
678 let strings = urls();
679 let table = SymbolTable::train(&borrow(&strings));
680 let mut bytes = Vec::new();
681 table.serialize(&mut bytes);
682 assert_eq!(bytes.len(), table.serialized_len());
683 let (read, consumed) = SymbolTable::deserialize(&bytes).unwrap();
684 assert_eq!(consumed, bytes.len());
685 assert_eq!(read.symbols, table.symbols);
686
687 let mut first = Vec::new();
690 let mut second = Vec::new();
691 table.compress(&strings[7], &mut first);
692 read.compress(&strings[7], &mut second);
693 assert_eq!(first, second);
694 }
695
696 #[test]
697 fn a_full_table_is_two_kilobytes_at_the_very_most() {
698 let strings = urls();
699 let table = SymbolTable::train(&borrow(&strings));
700 assert!(table.serialized_len() <= 1 + MAX_SYMBOLS * (1 + MAX_SYMBOL_LEN));
701 assert!(table.serialized_len() <= 2049);
702 }
703
704 #[test]
705 fn a_truncated_symbol_table_is_an_error() {
706 let strings = urls();
707 let table = SymbolTable::train(&borrow(&strings));
708 let mut bytes = Vec::new();
709 table.serialize(&mut bytes);
710 for len in 1..bytes.len().min(40) {
711 let error = SymbolTable::deserialize(&bytes[..len]).unwrap_err();
712 assert!(error.message().contains("ended in the middle"), "{error}");
713 }
714 }
715
716 #[test]
717 fn a_symbol_of_zero_bytes_is_an_error() {
718 let error = SymbolTable::deserialize(&[1, 0]).unwrap_err();
719 assert!(error.message().contains("is not a symbol"), "{error}");
720 }
721
722 #[test]
723 fn a_short_symbol_does_not_drag_the_rest_of_its_word_out_with_it() {
724 let table = SymbolTable::train(&[b"abcdabcdabcdabcd"]);
730 let mut compressed = Vec::new();
731 table.compress(b"abcdabcd", &mut compressed);
732 let mut out = Vec::new();
733 table.decompress(&compressed, &mut out).expect("decompresses");
734 table.decompress(&compressed, &mut out).expect("decompresses");
735 assert_eq!(out, b"abcdabcdabcdabcd");
736 }
737
738 #[test]
739 fn a_dangling_escape_is_an_error_and_not_a_panic() {
740 let table = SymbolTable::train(&[b"abcabcabc"]);
741 let error = table.decompress(&[ESCAPE], &mut Vec::new()).unwrap_err();
742 assert!(error.message().contains("escaped byte"), "{error}");
743 }
744
745 #[test]
746 fn a_code_the_table_does_not_have_is_an_error() {
747 let table = SymbolTable::train(&[b"abcabcabc"]);
748 let code = table.len() as u8;
749 let error = table.decompress(&[code], &mut Vec::new()).unwrap_err();
750 assert!(error.message().contains("not in the table"), "{error}");
751 }
752
753 #[test]
754 fn training_twice_on_the_same_sample_gives_the_same_table() {
755 let strings = urls();
757 let first = SymbolTable::train(&borrow(&strings));
758 let second = SymbolTable::train(&borrow(&strings));
759 assert_eq!(first.symbols, second.symbols);
760 }
761
762 #[test]
763 fn the_longest_match_wins_rather_than_the_first_one_found() {
764 let table = SymbolTable::build(vec![
765 Symbol::new(b"abc"),
766 Symbol::new(b"abcdef"),
767 Symbol::new(b"abcd"),
768 ]);
769 let mut out = Vec::new();
770 table.compress(b"abcdef", &mut out);
771 assert_eq!(out, vec![1]);
772 }
773
774 #[test]
775 fn a_symbol_longer_than_what_is_left_is_not_used() {
776 let table = SymbolTable::build(vec![Symbol::new(b"abcdef"), Symbol::new(b"ab")]);
777 let mut out = Vec::new();
778 table.compress(b"abcd", &mut out);
779 assert_eq!(out, vec![1, ESCAPE, b'c', ESCAPE, b'd']);
781 }
782
783 #[test]
784 fn concatenation_stops_at_eight_bytes() {
785 let long = Symbol::new(b"abcdef");
786 assert_eq!(long.concat(Symbol::new(b"ghijkl")).bytes(), b"abcdefgh");
787 assert_eq!(Symbol::new(b"ab").concat(Symbol::new(b"cd")).bytes(), b"abcd");
788 }
789
790 fn train_with_maps(samples: &[&[u8]]) -> SymbolTable {
793 use std::collections::HashMap;
794 let mut table = SymbolTable::empty();
795 for _ in 0..GENERATIONS {
796 let mut single = [0u32; IDS];
797 let mut pairs: HashMap<(u16, u16), u32> = HashMap::new();
798 for sample in samples {
799 let mut at = 0;
800 let mut previous: Option<u16> = None;
801 while at < sample.len() {
802 let (code, len) = table.match_at(sample, at);
803 let id =
804 if code == ESCAPE { 256 + u16::from(sample[at]) } else { u16::from(code) };
805 single[id as usize] += 1;
806 if let Some(previous) = previous {
807 *pairs.entry((previous, id)).or_insert(0) += 1;
808 }
809 previous = Some(id);
810 at += len;
811 }
812 }
813 let mut gains: HashMap<Symbol, u64> = HashMap::new();
814 for (id, count) in single.iter().enumerate().filter(|(_, count)| **count > 0) {
815 let symbol = symbol_of(&table, id as u16);
816 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
817 }
818 for ((first, second), count) in &pairs {
819 let symbol = symbol_of(&table, *first).concat(symbol_of(&table, *second));
820 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
821 }
822 let mut ranked: Vec<(Symbol, u64)> = gains.into_iter().collect();
823 ranked.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
824 ranked.truncate(MAX_SYMBOLS);
825 if ranked.is_empty() {
826 break;
827 }
828 table = SymbolTable::build(ranked.into_iter().map(|(symbol, _)| symbol).collect());
829 }
830 table
831 }
832
833 #[test]
834 fn flat_counts_train_the_same_table_as_hash_maps() {
835 let mut state = 0x9e37_79b9_7f4a_7c15u64;
836 let mut next = move || {
837 state ^= state << 13;
838 state ^= state >> 7;
839 state ^= state << 17;
840 state
841 };
842 let mut shapes: Vec<Vec<Vec<u8>>> = vec![urls(), Vec::new(), vec![Vec::new(); 3]];
843 shapes.push(
845 (0..400).map(|_| (0..12).map(|_| b"abc"[(next() % 3) as usize]).collect()).collect(),
846 );
847 shapes.push((0..300).map(|_| (0..40).map(|_| next() as u8).collect()).collect());
849 shapes.push(
851 (0..200)
852 .map(|index| format!("prefix-{}-suffix-{}", index % 7, index % 5).into_bytes())
853 .collect(),
854 );
855 for strings in &shapes {
856 let samples = borrow(strings);
857 assert_eq!(SymbolTable::train(&samples).symbols, train_with_maps(&samples).symbols);
858 }
859 }
860}