Add test harness for deterministic RNG in tests

itsibitzi · itsibitzi · commit 101314068ad1 · 2025-07-24T17:36:47.000+01:00
Fix issue where BitVec was required to be structurally identical rather
than just semantically equivalent. This caused issues when bits in the
LSB were toggled.

This adds a test harness that provides seeded random number generators
for tests, making failures reproducible via TEST_SEED env var.
diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs
@@ -95,16 +95,14 @@ impl<C: GeoConfig<Diff>> GeoDiffCount<'_, C> {
     /// having to construct another iterator with the remaining `BitChunk`s.
     fn from_bit_chunks<I: Iterator<Item = BitChunk>>(config: C, chunks: I) -> Self {
         let mut ones = iter_ones::<C::BucketType, _>(chunks.peekable());
-
         let mut msb = Vec::default();
         take_ref(&mut ones, config.max_msb_len() - 1).for_each(|bucket| {
             msb.push(bucket);
         });
         let smallest_msb = ones
             .next()
-            .map(|bucket| {
-                msb.push(bucket);
-                bucket
+            .inspect(|bucket| {
+                msb.push(*bucket);
             })
             .unwrap_or_default();
 
@@ -365,6 +363,7 @@ mod tests {
     use crate::{
         build_hasher::UnstableDefaultBuildHasher,
         config::{iter_ones, tests::test_estimate, FixedConfig},
+        test_rng::prng_test_harness,
     };
 
     use super::*;
@@ -517,28 +516,32 @@ mod tests {
 
     #[test]
     fn test_xor_plus_mask() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
-        let mask_size = 12;
-        let mask = 0b100001100000;
-        let mut a = GeoDiffCount7::default();
-        for _ in 0..10000 {
-            a.xor_bit(a.config.hash_to_bucket(rnd.next_u64()));
-        }
-        let mut expected = GeoDiffCount7::default();
-        let mut b = a.clone();
         for _ in 0..1000 {
-            let hash = rnd.next_u64();
-            b.xor_bit(b.config.hash_to_bucket(hash));
-            expected.xor_bit(expected.config.hash_to_bucket(hash));
-            assert_eq!(expected, xor(&a, &b));
-
-            let masked_a = masked(&a, mask, mask_size);
-            let masked_b = masked(&b, mask, mask_size);
-            let masked_expected = masked(&expected, mask, mask_size);
-            // FIXME: test failed once with:
-            // left: ~12.37563 (msb: [390, 334, 263, 242, 222, 215, 164, 148, 100, 97, 66, 36], |lsb|: 36)
-            // right: ~12.37563 (msb: [390, 334, 263, 242, 222, 215, 164, 148, 100, 97, 66, 36], |lsb|: 0)
-            assert_eq!(masked_expected, xor(&masked_a, &masked_b));
+            prng_test_harness(|mut rng| {
+                let mask_size = 12;
+                let mask = 0b100001100000;
+                let mut a = GeoDiffCount7::default();
+                for _ in 0..10000 {
+                    a.xor_bit(a.config.hash_to_bucket(rng.next_u64()));
+                }
+                let mut expected = GeoDiffCount7::default();
+                let mut b = a.clone();
+                for _ in 0..1000 {
+                    let hash = rng.next_u64();
+                    b.xor_bit(b.config.hash_to_bucket(hash));
+                    expected.xor_bit(expected.config.hash_to_bucket(hash));
+
+                    println!("a   -> {:?}", a);
+                    println!("b   -> {:?}", b);
+                    println!("exp -> {:?}", expected);
+
+                    assert_eq!(expected, xor(&a, &b));
+                    let masked_a = masked(&a, mask, mask_size);
+                    let masked_b = masked(&b, mask, mask_size);
+                    let masked_expected = masked(&expected, mask, mask_size);
+                    assert_eq!(masked_expected, xor(&masked_a, &masked_b));
+                }
+            });
         }
     }
 
diff --git a/crates/geo_filters/src/diff_count/bitvec.rs b/crates/geo_filters/src/diff_count/bitvec.rs
@@ -12,12 +12,20 @@ use crate::config::BITS_PER_BLOCK;
 /// bit consumes 1 byte). It only implements the minimum number of operations that we need for our
 /// GeoDiffCount implementation. In particular it supports xor-ing of two bit vectors and
 /// iterating through one bits.
-#[derive(Clone, Default, Debug, PartialEq, Eq)]
+#[derive(Clone, Default, Debug)]
 pub(crate) struct BitVec<'a> {
     num_bits: usize,
     blocks: Cow<'a, [u64]>,
 }
 
+impl PartialEq for BitVec<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        self.bit_chunks().eq(other.bit_chunks())
+    }
+}
+
+impl Eq for BitVec<'_> {}
+
 impl Ord for BitVec<'_> {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
         match self.num_bits.cmp(&other.num_bits) {
@@ -93,6 +101,8 @@ impl BitVec<'_> {
         assert!(index < self.num_bits);
         let (block_idx, bit_idx) = index.into_index_and_bit();
         self.blocks.to_mut()[block_idx] ^= bit_idx.into_block();
+
+        if self.blocks.to_mut()[block_idx] & bit_idx.into_block() == 0 {}
     }
 
     /// Returns an iterator over all blocks in reverse order.
diff --git a/crates/geo_filters/src/lib.rs b/crates/geo_filters/src/lib.rs
@@ -13,6 +13,8 @@ pub mod diff_count;
 pub mod distinct_count;
 #[cfg(feature = "evaluation")]
 pub mod evaluation;
+#[cfg(test)]
+mod test_rng;
 
 use std::hash::Hash;
 
diff --git a/crates/geo_filters/src/test_rng.rs b/crates/geo_filters/src/test_rng.rs
@@ -0,0 +1,24 @@
+use std::panic::{catch_unwind, resume_unwind, UnwindSafe};
+
+use rand::{rngs::StdRng, SeedableRng as _};
+
+/// Provides a seeded random number generator to tests which require some
+/// degree of randomization. If the test panics the harness will print the
+/// seed used for that run. You can then pass in this seed using the `TEST_SEED`
+/// environment variable when running your tests.
+pub fn prng_test_harness<F>(test_fn: F)
+where
+    F: Fn(StdRng) -> () + UnwindSafe,
+{
+    let seed = std::env::var("TEST_SEED")
+        .map(|s| s.parse::<u64>().expect("Parse TEST_SEED to u64"))
+        .unwrap_or_else(|_| rand::random());
+    let rng = StdRng::seed_from_u64(seed);
+    let maybe_panic = catch_unwind(move || {
+        test_fn(rng);
+    });
+    if let Err(panic_info) = maybe_panic {
+        eprintln!("Test failed! Reproduce with: TEST_SEED={}", seed);
+        resume_unwind(panic_info);
+    }
+}
