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 SymbolTable {
143 #[must_use]
146 pub fn empty() -> Self {
147 Self::build(Vec::new())
148 }
149
150 #[must_use]
156 pub fn train(samples: &[&[u8]]) -> Self {
157 let mut table = Self::empty();
158 for _ in 0..GENERATIONS {
159 let mut counts = Counts::new();
160 for sample in samples {
161 table.count(sample, &mut counts);
162 }
163 let next = counts.best(&table);
164 if next.is_empty() {
165 break;
166 }
167 table = Self::build(next);
168 }
169 table
170 }
171
172 #[must_use]
174 pub fn len(&self) -> usize {
175 self.symbols.len()
176 }
177
178 #[must_use]
180 pub fn is_empty(&self) -> bool {
181 self.symbols.is_empty()
182 }
183
184 #[must_use]
188 pub fn serialized_len(&self) -> usize {
189 1 + self.symbols.iter().map(|symbol| 1 + symbol.len()).sum::<usize>()
190 }
191
192 pub fn serialize(&self, out: &mut Vec<u8>) {
194 out.push(self.symbols.len() as u8);
195 for symbol in &self.symbols {
196 out.push(symbol.len);
197 out.extend_from_slice(&symbol.bytes());
198 }
199 }
200
201 pub fn deserialize(bytes: &[u8]) -> Result<(Self, usize)> {
207 let count = *bytes.first().ok_or_else(|| truncated("a symbol table header"))? as usize;
208 let mut at = 1;
209 let mut symbols = Vec::with_capacity(count);
210 for _ in 0..count {
211 let len = *bytes.get(at).ok_or_else(|| truncated("a symbol length"))? as usize;
212 if len == 0 || len > MAX_SYMBOL_LEN {
213 return Err(Error::internal(format!("a symbol of {len} bytes is not a symbol")));
214 }
215 at += 1;
216 let end = at + len;
217 if end > bytes.len() {
218 return Err(truncated("a symbol"));
219 }
220 symbols.push(Symbol::new(&bytes[at..end]));
221 at = end;
222 }
223 Ok((Self::build(symbols), at))
224 }
225
226 pub fn compress(&self, input: &[u8], out: &mut Vec<u8>) {
232 let mut at = 0;
233 while at < input.len() {
234 let (code, len) = self.match_at(input, at);
235 if code == ESCAPE {
236 out.push(ESCAPE);
237 out.push(input[at]);
238 } else {
239 out.push(code);
240 }
241 at += len;
242 }
243 }
244
245 pub fn decompress(&self, input: &[u8], out: &mut Vec<u8>) -> Result<()> {
251 let mut at = 0;
252 while at < input.len() {
253 let code = input[at];
254 at += 1;
255 if code == ESCAPE {
256 let literal = *input.get(at).ok_or_else(|| truncated("an escaped byte"))?;
257 out.push(literal);
258 at += 1;
259 } else {
260 let symbol = self
261 .symbols
262 .get(code as usize)
263 .ok_or_else(|| Error::internal(format!("code {code} is not in the table")))?;
264 out.extend_from_slice(&symbol.bytes());
265 }
266 }
267 Ok(())
268 }
269
270 fn match_at(&self, input: &[u8], at: usize) -> (u8, usize) {
272 let remaining = input.len() - at;
273 let word = load(input, at);
274 if remaining >= 3 {
277 if let Some((symbol, code)) = self.probe(word, remaining) {
278 return (code, symbol.len());
279 }
280 }
281 if remaining >= 2 {
282 let code = self.pair[(word & 0xffff) as usize];
283 if code != u16::MAX {
284 return (code as u8, 2);
285 }
286 }
287 let code = self.single[(word & 0xff) as usize];
288 if code == ESCAPE { (ESCAPE, 1) } else { (code, 1) }
289 }
290
291 fn probe(&self, word: u64, remaining: usize) -> Option<(Symbol, u8)> {
297 let mut slot = hash_of(word);
298 let mut best: Option<(Symbol, u8)> = None;
299 for _ in 0..PROBE {
300 match self.hash[slot] {
301 None => break,
302 Some((symbol, code)) => {
303 if symbol.len() <= remaining
304 && word & symbol.mask() == symbol.value
305 && best.is_none_or(|(found, _)| symbol.len() > found.len())
306 {
307 best = Some((symbol, code));
308 }
309 }
310 }
311 slot = (slot + 1) & (HASH_SLOTS - 1);
312 }
313 best
314 }
315
316 fn count(&self, input: &[u8], counts: &mut Counts) {
319 let mut at = 0;
320 let mut previous: Option<u16> = None;
321 while at < input.len() {
322 let (code, len) = self.match_at(input, at);
323 let id = if code == ESCAPE { 256 + u16::from(input[at]) } else { u16::from(code) };
324 counts.one(id);
325 if let Some(previous) = previous {
326 counts.two(previous, id);
327 }
328 previous = Some(id);
329 at += len;
330 }
331 }
332
333 fn build(symbols: Vec<Symbol>) -> Self {
334 let mut table = Self {
335 symbols,
336 single: vec![ESCAPE; 256],
337 pair: vec![u16::MAX; 65536],
338 hash: vec![None; HASH_SLOTS],
339 };
340 let mut order: Vec<(Symbol, u8)> =
343 table.symbols.iter().enumerate().map(|(code, symbol)| (*symbol, code as u8)).collect();
344 order.sort_by_key(|(symbol, code)| (std::cmp::Reverse(symbol.len()), *code));
345 for (symbol, code) in order {
346 match symbol.len() {
347 1 => {
348 let index = (symbol.value & 0xff) as usize;
349 if table.single[index] == ESCAPE {
350 table.single[index] = code;
351 }
352 }
353 2 => {
354 let index = (symbol.value & 0xffff) as usize;
355 if table.pair[index] == u16::MAX {
356 table.pair[index] = u16::from(code);
357 }
358 }
359 _ => {
360 let mut slot = hash_of(symbol.value);
361 for _ in 0..PROBE {
362 if table.hash[slot].is_none() {
363 table.hash[slot] = Some((symbol, code));
364 break;
365 }
366 slot = (slot + 1) & (HASH_SLOTS - 1);
367 }
368 }
369 }
370 }
371 table
372 }
373}
374
375fn load(input: &[u8], at: usize) -> u64 {
380 if at + 8 <= input.len() {
381 let bytes: [u8; 8] = input[at..at + 8].try_into().expect("eight bytes were checked");
382 u64::from_le_bytes(bytes)
383 } else {
384 let mut word = 0u64;
385 for (index, byte) in input[at..].iter().enumerate() {
386 word |= u64::from(*byte) << (8 * index);
387 }
388 word
389 }
390}
391
392fn hash_of(word: u64) -> usize {
396 let key = word & 0xff_ffff;
397 ((key.wrapping_mul(0x9e37_79b9_7f4a_7c15)) >> (64 - HASH_SLOTS.trailing_zeros())) as usize
398}
399
400struct Counts {
404 single: Vec<u32>,
405 pairs: HashMap<(u16, u16), u32>,
406}
407
408impl Counts {
409 fn new() -> Self {
410 Self { single: vec![0; 512], pairs: HashMap::new() }
411 }
412
413 fn one(&mut self, id: u16) {
414 self.single[id as usize] += 1;
415 }
416
417 fn two(&mut self, first: u16, second: u16) {
418 *self.pairs.entry((first, second)).or_insert(0) += 1;
419 }
420
421 fn best(&self, table: &SymbolTable) -> Vec<Symbol> {
428 let mut gains: HashMap<Symbol, u64> = HashMap::new();
429 for (id, count) in self.single.iter().enumerate() {
430 if *count == 0 {
431 continue;
432 }
433 let symbol = symbol_of(table, id as u16);
434 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
435 }
436 for ((first, second), count) in &self.pairs {
437 let symbol = symbol_of(table, *first).concat(symbol_of(table, *second));
438 *gains.entry(symbol).or_insert(0) += u64::from(*count) * symbol.len() as u64;
439 }
440 let mut ranked: Vec<(Symbol, u64)> = gains.into_iter().collect();
441 ranked.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
444 ranked.truncate(MAX_SYMBOLS);
445 ranked.into_iter().map(|(symbol, _)| symbol).collect()
446 }
447}
448
449fn symbol_of(table: &SymbolTable, id: u16) -> Symbol {
450 if (id as usize) < table.symbols.len() {
451 table.symbols[id as usize]
452 } else {
453 Symbol::single((id.saturating_sub(256)) as u8)
454 }
455}
456
457fn truncated(what: &str) -> Error {
458 Error::internal(format!("the input ended in the middle of {what}"))
459}
460
461#[cfg(test)]
462mod tests {
463 use super::*;
464
465 fn urls() -> Vec<Vec<u8>> {
469 let hosts = ["www.example.com", "shop.example.com", "news.other.example.org"];
470 let paths = ["/index.html", "/catalog/item", "/search", "/user/profile/settings"];
471 let mut out = Vec::new();
472 for index in 0..600 {
473 let host = hosts[index % hosts.len()];
474 let path = paths[(index / 3) % paths.len()];
475 out.push(
476 format!("http://{host}{path}?session={}&ref=google&page={}", index * 7, index % 20)
477 .into_bytes(),
478 );
479 }
480 out
481 }
482
483 fn borrow(strings: &[Vec<u8>]) -> Vec<&[u8]> {
484 strings.iter().map(Vec::as_slice).collect()
485 }
486
487 fn round_trip(table: &SymbolTable, strings: &[Vec<u8>]) -> (usize, usize) {
488 let mut raw = 0;
489 let mut compressed = 0;
490 for string in strings {
491 let mut bytes = Vec::new();
492 table.compress(string, &mut bytes);
493 let mut back = Vec::new();
494 table.decompress(&bytes, &mut back).unwrap();
495 assert_eq!(back, *string, "{}", String::from_utf8_lossy(string));
496 raw += string.len();
497 compressed += bytes.len();
498 }
499 (raw, compressed)
500 }
501
502 #[test]
503 fn urls_compress_by_more_than_half_and_come_back_unchanged() {
504 let strings = urls();
507 let table = SymbolTable::train(&borrow(&strings));
508 let (raw, compressed) = round_trip(&table, &strings);
509 let ratio = raw as f64 / compressed as f64;
510 assert!(ratio > 2.5, "{ratio:.2}x, {raw} to {compressed}");
511 assert!(table.len() > 100, "{} symbols", table.len());
512 }
513
514 #[test]
515 fn the_trainer_finds_the_long_repeated_pieces() {
516 let strings = urls();
517 let table = SymbolTable::train(&borrow(&strings));
518 let found: Vec<String> = (0..table.len())
519 .map(|code| String::from_utf8_lossy(&table.symbols[code].bytes()).into_owned())
520 .collect();
521 let long = found.iter().filter(|symbol| symbol.len() >= 6).count();
524 assert!(long > 60, "only {long} symbols of six bytes or more: {found:?}");
525 }
526
527 #[test]
528 fn english_text_round_trips_and_shrinks() {
529 let text: Vec<Vec<u8>> = "the quick brown fox jumps over the lazy dog while the other dog \
530 watches the fox and the dog and the fox go over the hill together"
531 .split(' ')
532 .map(|word| word.as_bytes().to_vec())
533 .collect();
534 let table = SymbolTable::train(&borrow(&text));
535 let (raw, compressed) = round_trip(&table, &text);
536 assert!(compressed < raw, "{raw} to {compressed}");
537 }
538
539 #[test]
540 fn incompressible_bytes_round_trip_and_cost_what_escaping_costs() {
541 let mut state = 0x1234_5678_9abc_def0u64;
545 let strings: Vec<Vec<u8>> = (0..100)
546 .map(|_| {
547 (0..64)
548 .map(|_| {
549 state ^= state << 13;
550 state ^= state >> 7;
551 state ^= state << 17;
552 state as u8
553 })
554 .collect()
555 })
556 .collect();
557 let table = SymbolTable::train(&borrow(&strings));
558 let (raw, compressed) = round_trip(&table, &strings);
559 assert!(compressed < raw * 2, "{raw} to {compressed}");
560 }
561
562 #[test]
563 fn an_empty_table_escapes_everything_and_still_round_trips() {
564 let table = SymbolTable::empty();
565 let strings = vec![b"hello".to_vec(), Vec::new(), b"x".to_vec()];
566 let (raw, compressed) = round_trip(&table, &strings);
567 assert_eq!(compressed, raw * 2);
568 }
569
570 #[test]
571 fn an_empty_string_compresses_to_nothing() {
572 let table = SymbolTable::train(&[b"abcabcabc"]);
573 let mut out = Vec::new();
574 table.compress(b"", &mut out);
575 assert!(out.is_empty());
576 let mut back = Vec::new();
577 table.decompress(&out, &mut back).unwrap();
578 assert!(back.is_empty());
579 }
580
581 #[test]
582 fn a_string_shorter_than_the_symbols_does_not_read_past_its_end() {
583 let table = SymbolTable::train(&[b"ab\0ab\0ab\0ab\0", b"abcdefgh"]);
586 for string in [b"a".to_vec(), b"ab".to_vec(), b"abc".to_vec()] {
587 let mut bytes = Vec::new();
588 table.compress(&string, &mut bytes);
589 let mut back = Vec::new();
590 table.decompress(&bytes, &mut back).unwrap();
591 assert_eq!(back, string);
592 }
593 }
594
595 #[test]
596 fn a_table_survives_being_written_and_read_back() {
597 let strings = urls();
598 let table = SymbolTable::train(&borrow(&strings));
599 let mut bytes = Vec::new();
600 table.serialize(&mut bytes);
601 assert_eq!(bytes.len(), table.serialized_len());
602 let (read, consumed) = SymbolTable::deserialize(&bytes).unwrap();
603 assert_eq!(consumed, bytes.len());
604 assert_eq!(read.symbols, table.symbols);
605
606 let mut first = Vec::new();
609 let mut second = Vec::new();
610 table.compress(&strings[7], &mut first);
611 read.compress(&strings[7], &mut second);
612 assert_eq!(first, second);
613 }
614
615 #[test]
616 fn a_full_table_is_two_kilobytes_at_the_very_most() {
617 let strings = urls();
618 let table = SymbolTable::train(&borrow(&strings));
619 assert!(table.serialized_len() <= 1 + MAX_SYMBOLS * (1 + MAX_SYMBOL_LEN));
620 assert!(table.serialized_len() <= 2049);
621 }
622
623 #[test]
624 fn a_truncated_symbol_table_is_an_error() {
625 let strings = urls();
626 let table = SymbolTable::train(&borrow(&strings));
627 let mut bytes = Vec::new();
628 table.serialize(&mut bytes);
629 for len in 1..bytes.len().min(40) {
630 let error = SymbolTable::deserialize(&bytes[..len]).unwrap_err();
631 assert!(error.message().contains("ended in the middle"), "{error}");
632 }
633 }
634
635 #[test]
636 fn a_symbol_of_zero_bytes_is_an_error() {
637 let error = SymbolTable::deserialize(&[1, 0]).unwrap_err();
638 assert!(error.message().contains("is not a symbol"), "{error}");
639 }
640
641 #[test]
642 fn a_dangling_escape_is_an_error_and_not_a_panic() {
643 let table = SymbolTable::train(&[b"abcabcabc"]);
644 let error = table.decompress(&[ESCAPE], &mut Vec::new()).unwrap_err();
645 assert!(error.message().contains("escaped byte"), "{error}");
646 }
647
648 #[test]
649 fn a_code_the_table_does_not_have_is_an_error() {
650 let table = SymbolTable::train(&[b"abcabcabc"]);
651 let code = table.len() as u8;
652 let error = table.decompress(&[code], &mut Vec::new()).unwrap_err();
653 assert!(error.message().contains("not in the table"), "{error}");
654 }
655
656 #[test]
657 fn training_twice_on_the_same_sample_gives_the_same_table() {
658 let strings = urls();
661 let first = SymbolTable::train(&borrow(&strings));
662 let second = SymbolTable::train(&borrow(&strings));
663 assert_eq!(first.symbols, second.symbols);
664 }
665
666 #[test]
667 fn the_longest_match_wins_rather_than_the_first_one_found() {
668 let table = SymbolTable::build(vec![
669 Symbol::new(b"abc"),
670 Symbol::new(b"abcdef"),
671 Symbol::new(b"abcd"),
672 ]);
673 let mut out = Vec::new();
674 table.compress(b"abcdef", &mut out);
675 assert_eq!(out, vec![1]);
676 }
677
678 #[test]
679 fn a_symbol_longer_than_what_is_left_is_not_used() {
680 let table = SymbolTable::build(vec![Symbol::new(b"abcdef"), Symbol::new(b"ab")]);
681 let mut out = Vec::new();
682 table.compress(b"abcd", &mut out);
683 assert_eq!(out, vec![1, ESCAPE, b'c', ESCAPE, b'd']);
685 }
686
687 #[test]
688 fn concatenation_stops_at_eight_bytes() {
689 let long = Symbol::new(b"abcdef");
690 assert_eq!(long.concat(Symbol::new(b"ghijkl")).bytes(), b"abcdefgh");
691 assert_eq!(Symbol::new(b"ab").concat(Symbol::new(b"cd")).bytes(), b"abcd");
692 }
693}