1use std::collections::HashMap;
46
47use rudb_common::{Error, Result};
48
49pub const ESCAPE: u8 = 255;
52
53pub const MAX_SYMBOLS: usize = 255;
55
56pub const MAX_SYMBOL_LEN: usize = 8;
59
60const GENERATIONS: usize = 5;
62
63const HASH_SLOTS: usize = 1024;
66
67const PROBE: usize = 8;
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
73struct Symbol {
74 value: u64,
75 len: u8,
76}
77
78impl Symbol {
79 fn new(bytes: &[u8]) -> Self {
80 let len = bytes.len().min(MAX_SYMBOL_LEN);
81 let mut value = 0u64;
82 for (index, byte) in bytes[..len].iter().enumerate() {
83 value |= u64::from(*byte) << (8 * index);
84 }
85 Self { value, len: len as u8 }
86 }
87
88 fn single(byte: u8) -> Self {
89 Self { value: u64::from(byte), len: 1 }
90 }
91
92 fn len(self) -> usize {
93 self.len as usize
94 }
95
96 fn mask(self) -> u64 {
97 mask_of(self.len())
98 }
99
100 fn bytes(self) -> Vec<u8> {
101 (0..self.len()).map(|index| (self.value >> (8 * index)) as u8).collect()
102 }
103
104 fn concat(self, other: Self) -> Self {
106 if self.len() >= MAX_SYMBOL_LEN {
107 return self;
108 }
109 let len = (self.len() + other.len()).min(MAX_SYMBOL_LEN);
110 let value = self.value | (other.value << (8 * self.len()));
111 Self { value: value & mask_of(len), len: len as u8 }
112 }
113}
114
115fn mask_of(len: usize) -> u64 {
116 if len >= 8 { u64::MAX } else { (1u64 << (8 * len)) - 1 }
117}
118
119pub struct SymbolTable {
121 symbols: Vec<Symbol>,
123 single: Vec<u8>,
125 pair: Vec<u16>,
127 hash: Vec<Option<(Symbol, u8)>>,
129}
130
131impl std::fmt::Debug for SymbolTable {
132 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133 formatter
135 .debug_struct("SymbolTable")
136 .field("symbols", &self.symbols.len())
137 .field("bytes", &self.serialized_len())
138 .finish()
139 }
140}
141
142impl PartialEq for SymbolTable {
149 fn eq(&self, other: &Self) -> bool {
150 self.symbols == other.symbols
151 }
152}
153
154impl Eq for SymbolTable {}
155
156impl SymbolTable {
157 #[must_use]
164 pub fn footprint(&self) -> usize {
165 size_of::<Self>()
166 + self.symbols.capacity() * size_of::<Symbol>()
167 + self.single.capacity()
168 + self.pair.capacity() * size_of::<u16>()
169 + self.hash.capacity() * size_of::<Option<(Symbol, u8)>>()
170 }
171
172 #[must_use]
175 pub fn empty() -> Self {
176 Self::build(Vec::new())
177 }
178
179 #[must_use]
185 pub fn train(samples: &[&[u8]]) -> Self {
186 let mut table = Self::empty();
187 for _ in 0..GENERATIONS {
188 let mut counts = Counts::new();
189 for sample in samples {
190 table.count(sample, &mut counts);
191 }
192 let next = counts.best(&table);
193 if next.is_empty() {
194 break;
195 }
196 table = Self::build(next);
197 }
198 table
199 }
200
201 #[must_use]
203 pub fn len(&self) -> usize {
204 self.symbols.len()
205 }
206
207 #[must_use]
209 pub fn is_empty(&self) -> bool {
210 self.symbols.is_empty()
211 }
212
213 #[must_use]
217 pub fn serialized_len(&self) -> usize {
218 1 + self.symbols.iter().map(|symbol| 1 + symbol.len()).sum::<usize>()
219 }
220
221 pub fn serialize(&self, out: &mut Vec<u8>) {
223 out.push(self.symbols.len() as u8);
224 for symbol in &self.symbols {
225 out.push(symbol.len);
226 out.extend_from_slice(&symbol.bytes());
227 }
228 }
229
230 pub fn deserialize(bytes: &[u8]) -> Result<(Self, usize)> {
236 let count = *bytes.first().ok_or_else(|| truncated("a symbol table header"))? as usize;
237 let mut at = 1;
238 let mut symbols = Vec::with_capacity(count);
239 for _ in 0..count {
240 let len = *bytes.get(at).ok_or_else(|| truncated("a symbol length"))? as usize;
241 if len == 0 || len > MAX_SYMBOL_LEN {
242 return Err(Error::internal(format!("a symbol of {len} bytes is not a symbol")));
243 }
244 at += 1;
245 let end = at + len;
246 if end > bytes.len() {
247 return Err(truncated("a symbol"));
248 }
249 symbols.push(Symbol::new(&bytes[at..end]));
250 at = end;
251 }
252 Ok((Self::build(symbols), at))
253 }
254
255 pub fn compress(&self, input: &[u8], out: &mut Vec<u8>) {
261 let mut at = 0;
262 while at < input.len() {
263 let (code, len) = self.match_at(input, at);
264 if code == ESCAPE {
265 out.push(ESCAPE);
266 out.push(input[at]);
267 } else {
268 out.push(code);
269 }
270 at += len;
271 }
272 }
273
274 pub fn decompress(&self, input: &[u8], out: &mut Vec<u8>) -> Result<()> {
291 out.reserve(input.len().saturating_mul(2).saturating_add(MAX_SYMBOL_LEN));
295 let mut at = 0;
296 while at < input.len() {
297 let code = input[at];
298 at += 1;
299 if code == ESCAPE {
300 let literal = *input.get(at).ok_or_else(|| truncated("an escaped byte"))?;
301 out.push(literal);
302 at += 1;
303 } else {
304 let symbol = *self
305 .symbols
306 .get(code as usize)
307 .ok_or_else(|| Error::internal(format!("code {code} is not in the table")))?;
308 out.extend_from_slice(&symbol.value.to_le_bytes());
309 out.truncate(out.len() - (MAX_SYMBOL_LEN - symbol.len()));
310 }
311 }
312 Ok(())
313 }
314
315 fn match_at(&self, input: &[u8], at: usize) -> (u8, usize) {
317 let remaining = input.len() - at;
318 let word = load(input, at);
319 if remaining >= 3 {
322 if let Some((symbol, code)) = self.probe(word, remaining) {
323 return (code, symbol.len());
324 }
325 }
326 if remaining >= 2 {
327 let code = self.pair[(word & 0xffff) as usize];
328 if code != u16::MAX {
329 return (code as u8, 2);
330 }
331 }
332 let code = self.single[(word & 0xff) as usize];
333 if code == ESCAPE { (ESCAPE, 1) } else { (code, 1) }
334 }
335
336 fn probe(&self, word: u64, remaining: usize) -> Option<(Symbol, u8)> {
342 let mut slot = hash_of(word);
343 let mut best: Option<(Symbol, u8)> = None;
344 for _ in 0..PROBE {
345 match self.hash[slot] {
346 None => break,
347 Some((symbol, code)) => {
348 if symbol.len() <= remaining
349 && word & symbol.mask() == symbol.value
350 && best.is_none_or(|(found, _)| symbol.len() > found.len())
351 {
352 best = Some((symbol, code));
353 }
354 }
355 }
356 slot = (slot + 1) & (HASH_SLOTS - 1);
357 }
358 best
359 }
360
361 fn count(&self, input: &[u8], counts: &mut Counts) {
364 let mut at = 0;
365 let mut previous: Option<u16> = None;
366 while at < input.len() {
367 let (code, len) = self.match_at(input, at);
368 let id = if code == ESCAPE { 256 + u16::from(input[at]) } else { u16::from(code) };
369 counts.one(id);
370 if let Some(previous) = previous {
371 counts.two(previous, id);
372 }
373 previous = Some(id);
374 at += len;
375 }
376 }
377
378 fn build(symbols: Vec<Symbol>) -> Self {
379 let mut table = Self {
380 symbols,
381 single: vec![ESCAPE; 256],
382 pair: vec![u16::MAX; 65536],
383 hash: vec![None; HASH_SLOTS],
384 };
385 let mut order: Vec<(Symbol, u8)> =
388 table.symbols.iter().enumerate().map(|(code, symbol)| (*symbol, code as u8)).collect();
389 order.sort_by_key(|(symbol, code)| (std::cmp::Reverse(symbol.len()), *code));
390 for (symbol, code) in order {
391 match symbol.len() {
392 1 => {
393 let index = (symbol.value & 0xff) as usize;
394 if table.single[index] == ESCAPE {
395 table.single[index] = code;
396 }
397 }
398 2 => {
399 let index = (symbol.value & 0xffff) as usize;
400 if table.pair[index] == u16::MAX {
401 table.pair[index] = u16::from(code);
402 }
403 }
404 _ => {
405 let mut slot = hash_of(symbol.value);
406 for _ in 0..PROBE {
407 if table.hash[slot].is_none() {
408 table.hash[slot] = Some((symbol, code));
409 break;
410 }
411 slot = (slot + 1) & (HASH_SLOTS - 1);
412 }
413 }
414 }
415 }
416 table
417 }
418}
419
420fn load(input: &[u8], at: usize) -> u64 {
425 if at + 8 <= input.len() {
426 let bytes: [u8; 8] = input[at..at + 8].try_into().expect("eight bytes were checked");
427 u64::from_le_bytes(bytes)
428 } else {
429 let mut word = 0u64;
430 for (index, byte) in input[at..].iter().enumerate() {
431 word |= u64::from(*byte) << (8 * index);
432 }
433 word
434 }
435}
436
437fn hash_of(word: u64) -> usize {
441 let key = word & 0xff_ffff;
442 ((key.wrapping_mul(0x9e37_79b9_7f4a_7c15)) >> (64 - HASH_SLOTS.trailing_zeros())) as usize
443}
444
445struct Counts {
449 single: Vec<u32>,
450 pairs: HashMap<(u16, u16), u32>,
451}
452
453impl Counts {
454 fn new() -> Self {
455 Self { single: vec![0; 512], pairs: HashMap::new() }
456 }
457
458 fn one(&mut self, id: u16) {
459 self.single[id as usize] += 1;
460 }
461
462 fn two(&mut self, first: u16, second: u16) {
463 *self.pairs.entry((first, second)).or_insert(0) += 1;
464 }
465
466 fn best(&self, table: &SymbolTable) -> Vec<Symbol> {
473 let mut gains: HashMap<Symbol, u64> = HashMap::new();
474 for (id, count) in self.single.iter().enumerate() {
475 if *count == 0 {
476 continue;
477 }
478 let symbol = symbol_of(table, id as u16);
479 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
480 }
481 for ((first, second), count) in &self.pairs {
482 let symbol = symbol_of(table, *first).concat(symbol_of(table, *second));
483 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
484 }
485 let mut ranked: Vec<(Symbol, u64)> = gains.into_iter().collect();
486 ranked.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
489 ranked.truncate(MAX_SYMBOLS);
490 ranked.into_iter().map(|(symbol, _)| symbol).collect()
491 }
492}
493
494fn symbol_of(table: &SymbolTable, id: u16) -> Symbol {
495 if (id as usize) < table.symbols.len() {
496 table.symbols[id as usize]
497 } else {
498 Symbol::single((id.saturating_sub(256)) as u8)
499 }
500}
501
502fn truncated(what: &str) -> Error {
503 Error::internal(format!("the input ended in the middle of {what}"))
504}
505
506#[cfg(test)]
507mod tests {
508 use super::*;
509
510 fn urls() -> Vec<Vec<u8>> {
514 let hosts = ["www.example.com", "shop.example.com", "news.other.example.org"];
515 let paths = ["/index.html", "/catalog/item", "/search", "/user/profile/settings"];
516 let mut out = Vec::new();
517 for index in 0..600 {
518 let host = hosts[index % hosts.len()];
519 let path = paths[(index / 3) % paths.len()];
520 out.push(
521 format!("http://{host}{path}?session={}&ref=google&page={}", index * 7, index % 20)
522 .into_bytes(),
523 );
524 }
525 out
526 }
527
528 fn borrow(strings: &[Vec<u8>]) -> Vec<&[u8]> {
529 strings.iter().map(Vec::as_slice).collect()
530 }
531
532 fn round_trip(table: &SymbolTable, strings: &[Vec<u8>]) -> (usize, usize) {
533 let mut raw = 0;
534 let mut compressed = 0;
535 for string in strings {
536 let mut bytes = Vec::new();
537 table.compress(string, &mut bytes);
538 let mut back = Vec::new();
539 table.decompress(&bytes, &mut back).unwrap();
540 assert_eq!(back, *string, "{}", String::from_utf8_lossy(string));
541 raw += string.len();
542 compressed += bytes.len();
543 }
544 (raw, compressed)
545 }
546
547 #[test]
548 fn urls_compress_by_more_than_half_and_come_back_unchanged() {
549 let strings = urls();
552 let table = SymbolTable::train(&borrow(&strings));
553 let (raw, compressed) = round_trip(&table, &strings);
554 let ratio = raw as f64 / compressed as f64;
555 assert!(ratio > 2.5, "{ratio:.2}x, {raw} to {compressed}");
556 assert!(table.len() > 100, "{} symbols", table.len());
557 }
558
559 #[test]
560 fn the_trainer_finds_the_long_repeated_pieces() {
561 let strings = urls();
562 let table = SymbolTable::train(&borrow(&strings));
563 let found: Vec<String> = (0..table.len())
564 .map(|code| String::from_utf8_lossy(&table.symbols[code].bytes()).into_owned())
565 .collect();
566 let long = found.iter().filter(|symbol| symbol.len() >= 6).count();
569 assert!(long > 60, "only {long} symbols of six bytes or more: {found:?}");
570 }
571
572 #[test]
573 fn english_text_round_trips_and_shrinks() {
574 let text: Vec<Vec<u8>> = "the quick brown fox jumps over the lazy dog while the other dog \
575 watches the fox and the dog and the fox go over the hill together"
576 .split(' ')
577 .map(|word| word.as_bytes().to_vec())
578 .collect();
579 let table = SymbolTable::train(&borrow(&text));
580 let (raw, compressed) = round_trip(&table, &text);
581 assert!(compressed < raw, "{raw} to {compressed}");
582 }
583
584 #[test]
585 fn incompressible_bytes_round_trip_and_cost_what_escaping_costs() {
586 let mut state = 0x1234_5678_9abc_def0u64;
590 let strings: Vec<Vec<u8>> = (0..100)
591 .map(|_| {
592 (0..64)
593 .map(|_| {
594 state ^= state << 13;
595 state ^= state >> 7;
596 state ^= state << 17;
597 state as u8
598 })
599 .collect()
600 })
601 .collect();
602 let table = SymbolTable::train(&borrow(&strings));
603 let (raw, compressed) = round_trip(&table, &strings);
604 assert!(compressed < raw * 2, "{raw} to {compressed}");
605 }
606
607 #[test]
608 fn an_empty_table_escapes_everything_and_still_round_trips() {
609 let table = SymbolTable::empty();
610 let strings = vec![b"hello".to_vec(), Vec::new(), b"x".to_vec()];
611 let (raw, compressed) = round_trip(&table, &strings);
612 assert_eq!(compressed, raw * 2);
613 }
614
615 #[test]
616 fn an_empty_string_compresses_to_nothing() {
617 let table = SymbolTable::train(&[b"abcabcabc"]);
618 let mut out = Vec::new();
619 table.compress(b"", &mut out);
620 assert!(out.is_empty());
621 let mut back = Vec::new();
622 table.decompress(&out, &mut back).unwrap();
623 assert!(back.is_empty());
624 }
625
626 #[test]
627 fn a_string_shorter_than_the_symbols_does_not_read_past_its_end() {
628 let table = SymbolTable::train(&[b"ab\0ab\0ab\0ab\0", b"abcdefgh"]);
631 for string in [b"a".to_vec(), b"ab".to_vec(), b"abc".to_vec()] {
632 let mut bytes = Vec::new();
633 table.compress(&string, &mut bytes);
634 let mut back = Vec::new();
635 table.decompress(&bytes, &mut back).unwrap();
636 assert_eq!(back, string);
637 }
638 }
639
640 #[test]
641 fn a_table_survives_being_written_and_read_back() {
642 let strings = urls();
643 let table = SymbolTable::train(&borrow(&strings));
644 let mut bytes = Vec::new();
645 table.serialize(&mut bytes);
646 assert_eq!(bytes.len(), table.serialized_len());
647 let (read, consumed) = SymbolTable::deserialize(&bytes).unwrap();
648 assert_eq!(consumed, bytes.len());
649 assert_eq!(read.symbols, table.symbols);
650
651 let mut first = Vec::new();
654 let mut second = Vec::new();
655 table.compress(&strings[7], &mut first);
656 read.compress(&strings[7], &mut second);
657 assert_eq!(first, second);
658 }
659
660 #[test]
661 fn a_full_table_is_two_kilobytes_at_the_very_most() {
662 let strings = urls();
663 let table = SymbolTable::train(&borrow(&strings));
664 assert!(table.serialized_len() <= 1 + MAX_SYMBOLS * (1 + MAX_SYMBOL_LEN));
665 assert!(table.serialized_len() <= 2049);
666 }
667
668 #[test]
669 fn a_truncated_symbol_table_is_an_error() {
670 let strings = urls();
671 let table = SymbolTable::train(&borrow(&strings));
672 let mut bytes = Vec::new();
673 table.serialize(&mut bytes);
674 for len in 1..bytes.len().min(40) {
675 let error = SymbolTable::deserialize(&bytes[..len]).unwrap_err();
676 assert!(error.message().contains("ended in the middle"), "{error}");
677 }
678 }
679
680 #[test]
681 fn a_symbol_of_zero_bytes_is_an_error() {
682 let error = SymbolTable::deserialize(&[1, 0]).unwrap_err();
683 assert!(error.message().contains("is not a symbol"), "{error}");
684 }
685
686 #[test]
687 fn a_short_symbol_does_not_drag_the_rest_of_its_word_out_with_it() {
688 let table = SymbolTable::train(&[b"abcdabcdabcdabcd"]);
694 let mut compressed = Vec::new();
695 table.compress(b"abcdabcd", &mut compressed);
696 let mut out = Vec::new();
697 table.decompress(&compressed, &mut out).expect("decompresses");
698 table.decompress(&compressed, &mut out).expect("decompresses");
699 assert_eq!(out, b"abcdabcdabcdabcd");
700 }
701
702 #[test]
703 fn a_dangling_escape_is_an_error_and_not_a_panic() {
704 let table = SymbolTable::train(&[b"abcabcabc"]);
705 let error = table.decompress(&[ESCAPE], &mut Vec::new()).unwrap_err();
706 assert!(error.message().contains("escaped byte"), "{error}");
707 }
708
709 #[test]
710 fn a_code_the_table_does_not_have_is_an_error() {
711 let table = SymbolTable::train(&[b"abcabcabc"]);
712 let code = table.len() as u8;
713 let error = table.decompress(&[code], &mut Vec::new()).unwrap_err();
714 assert!(error.message().contains("not in the table"), "{error}");
715 }
716
717 #[test]
718 fn training_twice_on_the_same_sample_gives_the_same_table() {
719 let strings = urls();
722 let first = SymbolTable::train(&borrow(&strings));
723 let second = SymbolTable::train(&borrow(&strings));
724 assert_eq!(first.symbols, second.symbols);
725 }
726
727 #[test]
728 fn the_longest_match_wins_rather_than_the_first_one_found() {
729 let table = SymbolTable::build(vec![
730 Symbol::new(b"abc"),
731 Symbol::new(b"abcdef"),
732 Symbol::new(b"abcd"),
733 ]);
734 let mut out = Vec::new();
735 table.compress(b"abcdef", &mut out);
736 assert_eq!(out, vec![1]);
737 }
738
739 #[test]
740 fn a_symbol_longer_than_what_is_left_is_not_used() {
741 let table = SymbolTable::build(vec![Symbol::new(b"abcdef"), Symbol::new(b"ab")]);
742 let mut out = Vec::new();
743 table.compress(b"abcd", &mut out);
744 assert_eq!(out, vec![1, ESCAPE, b'c', ESCAPE, b'd']);
746 }
747
748 #[test]
749 fn concatenation_stops_at_eight_bytes() {
750 let long = Symbol::new(b"abcdef");
751 assert_eq!(long.concat(Symbol::new(b"ghijkl")).bytes(), b"abcdefgh");
752 assert_eq!(Symbol::new(b"ab").concat(Symbol::new(b"cd")).bytes(), b"abcd");
753 }
754}