Skip to main content

rudb_vector/
fsst.rs

1//! FSST, the string encoding.
2//!
3//! Fast Static Symbol Table, from the 2020 paper by Boncz, Neumann and Leis. A table of at most 255
4//! symbols of one to eight bytes each, and compression is replacing the longest matching symbol at
5//! each position with its one byte code. A byte that no symbol covers is escaped, which costs two
6//! bytes, so the table has to be good or the output is larger than the input.
7//!
8//! ## Why this and not a general compressor
9//!
10//! `spec/06-compression.md` section 6.2 is blunt about it. FSST compresses text about 2x, which is
11//! worse than what zstd does to the same bytes, and the ratio is not why it is here. Two other
12//! properties are.
13//!
14//! The first is random access. Every string in a column is compressed independently against a
15//! shared table, so reading row 4,000,000 does not mean decompressing the four million before it. A
16//! block compressor gives up that property and gets it back by cutting the data into blocks, which
17//! means reading one string decompresses a block.
18//!
19//! The second is that a substring search can run against the compressed bytes. Compress the needle
20//! with the same symbol table and look for the compressed needle in the compressed haystack. That is
21//! what turns `URL LIKE '%google%'` from a decompress and scan into a scan, and section 6.7 says it
22//! is worth more on the ClickBench workload than any ratio improvement. It needs care, because the
23//! greedy match that compresses a needle standing alone can segment it differently from the way the
24//! same bytes were segmented inside a longer string, so a hit is a candidate and a miss is not a
25//! proof. The scan that uses it is M3 work and lives with the rest of encoded execution.
26//!
27//! ## Training
28//!
29//! The table is built from a sample rather than from the whole column, and the algorithm is the
30//! paper's: start with nothing, so every byte escapes, then repeat five times. Compress the sample
31//! with the table you have, count how often each symbol is used and how often each pair of adjacent
32//! symbols occurs, and build the next table from the best 255 of the symbols and the concatenations
33//! by gain, where gain is how many bytes of input the symbol accounts for. Five generations is what
34//! the paper found, and the shape of the thing is that the first generation learns single bytes, the
35//! second learns pairs, and the fifth is finding eight byte symbols like `https://`.
36//!
37//! ## Matching
38//!
39//! Three lookups in a fixed order, longest first. A hash table on the first three bytes for symbols
40//! of three bytes and up, a flat table indexed by the first two bytes, and a flat table indexed by
41//! the first one. The hash table probes eight slots and keeps the longest symbol that matches rather
42//! than the first, because several symbols share a three byte prefix and taking the first would make
43//! the ratio depend on insertion order.
44
45use rudb_common::{Error, Result};
46
47/// The code that means the next byte is a literal. 255 rather than 0 so that the 255 real codes are
48/// a contiguous range starting at zero and a code is its own index into the symbol table.
49pub const ESCAPE: u8 = 255;
50
51/// How many real symbols a table can hold.
52pub const MAX_SYMBOLS: usize = 255;
53
54/// The longest a symbol can be. Eight, so that a symbol is a `u64` and a match is a mask and a
55/// compare rather than a loop over bytes.
56pub const MAX_SYMBOL_LEN: usize = 8;
57
58/// How many generations the trainer runs. The paper's number.
59const GENERATIONS: usize = 5;
60
61/// Slots in the prefix hash table. A power of two, and four times the largest number of symbols that
62/// can be in it, which keeps the eight slot probe from filling up on a full table.
63const HASH_SLOTS: usize = 1024;
64
65/// How far a lookup probes before giving up. A miss here costs ratio and not correctness.
66const PROBE: usize = 8;
67
68/// One symbol. The bytes are in the low end of `value` in the order they appear, so that a match
69/// against the next eight bytes of input is one mask and one compare.
70#[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    /// The two symbols end to end, cut off at eight bytes.
103    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
117/// A trained symbol table, and everything needed to compress and decompress against it.
118pub struct SymbolTable {
119    /// Code to symbol. At most [`MAX_SYMBOLS`] long.
120    symbols: Vec<Symbol>,
121    /// First byte to code, or [`ESCAPE`] when no one byte symbol covers it.
122    single: Vec<u8>,
123    /// First two bytes to code, or `u16::MAX` when there is no two byte symbol for them.
124    pair: Vec<u16>,
125    /// Open addressed, keyed on the first three bytes, holding every symbol of three bytes or more.
126    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        // The lookup tables are 64k entries and printing them is never what anybody wanted.
132        formatter
133            .debug_struct("SymbolTable")
134            .field("symbols", &self.symbols.len())
135            .field("bytes", &self.serialized_len())
136            .finish()
137    }
138}
139
140/// Two tables are equal when they hold the same symbols in the same order.
141///
142/// The three lookup tables are built from the symbols when a table is built and hold nothing the
143/// symbols do not, so comparing them would be comparing the same information a second time over
144/// sixty five thousand entries. A vector in FSST form carries a table, and a vector is compared for
145/// equality all over the tests, so this is on a path that gets walked.
146impl 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    /// How many bytes of memory this table is holding.
156    ///
157    /// Mostly the hash table, which is sixty five thousand slots however few symbols are in it. That
158    /// is the number a vector in FSST form reports, and it is why the form is a decision about a
159    /// page rather than about a chunk: one table over a hundred chunks is nothing per chunk and one
160    /// table per chunk is a megabyte.
161    #[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    /// A table with no symbols, which escapes everything and doubles its input. The starting point
171    /// of training, and what a column of nothing but unique bytes ends up with.
172    #[must_use]
173    pub fn empty() -> Self {
174        Self::build(Vec::new())
175    }
176
177    /// Trains a table on a sample.
178    ///
179    /// The caller picks the sample. Section 6.3 says a systematic sample across the chunk rather
180    /// than the first N rows, because column data is frequently clustered, and that decision belongs
181    /// to whoever knows what the chunk is rather than to this function.
182    #[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    /// How many symbols are in the table.
201    #[must_use]
202    pub fn len(&self) -> usize {
203        self.symbols.len()
204    }
205
206    /// Whether the table has no symbols, in which case every byte of every string escapes.
207    #[must_use]
208    pub fn is_empty(&self) -> bool {
209        self.symbols.is_empty()
210    }
211
212    /// How many bytes [`serialize`](Self::serialize) writes. At most 2049 for a full table, and
213    /// that is the number section 6.4 is weighing when it says a shared symbol table is cheaper
214    /// than a shared dictionary.
215    #[must_use]
216    pub fn serialized_len(&self) -> usize {
217        1 + self.symbols.iter().map(|symbol| 1 + symbol.len()).sum::<usize>()
218    }
219
220    /// Writes the table itself, which has to travel with the data it compressed.
221    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    /// Reads back what [`serialize`](Self::serialize) wrote, and says how many bytes it consumed.
230    ///
231    /// # Errors
232    ///
233    /// If the bytes are truncated or describe a symbol of zero or more than eight bytes.
234    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    /// Compresses one string, appending to `out`.
255    ///
256    /// Strings are compressed one at a time against a shared table rather than as one stream,
257    /// because that is what keeps random access, which is the first of the two reasons this encoding
258    /// was chosen at all.
259    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    /// Decompresses one string, appending to `out`.
274    ///
275    /// A symbol goes out as all eight of the bytes its `u64` holds, and then the cursor steps back
276    /// over the ones that were not part of it. Eight is a length the compiler knows, so that is one
277    /// store. The length a symbol really has is only known at run time, so copying exactly that
278    /// many bytes is a call into `memcpy` for one to eight of them, and building a `Vec` to copy
279    /// them out of, which is what this used to do, is a heap allocation and a free on top.
280    ///
281    /// That mattered more than anything else in the engine. `SELECT COUNT(*) FROM hits WHERE URL
282    /// LIKE '%google%'` over ClickBench spends almost all of its time here, because the search
283    /// itself runs once per distinct URL and finding those means decompressing the column, and the
284    /// allocation, the free and the copy together were 41% of the query.
285    ///
286    /// # Errors
287    ///
288    /// If the input ends on an escape byte, or holds a code the table does not have.
289    pub fn decompress(&self, input: &[u8], out: &mut Vec<u8>) -> Result<()> {
290        // What the table is trained to reach, plus room for the tail of the last symbol, so the
291        // loop below mostly finds the space already there. Nothing here depends on the guess being
292        // right: too small and the growth happens where it always did.
293        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    /// The code and how many input bytes it covers. [`ESCAPE`] and 1 when nothing matches.
315    fn match_at(&self, input: &[u8], at: usize) -> (u8, usize) {
316        let remaining = input.len() - at;
317        let word = load(input, at);
318        // Written as a nested `if` rather than as a chained `if let` because the minimum supported
319        // Rust version is 1.85 and let chains landed in 1.88.
320        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    /// The longest symbol of three bytes or more matching here, if any.
336    ///
337    /// Longest rather than first, because several symbols share a three byte prefix and taking
338    /// whichever the probe reached first would make the compression ratio depend on the order the
339    /// table was built in.
340    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    /// Runs the matcher over a sample without producing output, recording what it used. This is the
361    /// counting half of a training generation.
362    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        // Longest first, so that a short symbol never displaces a long one out of the probe window
385        // and the flat tables get the lowest code for a duplicate.
386        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
419/// The next eight bytes as a little endian word, zero padded at the end of the input.
420///
421/// The padding is why every match checks the remaining length as well as the mask. Without that
422/// check a two byte symbol ending in a zero byte would match the last byte of a string.
423fn 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
436/// Hashes the first three bytes. The multiply and shift is the standard Fibonacci hash, which
437/// spreads a three byte key across the whole slot range where a mask of the low bits would put every
438/// symbol starting with the same letter in the same neighbourhood.
439fn 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
444/// What one training generation counts. Symbol ids below 256 are codes in the current table and ids
445/// from 256 up are escaped literal bytes, which is how a generation learns single bytes it does not
446/// have yet.
447struct Counts {
448    single: Vec<u32>,
449    /// Indexed by `first * IDS + second`. Flat rather than hashed because every id is under 512, so
450    /// every pair fits in a quarter million slots, and counting pairs is most of what training does.
451    pairs: Vec<u32>,
452    /// The pair slots that are not zero, so that reading and clearing them costs the pairs seen
453    /// rather than the whole array.
454    seen: Vec<u32>,
455}
456
457/// How many symbol ids there are: 256 codes and 256 escaped bytes.
458const 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    /// The 255 best symbols for the next generation.
485    ///
486    /// Gain is how many bytes of input a symbol accounts for, which is its length times how often it
487    /// was used. A concatenation is scored on the length it would have, so a pair of four byte
488    /// symbols scores as eight and a pair of six byte ones also scores as eight, because that is
489    /// what it would be cut down to.
490    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        // Different ids can spell the same symbol, a code and the pair it was learned from for one,
506        // so the gains are summed per symbol before anything is ranked.
507        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        // Gain first, then the symbol itself, so that two symbols with the same gain come out in the
516        // same order on every host and the table is a function of the sample and nothing else. The
517        // symbols are distinct by now, so the order is total and an unstable sort gives one answer.
518        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    /// A few hundred URLs in the shape ClickBench `hits` has them, which is the workload this
547    /// encoding was chosen for. Repetitive in the way real URLs are: a handful of hosts, a handful
548    /// of path shapes, and query strings that differ in a number.
549    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        // The number the paper reports on text is around 2x, and URLs are more repetitive than
586        // text. Anything under 2x here means the trainer is not finding the long symbols.
587        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        // Not a specific symbol, since which eight bytes win is a property of the sample, but the
603        // table has to be mostly long symbols or it has not learned anything.
604        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        // The worst case, and it has to be a correct worst case. Every byte escapes at two bytes
623        // each unless the trainer finds single byte symbols, which it will for the 255 most common
624        // of the 256 values.
625        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        // The load pads with zeros, so without the length check a two byte symbol whose second byte
665        // is zero would match the last byte of a string and swallow a byte that is not there.
666        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        // And the read back table compresses to the same bytes, which is the property that matters,
688        // since the lookup structures are rebuilt rather than stored.
689        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        // Decompression writes all eight bytes of a symbol and steps back over the ones that were
725        // not part of it, so a table of short symbols is where that would show. `ab` and `cd` are
726        // two bytes each and sit in a `u64` with six zero bytes above them, and if the step back
727        // were wrong those zeros would be in the answer. Appending twice checks it again at an
728        // offset, since the second write lands where the first one left the cursor.
729        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        // A table that differs run to run would make every size in the M1 report unreproducible.
756        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        // "ab" then two escapes, rather than a six byte symbol over four bytes of input.
780        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    /// The trainer the way it was written first, with the pairs and the gains in hash maps and one
791    /// stable sort over everything, kept here to check the flat counts pick the same symbols.
792    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        // Few letters, so that many pairs tie on gain and the tie break decides the table.
844        shapes.push(
845            (0..400).map(|_| (0..12).map(|_| b"abc"[(next() % 3) as usize]).collect()).collect(),
846        );
847        // Every byte value, so that escapes of all 256 bytes are counted.
848        shapes.push((0..300).map(|_| (0..40).map(|_| next() as u8).collect()).collect());
849        // Long repeats, so that concatenations reach eight bytes and get cut.
850        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}