Refactor geo_filters tests to use PRNG test harness

Author: itsibitzi

crates/geo_filters/src/config.rs

@@ -353,29 +353,30 @@ pub(crate) fn take_ref<I: Iterator>(iter: &mut I, n: usize) -> impl Iterator<Ite
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use rand::{RngCore, SeedableRng};
+    use rand::RngCore;
 
-    use crate::{Count, Method};
+    use crate::{test_rng::prng_test_harness, Count, Method};
 
     /// Runs estimation trials and returns the average precision and variance.
     pub(crate) fn test_estimate<M: Method, C: Count<M>>(f: impl Fn() -> C) -> (f32, f32) {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
-        let cnt = 10000usize;
-        let mut avg_precision = 0.0;
-        let mut avg_var = 0.0;
-        let trials = 500;
-        for _ in 0..trials {
-            let mut m = f();
-            // Insert cnt many random items.
-            for _ in 0..cnt {
-                m.push_hash(rnd.next_u64());
+        prng_test_harness(|rnd| {
+            let cnt = 10000usize;
+            let mut avg_precision = 0.0;
+            let mut avg_var = 0.0;
+            let trials = 500;
+            for _ in 0..trials {
+                let mut m = f();
+                // Insert cnt many random items.
+                for _ in 0..cnt {
+                    m.push_hash(rnd.next_u64());
+                }
+                // Compute the relative error between estimate and actually inserted items.
+                let high_precision = m.size() / cnt as f32 - 1.0;
+                // Take the average over trials many attempts.
+                avg_precision += high_precision / trials as f32;
+                avg_var += high_precision.powf(2.0) / trials as f32;
             }
-            // Compute the relative error between estimate and actually inserted items.
-            let high_precision = m.size() / cnt as f32 - 1.0;
-            // Take the average over trials many attempts.
-            avg_precision += high_precision / trials as f32;
-            avg_var += high_precision.powf(2.0) / trials as f32;
-        }
-        (avg_precision, avg_var)
+            (avg_precision, avg_var)
+        })
     }
 }

crates/geo_filters/src/config/lookup.rs

@@ -45,9 +45,12 @@ impl HashToBucketLookup {
 
 #[cfg(test)]
 mod tests {
-    use rand::{RngCore, SeedableRng};
+    use rand::RngCore;
 
-    use crate::config::{hash_to_bucket, phi_f64};
+    use crate::{
+        config::{hash_to_bucket, phi_f64},
+        test_rng::prng_test_harness,
+    };
 
     use super::HashToBucketLookup;
 
@@ -66,15 +69,17 @@ mod tests {
     fn lookup_random_hashes_variance<const B: usize>(n: u64) -> f64 {
         let phi = phi_f64(B);
         let buckets = HashToBucketLookup::new(B);
-        let mut var = 0.0;
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
-        for _ in 0..n {
-            let hash = rnd.next_u64();
-            let estimate = buckets.lookup(hash) as f64;
-            let real = hash_to_bucket(phi, hash) as f64;
-            let err = estimate - real; // assume the mean = 0.0
-            var += err.powf(2.0) / n as f64;
-        }
-        var
+
+        prng_test_harness(|rnd| {
+            let mut var = 0.0;
+            for _ in 0..n {
+                let hash = rnd.next_u64();
+                let estimate = buckets.lookup(hash) as f64;
+                let real = hash_to_bucket(phi, hash) as f64;
+                let err = estimate - real; // assume the mean = 0.0
+                var += err.powf(2.0) / n as f64;
+            }
+            var
+        })
     }
 }

crates/geo_filters/src/diff_count.rs

@@ -358,7 +358,7 @@ impl<C: GeoConfig<Diff>> Count<Diff> for GeoDiffCount<'_, C> {
 #[cfg(test)]
 mod tests {
     use itertools::Itertools;
-    use rand::{RngCore, SeedableRng};
+    use rand::RngCore;
 
     use crate::{
         build_hasher::UnstableDefaultBuildHasher,
@@ -462,29 +462,30 @@ mod tests {
 
     #[test]
     fn test_estimate_diff_size_fast() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
-        let mut a_p = GeoDiffCount7_50::default();
-        let mut a_hp = GeoDiffCount7::default();
-        let mut b_p = GeoDiffCount7_50::default();
-        let mut b_hp = GeoDiffCount7::default();
-        for _ in 0..10000 {
-            let hash = rnd.next_u64();
-            a_p.push_hash(hash);
-            a_hp.push_hash(hash);
-        }
-        for _ in 0..1000 {
-            let hash = rnd.next_u64();
-            b_p.push_hash(hash);
-            b_hp.push_hash(hash);
-        }
-        let c_p = xor(&a_p, &b_p);
-        let c_hp = xor(&a_hp, &b_hp);
+        prng_test_harness(|rnd| {
+            let mut a_p = GeoDiffCount7_50::default();
+            let mut a_hp = GeoDiffCount7::default();
+            let mut b_p = GeoDiffCount7_50::default();
+            let mut b_hp = GeoDiffCount7::default();
+            for _ in 0..10000 {
+                let hash = rnd.next_u64();
+                a_p.push_hash(hash);
+                a_hp.push_hash(hash);
+            }
+            for _ in 0..1000 {
+                let hash = rnd.next_u64();
+                b_p.push_hash(hash);
+                b_hp.push_hash(hash);
+            }
+            let c_p = xor(&a_p, &b_p);
+            let c_hp = xor(&a_hp, &b_hp);
 
-        assert_eq!(c_p.size(), a_p.size_with_sketch(&b_p));
-        assert_eq!(c_p.size(), b_p.size_with_sketch(&a_p));
+            assert_eq!(c_p.size(), a_p.size_with_sketch(&b_p));
+            assert_eq!(c_p.size(), b_p.size_with_sketch(&a_p));
 
-        assert_eq!(c_hp.size(), a_hp.size_with_sketch(&b_hp));
-        assert_eq!(c_hp.size(), b_hp.size_with_sketch(&a_hp));
+            assert_eq!(c_hp.size(), a_hp.size_with_sketch(&b_hp));
+            assert_eq!(c_hp.size(), b_hp.size_with_sketch(&a_hp));
+        });
     }
 
     #[test]
@@ -517,24 +518,19 @@ mod tests {
     #[test]
     fn test_xor_plus_mask() {
         for _ in 0..1000 {
-            prng_test_harness(|mut rng| {
+            prng_test_harness(|rnd| {
                 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()));
+                    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 = rng.next_u64();
+                    let hash = rnd.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);
@@ -547,17 +543,18 @@ mod tests {
 
     #[test]
     fn test_bit_chunks() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
         for _ in 0..100 {
-            let mut expected = GeoDiffCount7::default();
-            for _ in 0..1000 {
-                expected.push_hash(rnd.next_u64());
-            }
-            let actual = GeoDiffCount::from_bit_chunks(
-                expected.config.clone(),
-                expected.bit_chunks().peekable(),
-            );
-            assert_eq!(expected, actual);
+            prng_test_harness(|rnd| {
+                let mut expected = GeoDiffCount7::default();
+                for _ in 0..1000 {
+                    expected.push_hash(rnd.next_u64());
+                }
+                let actual = GeoDiffCount::from_bit_chunks(
+                    expected.config.clone(),
+                    expected.bit_chunks().peekable(),
+                );
+                assert_eq!(expected, actual);
+            });
         }
     }
 

crates/geo_filters/src/diff_count/bitvec.rs

@@ -101,8 +101,6 @@ 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.

crates/geo_filters/src/distinct_count.rs

@@ -232,11 +232,12 @@ fn or<C: GeoConfig<Distinct>>(
 #[cfg(test)]
 mod tests {
     use itertools::Itertools;
-    use rand::{RngCore, SeedableRng};
+    use rand::RngCore;
 
     use crate::build_hasher::UnstableDefaultBuildHasher;
     use crate::config::{iter_ones, tests::test_estimate, FixedConfig, VariableConfig};
     use crate::evaluation::simulation::simulate;
+    use crate::test_rng::prng_test_harness;
 
     use super::*;
 
@@ -321,19 +322,20 @@ mod tests {
 
     #[test]
     fn test_estimate_union_size_fast() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
-        let mut a = GeoDistinctCount7::default();
-        let mut b = GeoDistinctCount7::default();
-        for _ in 0..10000 {
-            a.push_hash(rnd.next_u64());
-        }
-        for _ in 0..1000 {
-            b.push_hash(rnd.next_u64());
-        }
-        let c = or(&a, &b);
+        prng_test_harness(|rnd| {
+            let mut a = GeoDistinctCount7::default();
+            let mut b = GeoDistinctCount7::default();
+            for _ in 0..10000 {
+                a.push_hash(rnd.next_u64());
+            }
+            for _ in 0..1000 {
+                b.push_hash(rnd.next_u64());
+            }
+            let c = or(&a, &b);
 
-        assert_eq!(c.size(), a.size_with_sketch(&b));
-        assert_eq!(c.size(), b.size_with_sketch(&a));
+            assert_eq!(c.size(), a.size_with_sketch(&b));
+            assert_eq!(c.size(), b.size_with_sketch(&a));
+        })
     }
 
     fn golden_section_min<F: Fn(f32) -> f32>(min: f32, max: f32, f: F) -> f32 {
@@ -392,15 +394,18 @@ mod tests {
 
     #[test]
     fn test_bit_chunks() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
         for _ in 0..100 {
-            let mut expected = GeoDistinctCount7::default();
-            for _ in 0..1000 {
-                expected.push_hash(rnd.next_u64());
-            }
-            let actual =
-                GeoDistinctCount::from_bit_chunks(expected.config.clone(), expected.bit_chunks());
-            assert_eq!(expected, actual);
+            prng_test_harness(|rnd| {
+                let mut expected = GeoDistinctCount7::default();
+                for _ in 0..1000 {
+                    expected.push_hash(rnd.next_u64());
+                }
+                let actual = GeoDistinctCount::from_bit_chunks(
+                    expected.config.clone(),
+                    expected.bit_chunks(),
+                );
+                assert_eq!(expected, actual);
+            })
         }
     }
 

crates/geo_filters/src/test_rng.rs

@@ -1,24 +1,25 @@
-use std::panic::{catch_unwind, resume_unwind, UnwindSafe};
+use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
 
 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)
+pub fn prng_test_harness<F, T>(test_fn: F) -> T
 where
-    F: Fn(StdRng) -> () + UnwindSafe,
+    F: Fn(&mut StdRng) -> T,
 {
     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);
+    let mut rng = StdRng::seed_from_u64(seed);
+    let maybe_panic = catch_unwind(AssertUnwindSafe(|| test_fn(&mut rng)));
+    match maybe_panic {
+        Ok(t) => t,
+        Err(panic_info) => {
+            eprintln!("Test failed! Reproduce with: TEST_SEED={}", seed);
+            resume_unwind(panic_info);
+        }
     }
 }