Speed up miri execution

.github/workflows/miri.yml

@@ -19,10 +19,17 @@ jobs:
       MIRIFLAGS: -Zmiri-strict-provenance -Zmiri-symbolic-alignment-check -Zmiri-backtrace=full
     steps:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7
-      - name: Run miri with nightly toolchain
+      - name: Install nightly toolchain with miri
         run: |
           rustup toolchain install nightly --component miri
           rustup override set nightly
+      - name: Rust Dependency Cache
+        uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2
+        with:
+          save-if: ${{ github.ref_name == 'develop' }}
+          shared-key: "miri-cache-${{ runner.os }}"
+      - name: Run miri
+        run: |
           cargo miri setup
           cargo miri test
 

src/builder.rs

@@ -279,9 +279,15 @@ impl CompressorBuilder {
             table.codes_one_byte.push(Code::new_escape(byte));
         }
 
-        // Fill codes_two_byte with pseudocode of first byte
-        for idx in 0..=65_535 {
-            table.codes_two_byte.push(Code::new_escape(idx as u8));
+        // Fill codes_two_byte with pseudocode of first byte.
+        //
+        // The pseudocode only depends on the low byte of the index, so the 65,536 entries are
+        // 256 copies of the 256-entry `codes_one_byte` row we just built. Copying the row is
+        // considerably cheaper than pushing every element individually.
+        for _ in 0..256 {
+            table
+                .codes_two_byte
+                .extend_from_slice(&table.codes_one_byte);
         }
 
         table
@@ -415,6 +421,13 @@ impl CompressorBuilder {
 
         let mut symbol_lens = [0u8; FSST_CODE_BASE as usize];
 
+        // The `codes_two_byte` slots claimed by a two-byte symbol, as (slot, new code) pairs
+        // recorded in old-code (i.e. insertion) order. Replaying them in that order below
+        // reproduces the last-write-wins behaviour of `insert` in the case where the same
+        // two-byte symbol was inserted more than once.
+        let mut two_byte_slots = [(0u16, 0u8); FSST_CODE_BASE as usize];
+        let mut n_two_byte_slots = 0usize;
+
         for i in 0..(self.n_symbols as usize) {
             let symbol = self.symbols[256 + i];
             let len = symbol.len();
@@ -436,6 +449,9 @@ impl CompressorBuilder {
                     new_codes[i] = no_suffix_code;
                     no_suffix_code += 1;
                 }
+
+                two_byte_slots[n_two_byte_slots] = (symbol.first2(), new_codes[i]);
+                n_two_byte_slots += 1;
             } else {
                 // Assign new code based on the next code available for the given length symbol
                 new_codes[i] = codes_by_length[len - 1];
@@ -466,15 +482,17 @@ impl CompressorBuilder {
 
         // Rewrite the codes_two_byte table to point at the new code values.
         // Replace pseudocodes with escapes.
-        for two_bytes in 0..=65_535 {
-            let two_byte = self.codes_two_byte[two_bytes];
-            if two_byte.extended_code() >= FSST_CODE_BASE {
-                let new_code = new_codes[two_byte.code() as usize];
-                self.codes_two_byte[two_bytes] = Code::new_symbol(new_code, 2);
-            } else {
-                // The one-byte code for the given code number here...
-                self.codes_two_byte[two_bytes] = self.codes_one_byte[two_bytes & 0xFF];
-            }
+        //
+        // A slot holds a symbol code if and only if some two-byte symbol was inserted with that
+        // slot as its `first2()`; every other slot falls back to the one-byte code for its low
+        // byte. Rather than visiting all 65,536 slots, bulk-copy the (already rewritten)
+        // `codes_one_byte` row 256 times and then patch the at most 255 symbol slots.
+        self.codes_two_byte.clear();
+        for _ in 0..256 {
+            self.codes_two_byte.extend_from_slice(&self.codes_one_byte);
+        }
+        for &(slot, new_code) in &two_byte_slots[..n_two_byte_slots] {
+            self.codes_two_byte[slot as usize] = Code::new_symbol(new_code, 2);
         }
 
         // Reset values in the hash table as well.
@@ -513,10 +531,7 @@ impl CompressorBuilder {
 /// The number of generations used for training. This is taken from the [FSST paper].
 ///
 /// [FSST paper]: https://www.vldb.org/pvldb/vol13/p2649-boncz.pdf
-#[cfg(not(miri))]
 const GENERATIONS: [usize; 5] = [8usize, 38, 68, 98, 128];
-#[cfg(miri)]
-const GENERATIONS: [usize; 3] = [8usize, 38, 128];
 
 const FSST_SAMPLETARGET: usize = 1 << 14;
 const FSST_SAMPLELINE: usize = 512;

tests/correctness.rs

@@ -15,6 +15,14 @@ static DECLARATION: &str = include_str!("./fixtures/declaration.txt");
 
 static ART_OF_WAR: &str = include_str!("./fixtures/art_of_war.txt");
 
+/// Miri interprets every memory access, so the cost of these tests scales with the size of the
+/// corpus rather than being dominated by fixed setup. The unsafe code paths under test are hit
+/// just as thoroughly by a small corpus, so the size-sensitive tests below scale their inputs
+/// down under miri.
+const fn scaled(full: usize, under_miri: usize) -> usize {
+    if cfg!(miri) { under_miri } else { full }
+}
+
 #[test]
 fn test_basic() {
     // Roundtrip the declaration
@@ -74,12 +82,14 @@ fn test_large() {
 
 #[test]
 fn test_chinese() {
-    let trained = Compressor::train(&vec![ART_OF_WAR.as_bytes()]);
+    // A byte prefix is enough under miri: this is a byte-level roundtrip, and the multi-byte
+    // UTF-8 sequences that make this case interesting are dense from the very first byte.
+    let corpus = &ART_OF_WAR.as_bytes()[..scaled(ART_OF_WAR.len(), 1_024)];
+
+    let trained = Compressor::train(&vec![corpus]);
     assert_eq!(
-        ART_OF_WAR.as_bytes(),
-        trained
-            .decompressor()
-            .decompress(&trained.compress(ART_OF_WAR.as_bytes()))
+        corpus,
+        trained.decompressor().decompress(&trained.compress(corpus))
     );
 }
 
@@ -90,7 +100,7 @@ fn test_all_escape_roundtrip() {
     let decompressor = compressor.decompressor();
 
     // Large enough to exercise the 8-byte block loop in decompress_into.
-    let input: Vec<u8> = (0..=255u8).cycle().take(4096).collect();
+    let input: Vec<u8> = (0..=255u8).cycle().take(scaled(4096, 512)).collect();
     let compressed = compressor.compress(&input);
     // All-escape compressed size should be exactly 2x input.
     assert_eq!(compressed.len(), input.len() * 2);
@@ -120,22 +130,34 @@ fn test_invalid_tail_code_not_in_symbol_table_panics() {
 
 #[test]
 fn test_large_with_rebuild() {
-    let corpus: Vec<u8> = DECLARATION.bytes().cycle().take(10_240).collect();
+    let corpus: Vec<u8> = DECLARATION
+        .bytes()
+        .cycle()
+        .take(scaled(10_240, 1_024))
+        .collect();
+    // `DECLARATION` is pure ASCII, so slicing it at a byte index stays on a char boundary.
+    let text = &DECLARATION[..scaled(DECLARATION.len(), 1_024)];
 
     let trained = Compressor::train(&vec![&corpus]);
-    let compressed = trained.compress(DECLARATION.as_bytes());
-
-    // let compressed = trained.compress(&massive);
-    let rebuilt = Compressor::rebuild_from(trained.symbol_table(), trained.symbol_lengths());
-    let recompressed = rebuilt.compress(DECLARATION.as_bytes());
+    let compressed = trained.compress(text.as_bytes());
+
+    // `symbol_table()` / `symbol_lengths()` are padded to 255 entries, so they have to be sliced
+    // to `n_symbols()` before being handed back to `rebuild_from`. Passing the padded arrays
+    // makes the rebuilt table disagree with the trained one whenever the table is not full.
+    let n_symbols = trained.n_symbols();
+    let rebuilt = Compressor::rebuild_from(
+        &trained.symbol_table()[..n_symbols],
+        &trained.symbol_lengths()[..n_symbols],
+    );
+    let recompressed = rebuilt.compress(text.as_bytes());
 
     assert_eq!(compressed, recompressed);
 
     // Ensure round-trip after rebuilding the compressor
     let decompressed = rebuilt.decompressor().decompress(&recompressed);
     assert_eq!(
         unsafe { std::str::from_utf8_unchecked(&decompressed) },
-        DECLARATION,
+        text,
     );
 }
 
@@ -158,10 +180,6 @@ fn test_pruning_small_input() {
 
     // 0xFF (count=3) survives: pruning lowers the count threshold to 1,
     // and saves (3) > cost (2). Bytes 200..210 (count=1) are pruned.
-    //
-    // Under miri, only 3 generations run (vs 5 normally), so the longest
-    // merged symbol reaches 4 bytes instead of 8.
-    #[cfg(not(miri))]
     assert_eq!(
         &compressor.symbol_table()[0..compressor.n_symbols()],
         &[
@@ -171,16 +189,6 @@ fn test_pruning_small_input() {
             Symbol::from_u8(0xFF),
         ],
     );
-    #[cfg(miri)]
-    assert_eq!(
-        &compressor.symbol_table()[0..compressor.n_symbols()],
-        &[
-            Symbol::from_slice(b"aa\0\0\0\0\0\0"),
-            Symbol::from_slice(b"aaaa\0\0\0\0"),
-            Symbol::from_u8(b'a'),
-            Symbol::from_u8(0xFF),
-        ],
-    );
 
     let compressed = compressor.compress(&corpus);
     assert_eq!(compressor.decompressor().decompress(&compressed), corpus);

tests/exact_capacity.rs

@@ -98,8 +98,10 @@ fn test_decompress_exact_capacity() {
         builder.build()
     };
 
-    // Create a large, highly compressible string
-    let plaintext = vec![b'a'; 100_000];
+    // Create a large, highly compressible string. Under miri a much smaller buffer still
+    // exercises both the 8-byte block loop and the byte-by-byte tail fallback this test targets,
+    // and miri's cost here is proportional to the buffer size.
+    let plaintext = vec![b'a'; if cfg!(miri) { 2_000 } else { 100_000 }];
 
     // Compress it into an over-allocated buffer to avoid any compression bugs
     let mut compressed = Vec::with_capacity(plaintext.len() * 2);