add fuzzer, fix bug

.github/workflows/fuzz.yml

@@ -0,0 +1,28 @@
+name: Fuzz
+
+on:
+  schedule:
+    - cron: "0 0 * * *" # daily
+  workflow_dispatch:
+
+jobs:
+  fuzz:
+    name: "fuzz"
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install cargo fuzz
+        run: cargo install cargo-fuzz
+      - name: Run fuzzing target
+        run: cargo fuzz run fuzz_compress -- -max_total_time=600
+        continue-on-error: true
+      - name: Archive crash artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: fuzzing-crash-artifacts
+          path: fuzz/artifacts
+      - name: Archive fuzzing corpus
+        uses: actions/upload-artifact@v4
+        with:
+          name: fuzzing-corpus
+          path: fuzz/corpus

fuzz/.gitignore

@@ -0,0 +1,4 @@
+target
+corpus
+artifacts
+coverage

fuzz/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "fsst-rs-fuzz"
+version = "0.0.0"
+publish = false
+edition = "2021"
+
+[package.metadata]
+cargo-fuzz = true
+
+[dependencies]
+libfuzzer-sys = "0.4"
+
+[dependencies.fsst-rs]
+path = ".."
+
+[[bin]]
+name = "fuzz_train"
+path = "fuzz_targets/fuzz_train.rs"
+test = false
+doc = false
+bench = false
+
+[[bin]]
+name = "fuzz_compress"
+path = "fuzz_targets/fuzz_compress.rs"
+test = false
+doc = false
+bench = false

fuzz/fuzz_targets/fuzz_compress.rs

@@ -0,0 +1,8 @@
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+
+fuzz_target!(|data: &[u8]| {
+    let table = fsst_rs::train("the quick brown fox jumped over the lazy dog".as_bytes());
+    let _ = table.compress(data);
+});

fuzz/fuzz_targets/fuzz_train.rs

@@ -0,0 +1,7 @@
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+
+fuzz_target!(|data: &[u8]| {
+    let _ = fsst_rs::train(data);
+});

src/lib.rs

@@ -386,10 +386,25 @@ impl SymbolTable {
             remaining_bytes.is_positive(),
             "in_ptr exceeded in_end, should not be possible"
         );
-
-        // Shift off the remaining bytes
-        let mut last_word = unsafe { (in_ptr as *const u64).read_unaligned() };
-        last_word = mask_prefix(last_word, remaining_bytes as usize);
+        let remaining_bytes = remaining_bytes as usize;
+
+        // Load the last `remaining_byte`s of data into a final world. We then replicate the loop above,
+        // but shift data out of this word rather than advancing an input pointer and potentially reading
+        // unowned memory.
+        let mut last_word = unsafe {
+            match remaining_bytes {
+                0 => 0,
+                1 => extract_u64::<1>(in_ptr),
+                2 => extract_u64::<2>(in_ptr),
+                3 => extract_u64::<3>(in_ptr),
+                4 => extract_u64::<4>(in_ptr),
+                5 => extract_u64::<5>(in_ptr),
+                6 => extract_u64::<6>(in_ptr),
+                7 => extract_u64::<7>(in_ptr),
+                8 => extract_u64::<8>(in_ptr),
+                _ => unreachable!("remaining bytes must be <= 8"),
+            }
+        };
 
         while in_ptr < in_end && out_ptr < out_end {
             unsafe {
@@ -466,17 +481,6 @@ impl SymbolTable {
     }
 }
 
-/// Mask the word, keeping only the `prefix_bytes` front.
-fn mask_prefix(word: u64, prefix_bytes: usize) -> u64 {
-    let mask = if prefix_bytes == 0 {
-        0
-    } else {
-        u64::MAX >> (8 * (8 - prefix_bytes))
-    };
-
-    word & mask
-}
-
 fn advance_8byte_word(word: u64, bytes: usize) -> u64 {
     // shift the word off the right-end, because little endian means the first
     // char is stored in the LSB.
@@ -499,3 +503,37 @@ fn compare_masked(left: u64, right: u64, ignored_bits: u16) -> bool {
 
     (left & mask) == right
 }
+
+/// This is a function that will get monomorphized based on the value of `N` to do
+/// a load of `N` values from the pointer in a minimum number of instructions into
+/// an output `u64`.
+unsafe fn extract_u64<const N: usize>(ptr: *const u8) -> u64 {
+    match N {
+        1 => ptr.read() as u64,
+        2 => (ptr as *const u16).read_unaligned() as u64,
+        3 => {
+            let low = ptr.read() as u64;
+            let high = (ptr.byte_add(1) as *const u16).read_unaligned() as u64;
+            high << 8 | low
+        }
+        4 => (ptr as *const u32).read_unaligned() as u64,
+        5 => {
+            let low = (ptr as *const u32).read_unaligned() as u64;
+            let high = ptr.byte_add(4).read() as u64;
+            high << 32 | low
+        }
+        6 => {
+            let low = (ptr as *const u32).read_unaligned() as u64;
+            let high = (ptr.byte_add(4) as *const u16).read_unaligned() as u64;
+            high << 32 | low
+        }
+        7 => {
+            let low = (ptr as *const u32).read_unaligned() as u64;
+            let mid = (ptr.byte_add(4) as *const u16).read_unaligned() as u64;
+            let high = ptr.byte_add(6).read() as u64;
+            (high << 48) | (mid << 32) | low
+        }
+        8 => (ptr as *const u64).read_unaligned(),
+        _ => unreachable!("N must be <= 8"),
+    }
+}

tests/correctness.rs

@@ -48,7 +48,7 @@ fn test_one_byte() {
 #[test]
 fn test_zeros() {
     println!("training zeros");
-    let training_data: Vec<u8> = vec![0, 1, 2, 3, 4];
+    let training_data: Vec<u8> = vec![0, 1, 2, 3, 4, 0];
     let trained = fsst_rs::train(&training_data);
     println!("compressing with zeros");
     let compressed = trained.compress(&[0, 4]);