diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 00000000..00c55758 Binary files /dev/null and b/.DS_Store differ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e5a4bc5f..addfffc8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,6 +41,9 @@ jobs: - name: Test run: cargo test --verbose + - name: Test with parallel + run: cargo test --verbose --features parallel + lint: name: Formatting and Clippy runs-on: ubuntu-latest @@ -74,3 +77,57 @@ jobs: - name: Format run: cargo +nightly fmt --all -- --check + check_crates: + name: Check Crates + runs-on: ubuntu-latest + if: "! contains(toJSON(github.event.commits.*.message), '[skip-ci]')" + + steps: + - uses: actions/checkout@v4 + + - uses: dtolnay/rust-toolchain@stable + id: rs-stable + + - uses: actions/cache@v3 + with: + path: | + ~/.cargo/bin/ + ~/.cargo/registry/index/ + ~/.cargo/registry/cache/ + ~/.cargo/git/db/ + target/ + key: rust-${{ steps.rs-stable.outputs.rustc_hash }}-${{ hashFiles('**/Cargo.toml', '**/Cargo.lock') }} + + - name: Cargo Check Crates + run: | + cargo check --verbose --package p3-air + cargo check --verbose --package p3-baby-bear + cargo check --verbose --package p3-blake3 + cargo check --verbose --package p3-brakedown + cargo check --verbose --package p3-challenger + cargo check --verbose --package p3-code + cargo check --verbose --package p3-commit + cargo check --verbose --package p3-dft + cargo check --verbose --package p3-field + cargo check --verbose --package p3-field-testing + cargo check --verbose --package p3-fri + cargo check --verbose --package p3-goldilocks + cargo check --verbose --package p3-interpolation + cargo check --verbose --package p3-keccak + cargo check --verbose --package p3-keccak-air + cargo check --verbose --package p3-lde + cargo check --verbose --package p3-matrix + cargo check --verbose --package p3-maybe-rayon + cargo check --verbose --package p3-mds + cargo check --verbose --package p3-merkle-tree + cargo check --verbose --package p3-mersenne-31 + cargo check --verbose --package p3-monolith + cargo check --verbose --package p3-multi-stark + cargo check --verbose --package p3-poseidon + cargo check --verbose --package p3-poseidon2 + cargo check --verbose --package p3-reed-solomon + cargo check --verbose --package p3-rescue + cargo check --verbose --package p3-symmetric + cargo check --verbose --package p3-tensor-pcs + cargo check --verbose --package p3-uni-stark + cargo check --verbose --package p3-util \ No newline at end of file diff --git a/air/src/virtual_column.rs b/air/src/virtual_column.rs index 083b3838..0ccc9835 100644 --- a/air/src/virtual_column.rs +++ b/air/src/virtual_column.rs @@ -1,17 +1,18 @@ use alloc::vec; use alloc::vec::Vec; -use core::ops::{Add, Mul}; +use core::ops::Mul; use p3_field::{AbstractField, Field}; /// An affine function over columns in a PAIR. +#[derive(Clone, Debug)] pub struct VirtualPairCol { column_weights: Vec<(PairCol, F)>, constant: F, } /// A column in a PAIR, i.e. either a preprocessed column or a main trace column. -#[derive(Debug, Clone, Copy)] +#[derive(Clone, Copy, Debug)] pub enum PairCol { Preprocessed(usize), Main(usize), @@ -111,10 +112,11 @@ impl VirtualPairCol { pub fn apply(&self, preprocessed: &[Var], main: &[Var]) -> Expr where - Expr: AbstractField + Mul + Add, + F: Into, + Expr: AbstractField + Mul, Var: Into + Copy, { - let mut result = Expr::zero() + self.constant; + let mut result = self.constant.into(); for (column, weight) in &self.column_weights { result += column.get(preprocessed, main).into() * *weight; } diff --git a/baby-bear/Cargo.toml b/baby-bear/Cargo.toml index b9b665fa..658b4889 100644 --- a/baby-bear/Cargo.toml +++ b/baby-bear/Cargo.toml @@ -13,6 +13,7 @@ serde = { version = "1.0", default-features = false, features = ["derive"] } p3-field-testing = { path = "../field-testing" } criterion = "0.5.1" rand_chacha = "0.3.1" +serde_json = "1.0.113" [[bench]] name = "inverse" diff --git a/baby-bear/src/baby_bear.rs b/baby-bear/src/baby_bear.rs index 9b9e71ed..0e73288f 100644 --- a/baby-bear/src/baby_bear.rs +++ b/baby-bear/src/baby_bear.rs @@ -8,7 +8,7 @@ use p3_field::{ }; use rand::distributions::{Distribution, Standard}; use rand::Rng; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; /// The Baby Bear prime const P: u32 = 0x78000001; @@ -34,7 +34,7 @@ const MONTY_MU: u32 = if cfg!(all(target_arch = "aarch64", target_feature = "neo const MONTY_MASK: u32 = ((1u64 << MONTY_BITS) - 1) as u32; /// The prime field `2^31 - 2^27 + 1`, a.k.a. the Baby Bear field. -#[derive(Copy, Clone, Default, Eq, Hash, PartialEq, Serialize, Deserialize)] +#[derive(Copy, Clone, Default, Eq, Hash, PartialEq)] #[repr(transparent)] // `PackedBabyBearNeon` relies on this! pub struct BabyBear { // This is `pub(crate)` just for tests. If you're accessing `value` outside of those, you're @@ -89,6 +89,19 @@ impl Distribution for Standard { } } +impl Serialize for BabyBear { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_u32(self.as_canonical_u32()) + } +} + +impl<'de> Deserialize<'de> for BabyBear { + fn deserialize>(d: D) -> Result { + let val = u32::deserialize(d)?; + Ok(BabyBear::from_canonical_u32(val)) + } +} + const MONTY_ZERO: u32 = to_monty(0); const MONTY_ONE: u32 = to_monty(1); const MONTY_TWO: u32 = to_monty(2); @@ -493,6 +506,37 @@ mod tests { assert_eq!(m1.exp_u64(1725656503).exp_const_u64::<7>(), m1); assert_eq!(m2.exp_u64(1725656503).exp_const_u64::<7>(), m2); assert_eq!(f_2.exp_u64(1725656503).exp_const_u64::<7>(), f_2); + + let f_serialized = serde_json::to_string(&f).unwrap(); + let f_deserialized: F = serde_json::from_str(&f_serialized).unwrap(); + assert_eq!(f, f_deserialized); + + let f_1_serialized = serde_json::to_string(&f_1).unwrap(); + let f_1_deserialized: F = serde_json::from_str(&f_1_serialized).unwrap(); + let f_1_serialized_again = serde_json::to_string(&f_1_deserialized).unwrap(); + let f_1_deserialized_again: F = serde_json::from_str(&f_1_serialized_again).unwrap(); + assert_eq!(f_1, f_1_deserialized); + assert_eq!(f_1, f_1_deserialized_again); + + let f_2_serialized = serde_json::to_string(&f_2).unwrap(); + let f_2_deserialized: F = serde_json::from_str(&f_2_serialized).unwrap(); + assert_eq!(f_2, f_2_deserialized); + + let f_p_minus_1_serialized = serde_json::to_string(&f_p_minus_1).unwrap(); + let f_p_minus_1_deserialized: F = serde_json::from_str(&f_p_minus_1_serialized).unwrap(); + assert_eq!(f_p_minus_1, f_p_minus_1_deserialized); + + let f_p_minus_2_serialized = serde_json::to_string(&f_p_minus_2).unwrap(); + let f_p_minus_2_deserialized: F = serde_json::from_str(&f_p_minus_2_serialized).unwrap(); + assert_eq!(f_p_minus_2, f_p_minus_2_deserialized); + + let m1_serialized = serde_json::to_string(&m1).unwrap(); + let m1_deserialized: F = serde_json::from_str(&m1_serialized).unwrap(); + assert_eq!(m1, m1_deserialized); + + let m2_serialized = serde_json::to_string(&m2).unwrap(); + let m2_deserialized: F = serde_json::from_str(&m2_serialized).unwrap(); + assert_eq!(m2, m2_deserialized); } test_field!(crate::BabyBear); diff --git a/baby-bear/src/x86_64_avx2.rs b/baby-bear/src/x86_64_avx2.rs index 22631318..a208a5f4 100644 --- a/baby-bear/src/x86_64_avx2.rs +++ b/baby-bear/src/x86_64_avx2.rs @@ -4,6 +4,8 @@ use core::mem::transmute; use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign}; use p3_field::{AbstractField, Field, PackedField}; +use rand::distributions::{Distribution, Standard}; +use rand::Rng; use crate::BabyBear; @@ -500,6 +502,13 @@ impl Sub for BabyBear { } } +impl Distribution for Standard { + #[inline] + fn sample(&self, rng: &mut R) -> PackedBabyBearAVX2 { + PackedBabyBearAVX2(rng.gen()) + } +} + #[inline] #[must_use] fn interleave1(a: __m256i, b: __m256i) -> (__m256i, __m256i) { diff --git a/commit/src/pcs.rs b/commit/src/pcs.rs index 8a56d50c..5bdb6df0 100644 --- a/commit/src/pcs.rs +++ b/commit/src/pcs.rs @@ -89,7 +89,7 @@ where fn commit_shifted_batches( &self, polynomials: Vec, - coset_shifts: &[Val], + coset_shift: &[Val], ) -> (Self::Commitment, Self::ProverData); fn commit_shifted_batch( diff --git a/field/src/extension/binomial_extension.rs b/field/src/extension/binomial_extension.rs index 39551332..458126dc 100644 --- a/field/src/extension/binomial_extension.rs +++ b/field/src/extension/binomial_extension.rs @@ -43,6 +43,7 @@ impl From for BinomialExtensionField, const D: usize> ExtensionField for BinomialExtensionField { + type ExtensionPacking = BinomialExtensionField; } impl, const D: usize> HasFrobenius for BinomialExtensionField { diff --git a/field/src/field.rs b/field/src/field.rs index 32e5b248..0b4ffc79 100644 --- a/field/src/field.rs +++ b/field/src/field.rs @@ -294,6 +294,12 @@ pub trait AbstractExtensionField: } pub trait ExtensionField: Field + AbstractExtensionField { + type ExtensionPacking: AbstractExtensionField + + 'static + + Copy + + Send + + Sync; + fn is_in_basefield(&self) -> bool { self.as_base_slice()[1..].iter().all(Field::is_zero) } @@ -306,7 +312,9 @@ pub trait ExtensionField: Field + AbstractExtensionField { } } -impl ExtensionField for F {} +impl ExtensionField for F { + type ExtensionPacking = F::Packing; +} impl AbstractExtensionField for AF { const D: usize = 1; diff --git a/fri/src/verifier.rs b/fri/src/verifier.rs index d7158d2e..3aee9044 100644 --- a/fri/src/verifier.rs +++ b/fri/src/verifier.rs @@ -153,8 +153,8 @@ where x = x.square(); } - debug_assert!(index == 0 || index == 1); - debug_assert!(x.is_one() || x == F::two_adic_generator(1)); + debug_assert!(index < config.blowup(), "index was {}", index); + debug_assert_eq!(x.exp_power_of_2(config.log_blowup), F::one()); Ok(folded_eval) } diff --git a/keccak-air/Cargo.toml b/keccak-air/Cargo.toml index f3474019..c1b5f97d 100644 --- a/keccak-air/Cargo.toml +++ b/keccak-air/Cargo.toml @@ -39,4 +39,7 @@ name = "prove_baby_bear_poseidon2" name = "prove_goldilocks_keccak" [features] +# TODO: Consider removing, at least when this gets split off into another repository. +# We should be able to enable p3-maybe-rayon/parallel directly; this just doesn't +# seem to work when using cargo with the -p or --package option. parallel = ["p3-maybe-rayon/parallel"] diff --git a/keccak-air/examples/prove_baby_bear_keccak.rs b/keccak-air/examples/prove_baby_bear_keccak.rs index ea43473b..c86eed21 100644 --- a/keccak-air/examples/prove_baby_bear_keccak.rs +++ b/keccak-air/examples/prove_baby_bear_keccak.rs @@ -3,7 +3,6 @@ use p3_challenger::DuplexChallenger; use p3_commit::ExtensionMmcs; use p3_dft::Radix2DitParallel; use p3_field::extension::BinomialExtensionField; -use p3_field::Field; use p3_fri::{FriConfig, TwoAdicFriPcs, TwoAdicFriPcsConfig}; use p3_keccak::Keccak256Hash; use p3_keccak_air::{generate_trace_rows, KeccakAir}; @@ -31,9 +30,7 @@ fn main() -> Result<(), VerificationError> { .init(); type Val = BabyBear; - type Domain = Val; type Challenge = BinomialExtensionField; - type PackedChallenge = BinomialExtensionField<::Packing, 4>; type Perm = Poseidon2; let perm = Perm::new_from_rng(8, 22, DiffusionMatrixBabybear, &mut thread_rng()); @@ -65,7 +62,7 @@ fn main() -> Result<(), VerificationError> { TwoAdicFriPcs>; let pcs = Pcs::new(fri_config, dft, val_mmcs); - type MyConfig = StarkConfig; + type MyConfig = StarkConfig; let config = StarkConfig::new(pcs); let mut challenger = Challenger::new(perm.clone()); diff --git a/keccak-air/examples/prove_baby_bear_poseidon2.rs b/keccak-air/examples/prove_baby_bear_poseidon2.rs index 7308eb34..4300d5fe 100644 --- a/keccak-air/examples/prove_baby_bear_poseidon2.rs +++ b/keccak-air/examples/prove_baby_bear_poseidon2.rs @@ -30,9 +30,7 @@ fn main() -> Result<(), VerificationError> { .init(); type Val = BabyBear; - type Domain = Val; type Challenge = BinomialExtensionField; - type PackedChallenge = BinomialExtensionField<::Packing, 4>; type Perm = Poseidon2; let perm = Perm::new_from_rng(8, 22, DiffusionMatrixBabybear, &mut thread_rng()); @@ -64,7 +62,7 @@ fn main() -> Result<(), VerificationError> { TwoAdicFriPcs>; let pcs = Pcs::new(fri_config, dft, val_mmcs); - type MyConfig = StarkConfig; + type MyConfig = StarkConfig; let config = StarkConfig::new(pcs); let mut challenger = Challenger::new(perm.clone()); diff --git a/keccak-air/examples/prove_goldilocks_keccak.rs b/keccak-air/examples/prove_goldilocks_keccak.rs index fa8fe9ac..af69d08e 100644 --- a/keccak-air/examples/prove_goldilocks_keccak.rs +++ b/keccak-air/examples/prove_goldilocks_keccak.rs @@ -2,7 +2,6 @@ use p3_challenger::DuplexChallenger; use p3_commit::ExtensionMmcs; use p3_dft::Radix2DitParallel; use p3_field::extension::BinomialExtensionField; -use p3_field::Field; use p3_fri::{FriConfig, TwoAdicFriPcs, TwoAdicFriPcsConfig}; use p3_goldilocks::Goldilocks; use p3_keccak::Keccak256Hash; @@ -31,9 +30,7 @@ fn main() -> Result<(), VerificationError> { .init(); type Val = Goldilocks; - type Domain = Val; type Challenge = BinomialExtensionField; - type PackedChallenge = BinomialExtensionField<::Packing, 2>; type Perm = Poseidon2; let perm = Perm::new_from_rng(8, 22, DiffusionMatrixGoldilocks, &mut thread_rng()); @@ -64,7 +61,7 @@ fn main() -> Result<(), VerificationError> { TwoAdicFriPcs>; let pcs = Pcs::new(fri_config, dft, val_mmcs); - type MyConfig = StarkConfig; + type MyConfig = StarkConfig; let config = StarkConfig::new(pcs); let mut challenger = Challenger::new(perm.clone()); diff --git a/keccak-air/src/constants.rs b/keccak-air/src/constants.rs index 3ca7edd6..dac6efac 100644 --- a/keccak-air/src/constants.rs +++ b/keccak-air/src/constants.rs @@ -1,6 +1,6 @@ use crate::BITS_PER_LIMB; -pub(crate) const R: [[u8; 5]; 5] = [ +pub const R: [[u8; 5]; 5] = [ [0, 36, 3, 41, 18], [1, 44, 10, 45, 2], [62, 6, 43, 15, 61], diff --git a/keccak-air/src/lib.rs b/keccak-air/src/lib.rs index 4214ec81..f9d5d278 100644 --- a/keccak-air/src/lib.rs +++ b/keccak-air/src/lib.rs @@ -13,7 +13,7 @@ mod round_flags; pub use air::*; pub use columns::*; -pub use constants::RC; +pub use constants::*; pub use generation::*; pub const NUM_ROUNDS: usize = 24; diff --git a/merkle-tree/src/mmcs.rs b/merkle-tree/src/mmcs.rs index 9c0243ca..f7268cd5 100644 --- a/merkle-tree/src/mmcs.rs +++ b/merkle-tree/src/mmcs.rs @@ -192,7 +192,7 @@ mod tests { type F = BabyBear; - type Perm = Poseidon2; + type Perm = Poseidon2; type MyHash = PaddingFreeSponge; type MyCompress = TruncatedPermutation; type MyMmcs = FieldMerkleTreeMmcs<::Packing, MyHash, MyCompress, 8>; diff --git a/mersenne-31/src/aarch64_neon.rs b/mersenne-31/src/aarch64_neon.rs index 0df29a5a..e05a67c5 100644 --- a/mersenne-31/src/aarch64_neon.rs +++ b/mersenne-31/src/aarch64_neon.rs @@ -4,6 +4,8 @@ use core::mem::transmute; use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign}; use p3_field::{AbstractField, Field, PackedField}; +use rand::distributions::{Distribution, Standard}; +use rand::Rng; use crate::Mersenne31; @@ -470,6 +472,13 @@ impl Sub for Mersenne31 { } } +impl Distribution for Standard { + #[inline] + fn sample(&self, rng: &mut R) -> PackedMersenne31Neon { + PackedMersenne31Neon(rng.gen()) + } +} + #[inline] #[must_use] fn interleave1(v0: uint32x4_t, v1: uint32x4_t) -> (uint32x4_t, uint32x4_t) { diff --git a/mersenne-31/src/x86_64_avx2.rs b/mersenne-31/src/x86_64_avx2.rs index c48a0e9f..6cc0d7e6 100644 --- a/mersenne-31/src/x86_64_avx2.rs +++ b/mersenne-31/src/x86_64_avx2.rs @@ -4,6 +4,8 @@ use core::mem::transmute; use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign}; use p3_field::{AbstractField, Field, PackedField}; +use rand::distributions::{Distribution, Standard}; +use rand::Rng; use crate::Mersenne31; @@ -474,6 +476,13 @@ impl Sub for Mersenne31 { } } +impl Distribution for Standard { + #[inline] + fn sample(&self, rng: &mut R) -> PackedMersenne31AVX2 { + PackedMersenne31AVX2(rng.gen()) + } +} + #[inline] #[must_use] fn interleave1(a: __m256i, b: __m256i) -> (__m256i, __m256i) { diff --git a/poseidon2/Cargo.toml b/poseidon2/Cargo.toml index 564fc52c..3f60c247 100644 --- a/poseidon2/Cargo.toml +++ b/poseidon2/Cargo.toml @@ -14,6 +14,8 @@ p3-symmetric = { path = "../symmetric" } rand = "0.8.5" [dev-dependencies] +ark-ff = { version = "^0.4.0", default-features = false } +zkhash = { git = "https://github.com/HorizenLabs/poseidon2" } criterion = "0.5.1" [[bench]] diff --git a/poseidon2/src/external.rs b/poseidon2/src/external.rs deleted file mode 100644 index 239955a4..00000000 --- a/poseidon2/src/external.rs +++ /dev/null @@ -1,56 +0,0 @@ -use p3_field::AbstractField; -use p3_mds::m4::M4Mds; -use p3_symmetric::Permutation; - -pub fn matmul_external_mut( - input: &mut [AF; WIDTH], - m4: &M4Mds, -) { - match WIDTH { - 2 => { - // Matrix circ(2, 1) - let mut sum = input[0].clone(); - sum += input[1].clone(); - input[0] += sum.clone(); - input[1] += sum.clone(); - } - 3 => { - // Matrix circ(2, 1, 1) - let mut sum = input[0].clone(); - sum += input[1].clone(); - sum += input[2].clone(); - input[0] += sum.clone(); - input[1] += sum.clone(); - input[2] += sum.clone(); - } - 4 => { - // Applying cheap 4x4 MDS matrix to each 4-element part of the state - matmul_m4(input, m4) - } - 8 | 12 | 16 | 20 | 24 => { - // First, apply Diag(M4, ... , M4) - matmul_m4(input, m4); - - // Apply the circulant matrix ( - let t4 = WIDTH / 4; - let mut stored = [AF::zero(), AF::zero(), AF::zero(), AF::zero()]; - for l in 0..4 { - stored[l] = input[l].clone(); - for j in 1..t4 { - stored[l] += input[4 * j + l].clone(); - } - } - for i in 0..input.len() { - input[i] += stored[i % 4].clone(); - } - } - _ => panic!("Unsupported width"), - } -} - -#[inline] -fn matmul_m4(input: &mut [AF; WIDTH], m4: &M4Mds) { - input - .chunks_exact_mut(4) - .for_each(|x| m4.permute_mut(x.try_into().unwrap())); -} diff --git a/poseidon2/src/lib.rs b/poseidon2/src/lib.rs index d7c6c43b..5db151ac 100644 --- a/poseidon2/src/lib.rs +++ b/poseidon2/src/lib.rs @@ -10,16 +10,15 @@ extern crate alloc; mod babybear; mod diffusion; -mod external; mod goldilocks; +mod matrix; use alloc::vec::Vec; pub use babybear::DiffusionMatrixBabybear; pub use diffusion::DiffusionPermutation; -pub use external::*; pub use goldilocks::DiffusionMatrixGoldilocks; +use matrix::Poseidon2MEMatrix; use p3_field::{AbstractField, PrimeField}; -use p3_mds::m4::M4Mds; use p3_symmetric::{CryptographicPermutation, Permutation}; use rand::distributions::Standard; use rand::prelude::Distribution; @@ -41,9 +40,6 @@ pub struct Poseidon2 { /// The linear layer used in internal rounds (only needs diffusion property, not MDS). internal_linear_layer: Diffusion, - - /// The matrix `M4` used in the external linear layer. - m4: M4Mds, } impl Poseidon2 @@ -63,7 +59,6 @@ where rounds_p, constants, internal_linear_layer, - m4: M4Mds, } } @@ -88,7 +83,6 @@ where rounds_p, constants, internal_linear_layer: internal_mds, - m4: M4Mds, } } @@ -118,14 +112,6 @@ where { state.iter_mut().for_each(|a| *a = self.sbox_p(a)); } - - #[inline] - fn external_linear_permute_mut(&self, state: &mut [AF; WIDTH]) - where - AF: AbstractField, - { - matmul_external_mut(state, &self.m4); - } } impl Permutation<[AF; WIDTH]> @@ -136,21 +122,23 @@ where Diffusion: DiffusionPermutation, { fn permute_mut(&self, state: &mut [AF; WIDTH]) { + let external_linear_layer = Poseidon2MEMatrix::; + // The initial linear layer. - self.external_linear_permute_mut(state); + external_linear_layer.permute_mut(state); // The first half of the external rounds. let rounds = self.rounds_f + self.rounds_p; - let rounds_f_beggining = self.rounds_f / 2; - for r in 0..rounds_f_beggining { + let rounds_f_beginning = self.rounds_f / 2; + for r in 0..rounds_f_beginning { self.add_rc(state, &self.constants[r]); self.sbox(state); - self.external_linear_permute_mut(state); + external_linear_layer.permute_mut(state); } // The internal rounds. - let p_end = rounds_f_beggining + self.rounds_p; - for r in self.rounds_f..p_end { + let p_end = rounds_f_beginning + self.rounds_p; + for r in rounds_f_beginning..p_end { state[0] += AF::from_f(self.constants[r][0]); state[0] = self.sbox_p(&state[0]); self.internal_linear_layer.permute_mut(state); @@ -160,7 +148,7 @@ where for r in p_end..rounds { self.add_rc(state, &self.constants[r]); self.sbox(state); - self.external_linear_permute_mut(state); + external_linear_layer.permute_mut(state); } } } @@ -173,3 +161,238 @@ where Diffusion: DiffusionPermutation, { } + +#[cfg(test)] +mod tests { + use alloc::vec::Vec; + + use ark_ff::{BigInteger, PrimeField}; + use p3_baby_bear::BabyBear; + use p3_field::AbstractField; + use p3_goldilocks::Goldilocks; + use p3_symmetric::Permutation; + use rand::Rng; + use zkhash::fields::babybear::FpBabyBear; + use zkhash::fields::goldilocks::FpGoldiLocks; + use zkhash::poseidon2::poseidon2::Poseidon2 as Poseidon2Ref; + use zkhash::poseidon2::poseidon2_instance_babybear::{POSEIDON2_BABYBEAR_16_PARAMS, RC16}; + use zkhash::poseidon2::poseidon2_instance_goldilocks::{ + POSEIDON2_GOLDILOCKS_12_PARAMS, POSEIDON2_GOLDILOCKS_8_PARAMS, RC12, RC8, + }; + + use crate::goldilocks::DiffusionMatrixGoldilocks; + use crate::{DiffusionMatrixBabybear, Poseidon2}; + + fn goldilocks_from_ark_ff(input: FpGoldiLocks) -> Goldilocks { + let as_bigint = input.into_bigint(); + let mut as_bytes = as_bigint.to_bytes_le(); + as_bytes.resize(8, 0); + let as_u64 = u64::from_le_bytes(as_bytes[0..8].try_into().unwrap()); + Goldilocks::from_wrapped_u64(as_u64) + } + + fn babybear_from_ark_ff(input: FpBabyBear) -> BabyBear { + let as_bigint = input.into_bigint(); + let mut as_bytes = as_bigint.to_bytes_le(); + as_bytes.resize(4, 0); + let as_u32 = u32::from_le_bytes(as_bytes[0..4].try_into().unwrap()); + BabyBear::from_wrapped_u32(as_u32) + } + + #[test] + fn test_poseidon2_goldilocks_width_8() { + const WIDTH: usize = 8; + const D: u64 = 7; + const ROUNDS_F: usize = 8; + const ROUNDS_P: usize = 22; + + type F = Goldilocks; + + let mut rng = rand::thread_rng(); + + // Poiseidon2 reference implementation from zkhash repo. + let poseidon2_ref = Poseidon2Ref::new(&POSEIDON2_GOLDILOCKS_8_PARAMS); + + // Copy over round constants from zkhash. + let round_constants: Vec<[F; WIDTH]> = RC8 + .iter() + .map(|vec| { + vec.iter() + .cloned() + .map(goldilocks_from_ark_ff) + .collect::>() + .try_into() + .unwrap() + }) + .collect(); + + // Our Poseidon2 implementation. + let poseidon2: Poseidon2 = Poseidon2::new( + ROUNDS_F, + ROUNDS_P, + round_constants, + DiffusionMatrixGoldilocks, + ); + + // Generate random input and convert to both Goldilocks field formats. + let input_u64 = rng.gen::<[u64; WIDTH]>(); + let input_ref = input_u64 + .iter() + .cloned() + .map(FpGoldiLocks::from) + .collect::>(); + let input = input_u64.map(F::from_wrapped_u64); + + // Check that the conversion is correct. + assert!(input_ref + .iter() + .zip(input.iter()) + .all(|(a, b)| goldilocks_from_ark_ff(*a) == *b)); + + // Run reference implementation. + let output_ref = poseidon2_ref.permutation(&input_ref); + let expected: [F; WIDTH] = output_ref + .iter() + .cloned() + .map(goldilocks_from_ark_ff) + .collect::>() + .try_into() + .unwrap(); + + // Run our implementation. + let mut output = input; + poseidon2.permute_mut(&mut output); + + assert_eq!(output, expected); + } + + #[test] + fn test_poseidon2_goldilocks_width_12() { + const WIDTH: usize = 12; + const D: u64 = 7; + const ROUNDS_F: usize = 8; + const ROUNDS_P: usize = 22; + + type F = Goldilocks; + + let mut rng = rand::thread_rng(); + + // Poiseidon2 reference implementation from zkhash repo. + let poseidon2_ref = Poseidon2Ref::new(&POSEIDON2_GOLDILOCKS_12_PARAMS); + + // Copy over round constants from zkhash. + let round_constants: Vec<[F; WIDTH]> = RC12 + .iter() + .map(|vec| { + vec.iter() + .cloned() + .map(goldilocks_from_ark_ff) + .collect::>() + .try_into() + .unwrap() + }) + .collect(); + + // Our Poseidon2 implementation. + let poseidon2: Poseidon2 = Poseidon2::new( + ROUNDS_F, + ROUNDS_P, + round_constants, + DiffusionMatrixGoldilocks, + ); + + // Generate random input and convert to both Goldilocks field formats. + let input_u64 = rng.gen::<[u64; WIDTH]>(); + let input_ref = input_u64 + .iter() + .cloned() + .map(FpGoldiLocks::from) + .collect::>(); + let input = input_u64.map(F::from_wrapped_u64); + + // Check that the conversion is correct. + assert!(input_ref + .iter() + .zip(input.iter()) + .all(|(a, b)| goldilocks_from_ark_ff(*a) == *b)); + + // Run reference implementation. + let output_ref = poseidon2_ref.permutation(&input_ref); + let expected: [F; WIDTH] = output_ref + .iter() + .cloned() + .map(goldilocks_from_ark_ff) + .collect::>() + .try_into() + .unwrap(); + + // Run our implementation. + let mut output = input; + poseidon2.permute_mut(&mut output); + + assert_eq!(output, expected); + } + + #[test] + fn test_poseidon2_babybear_width_16() { + const WIDTH: usize = 16; + const D: u64 = 7; + const ROUNDS_F: usize = 8; + const ROUNDS_P: usize = 13; + + type F = BabyBear; + + let mut rng = rand::thread_rng(); + + // Poiseidon2 reference implementation from zkhash repo. + let poseidon2_ref = Poseidon2Ref::new(&POSEIDON2_BABYBEAR_16_PARAMS); + + // Copy over round constants from zkhash. + let round_constants: Vec<[F; WIDTH]> = RC16 + .iter() + .map(|vec| { + vec.iter() + .cloned() + .map(babybear_from_ark_ff) + .collect::>() + .try_into() + .unwrap() + }) + .collect(); + + // Our Poseidon2 implementation. + let poseidon2: Poseidon2 = + Poseidon2::new(ROUNDS_F, ROUNDS_P, round_constants, DiffusionMatrixBabybear); + + // Generate random input and convert to both BabyBear field formats. + let input_u32 = rng.gen::<[u32; WIDTH]>(); + let input_ref = input_u32 + .iter() + .cloned() + .map(FpBabyBear::from) + .collect::>(); + let input = input_u32.map(F::from_wrapped_u32); + + // Check that the conversion is correct. + assert!(input_ref + .iter() + .zip(input.iter()) + .all(|(a, b)| babybear_from_ark_ff(*a) == *b)); + + // Run reference implementation. + let output_ref = poseidon2_ref.permutation(&input_ref); + let expected: [F; WIDTH] = output_ref + .iter() + .cloned() + .map(babybear_from_ark_ff) + .collect::>() + .try_into() + .unwrap(); + + // Run our implementation. + let mut output = input; + poseidon2.permute_mut(&mut output); + + assert_eq!(output, expected); + } +} diff --git a/poseidon2/src/matrix.rs b/poseidon2/src/matrix.rs new file mode 100644 index 00000000..f2a389e2 --- /dev/null +++ b/poseidon2/src/matrix.rs @@ -0,0 +1,64 @@ +use p3_field::{AbstractField, PrimeField}; +use p3_symmetric::Permutation; + +extern crate alloc; + +// The t x t matrix M_E := circ(2M_4, M_4, ..., M_4), where M_4 is the 4 x 4 matrix +// [ 5 7 1 3 ] +// [ 4 6 1 1 ] +// [ 1 3 5 7 ] +// [ 1 1 4 6 ]. +// The permutation calculation is based on Appendix B from the Poseidon2 paper. +#[derive(Copy, Clone, Default)] +pub struct Poseidon2MEMatrix; + +// Multiply a 4-element vector x by M_4, in place. +// This uses the formula from the start of Appendix B, with multiplications unrolled into additions. +fn apply_m_4(x: &mut [AF]) +where + AF: AbstractField, + AF::F: PrimeField, +{ + let t0 = x[0].clone() + x[1].clone(); + let t1 = x[2].clone() + x[3].clone(); + let t2 = x[1].clone() + x[1].clone() + t1.clone(); + let t3 = x[3].clone() + x[3].clone() + t0.clone(); + let t4 = t1.clone() + t1.clone() + t1.clone() + t1 + t3.clone(); + let t5 = t0.clone() + t0.clone() + t0.clone() + t0 + t2.clone(); + let t6 = t3 + t5.clone(); + let t7 = t2 + t4.clone(); + x[0] = t6; + x[1] = t5; + x[2] = t7; + x[3] = t4; +} + +impl Permutation<[AF; WIDTH]> for Poseidon2MEMatrix +where + AF: AbstractField, + AF::F: PrimeField, +{ + fn permute_mut(&self, state: &mut [AF; WIDTH]) { + // First, we apply M_4 to each consecutive four elements of the state. + // In Appendix B's terminology, this replaces each x_i with x_i'. + for i in (0..WIDTH).step_by(4) { + apply_m_4(&mut state[i..i + 4]); + } + + // Now, we apply the outer circulant matrix (to compute the y_i values). + + // We first precompute the four sums of every four elements. + let sums: [AF; 4] = core::array::from_fn(|k| { + (0..WIDTH) + .step_by(4) + .map(|j| state[j + k].clone()) + .sum::() + }); + + // The formula for each y_i involves 2x_i' term and x_j' terms for each j that equals i mod 4. + // In other words, we can add a single copy of x_i' to the appropriate one of our precomputed sums + for i in 0..WIDTH { + state[i] += sums[i % 4].clone(); + } + } +} diff --git a/tensor-pcs/Cargo.toml b/tensor-pcs/Cargo.toml index 20c4a9e5..fde22e7a 100644 --- a/tensor-pcs/Cargo.toml +++ b/tensor-pcs/Cargo.toml @@ -11,3 +11,4 @@ p3-commit = { path = "../commit" } p3-field = { path = "../field" } p3-matrix = { path = "../matrix" } p3-util = { path = "../util" } +serde = "1.0.196" diff --git a/uni-stark/src/config.rs b/uni-stark/src/config.rs index a973ddbb..960ff273 100644 --- a/uni-stark/src/config.rs +++ b/uni-stark/src/config.rs @@ -2,17 +2,20 @@ use core::marker::PhantomData; use p3_challenger::{CanObserve, FieldChallenger}; use p3_commit::{Pcs, UnivariatePcsWithLde}; -use p3_field::{AbstractExtensionField, ExtensionField, PackedField, TwoAdicField}; +use p3_field::{ExtensionField, Field, TwoAdicField}; use p3_matrix::dense::RowMajorMatrix; +pub type PackedVal = <::Val as Field>::Packing; +pub type PackedChallenge = <::Challenge as ExtensionField< + ::Val, +>>::ExtensionPacking; + pub trait StarkGenericConfig { /// The field over which trace data is encoded. type Val: TwoAdicField; - type PackedVal: PackedField; /// The field from which most random challenges are drawn. type Challenge: ExtensionField + TwoAdicField; - type PackedChallenge: AbstractExtensionField + Copy; /// The PCS used to commit to trace polynomials. type Pcs: UnivariatePcsWithLde< @@ -29,14 +32,12 @@ pub trait StarkGenericConfig { fn pcs(&self) -> &Self::Pcs; } -pub struct StarkConfig { +pub struct StarkConfig { pcs: Pcs, - _phantom: PhantomData<(Val, Challenge, PackedChallenge, Challenger)>, + _phantom: PhantomData<(Val, Challenge, Challenger)>, } -impl - StarkConfig -{ +impl StarkConfig { pub fn new(pcs: Pcs) -> Self { Self { pcs, @@ -45,20 +46,17 @@ impl } } -impl StarkGenericConfig - for StarkConfig +impl StarkGenericConfig + for StarkConfig where Val: TwoAdicField, Challenge: ExtensionField + TwoAdicField, - PackedChallenge: AbstractExtensionField + Copy, Pcs: UnivariatePcsWithLde, Challenger>, Challenger: FieldChallenger + CanObserve<>>::Commitment>, { type Val = Val; - type PackedVal = Val::Packing; type Challenge = Challenge; - type PackedChallenge = PackedChallenge; type Pcs = Pcs; type Challenger = Challenger; diff --git a/uni-stark/src/folder.rs b/uni-stark/src/folder.rs index 8a5de17f..25e5beca 100644 --- a/uni-stark/src/folder.rs +++ b/uni-stark/src/folder.rs @@ -1,15 +1,15 @@ use p3_air::{AirBuilder, TwoRowMatrixView}; use p3_field::{AbstractField, Field}; -use crate::StarkGenericConfig; +use crate::{PackedChallenge, PackedVal, StarkGenericConfig}; pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> { - pub main: TwoRowMatrixView<'a, SC::PackedVal>, - pub is_first_row: SC::PackedVal, - pub is_last_row: SC::PackedVal, - pub is_transition: SC::PackedVal, + pub main: TwoRowMatrixView<'a, PackedVal>, + pub is_first_row: PackedVal, + pub is_last_row: PackedVal, + pub is_transition: PackedVal, pub alpha: SC::Challenge, - pub accumulator: SC::PackedChallenge, + pub accumulator: PackedChallenge, } pub struct VerifierConstraintFolder<'a, Challenge> { @@ -23,9 +23,9 @@ pub struct VerifierConstraintFolder<'a, Challenge> { impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> { type F = SC::Val; - type Expr = SC::PackedVal; - type Var = SC::PackedVal; - type M = TwoRowMatrixView<'a, SC::PackedVal>; + type Expr = PackedVal; + type Var = PackedVal; + type M = TwoRowMatrixView<'a, PackedVal>; fn main(&self) -> Self::M { self.main @@ -48,8 +48,8 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> { } fn assert_zero>(&mut self, x: I) { - let x: SC::PackedVal = x.into(); - self.accumulator *= SC::PackedChallenge::from_f(self.alpha); + let x: PackedVal = x.into(); + self.accumulator *= PackedChallenge::::from_f(self.alpha); self.accumulator += x; } } diff --git a/uni-stark/src/lib.rs b/uni-stark/src/lib.rs index 21135165..4132539c 100644 --- a/uni-stark/src/lib.rs +++ b/uni-stark/src/lib.rs @@ -15,8 +15,9 @@ mod symbolic_variable; mod verifier; mod zerofier_coset; -pub mod check_constraints; +mod check_constraints; +pub use check_constraints::*; pub use config::*; pub use decompose::*; pub use folder::*; diff --git a/uni-stark/src/prover.rs b/uni-stark/src/prover.rs index 5c399224..399f8eb0 100644 --- a/uni-stark/src/prover.rs +++ b/uni-stark/src/prover.rs @@ -17,12 +17,16 @@ use tracing::{info_span, instrument}; use crate::symbolic_builder::{get_log_quotient_degree, SymbolicAirBuilder}; use crate::{ - decompose_and_flatten, Commitments, OpenedValues, Proof, ProverConstraintFolder, - StarkGenericConfig, ZerofierOnCoset, + decompose_and_flatten, Commitments, OpenedValues, PackedChallenge, PackedVal, Proof, + ProverConstraintFolder, StarkGenericConfig, ZerofierOnCoset, }; #[instrument(skip_all)] -pub fn prove Air>>( +pub fn prove< + SC, + #[cfg(debug_assertions)] A: for<'a> Air>, + #[cfg(not(debug_assertions))] A, +>( config: &SC, air: &A, challenger: &mut SC::Challenger, @@ -32,7 +36,7 @@ where SC: StarkGenericConfig, A: Air> + for<'a> Air>, { - // #[cfg(debug_assertions)] + #[cfg(debug_assertions)] crate::check_constraints::check_constraints(air, &trace); let degree = trace.height(); @@ -146,7 +150,7 @@ where // length `WIDTH`. In the edge case where `quotient_size < WIDTH`, we need to pad those vectors // in order for the slices to exist. The entries beyond quotient_size will be ignored, so we can // just use default values. - for _ in quotient_size..SC::PackedVal::WIDTH { + for _ in quotient_size..PackedVal::::WIDTH { coset.push(SC::Val::default()); lagrange_first_evals.push(SC::Val::default()); lagrange_last_evals.push(SC::Val::default()); @@ -154,20 +158,20 @@ where (0..quotient_size) .into_par_iter() - .step_by(SC::PackedVal::WIDTH) + .step_by(PackedVal::::WIDTH) .flat_map_iter(|i_local_start| { let wrap = |i| i % quotient_size; let i_next_start = wrap(i_local_start + next_step); - let i_range = i_local_start..i_local_start + SC::PackedVal::WIDTH; + let i_range = i_local_start..i_local_start + PackedVal::::WIDTH; - let x = *SC::PackedVal::from_slice(&coset[i_range.clone()]); + let x = *PackedVal::::from_slice(&coset[i_range.clone()]); let is_transition = x - subgroup_last; - let is_first_row = *SC::PackedVal::from_slice(&lagrange_first_evals[i_range.clone()]); - let is_last_row = *SC::PackedVal::from_slice(&lagrange_last_evals[i_range]); + let is_first_row = *PackedVal::::from_slice(&lagrange_first_evals[i_range.clone()]); + let is_last_row = *PackedVal::::from_slice(&lagrange_last_evals[i_range]); let local: Vec<_> = (0..trace_lde.width()) .map(|col| { - SC::PackedVal::from_fn(|offset| { + PackedVal::::from_fn(|offset| { let row = wrap(i_local_start + offset); trace_lde.get(row, col) }) @@ -175,14 +179,14 @@ where .collect(); let next: Vec<_> = (0..trace_lde.width()) .map(|col| { - SC::PackedVal::from_fn(|offset| { + PackedVal::::from_fn(|offset| { let row = wrap(i_next_start + offset); trace_lde.get(row, col) }) }) .collect(); - let accumulator = SC::PackedChallenge::zero(); + let accumulator = PackedChallenge::::zero(); let mut folder = ProverConstraintFolder { main: TwoRowMatrixView { local: &local, @@ -197,11 +201,11 @@ where air.eval(&mut folder); // quotient(x) = constraints(x) / Z_H(x) - let zerofier_inv: SC::PackedVal = zerofier_on_coset.eval_inverse_packed(i_local_start); + let zerofier_inv: PackedVal = zerofier_on_coset.eval_inverse_packed(i_local_start); let quotient = folder.accumulator * zerofier_inv; // "Transpose" D packed base coefficients into WIDTH scalar extension coefficients. - let limit = SC::PackedVal::WIDTH.min(quotient_size); + let limit = PackedVal::::WIDTH.min(quotient_size); (0..limit).map(move |idx_in_packing| { let quotient_value = (0..>::D) .map(|coeff_idx| quotient.as_base_slice()[coeff_idx].as_slice()[idx_in_packing]) diff --git a/uni-stark/src/symbolic_expression.rs b/uni-stark/src/symbolic_expression.rs index 03bbb614..ee281235 100644 --- a/uni-stark/src/symbolic_expression.rs +++ b/uni-stark/src/symbolic_expression.rs @@ -1,5 +1,5 @@ use alloc::rc::Rc; -use core::fmt::{Debug, Display, Formatter}; +use core::fmt::Debug; use core::iter::{Product, Sum}; use core::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; @@ -266,19 +266,3 @@ impl Product for SymbolicExpression { iter.map(|x| Self::from(x)).product() } } - -impl Display for SymbolicExpression { - fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { - match self { - SymbolicExpression::Variable(v) => write!(f, "{}", v), - SymbolicExpression::IsFirstRow => write!(f, "IsFirstRow"), - SymbolicExpression::IsLastRow => write!(f, "IsLastRow"), - SymbolicExpression::IsTransition => write!(f, "IsTransition"), - SymbolicExpression::Constant(c) => write!(f, "{}", c), - SymbolicExpression::Add { x, y, .. } => write!(f, "({} + {})", x, y), - SymbolicExpression::Sub { x, y, .. } => write!(f, "({} - {})", x, y), - SymbolicExpression::Neg { x, .. } => write!(f, "-{}", x), - SymbolicExpression::Mul { x, y, .. } => write!(f, "({} * {})", x, y), - } - } -} diff --git a/uni-stark/src/symbolic_variable.rs b/uni-stark/src/symbolic_variable.rs index 6a402779..ea31d20e 100644 --- a/uni-stark/src/symbolic_variable.rs +++ b/uni-stark/src/symbolic_variable.rs @@ -1,4 +1,3 @@ -use core::fmt::{Display, Formatter}; use core::marker::PhantomData; use core::ops::{Add, Mul, Sub}; @@ -125,13 +124,3 @@ impl Mul> for SymbolicExpression { self * Self::from(rhs) } } - -impl Display for SymbolicVariable { - fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { - if self.is_next { - write!(f, "Next({})", self.column) - } else { - write!(f, "Local({})", self.column) - } - } -} diff --git a/uni-stark/tests/mul_air.rs b/uni-stark/tests/mul_air.rs index e14db7b0..b33894c8 100644 --- a/uni-stark/tests/mul_air.rs +++ b/uni-stark/tests/mul_air.rs @@ -71,9 +71,7 @@ fn test_prove_baby_bear() -> Result<(), VerificationError> { const HEIGHT: usize = 1 << 6; type Val = BabyBear; - type Domain = Val; type Challenge = BinomialExtensionField; - type PackedChallenge = BinomialExtensionField<::Packing, 4>; type Perm = Poseidon2; let perm = Perm::new_from_rng(8, 22, DiffusionMatrixBabybear, &mut thread_rng()); @@ -105,7 +103,7 @@ fn test_prove_baby_bear() -> Result<(), VerificationError> { TwoAdicFriPcs>; let pcs = Pcs::new(fri_config, dft, val_mmcs); - type MyConfig = StarkConfig; + type MyConfig = StarkConfig; let config = StarkConfig::new(pcs); let mut challenger = Challenger::new(perm.clone());