Add Arbitrary implementation for other types in fp

ext/crates/fp/src/field/field_internal.rs

@@ -47,7 +47,7 @@ macro_rules! normal_from_assign {
 /// operations to the outside world.
 #[allow(private_bounds)]
 pub trait FieldInternal:
-    std::fmt::Debug + Copy + PartialEq + Eq + Hash + Sized + crate::MaybeArbitrary<()>
+    std::fmt::Debug + Copy + PartialEq + Eq + Hash + Sized + crate::MaybeArbitrary<()> + 'static
 {
     /// The internal representation of a field element.
     type ElementContainer: FieldElementContainer;

ext/crates/fp/src/matrix/matrix_inner.rs

@@ -962,6 +962,120 @@ impl std::ops::AddAssign<&Self> for Matrix {
     }
 }
 
+#[cfg(feature = "proptest")]
+pub mod arbitrary {
+    use proptest::prelude::*;
+
+    use super::*;
+    use crate::{
+        field::Fp,
+        vector::{arbitrary::FqVectorArbParams, FqVector},
+    };
+
+    pub const MAX_ROWS: usize = 100;
+    pub const MAX_COLUMNS: usize = 100;
+
+    #[derive(Debug, Clone)]
+    pub struct MatrixArbParams {
+        pub p: Option<ValidPrime>,
+        pub rows: BoxedStrategy<usize>,
+        pub columns: BoxedStrategy<usize>,
+    }
+
+    impl Default for MatrixArbParams {
+        fn default() -> Self {
+            Self {
+                p: None,
+                rows: (1..=MAX_ROWS).boxed(),
+                columns: (1..=MAX_COLUMNS).boxed(),
+            }
+        }
+    }
+
+    impl Arbitrary for Matrix {
+        type Parameters = MatrixArbParams;
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+            let p = match args.p {
+                Some(p) => Just(p).boxed(),
+                None => any::<ValidPrime>().boxed(),
+            };
+
+            (p, args.rows, args.columns)
+                .prop_flat_map(|(p, rows, columns)| {
+                    let row_strategy = any_with::<FqVector<_>>(FqVectorArbParams {
+                        fq: Some(Fp::new(p)),
+                        len: Just(columns).boxed(),
+                    })
+                    .prop_map(|v| -> FpVector { v.into() });
+
+                    let rows = proptest::collection::vec(row_strategy, rows);
+                    (Just(p), rows, Just(columns))
+                })
+                .prop_map(|(p, rows, columns)| Self::from_rows(p, rows, columns))
+                .boxed()
+        }
+    }
+
+    impl Matrix {
+        /// Generate an arbitrary row-reduced matrix.
+        ///
+        /// This is more interesting than just generating an arbitrary matrix and row-reducing. If
+        /// we pick a matrix uniformly at random in the space of all $n \times m$ matrices, it has a
+        /// very high probability of having full rank with all its pivots in the first $n$ columns.
+        /// This implies that, after projecting to the space of row-reduced matrices, the output is
+        /// very likely to be an identity matrix augmented by a random matrix. If $m$ is
+        /// significantly larger than $n$, this is only a tiny subspace of the space of all
+        /// row-reduced matrices.
+        ///
+        /// While a search through *all* $n \times m$ matrices will also cover all row-reduced
+        /// matrices, in practice this space is so large that we only test a vanishingly small
+        /// fraction of it. Therefore, if a method that is sensitive to the pivot structure of the
+        /// input matrix is proptested using `arbitrary_with`, it is unlikely that the tests will
+        /// cover many matrices with interesting pivots, while those are the most likely to cause
+        /// bugs. This function attempts to generate a matrix that is chosen uniformly at random
+        /// directly in the space of all row-reduced matrices.
+        ///
+        /// In practice, this is not quite right. There is no randomness in the code; instead we
+        /// generate a `Strategy` that samples from only the space of row-reduced matrices. Also,
+        /// depending on the parameters, the strategy may output matrices that are not all of the
+        /// same size or even over the same ground field, so using the word "space" is slightly
+        /// improper, mathematically speaking.
+        pub fn arbitrary_rref_with(args: MatrixArbParams) -> impl Strategy<Value = Self> {
+            Self::arbitrary_with(args)
+                .prop_flat_map(|m| {
+                    let column_vec = (0..m.columns()).collect::<Vec<_>>();
+                    let smallest_dim = std::cmp::min(m.rows(), m.columns());
+                    let pivot_cols = proptest::sample::subsequence(column_vec, 0..=smallest_dim);
+                    (Just(m), pivot_cols)
+                })
+                .prop_map(|(mut m, pivot_cols)| {
+                    // Ensure rows start with 0s followed by a 1 in their pivot column
+                    for (row_idx, row) in m.iter_mut().enumerate() {
+                        if let Some(&col_idx) = pivot_cols.get(row_idx) {
+                            row.slice_mut(0, col_idx).set_to_zero();
+                            row.set_entry(col_idx, 1);
+                        } else {
+                            row.set_to_zero();
+                        }
+                    }
+                    // Set all other entries in the pivot columns to 0
+                    for (row_idx, &col_idx) in pivot_cols.iter().enumerate() {
+                        for row in m.iter_mut().take(row_idx) {
+                            row.set_entry(col_idx, 0);
+                        }
+                    }
+                    m
+                })
+        }
+
+        pub fn arbitrary_rref() -> impl Strategy<Value = Self> {
+            Self::arbitrary_rref_with(MatrixArbParams::default())
+        }
+    }
+}
+
 /// This models an augmented matrix.
 ///
 /// In an ideal world, this will have no public fields. The inner matrix
@@ -1223,6 +1337,8 @@ impl<'a> MatrixSliceMut<'a> {
 
 #[cfg(test)]
 mod tests {
+    use proptest::prelude::*;
+
     use super::*;
 
     #[test]
@@ -1278,4 +1394,14 @@ mod tests {
             assert_eq!(m.pivots(), &goal_pivots)
         }
     }
+
+    proptest! {
+        // Test that `arbitrary_rref` generates matrices in rref.
+        #[test]
+        fn test_arbitrary_rref(m in Matrix::arbitrary_rref()) {
+            let mut m_red = m.clone();
+            m_red.row_reduce();
+            prop_assert_eq!(m, m_red);
+        }
+    }
 }

ext/crates/fp/src/matrix/mod.rs

@@ -13,3 +13,5 @@ pub use matrix_inner::{AugmentedMatrix, Matrix, MatrixSliceMut};
 pub use quasi_inverse::QuasiInverse;
 pub use subquotient::Subquotient;
 pub use subspace::Subspace;
+#[cfg(feature = "proptest")]
+pub use {matrix_inner::arbitrary::*, subquotient::arbitrary::*, subspace::arbitrary::*};

ext/crates/fp/src/matrix/subquotient.rs

@@ -96,10 +96,8 @@ impl Subquotient {
 
     /// The pivot columns of the complement to the subspace
     pub fn complement_pivots(&self) -> impl Iterator<Item = usize> + '_ {
-        (0..self.ambient_dimension()).filter(|&i| {
-            !self.quotient.pivots().contains(&(i as isize))
-                && !self.gens.pivots().contains(&(i as isize))
-        })
+        (0..self.ambient_dimension())
+            .filter(|&i| self.quotient.pivots()[i] < 0 && self.gens.pivots()[i] < 0)
     }
 
     pub fn quotient(&mut self, elt: FpSlice) {
@@ -184,9 +182,62 @@ impl Subquotient {
     }
 }
 
+#[cfg(feature = "proptest")]
+pub mod arbitrary {
+    use proptest::prelude::*;
+
+    use super::*;
+    use crate::matrix::subspace::arbitrary::SubspaceArbParams;
+    pub use crate::matrix::subspace::arbitrary::MAX_DIM;
+
+    #[derive(Debug, Clone)]
+    pub struct SubquotientArbParams {
+        pub p: Option<ValidPrime>,
+        pub dim: BoxedStrategy<usize>,
+    }
+
+    impl Default for SubquotientArbParams {
+        fn default() -> Self {
+            Self {
+                p: None,
+                dim: (0..=MAX_DIM).boxed(),
+            }
+        }
+    }
+
+    impl Arbitrary for Subquotient {
+        type Parameters = SubquotientArbParams;
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+            let p = match args.p {
+                Some(p) => Just(p).boxed(),
+                None => any::<ValidPrime>().boxed(),
+            };
+
+            (p, args.dim)
+                .prop_flat_map(|(p, dim)| {
+                    let sub = Subspace::arbitrary_with(SubspaceArbParams {
+                        p: Some(p),
+                        dim: Just(dim).boxed(),
+                    });
+                    let quotient = Subspace::arbitrary_with(SubspaceArbParams {
+                        p: Some(p),
+                        dim: Just(dim).boxed(),
+                    });
+
+                    (sub, quotient)
+                })
+                .prop_map(|(sub, quotient)| Self::from_parts(sub, quotient))
+                .boxed()
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use expect_test::expect;
+    use proptest::prelude::*;
 
     use super::*;
 
@@ -220,4 +271,14 @@ mod tests {
 
         assert_eq!(sq.gens().count(), 1);
     }
+
+    proptest! {
+        #[test]
+        fn test_sum_quotient_gens_complement_is_ambient(sq: Subquotient) {
+            let quotient_dim = sq.zeros().dimension();
+            let gens_dim = sq.gens().count();
+            let complement_dim = sq.complement_pivots().count();
+            assert_eq!(quotient_dim + gens_dim + complement_dim, sq.ambient_dimension());
+        }
+    }
 }

ext/crates/fp/src/matrix/subspace.rs

@@ -298,3 +298,50 @@ impl std::fmt::Display for Subspace {
         Ok(())
     }
 }
+
+#[cfg(feature = "proptest")]
+pub mod arbitrary {
+    use proptest::prelude::*;
+
+    use super::*;
+    use crate::matrix::matrix_inner::arbitrary::MatrixArbParams;
+    pub use crate::matrix::matrix_inner::arbitrary::MAX_COLUMNS as MAX_DIM;
+
+    #[derive(Debug, Clone)]
+    pub struct SubspaceArbParams {
+        pub p: Option<ValidPrime>,
+        pub dim: BoxedStrategy<usize>,
+    }
+
+    impl Default for SubspaceArbParams {
+        fn default() -> Self {
+            Self {
+                p: None,
+                dim: (0..=MAX_DIM).boxed(),
+            }
+        }
+    }
+
+    impl Arbitrary for Subspace {
+        type Parameters = SubspaceArbParams;
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+            let p = match args.p {
+                Some(p) => Just(p).boxed(),
+                None => any::<ValidPrime>().boxed(),
+            };
+
+            (p, args.dim)
+                .prop_flat_map(move |(p, dim)| {
+                    Matrix::arbitrary_rref_with(MatrixArbParams {
+                        p: Some(p),
+                        rows: (0..=dim + 1).boxed(),
+                        columns: Just(dim).boxed(),
+                    })
+                })
+                .prop_map(Self::from_matrix)
+                .boxed()
+        }
+    }
+}

ext/crates/fp/src/vector/impl_fqvector.rs

@@ -408,3 +408,45 @@ impl<F: Field> std::fmt::Display for FqVector<F> {
         self.as_slice().fmt(f)
     }
 }
+
+#[cfg(feature = "proptest")]
+pub mod arbitrary {
+    use proptest::prelude::*;
+
+    use super::*;
+
+    pub const MAX_LEN: usize = 10_000;
+
+    #[derive(Debug, Clone)]
+    pub struct FqVectorArbParams<F> {
+        pub fq: Option<F>,
+        pub len: BoxedStrategy<usize>,
+    }
+
+    impl<F> Default for FqVectorArbParams<F> {
+        fn default() -> Self {
+            Self {
+                fq: None,
+                len: (0..=MAX_LEN).boxed(),
+            }
+        }
+    }
+
+    impl<F: Field> Arbitrary for FqVector<F> {
+        type Parameters = FqVectorArbParams<F>;
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+            let fq = match args.fq {
+                Some(fq) => Just(fq).boxed(),
+                None => any::<F>().boxed(),
+            };
+            (fq, args.len)
+                .prop_flat_map(|(fq, len)| {
+                    (Just(fq), proptest::collection::vec(fq.arb_element(), len))
+                })
+                .prop_map(|(fq, v)| Self::from_slice(fq, &v))
+                .boxed()
+        }
+    }
+}

ext/crates/fp/src/vector/mod.rs

@@ -7,21 +7,21 @@ mod impl_fqvector;
 mod iter;
 
 pub use fp_wrapper::*;
+#[cfg(feature = "proptest")]
+pub use impl_fqvector::arbitrary;
 pub use inner::*;
 
 #[cfg(test)]
 pub(super) mod tests {
     use itertools::Itertools;
     use proptest::prelude::*;
 
-    use super::inner::FqVector;
+    use super::{arbitrary::MAX_LEN as MAX_TEST_VEC_LEN, inner::FqVector};
     use crate::{
         field::{element::FieldElement, fp::F2, Field},
         limb,
     };
 
-    pub const MAX_TEST_VEC_LEN: usize = 10_000;
-
     pub struct VectorDiffEntry<F: Field> {
         pub index: usize,
         pub left: FieldElement<F>,