Address reworks in geo_filters work

- Pass in RNG to helper functions
- Add iterations to test harness

Author: itsibitzi

crates/geo_filters/src/config.rs

@@ -353,30 +353,31 @@ pub(crate) fn take_ref<I: Iterator>(iter: &mut I, n: usize) -> impl Iterator<Ite
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use rand::RngCore;
+    use rand::{rngs::StdRng, RngCore};
 
-    use crate::{test_rng::prng_test_harness, Count, Method};
+    use crate::{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) {
-        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;
+    pub(crate) fn test_estimate<M: Method, C: Count<M>>(
+        rnd: &mut StdRng,
+        f: impl Fn() -> C,
+    ) -> (f32, f32) {
+        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());
             }
-            (avg_precision, avg_var)
-        })
+            // 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)
     }
 }

crates/geo_filters/src/config/lookup.rs

@@ -45,7 +45,7 @@ impl HashToBucketLookup {
 
 #[cfg(test)]
 mod tests {
-    use rand::RngCore;
+    use rand::{rngs::StdRng, RngCore};
 
     use crate::{
         config::{hash_to_bucket, phi_f64},
@@ -56,30 +56,32 @@ mod tests {
 
     #[test]
     fn test_lookup_7() {
-        let var = lookup_random_hashes_variance::<7>(1 << 16);
-        assert!(var < 1e-4, "variance {var} too large");
+        prng_test_harness(1, |rnd| {
+            let var = lookup_random_hashes_variance::<7>(rnd, 1 << 16);
+            assert!(var < 1e-4, "variance {var} too large");
+        });
     }
 
     #[test]
     fn test_lookup_13() {
-        let var = lookup_random_hashes_variance::<13>(1 << 16);
-        assert!(var < 1e-4, "variance {var} too large");
+        prng_test_harness(1, |rnd| {
+            let var = lookup_random_hashes_variance::<13>(rnd, 1 << 16);
+            assert!(var < 1e-4, "variance {var} too large");
+        });
     }
 
-    fn lookup_random_hashes_variance<const B: usize>(n: u64) -> f64 {
+    fn lookup_random_hashes_variance<const B: usize>(rnd: &mut StdRng, n: u64) -> f64 {
         let phi = phi_f64(B);
         let buckets = HashToBucketLookup::new(B);
 
-        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
-        })
+        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

@@ -432,35 +432,39 @@ mod tests {
 
     #[test]
     fn test_estimate_fast() {
-        let (avg_precision, avg_var) = test_estimate(GeoDiffCount7::default);
-        println!(
-            "avg precision: {} with standard deviation: {}",
-            avg_precision,
-            avg_var.sqrt(),
-        );
-        // Make sure that the estimate converges to the correct value.
-        assert!(avg_precision.abs() < 0.04);
-        // We should theoretically achieve a standard deviation of about 0.12
-        assert!(avg_var.sqrt() < 0.14);
+        prng_test_harness(1, |rnd| {
+            let (avg_precision, avg_var) = test_estimate(rnd, GeoDiffCount7::default);
+            println!(
+                "avg precision: {} with standard deviation: {}",
+                avg_precision,
+                avg_var.sqrt(),
+            );
+            // Make sure that the estimate converges to the correct value.
+            assert!(avg_precision.abs() < 0.04);
+            // We should theoretically achieve a standard deviation of about 0.12
+            assert!(avg_var.sqrt() < 0.14);
+        })
     }
 
     #[test]
     fn test_estimate_fast_low_precision() {
-        let (avg_precision, avg_var) = test_estimate(GeoDiffCount7_50::default);
-        println!(
-            "avg precision: {} with standard deviation: {}",
-            avg_precision,
-            avg_var.sqrt(),
-        );
-        // Make sure that the estimate converges to the correct value.
-        assert!(avg_precision.abs() < 0.15);
-        // We should theoretically achieve a standard deviation of about 0.25
-        assert!(avg_var.sqrt() < 0.4);
+        prng_test_harness(1, |rnd| {
+            let (avg_precision, avg_var) = test_estimate(rnd, GeoDiffCount7_50::default);
+            println!(
+                "avg precision: {} with standard deviation: {}",
+                avg_precision,
+                avg_var.sqrt(),
+            );
+            // Make sure that the estimate converges to the correct value.
+            assert!(avg_precision.abs() < 0.15);
+            // We should theoretically achieve a standard deviation of about 0.25
+            assert!(avg_var.sqrt() < 0.4);
+        });
     }
 
     #[test]
     fn test_estimate_diff_size_fast() {
-        prng_test_harness(|rnd| {
+        prng_test_harness(1, |rnd| {
             let mut a_p = GeoDiffCount7_50::default();
             let mut a_hp = GeoDiffCount7::default();
             let mut b_p = GeoDiffCount7_50::default();
@@ -515,45 +519,41 @@ mod tests {
 
     #[test]
     fn test_xor_plus_mask() {
-        for _ in 0..1000 {
-            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(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);
-                    assert_eq!(masked_expected, xor(&masked_a, &masked_b));
-                }
-            });
-        }
+        prng_test_harness(1000, |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(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);
+                assert_eq!(masked_expected, xor(&masked_a, &masked_b));
+            }
+        });
     }
 
     #[test]
     fn test_bit_chunks() {
-        for _ in 0..100 {
-            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);
-            });
-        }
+        prng_test_harness(100, |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);
+        });
     }
 
     #[test]

crates/geo_filters/src/distinct_count.rs

@@ -308,21 +308,23 @@ mod tests {
 
     #[test]
     fn test_estimate_fast() {
-        let (avg_precision, avg_var) = test_estimate(GeoDistinctCount7::default);
-        println!(
-            "avg precision: {} with standard deviation: {}",
-            avg_precision,
-            avg_var.sqrt(),
-        );
-        // Make sure that the estimate converges to the correct value.
-        assert!(avg_precision.abs() < 0.04);
-        // We should theoretically achieve a standard deviation of about 0.065
-        assert!(avg_var.sqrt() < 0.08);
+        prng_test_harness(1, |rnd| {
+            let (avg_precision, avg_var) = test_estimate(rnd, GeoDistinctCount7::default);
+            println!(
+                "avg precision: {} with standard deviation: {}",
+                avg_precision,
+                avg_var.sqrt(),
+            );
+            // Make sure that the estimate converges to the correct value.
+            assert!(avg_precision.abs() < 0.04);
+            // We should theoretically achieve a standard deviation of about 0.065
+            assert!(avg_var.sqrt() < 0.08);
+        })
     }
 
     #[test]
     fn test_estimate_union_size_fast() {
-        prng_test_harness(|rnd| {
+        prng_test_harness(1, |rnd| {
             let mut a = GeoDistinctCount7::default();
             let mut b = GeoDistinctCount7::default();
             for _ in 0..10000 {
@@ -394,19 +396,15 @@ mod tests {
 
     #[test]
     fn test_bit_chunks() {
-        for _ in 0..100 {
-            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);
-            })
-        }
+        prng_test_harness(100, |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);
+        })
     }
 
     #[test]

crates/geo_filters/src/test_rng.rs

@@ -6,15 +6,31 @@ use rand::{rngs::StdRng, SeedableRng as _};
 /// 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, T>(test_fn: F) -> T
+///
+/// You can provide a number of `iterations` this harness will run with randomly
+/// generated seeds. If a manual seed is provided via the environment then the test
+/// is only ran once with this seed.
+pub fn prng_test_harness<F>(iterations: usize, mut test_fn: F)
 where
-    F: Fn(&mut StdRng) -> T,
+    F: FnMut(&mut StdRng) -> (),
 {
-    let seed = std::env::var("TEST_SEED")
+    let maybe_manual_seed = std::env::var("TEST_SEED")
         .map(|s| s.parse::<u64>().expect("Parse TEST_SEED to u64"))
-        .unwrap_or_else(|_| rand::random());
-    let mut rng = StdRng::seed_from_u64(seed);
-    let maybe_panic = catch_unwind(AssertUnwindSafe(|| test_fn(&mut rng)));
+        .ok();
+    let mut seed = 0;
+    let maybe_panic = catch_unwind(AssertUnwindSafe(|| {
+        if let Some(manual_seed) = maybe_manual_seed {
+            seed = manual_seed;
+            let mut rng = StdRng::seed_from_u64(seed);
+            test_fn(&mut rng);
+        } else {
+            for _ in 0..iterations {
+                seed = rand::random();
+                let mut rng = StdRng::seed_from_u64(seed);
+                test_fn(&mut rng);
+            }
+        }
+    }));
     match maybe_panic {
         Ok(t) => t,
         Err(panic_info) => {