Improved BA/GF iteration tests

Closes #865
private-attribution · Jan 26, 2024 · 14b87da · 14b87da
1 parent 0d99f0d
commit 14b87da
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 48 deletions.
diff --git a/ipa-core/src/ff/boolean_array.rs b/ipa-core/src/ff/boolean_array.rs
@@ -284,6 +284,11 @@ macro_rules! boolean_array_impl {
             #[derive(Clone, Copy, PartialEq, Eq, Debug)]
             pub struct $name(pub(super) Store);
 
+            impl $name {
+                #[cfg(all(test, unit_test))]
+                const STORE_LEN: usize = bitvec::mem::elts::<u8>($bits);
+            }
+
             impl ArrayAccess for $name {
                 type Output = Boolean;
                 type Iter<'a> = BAIterator<'a>;
@@ -427,12 +432,26 @@ macro_rules! boolean_array_impl {
 
             #[cfg(all(test, unit_test))]
             mod tests {
+                use proptest::{
+                    prelude::{prop, Arbitrary, Strategy},
+                    proptest,
+                };
                 use rand::{thread_rng, Rng};
 
                 use super::*;
 
-                // Only small BAs expose this via `Field`.
-                const ONE: $name = $name(bitarr_one!($bits));
+                impl Arbitrary for $name {
+                    type Parameters = <[u8; $name::STORE_LEN] as Arbitrary>::Parameters;
+                    type Strategy = prop::strategy::Map<
+                        <[u8; $name::STORE_LEN] as Arbitrary>::Strategy,
+                        fn([u8; $name::STORE_LEN]) -> Self,
+                    >;
+
+                    fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+                        <[u8; $name::STORE_LEN]>::arbitrary_with(args)
+                            .prop_map(|arr| $name(Store::from(arr)))
+                    }
+                }
 
                 #[test]
                 fn set_boolean_array() {
@@ -444,46 +463,31 @@ macro_rules! boolean_array_impl {
                     assert_eq!(ba.get(i), Some(a));
                 }
 
-                #[test]
-                fn iterate_boolean_array() {
-                    let bits = ONE;
-                    let iter = bits.iter();
-                    for (i, j) in iter.enumerate() {
-                        if i == 0 {
-                            assert_eq!(j, Boolean::ONE);
-                        } else {
-                            assert_eq!(j, Boolean::ZERO);
+                proptest! {
+                    #[test]
+                    fn iterate_boolean_array(a: $name) {
+                        let mut iter = a.iter().enumerate();
+                        assert_eq!(iter.len(), $bits);
+                        while let Some((i, b)) = iter.next() {
+                            assert_eq!(bool::from(b), a.0[i]);
+                            assert_eq!(iter.len(), $bits - 1 - i);
                         }
                     }
-                }
 
-                #[test]
-                fn iterate_secret_shared_boolean_array() {
+                    #[test]
+                    fn iterate_secret_shared_boolean_array(a: AdditiveShare<$name>) {
                     use crate::secret_sharing::replicated::ReplicatedSecretSharing;
-                    let bits = AdditiveShare::new(ONE, ONE);
-                    let iter = bits.into_iter();
-                    assert_eq!(iter.len(), $bits);
-                    for (i, j) in iter.enumerate() {
-                        if i == 0 {
-                            assert_eq!(j, AdditiveShare::new(Boolean::ONE, Boolean::ONE));
-                        } else {
-                            assert_eq!(j, AdditiveShare::<Boolean>::ZERO);
+                        let mut iter = a.iter().enumerate();
+                        assert_eq!(iter.len(), $bits);
+                        while let Some((i, sb)) = iter.next() {
+                            let left = Boolean::from(a.left().0[i]);
+                            let right = Boolean::from(a.right().0[i]);
+                            assert_eq!(sb, AdditiveShare::new(left, right));
+                            assert_eq!(iter.len(), $bits - 1 - i);
                         }
                     }
                 }
 
-                #[test]
-                fn iterate_secret_shared_boolean_array_len() {
-                    use crate::secret_sharing::replicated::ReplicatedSecretSharing;
-                    let bits = AdditiveShare::new(ONE, ONE);
-                    let mut iter = bits.into_iter();
-                    assert_eq!(iter.len(), $bits);
-                    for b in (0..$bits).rev() {
-                        iter.next().unwrap();
-                        assert_eq!(iter.len(), b);
-                    }
-                }
-
                 #[test]
                 fn serde() {
                     let ba = thread_rng().gen::<$name>();

diff --git a/ipa-core/src/ff/galois_field.rs b/ipa-core/src/ff/galois_field.rs
@@ -624,21 +624,12 @@ macro_rules! bit_array_impl {
                     }
 
                     #[test]
-                    fn arrayacces_iter(a: $name) {
-                        let mut val = 0u128;
-                        for (i, b) in a.iter().enumerate() {
-                            val |= u128::from(b) << i;
-                        }
-                        assert_eq!(val, a.as_u128());
-                    }
-
-                    #[test]
-                    fn arrayacces_iter_len(a: $name) {
-                        let mut iter = a.iter();
+                    fn arrayaccess_iter(a: $name) {
+                        let mut iter = a.iter().enumerate();
                         assert_eq!(iter.len(), $bits);
-                        for i in (0..$bits).rev() {
-                            iter.next().unwrap();
-                            assert_eq!(iter.len(), i);
+                        while let Some((i, b)) = iter.next() {
+                            assert_eq!(u128::from(b), (a.as_u128() >> i) & 1);
+                            assert_eq!(iter.len(), $bits - 1 - i);
                         }
                     }
                 }