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<()> {
280 let mut at = 0;
281 while at < input.len() {
282 let code = input[at];
283 at += 1;
284 if code == ESCAPE {
285 let literal = *input.get(at).ok_or_else(|| truncated("an escaped byte"))?;
286 out.push(literal);
287 at += 1;
288 } else {
289 let symbol = self
290 .symbols
291 .get(code as usize)
292 .ok_or_else(|| Error::internal(format!("code {code} is not in the table")))?;
293 out.extend_from_slice(&symbol.bytes());
294 }
295 }
296 Ok(())
297 }
298
299 fn match_at(&self, input: &[u8], at: usize) -> (u8, usize) {
301 let remaining = input.len() - at;
302 let word = load(input, at);
303 if remaining >= 3 {
306 if let Some((symbol, code)) = self.probe(word, remaining) {
307 return (code, symbol.len());
308 }
309 }
310 if remaining >= 2 {
311 let code = self.pair[(word & 0xffff) as usize];
312 if code != u16::MAX {
313 return (code as u8, 2);
314 }
315 }
316 let code = self.single[(word & 0xff) as usize];
317 if code == ESCAPE { (ESCAPE, 1) } else { (code, 1) }
318 }
319
320 fn probe(&self, word: u64, remaining: usize) -> Option<(Symbol, u8)> {
326 let mut slot = hash_of(word);
327 let mut best: Option<(Symbol, u8)> = None;
328 for _ in 0..PROBE {
329 match self.hash[slot] {
330 None => break,
331 Some((symbol, code)) => {
332 if symbol.len() <= remaining
333 && word & symbol.mask() == symbol.value
334 && best.is_none_or(|(found, _)| symbol.len() > found.len())
335 {
336 best = Some((symbol, code));
337 }
338 }
339 }
340 slot = (slot + 1) & (HASH_SLOTS - 1);
341 }
342 best
343 }
344
345 fn count(&self, input: &[u8], counts: &mut Counts) {
348 let mut at = 0;
349 let mut previous: Option<u16> = None;
350 while at < input.len() {
351 let (code, len) = self.match_at(input, at);
352 let id = if code == ESCAPE { 256 + u16::from(input[at]) } else { u16::from(code) };
353 counts.one(id);
354 if let Some(previous) = previous {
355 counts.two(previous, id);
356 }
357 previous = Some(id);
358 at += len;
359 }
360 }
361
362 fn build(symbols: Vec<Symbol>) -> Self {
363 let mut table = Self {
364 symbols,
365 single: vec![ESCAPE; 256],
366 pair: vec![u16::MAX; 65536],
367 hash: vec![None; HASH_SLOTS],
368 };
369 let mut order: Vec<(Symbol, u8)> =
372 table.symbols.iter().enumerate().map(|(code, symbol)| (*symbol, code as u8)).collect();
373 order.sort_by_key(|(symbol, code)| (std::cmp::Reverse(symbol.len()), *code));
374 for (symbol, code) in order {
375 match symbol.len() {
376 1 => {
377 let index = (symbol.value & 0xff) as usize;
378 if table.single[index] == ESCAPE {
379 table.single[index] = code;
380 }
381 }
382 2 => {
383 let index = (symbol.value & 0xffff) as usize;
384 if table.pair[index] == u16::MAX {
385 table.pair[index] = u16::from(code);
386 }
387 }
388 _ => {
389 let mut slot = hash_of(symbol.value);
390 for _ in 0..PROBE {
391 if table.hash[slot].is_none() {
392 table.hash[slot] = Some((symbol, code));
393 break;
394 }
395 slot = (slot + 1) & (HASH_SLOTS - 1);
396 }
397 }
398 }
399 }
400 table
401 }
402}
403
404fn load(input: &[u8], at: usize) -> u64 {
409 if at + 8 <= input.len() {
410 let bytes: [u8; 8] = input[at..at + 8].try_into().expect("eight bytes were checked");
411 u64::from_le_bytes(bytes)
412 } else {
413 let mut word = 0u64;
414 for (index, byte) in input[at..].iter().enumerate() {
415 word |= u64::from(*byte) << (8 * index);
416 }
417 word
418 }
419}
420
421fn hash_of(word: u64) -> usize {
425 let key = word & 0xff_ffff;
426 ((key.wrapping_mul(0x9e37_79b9_7f4a_7c15)) >> (64 - HASH_SLOTS.trailing_zeros())) as usize
427}
428
429struct Counts {
433 single: Vec<u32>,
434 pairs: HashMap<(u16, u16), u32>,
435}
436
437impl Counts {
438 fn new() -> Self {
439 Self { single: vec![0; 512], pairs: HashMap::new() }
440 }
441
442 fn one(&mut self, id: u16) {
443 self.single[id as usize] += 1;
444 }
445
446 fn two(&mut self, first: u16, second: u16) {
447 *self.pairs.entry((first, second)).or_insert(0) += 1;
448 }
449
450 fn best(&self, table: &SymbolTable) -> Vec<Symbol> {
457 let mut gains: HashMap<Symbol, u64> = HashMap::new();
458 for (id, count) in self.single.iter().enumerate() {
459 if *count == 0 {
460 continue;
461 }
462 let symbol = symbol_of(table, id as u16);
463 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
464 }
465 for ((first, second), count) in &self.pairs {
466 let symbol = symbol_of(table, *first).concat(symbol_of(table, *second));
467 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
468 }
469 let mut ranked: Vec<(Symbol, u64)> = gains.into_iter().collect();
470 ranked.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
473 ranked.truncate(MAX_SYMBOLS);
474 ranked.into_iter().map(|(symbol, _)| symbol).collect()
475 }
476}
477
478fn symbol_of(table: &SymbolTable, id: u16) -> Symbol {
479 if (id as usize) < table.symbols.len() {
480 table.symbols[id as usize]
481 } else {
482 Symbol::single((id.saturating_sub(256)) as u8)
483 }
484}
485
486fn truncated(what: &str) -> Error {
487 Error::internal(format!("the input ended in the middle of {what}"))
488}
489
490#[cfg(test)]
491mod tests {
492 use super::*;
493
494 fn urls() -> Vec<Vec<u8>> {
498 let hosts = ["www.example.com", "shop.example.com", "news.other.example.org"];
499 let paths = ["/index.html", "/catalog/item", "/search", "/user/profile/settings"];
500 let mut out = Vec::new();
501 for index in 0..600 {
502 let host = hosts[index % hosts.len()];
503 let path = paths[(index / 3) % paths.len()];
504 out.push(
505 format!("http://{host}{path}?session={}&ref=google&page={}", index * 7, index % 20)
506 .into_bytes(),
507 );
508 }
509 out
510 }
511
512 fn borrow(strings: &[Vec<u8>]) -> Vec<&[u8]> {
513 strings.iter().map(Vec::as_slice).collect()
514 }
515
516 fn round_trip(table: &SymbolTable, strings: &[Vec<u8>]) -> (usize, usize) {
517 let mut raw = 0;
518 let mut compressed = 0;
519 for string in strings {
520 let mut bytes = Vec::new();
521 table.compress(string, &mut bytes);
522 let mut back = Vec::new();
523 table.decompress(&bytes, &mut back).unwrap();
524 assert_eq!(back, *string, "{}", String::from_utf8_lossy(string));
525 raw += string.len();
526 compressed += bytes.len();
527 }
528 (raw, compressed)
529 }
530
531 #[test]
532 fn urls_compress_by_more_than_half_and_come_back_unchanged() {
533 let strings = urls();
536 let table = SymbolTable::train(&borrow(&strings));
537 let (raw, compressed) = round_trip(&table, &strings);
538 let ratio = raw as f64 / compressed as f64;
539 assert!(ratio > 2.5, "{ratio:.2}x, {raw} to {compressed}");
540 assert!(table.len() > 100, "{} symbols", table.len());
541 }
542
543 #[test]
544 fn the_trainer_finds_the_long_repeated_pieces() {
545 let strings = urls();
546 let table = SymbolTable::train(&borrow(&strings));
547 let found: Vec<String> = (0..table.len())
548 .map(|code| String::from_utf8_lossy(&table.symbols[code].bytes()).into_owned())
549 .collect();
550 let long = found.iter().filter(|symbol| symbol.len() >= 6).count();
553 assert!(long > 60, "only {long} symbols of six bytes or more: {found:?}");
554 }
555
556 #[test]
557 fn english_text_round_trips_and_shrinks() {
558 let text: Vec<Vec<u8>> = "the quick brown fox jumps over the lazy dog while the other dog \
559 watches the fox and the dog and the fox go over the hill together"
560 .split(' ')
561 .map(|word| word.as_bytes().to_vec())
562 .collect();
563 let table = SymbolTable::train(&borrow(&text));
564 let (raw, compressed) = round_trip(&table, &text);
565 assert!(compressed < raw, "{raw} to {compressed}");
566 }
567
568 #[test]
569 fn incompressible_bytes_round_trip_and_cost_what_escaping_costs() {
570 let mut state = 0x1234_5678_9abc_def0u64;
574 let strings: Vec<Vec<u8>> = (0..100)
575 .map(|_| {
576 (0..64)
577 .map(|_| {
578 state ^= state << 13;
579 state ^= state >> 7;
580 state ^= state << 17;
581 state as u8
582 })
583 .collect()
584 })
585 .collect();
586 let table = SymbolTable::train(&borrow(&strings));
587 let (raw, compressed) = round_trip(&table, &strings);
588 assert!(compressed < raw * 2, "{raw} to {compressed}");
589 }
590
591 #[test]
592 fn an_empty_table_escapes_everything_and_still_round_trips() {
593 let table = SymbolTable::empty();
594 let strings = vec![b"hello".to_vec(), Vec::new(), b"x".to_vec()];
595 let (raw, compressed) = round_trip(&table, &strings);
596 assert_eq!(compressed, raw * 2);
597 }
598
599 #[test]
600 fn an_empty_string_compresses_to_nothing() {
601 let table = SymbolTable::train(&[b"abcabcabc"]);
602 let mut out = Vec::new();
603 table.compress(b"", &mut out);
604 assert!(out.is_empty());
605 let mut back = Vec::new();
606 table.decompress(&out, &mut back).unwrap();
607 assert!(back.is_empty());
608 }
609
610 #[test]
611 fn a_string_shorter_than_the_symbols_does_not_read_past_its_end() {
612 let table = SymbolTable::train(&[b"ab\0ab\0ab\0ab\0", b"abcdefgh"]);
615 for string in [b"a".to_vec(), b"ab".to_vec(), b"abc".to_vec()] {
616 let mut bytes = Vec::new();
617 table.compress(&string, &mut bytes);
618 let mut back = Vec::new();
619 table.decompress(&bytes, &mut back).unwrap();
620 assert_eq!(back, string);
621 }
622 }
623
624 #[test]
625 fn a_table_survives_being_written_and_read_back() {
626 let strings = urls();
627 let table = SymbolTable::train(&borrow(&strings));
628 let mut bytes = Vec::new();
629 table.serialize(&mut bytes);
630 assert_eq!(bytes.len(), table.serialized_len());
631 let (read, consumed) = SymbolTable::deserialize(&bytes).unwrap();
632 assert_eq!(consumed, bytes.len());
633 assert_eq!(read.symbols, table.symbols);
634
635 let mut first = Vec::new();
638 let mut second = Vec::new();
639 table.compress(&strings[7], &mut first);
640 read.compress(&strings[7], &mut second);
641 assert_eq!(first, second);
642 }
643
644 #[test]
645 fn a_full_table_is_two_kilobytes_at_the_very_most() {
646 let strings = urls();
647 let table = SymbolTable::train(&borrow(&strings));
648 assert!(table.serialized_len() <= 1 + MAX_SYMBOLS * (1 + MAX_SYMBOL_LEN));
649 assert!(table.serialized_len() <= 2049);
650 }
651
652 #[test]
653 fn a_truncated_symbol_table_is_an_error() {
654 let strings = urls();
655 let table = SymbolTable::train(&borrow(&strings));
656 let mut bytes = Vec::new();
657 table.serialize(&mut bytes);
658 for len in 1..bytes.len().min(40) {
659 let error = SymbolTable::deserialize(&bytes[..len]).unwrap_err();
660 assert!(error.message().contains("ended in the middle"), "{error}");
661 }
662 }
663
664 #[test]
665 fn a_symbol_of_zero_bytes_is_an_error() {
666 let error = SymbolTable::deserialize(&[1, 0]).unwrap_err();
667 assert!(error.message().contains("is not a symbol"), "{error}");
668 }
669
670 #[test]
671 fn a_dangling_escape_is_an_error_and_not_a_panic() {
672 let table = SymbolTable::train(&[b"abcabcabc"]);
673 let error = table.decompress(&[ESCAPE], &mut Vec::new()).unwrap_err();
674 assert!(error.message().contains("escaped byte"), "{error}");
675 }
676
677 #[test]
678 fn a_code_the_table_does_not_have_is_an_error() {
679 let table = SymbolTable::train(&[b"abcabcabc"]);
680 let code = table.len() as u8;
681 let error = table.decompress(&[code], &mut Vec::new()).unwrap_err();
682 assert!(error.message().contains("not in the table"), "{error}");
683 }
684
685 #[test]
686 fn training_twice_on_the_same_sample_gives_the_same_table() {
687 let strings = urls();
690 let first = SymbolTable::train(&borrow(&strings));
691 let second = SymbolTable::train(&borrow(&strings));
692 assert_eq!(first.symbols, second.symbols);
693 }
694
695 #[test]
696 fn the_longest_match_wins_rather_than_the_first_one_found() {
697 let table = SymbolTable::build(vec![
698 Symbol::new(b"abc"),
699 Symbol::new(b"abcdef"),
700 Symbol::new(b"abcd"),
701 ]);
702 let mut out = Vec::new();
703 table.compress(b"abcdef", &mut out);
704 assert_eq!(out, vec![1]);
705 }
706
707 #[test]
708 fn a_symbol_longer_than_what_is_left_is_not_used() {
709 let table = SymbolTable::build(vec![Symbol::new(b"abcdef"), Symbol::new(b"ab")]);
710 let mut out = Vec::new();
711 table.compress(b"abcd", &mut out);
712 assert_eq!(out, vec![1, ESCAPE, b'c', ESCAPE, b'd']);
714 }
715
716 #[test]
717 fn concatenation_stops_at_eight_bytes() {
718 let long = Symbol::new(b"abcdef");
719 assert_eq!(long.concat(Symbol::new(b"ghijkl")).bytes(), b"abcdefgh");
720 assert_eq!(Symbol::new(b"ab").concat(Symbol::new(b"cd")).bytes(), b"abcd");
721 }
722}