From feccc3ac199fdbd8f741c2faf006f58123f7e692 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 31 Jul 2026 16:53:43 +0100 Subject: [PATCH] Speed up miri execution Signed-off-by: Robert Kruszewski --- .github/workflows/miri.yml | 9 +++++- src/builder.rs | 45 ++++++++++++++++++--------- tests/correctness.rs | 62 +++++++++++++++++++++----------------- tests/exact_capacity.rs | 6 ++-- 4 files changed, 77 insertions(+), 45 deletions(-) diff --git a/.github/workflows/miri.yml b/.github/workflows/miri.yml index 252d48e..a87beb8 100644 --- a/.github/workflows/miri.yml +++ b/.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 diff --git a/src/builder.rs b/src/builder.rs index 6dca204..556cf4c 100644 --- a/src/builder.rs +++ b/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; diff --git a/tests/correctness.rs b/tests/correctness.rs index dd8bb77..57da09a 100644 --- a/tests/correctness.rs +++ b/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 = (0..=255u8).cycle().take(4096).collect(); + let input: Vec = (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,14 +130,26 @@ fn test_invalid_tail_code_not_in_symbol_table_panics() { #[test] fn test_large_with_rebuild() { - let corpus: Vec = DECLARATION.bytes().cycle().take(10_240).collect(); + let corpus: Vec = 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); @@ -135,7 +157,7 @@ fn test_large_with_rebuild() { 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); diff --git a/tests/exact_capacity.rs b/tests/exact_capacity.rs index 591b7ba..cec5c94 100644 --- a/tests/exact_capacity.rs +++ b/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);