Merge branch 'main' into feat/fri-higher-arity-2

Author: OsamaAlkhodairy

README.md

@@ -63,14 +63,14 @@ Many variations are possible, with different fields, hashes and so forth, which
 
 For example, to prove 2^20 Poseidon2 permutations of width 16, using the `KoalaBear` field, `Radix2DitParallel` DFT and `KeccakF` as the Merkle tree hash:
 ```
-RUSTFLAGS="-Ctarget-cpu=native" cargo run --example prove_monty_31 --release --features parallel -- --field koala-bear --objective poseidon-2-permutations --log-trace-length 17 --discrete-fourier-transform radix-2-dit-parallel --merkle-hash keccak-f
+RUSTFLAGS="-Ctarget-cpu=native" cargo run --example prove_prime_field_31 --release --features parallel -- --field koala-bear --objective poseidon-2-permutations --log-trace-length 17 --discrete-fourier-transform radix-2-dit-parallel --merkle-hash keccak-f
 ```
 
 Currently the options for the command line arguments are:
-- `--field` (`-f`): `koala-bear` or `baby-bear`.
+- `--field` (`-f`): `mersenne-31` or `koala-bear` or `baby-bear`.
 - `--objective` (`-o`): `blake-3-permutations, poseidon-2-permutations, keccak-f-permutations`.
 - `--log-trace-length` (`-l`): Accepts any integer between `0` and `255`. The number of permutations proven is `trace_length, 8*trace_length` and `trace_length/24` for `blake3, poseidon2` and `keccakf` respectively. 
-- `--discrete-fourier-transform` (`-d`): `radix-2-dit-parallel, recursive-dft`.
+- `--discrete-fourier-transform` (`-d`): `radix-2-dit-parallel, recursive-dft`. This option should be omitted if the field choice is `mersenne-31` as the circle stark currently only supports a single discrete fourier transform.
 - `--merkle-hash` (`-m`): `poseidon-2, keccak-f`.
 
 Extra speedups may be possible with some configuration changes:

air/src/utils.rs

@@ -76,15 +76,23 @@ pub fn checked_andn<F: Field>(x: F, y: F) -> F {
 /// The output array is in little-endian order.
 #[inline]
 pub fn u32_to_bits_le<FA: FieldAlgebra>(val: u32) -> [FA; 32] {
-    // We do this over F::from_canonical_u32 as from_canonical_u32 can be slow
-    // like in the case of monty field.
-    array::from_fn(|i| {
-        if val & (1 << i) != 0 {
-            FA::ONE
-        } else {
-            FA::ZERO
-        }
-    })
+    array::from_fn(|i| FA::from_bool(val & (1 << i) != 0))
+}
+
+/// Convert a 64-bit integer into an array of 64 0 or 1 field elements.
+///
+/// The output array is in little-endian order.
+#[inline]
+pub fn u64_to_bits_le<FA: FieldAlgebra>(val: u64) -> [FA; 64] {
+    array::from_fn(|i| FA::from_bool(val & (1 << i) != 0))
+}
+
+/// Convert a 64-bit integer into an array of four field elements representing the 16 bit limb decomposition.
+///
+/// The output array is in little-endian order.
+#[inline]
+pub fn u64_to_16_bit_limbs<FA: FieldAlgebra>(val: u64) -> [FA; 4] {
+    array::from_fn(|i| FA::from_canonical_u16((val >> (16 * i)) as u16))
 }
 
 /// Verify that `a = b + c + d mod 2^32`

blake3-air/examples/prove_blake3_m31_keccak.rs

@@ -18,9 +18,13 @@ use crate::{generate_trace_rows, Blake3State, FullRound, QuarterRound};
 pub struct Blake3Air {}
 
 impl Blake3Air {
-    pub fn generate_trace_rows<F: PrimeField64>(&self, num_hashes: usize) -> RowMajorMatrix<F> {
+    pub fn generate_trace_rows<F: PrimeField64>(
+        &self,
+        num_hashes: usize,
+        extra_capacity_bits: usize,
+    ) -> RowMajorMatrix<F> {
         let inputs = (0..num_hashes).map(|_| random()).collect::<Vec<_>>();
-        generate_trace_rows(inputs)
+        generate_trace_rows(inputs, extra_capacity_bits)
     }
 
     /// Verify that the quarter round function has been correctly computed.

blake3-air/examples/prove_blake3_m31_poseidon2.rs

@@ -13,14 +13,23 @@ use crate::{Blake3State, FullRound};
 
 // TODO: Take generic iterable
 #[instrument(name = "generate Blake3 trace", skip_all)]
-pub fn generate_trace_rows<F: PrimeField64>(inputs: Vec<[u32; 24]>) -> RowMajorMatrix<F> {
+pub fn generate_trace_rows<F: PrimeField64>(
+    inputs: Vec<[u32; 24]>,
+    extra_capacity_bits: usize,
+) -> RowMajorMatrix<F> {
     let num_rows = inputs.len();
     assert!(
         num_rows.is_power_of_two(),
         "Callers expected to pad inputs to VECTOR_LEN times a power of two"
     );
 
-    let mut trace = RowMajorMatrix::new(F::zero_vec(num_rows * NUM_BLAKE3_COLS), NUM_BLAKE3_COLS);
+    let trace_length = num_rows * NUM_BLAKE3_COLS;
+
+    // We allocate extra_capacity_bits now as this will be needed by the dft.
+    let mut long_trace = F::zero_vec(trace_length << extra_capacity_bits);
+    long_trace.truncate(trace_length);
+
+    let mut trace = RowMajorMatrix::new(long_trace, NUM_BLAKE3_COLS);
     let (prefix, rows, suffix) = unsafe { trace.values.align_to_mut::<Blake3Cols<F>>() };
     assert!(prefix.is_empty(), "Alignment should match");
     assert!(suffix.is_empty(), "Alignment should match");
@@ -205,14 +214,14 @@ fn verifiable_half_round(
     let (rot_1, rot_2) = if flag { (8, 7) } else { (16, 12) };
 
     // The first summation:
-    a += b;
-    a += m;
+    a = a.wrapping_add(b);
+    a = a.wrapping_add(m);
 
     // The first xor:
     d = (d ^ a).rotate_right(rot_1);
 
     // The second summation:
-    c += d;
+    c = c.wrapping_add(d);
 
     // The second xor:
     b = (b ^ c).rotate_right(rot_2);

blake3-air/src/air.rs

@@ -109,10 +109,6 @@ impl FieldAlgebra for Bn254Fr {
         f
     }
 
-    fn from_bool(b: bool) -> Self {
-        Self::new(FFBn254Fr::from(b as u64))
-    }
-
     fn from_canonical_u8(n: u8) -> Self {
         Self::new(FFBn254Fr::from(n as u64))
     }

blake3-air/src/generation.rs

@@ -13,3 +13,4 @@ tracing.workspace = true
 
 [dev-dependencies]
 p3-goldilocks.workspace = true
+p3-baby-bear.workspace = true

bn254-fr/src/lib.rs

@@ -148,8 +148,8 @@ where
     P: CryptographicPermutation<[F; WIDTH]>,
 {
     fn sample_bits(&mut self, bits: usize) -> usize {
-        debug_assert!(bits < (usize::BITS as usize));
-        debug_assert!((1 << bits) < F::ORDER_U64);
+        assert!(bits < (usize::BITS as usize));
+        assert!((1 << bits) < F::ORDER_U64);
         let rand_f: F = self.sample();
         let rand_usize = rand_f.as_canonical_u64() as usize;
         rand_usize & ((1 << bits) - 1)
@@ -160,45 +160,88 @@ where
 mod tests {
     use core::iter;
 
+    use p3_baby_bear::BabyBear;
     use p3_field::FieldAlgebra;
     use p3_goldilocks::Goldilocks;
     use p3_symmetric::Permutation;
 
     use super::*;
+    use crate::grinding_challenger::GrindingChallenger;
 
     const WIDTH: usize = 24;
     const RATE: usize = 16;
 
-    type TestArray = [F; WIDTH];
-    type F = Goldilocks;
+    type G = Goldilocks;
+    type BB = BabyBear;
 
     #[derive(Clone)]
     struct TestPermutation {}
 
-    impl Permutation<TestArray> for TestPermutation {
-        fn permute_mut(&self, input: &mut TestArray) {
+    impl<F: Clone> Permutation<[F; WIDTH]> for TestPermutation {
+        fn permute_mut(&self, input: &mut [F; WIDTH]) {
             input.reverse()
         }
     }
 
-    impl CryptographicPermutation<TestArray> for TestPermutation {}
+    impl<F: Clone> CryptographicPermutation<[F; WIDTH]> for TestPermutation {}
 
     #[test]
     fn test_duplex_challenger() {
-        type Chal = DuplexChallenger<F, TestPermutation, WIDTH, RATE>;
+        type Chal = DuplexChallenger<G, TestPermutation, WIDTH, RATE>;
         let permutation = TestPermutation {};
         let mut duplex_challenger = DuplexChallenger::new(permutation);
 
         // Observe 12 elements.
-        (0..12).for_each(|element| duplex_challenger.observe(F::from_canonical_u8(element as u8)));
+        (0..12).for_each(|element| duplex_challenger.observe(G::from_canonical_u8(element as u8)));
 
-        let state_after_duplexing: Vec<_> = iter::repeat(F::ZERO)
+        let state_after_duplexing: Vec<_> = iter::repeat(G::ZERO)
             .take(12)
-            .chain((0..12).map(F::from_canonical_u8).rev())
+            .chain((0..12).map(G::from_canonical_u8).rev())
             .collect();
 
-        let expected_samples: Vec<F> = state_after_duplexing[..16].iter().copied().rev().collect();
-        let samples = <Chal as CanSample<F>>::sample_vec(&mut duplex_challenger, 16);
+        let expected_samples: Vec<G> = state_after_duplexing[..16].iter().copied().rev().collect();
+        let samples = <Chal as CanSample<G>>::sample_vec(&mut duplex_challenger, 16);
         assert_eq!(samples, expected_samples);
     }
+
+    #[test]
+    #[should_panic]
+    fn test_duplex_challenger_sample_bits_security() {
+        type GoldilocksChal = DuplexChallenger<G, TestPermutation, WIDTH, RATE>;
+        let permutation = TestPermutation {};
+        let mut duplex_challenger = GoldilocksChal::new(permutation);
+
+        for _ in 0..100 {
+            assert!(duplex_challenger.sample_bits(129) < 4);
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_duplex_challenger_sample_bits_security_small_field() {
+        type BabyBearChal = DuplexChallenger<BB, TestPermutation, WIDTH, RATE>;
+        let permutation = TestPermutation {};
+        let mut duplex_challenger = BabyBearChal::new(permutation);
+
+        for _ in 0..100 {
+            assert!(duplex_challenger.sample_bits(40) < 1 << 31);
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_duplex_challenger_grind_security() {
+        type GoldilocksChal = DuplexChallenger<G, TestPermutation, WIDTH, RATE>;
+        let permutation = TestPermutation {};
+        let mut duplex_challenger = GoldilocksChal::new(permutation);
+
+        // This should cause sample_bits (and hence grind and check_witness) to
+        // panic. If bit sizes were not constrained correctly inside the
+        // challenger, (1 << too_many_bits) would loop around, incorrectly
+        // grinding and accepting a 1-bit PoW.
+        let too_many_bits = usize::BITS as usize;
+
+        let witness = duplex_challenger.grind(too_many_bits);
+        assert!(duplex_challenger.check_witness(too_many_bits, witness));
+    }
 }