Merge branch 'main' into sc-20250723-serialize

Author: itsibitzi

crates/geo_filters/src/config.rs

@@ -353,13 +353,15 @@ 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::{rngs::StdRng, RngCore};
 
     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) {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
+    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;

crates/geo_filters/src/config/lookup.rs

@@ -45,29 +45,36 @@ impl HashToBucketLookup {
 
 #[cfg(test)]
 mod tests {
-    use rand::{RngCore, SeedableRng};
+    use rand::{rngs::StdRng, 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;
 
     #[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);
+
         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;

crates/geo_filters/src/diff_count.rs

@@ -96,21 +96,17 @@ impl<'a, C: GeoConfig<Diff>> GeoDiffCount<'a, 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();
-
         let lsb = BitVec::from_bit_chunks(ones.into_bitchunks(), smallest_msb.into_usize());
-
         let result = Self {
             config,
             msb: Cow::from(msb),
@@ -217,17 +213,16 @@ impl<'a, C: GeoConfig<Diff>> GeoDiffCount<'a, C> {
             match msb.binary_search_by(|k| bucket.cmp(k)) {
                 Ok(idx) => {
                     msb.remove(idx);
-                    let (first, second) = {
+                    let first = {
                         let mut lsb = iter_ones(self.lsb.bit_chunks().peekable());
-                        (lsb.next(), lsb.next())
+                        lsb.next()
                     };
-                    let new_smallest = if let Some(smallest) = first {
+                    if let Some(smallest) = first {
                         msb.push(C::BucketType::from_usize(smallest));
-                        second.map(|_| smallest).unwrap_or(0)
+                        self.lsb.resize(smallest);
                     } else {
-                        0
+                        self.lsb.resize(0);
                     };
-                    self.lsb.resize(new_smallest);
                 }
                 Err(idx) => {
                     msb.insert(idx, bucket);
@@ -245,6 +240,12 @@ impl<'a, C: GeoConfig<Diff>> GeoDiffCount<'a, C> {
                         // ensure LSB bit vector has the space for `smallest`
                         self.lsb.resize(new_smallest);
                         self.lsb.toggle(smallest);
+                    } else if msb.len() == self.config.max_msb_len() {
+                        let smallest = msb
+                            .last()
+                            .expect("should have at least one element")
+                            .into_usize();
+                        self.lsb.resize(smallest);
                     }
                 }
             }
@@ -418,11 +419,12 @@ mod tests {
     use std::io::Write;
 
     use itertools::Itertools;
-    use rand::{seq::IteratorRandom, RngCore, SeedableRng};
+    use rand::{rngs::StdRng, seq::IteratorRandom, RngCore};
 
     use crate::{
         build_hasher::UnstableDefaultBuildHasher,
         config::{iter_ones, tests::test_estimate, FixedConfig},
+        test_rng::prng_test_harness,
     };
 
     use super::*;
@@ -493,57 +495,62 @@ 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() {
-        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(1, |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]
@@ -575,45 +582,39 @@ 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(10, |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() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
-        for _ in 0..100 {
+        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(),
-            );
+            let actual =
+                GeoDiffCount::from_bit_chunks(expected.config.clone(), expected.bit_chunks());
             assert_eq!(expected, actual);
-        }
+        });
     }
 
     #[test]
@@ -656,44 +657,44 @@ mod tests {
     // This helper exists in order to easily test serializing types with different
     // bucket types in the MSB sparse bit field representation. See tests below.
     #[cfg(target_endian = "little")]
-    fn serialization_round_trip<C: GeoConfig<Diff> + Default>() {
-        let mut rnd = rand::rngs::StdRng::from_os_rng();
+    fn serialization_round_trip<C: GeoConfig<Diff> + Default>(rnd: &mut StdRng) {
         // Run 100 simulations of random values being put into
         // a diff counter. "Serializing" to a vector to emulate
         // writing to a disk, and then deserializing and asserting
         // the filters are equal.
-        for _ in 0..100 {
-            let mut before = GeoDiffCount::<'_, C>::default();
-            // Select a random number of items to insert.
-            let items = (1..1000).choose(&mut rnd).unwrap();
-            for _ in 0..items {
-                before.push_hash(rnd.next_u64());
-            }
-            let mut writer = vec![];
-            // Insert some padding to emulate alignment issues with the slices.
-            // A previous version of this test never panicked even though we were
-            // violating the alignment preconditions for the `from_raw_parts` function.
-            let padding = [0_u8; 8];
-            let pad_amount = (0..8).choose(&mut rnd).unwrap();
-            writer.write_all(&padding[..pad_amount]).unwrap();
-            before.write(&mut writer).unwrap();
-            let after =
-                GeoDiffCount::<'_, C>::from_bytes(before.config.clone(), &writer[pad_amount..]);
-            assert_eq!(before, after);
+        let mut before = GeoDiffCount::<'_, C>::default();
+        // Select a random number of items to insert.
+        let items = (1..1000).choose(rnd).unwrap();
+        for _ in 0..items {
+            before.push_hash(rnd.next_u64());
         }
+        let mut writer = vec![];
+        // Insert some padding to emulate alignment issues with the slices.
+        // A previous version of this test never panicked even though we were
+        // violating the alignment preconditions for the `from_raw_parts` function.
+        let padding = [0_u8; 8];
+        let pad_amount = (0..8).choose(rnd).unwrap();
+        writer.write_all(&padding[..pad_amount]).unwrap();
+        before.write(&mut writer).unwrap();
+        let after = GeoDiffCount::<'_, C>::from_bytes(before.config.clone(), &writer[pad_amount..]);
+        assert_eq!(before, after);
     }
 
     #[test]
     #[cfg(target_endian = "little")]
     fn test_serialization_round_trip_7() {
-        // Uses a u16 for MSB buckets.
-        serialization_round_trip::<GeoDiffConfig7>();
+        prng_test_harness(100, |rnd| {
+            // Uses a u16 for MSB buckets.
+            serialization_round_trip::<GeoDiffConfig7>(rnd);
+        });
     }
 
     #[test]
     #[cfg(target_endian = "little")]
     fn test_serialization_round_trip_13() {
-        // Uses a u32 for MSB buckets.
-        serialization_round_trip::<GeoDiffConfig13>();
+        prng_test_harness(100, |rnd| {
+            // Uses a u32 for MSB buckets.
+            serialization_round_trip::<GeoDiffConfig13>(rnd);
+        });
     }
 }