From f9df5e9b348a4e85195285300ea38a7975ed192c Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Mon, 10 Nov 2025 15:22:18 -0500 Subject: [PATCH 01/26] everything works now except for cyclic recursion --- plonky2/Cargo.toml | 35 +- plonky2/benches/hashing.rs | 57 ++- plonky2/src/gadgets/interpolation.rs | 2 +- plonky2/src/hash/mod.rs | 1 + plonky2/src/hash/poseidon2._rs | 342 +++++++++++++ plonky2/src/hash/poseidon2/config.rs | 167 ++++++ plonky2/src/hash/poseidon2/gate.rs | 532 +++++++++++++++++++ plonky2/src/hash/poseidon2/hash.rs | 589 ++++++++++++++++++++++ plonky2/src/hash/poseidon2/mod.rs | 7 + plonky2/src/hash/poseidon2/p3.rs | 67 +++ plonky2/src/hash/poseidon2/pure.rs | 144 ++++++ plonky2/src/plonk/config.rs | 5 +- plonky2/src/recursion/cyclic_recursion.rs | 236 ++++----- 13 files changed, 2047 insertions(+), 137 deletions(-) create mode 100644 plonky2/src/hash/poseidon2._rs create mode 100644 plonky2/src/hash/poseidon2/config.rs create mode 100644 plonky2/src/hash/poseidon2/gate.rs create mode 100644 plonky2/src/hash/poseidon2/hash.rs create mode 100644 plonky2/src/hash/poseidon2/mod.rs create mode 100644 plonky2/src/hash/poseidon2/p3.rs create mode 100644 plonky2/src/hash/poseidon2/pure.rs diff --git a/plonky2/Cargo.toml b/plonky2/Cargo.toml index 83ff08519..37bd6b25b 100644 --- a/plonky2/Cargo.toml +++ b/plonky2/Cargo.toml @@ -38,6 +38,11 @@ plonky2_field = { version = "1.0.0", path = "../field", default-features = false plonky2_maybe_rayon = { version = "1.0.0", path = "../maybe_rayon", default-features = false } plonky2_util = { version = "1.0.0", path = "../util", default-features = false } +# Plonky3 dependencies for Poseidon2 +p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" } +p3-goldilocks = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" } +p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" } +p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" } [target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies] getrandom = { version = "0.2", default-features = false, features = ["js"] } @@ -60,29 +65,29 @@ jemallocator = "0.5.0" name = "generate_constants" required-features = ["rand_chacha"] -[[bench]] -name = "field_arithmetic" -harness = false +# [[bench]] +# name = "field_arithmetic" +# harness = false -[[bench]] -name = "ffts" -harness = false +# [[bench]] +# name = "ffts" +# harness = false [[bench]] name = "hashing" harness = false -[[bench]] -name = "merkle" -harness = false +# [[bench]] +# name = "merkle" +# harness = false -[[bench]] -name = "transpose" -harness = false +# [[bench]] +# name = "transpose" +# harness = false -[[bench]] -name = "reverse_index_bits" -harness = false +# [[bench]] +# name = "reverse_index_bits" +# harness = false # Display math equations properly in documentation [package.metadata.docs.rs] diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs index cf64b764a..acbe5cb74 100644 --- a/plonky2/benches/hashing.rs +++ b/plonky2/benches/hashing.rs @@ -3,9 +3,10 @@ mod allocator; use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; use plonky2::field::goldilocks_field::GoldilocksField; use plonky2::field::types::Sample; -use plonky2::hash::hash_types::{BytesHash, RichField}; +use plonky2::hash::hash_types::{BytesHash, HashOut, RichField}; use plonky2::hash::keccak::KeccakHash; use plonky2::hash::poseidon::{Poseidon, SPONGE_WIDTH}; +use plonky2::hash::poseidon2::{Poseidon2Hash, POSEIDON2_WIDTH}; use plonky2::plonk::config::Hasher; use tynm::type_name; @@ -32,8 +33,62 @@ pub(crate) fn bench_poseidon(c: &mut Criterion) { ); } +pub(crate) fn bench_poseidon2(c: &mut Criterion) { + c.bench_function( + &format!( + "poseidon2<{}, {POSEIDON2_WIDTH}>", + type_name::() + ), + |b| { + b.iter_batched( + || { + let input: Vec = (0..POSEIDON2_WIDTH) + .map(|_| GoldilocksField::rand()) + .collect(); + input + }, + |input| Poseidon2Hash::hash_no_pad(&input), + BatchSize::SmallInput, + ) + }, + ); +} + +pub(crate) fn bench_poseidon2_two_to_one(c: &mut Criterion) { + c.bench_function( + &format!("poseidon2_two_to_one<{}>", type_name::()), + |b| { + b.iter_batched( + || { + let left = HashOut { + elements: [ + GoldilocksField::rand(), + GoldilocksField::rand(), + GoldilocksField::rand(), + GoldilocksField::rand(), + ], + }; + let right = HashOut { + elements: [ + GoldilocksField::rand(), + GoldilocksField::rand(), + GoldilocksField::rand(), + GoldilocksField::rand(), + ], + }; + (left, right) + }, + |(left, right)| Poseidon2Hash::two_to_one(left, right), + BatchSize::SmallInput, + ) + }, + ); +} + fn criterion_benchmark(c: &mut Criterion) { bench_poseidon::(c); + bench_poseidon2(c); + bench_poseidon2_two_to_one(c); bench_keccak::(c); } diff --git a/plonky2/src/gadgets/interpolation.rs b/plonky2/src/gadgets/interpolation.rs index 39b048af4..9aedf7143 100644 --- a/plonky2/src/gadgets/interpolation.rs +++ b/plonky2/src/gadgets/interpolation.rs @@ -86,7 +86,7 @@ mod tests { let value_targets = values .iter() - .map(|&v| (builder.constant_extension(v))) + .map(|&v| builder.constant_extension(v)) .collect::>(); let zt = builder.constant_extension(z); diff --git a/plonky2/src/hash/mod.rs b/plonky2/src/hash/mod.rs index 0e4bb8a59..0d3436973 100644 --- a/plonky2/src/hash/mod.rs +++ b/plonky2/src/hash/mod.rs @@ -10,4 +10,5 @@ pub mod merkle_proofs; pub mod merkle_tree; pub mod path_compression; pub mod poseidon; +pub mod poseidon2; pub mod poseidon_goldilocks; diff --git a/plonky2/src/hash/poseidon2._rs b/plonky2/src/hash/poseidon2._rs new file mode 100644 index 000000000..9867a2f64 --- /dev/null +++ b/plonky2/src/hash/poseidon2._rs @@ -0,0 +1,342 @@ +//! Poseidon2 hash function implementation using Plonky3's Poseidon2. +//! +//! This module provides a wrapper around Plonky3's Poseidon2 implementation, +//! making it compatible with Plonky2's hash interface. + +#[cfg(not(feature = "std"))] +use alloc::vec::Vec; +use core::fmt::Debug; + +use p3_field::{PrimeCharacteristicRing, PrimeField64}; +use p3_goldilocks::{ + Goldilocks as GoldilocksP3, Poseidon2GoldilocksHL as Poseidon2P3Impl, + HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS, HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS, +}; +use p3_poseidon2::{ExternalLayerConstants, Poseidon2}; +use p3_symmetric::Permutation; + +use crate::field::goldilocks_field::GoldilocksField; +use crate::field::types::{Field, PrimeField64 as Plonky2PrimeField64}; +use crate::hash::hash_types::{HashOut, RichField, NUM_HASH_OUT_ELTS}; +use crate::hash::hashing::{compress, hash_n_to_hash_no_pad, PlonkyPermutation}; +use crate::iop::target::Target; +use crate::plonk::config::{AlgebraicHasher, Hasher}; + +/// Width of the Poseidon2 permutation (rate + capacity) +pub const POSEIDON2_WIDTH: usize = 12; +/// Rate of the sponge construction +pub const POSEIDON2_RATE: usize = 8; +/// Number of rounds for Poseidon2 (D parameter for x^D S-box) +pub const POSEIDON2_D: u64 = 7; + +/// Wrapper around Plonky3's Poseidon2 for Goldilocks field +pub struct Poseidon2Goldilocks { + /// The underlying Plonky3 Poseidon2 instance + permutation: Poseidon2P3Impl, +} + +// eraseamusegavelreign + +impl Debug for Poseidon2Goldilocks { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("Poseidon2Goldilocks").finish() + } +} + +impl Default for Poseidon2Goldilocks { + fn default() -> Self { + Self::new() + } +} + +impl Poseidon2Goldilocks { + /// Create a new Poseidon2 instance with default parameters + /// Uses Plonky3's Poseidon2 implementation for Goldilocks with precomputed constants + pub fn new() -> Self { + // Use Plonky3's Poseidon2 for Goldilocks with precomputed constants + // Convert the precomputed constants to the proper format + + // External round constants - convert 2D array of u64 to Vec<[Goldilocks; 8]> + let external_initial: Vec<[GoldilocksP3; 8]> = HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS[0] + .iter() + .map(|&round| round.map(GoldilocksP3::from_u64)) + .collect(); + let external_terminal: Vec<[GoldilocksP3; 8]> = HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS[1] + .iter() + .map(|&round| round.map(GoldilocksP3::from_u64)) + .collect(); + + // Internal round constants - convert Vec to Vec + let internal_constants: Vec = HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS + .iter() + .map(|&c| GoldilocksP3::from_u64(c)) + .collect(); + + let external_constants = ExternalLayerConstants::new(external_initial, external_terminal); + let permutation = Poseidon2::new(external_constants, internal_constants); + Self { permutation } + } + + /// Permute the state using Poseidon2 + pub fn permute(&self, state: [GoldilocksField; POSEIDON2_WIDTH]) -> [GoldilocksField; POSEIDON2_WIDTH] { + // Convert plonky2 fields to plonky3 fields + let p3_state: [GoldilocksP3; POSEIDON2_WIDTH] = state.map(|f| { + let val = Plonky2PrimeField64::to_canonical_u64(&f); + GoldilocksP3::from_u64(val) + }); + + // Apply permutation + let output = self.permutation.permute(p3_state); + + // Convert back to plonky2 fields + output.map(|f| GoldilocksField::from_canonical_u64(PrimeField64::as_canonical_u64(&f))) + } +} + +/// Permutation for Poseidon2 +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub struct Poseidon2Permutation([T; POSEIDON2_WIDTH]); + +impl Default for Poseidon2Permutation { + fn default() -> Self { + Self([T::default(); POSEIDON2_WIDTH]) + } +} + +impl AsRef<[T]> for Poseidon2Permutation { + fn as_ref(&self) -> &[T] { + &self.0 + } +} + +impl PlonkyPermutation for Poseidon2Permutation { + const RATE: usize = POSEIDON2_RATE; + const WIDTH: usize = POSEIDON2_WIDTH; + + fn new>(iter: I) -> Self { + let mut state = [GoldilocksField::ZERO; POSEIDON2_WIDTH]; + for (i, item) in iter.into_iter().take(POSEIDON2_WIDTH).enumerate() { + state[i] = item; + } + Self(state) + } + + fn set_elt(&mut self, elt: GoldilocksField, idx: usize) { + assert!(idx < POSEIDON2_WIDTH); + self.0[idx] = elt; + } + + fn set_from_iter>(&mut self, elts: I, start_idx: usize) { + assert!(start_idx <= POSEIDON2_WIDTH); + for (i, elt) in elts.into_iter().take(POSEIDON2_WIDTH - start_idx).enumerate() { + self.0[start_idx + i] = elt; + } + } + + fn set_from_slice(&mut self, elts: &[GoldilocksField], start_idx: usize) { + assert!(start_idx <= POSEIDON2_WIDTH); + let len = core::cmp::min(elts.len(), POSEIDON2_WIDTH - start_idx); + self.0[start_idx..start_idx + len].copy_from_slice(&elts[..len]); + } + + fn permute(&mut self) { + let poseidon2 = Poseidon2Goldilocks::new(); + self.0 = poseidon2.permute(self.0); + } + + fn squeeze(&self) -> &[GoldilocksField] { + &self.0[..POSEIDON2_RATE] + } +} + +impl PlonkyPermutation for Poseidon2Permutation { + const RATE: usize = POSEIDON2_RATE; + const WIDTH: usize = POSEIDON2_WIDTH; + + fn new>(iter: I) -> Self { + let mut state = [Target::default(); POSEIDON2_WIDTH]; + for (i, item) in iter.into_iter().take(POSEIDON2_WIDTH).enumerate() { + state[i] = item; + } + Self(state) + } + + fn set_elt(&mut self, elt: Target, idx: usize) { + assert!(idx < POSEIDON2_WIDTH); + self.0[idx] = elt; + } + + fn set_from_iter>(&mut self, elts: I, start_idx: usize) { + assert!(start_idx <= POSEIDON2_WIDTH); + for (i, elt) in elts.into_iter().take(POSEIDON2_WIDTH - start_idx).enumerate() { + self.0[start_idx + i] = elt; + } + } + + fn set_from_slice(&mut self, elts: &[Target], start_idx: usize) { + assert!(start_idx <= POSEIDON2_WIDTH); + let len = core::cmp::min(elts.len(), POSEIDON2_WIDTH - start_idx); + self.0[start_idx..start_idx + len].copy_from_slice(&elts[..len]); + } + + fn permute(&mut self) { + // For circuit targets, this is handled by the circuit builder + // The actual permutation is applied through gates + } + + fn squeeze(&self) -> &[Target] { + &self.0[..POSEIDON2_RATE] + } +} + +/// Poseidon2 hash function +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub struct Poseidon2Hash; + +impl Hasher for Poseidon2Hash { + const HASH_SIZE: usize = 4 * 8; // 4 field elements, 8 bytes each + type Hash = HashOut; + type Permutation = Poseidon2Permutation; + + fn hash_no_pad(input: &[GoldilocksField]) -> Self::Hash { + hash_n_to_hash_no_pad::(input) + } + + fn two_to_one(left: Self::Hash, right: Self::Hash) -> Self::Hash { + compress::(left, right) + } +} + +impl AlgebraicHasher for Poseidon2Hash { + type AlgebraicPermutation = Poseidon2Permutation; + + fn permute_swapped( + mut inputs: Self::AlgebraicPermutation, + swap: crate::iop::target::BoolTarget, + builder: &mut crate::plonk::circuit_builder::CircuitBuilder, + ) -> Self::AlgebraicPermutation + where + GoldilocksField: RichField + crate::field::extension::Extendable, + { + // Swap the first two chunks (first 4 elements with next 4 elements) if swap is true + // This is used for Merkle proof verification + let chunk_size = NUM_HASH_OUT_ELTS; + + for i in 0..chunk_size { + let left = inputs.0[i]; + let right = inputs.0[chunk_size + i]; + let (new_left, new_right) = ( + builder.select(swap, right, left), + builder.select(swap, left, right), + ); + inputs.0[i] = new_left; + inputs.0[chunk_size + i] = new_right; + } + + // Apply the permutation + // For now, we use witness generation since we don't have custom Poseidon2 gates yet + // In production, you would want to implement custom gates for efficiency + let output = builder.add_virtual_targets(POSEIDON2_WIDTH); + + // Add a constraint that will be checked during witness generation + // This is a placeholder - a full implementation would use custom gates + for _i in 0..POSEIDON2_WIDTH { + // The actual permutation will be computed during witness generation + // and verified by the constraints + let _ = builder.add_virtual_target(); + } + + Poseidon2Permutation(output.try_into().unwrap()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::field::types::{Field, Field64}; + + #[test] + fn test_poseidon2_permutation() { + let mut state = Poseidon2Permutation::::new( + (0..POSEIDON2_WIDTH).map(|i| GoldilocksField::from_canonical_u64(i as u64)) + ); + + let original_state = state.0; + state.permute(); + let permuted_state = state.0; + + // Check that permutation changes the state + assert_ne!(original_state, permuted_state); + + // Check that all outputs are valid field elements + for elem in &permuted_state { + assert!(elem.to_canonical_u64() < GoldilocksField::ORDER); + } + } + + #[test] + fn test_poseidon2_hash() { + let input: Vec = (0..8) + .map(GoldilocksField::from_canonical_u64) + .collect(); + + let hash = Poseidon2Hash::hash_no_pad(&input); + + // Verify hash has correct number of elements + assert_eq!(hash.elements.len(), NUM_HASH_OUT_ELTS); + + // Verify hash is deterministic + let hash2 = Poseidon2Hash::hash_no_pad(&input); + assert_eq!(hash, hash2); + } + + #[test] + fn test_poseidon2_two_to_one() { + let left = HashOut { + elements: [ + GoldilocksField::from_canonical_u64(1), + GoldilocksField::from_canonical_u64(2), + GoldilocksField::from_canonical_u64(3), + GoldilocksField::from_canonical_u64(4), + ], + }; + let right = HashOut { + elements: [ + GoldilocksField::from_canonical_u64(5), + GoldilocksField::from_canonical_u64(6), + GoldilocksField::from_canonical_u64(7), + GoldilocksField::from_canonical_u64(8), + ], + }; + + let hash = Poseidon2Hash::two_to_one(left, right); + + // Verify hash has correct number of elements + assert_eq!(hash.elements.len(), NUM_HASH_OUT_ELTS); + + // Verify hash is deterministic + let hash2 = Poseidon2Hash::two_to_one(left, right); + assert_eq!(hash, hash2); + + // Verify hash is different from inputs + assert_ne!(hash, left); + assert_ne!(hash, right); + } + + #[test] + fn test_sponge_construction() { + // Test that the sponge construction works correctly with multiple absorb/squeeze cycles + let input1: Vec = (0..16) + .map(GoldilocksField::from_canonical_u64) + .collect(); + let input2: Vec = (16..32) + .map(GoldilocksField::from_canonical_u64) + .collect(); + + let hash1 = Poseidon2Hash::hash_no_pad(&input1); + let hash2 = Poseidon2Hash::hash_no_pad(&input2); + + // Different inputs should produce different hashes + assert_ne!(hash1, hash2); + } +} diff --git a/plonky2/src/hash/poseidon2/config.rs b/plonky2/src/hash/poseidon2/config.rs new file mode 100644 index 000000000..96f25099b --- /dev/null +++ b/plonky2/src/hash/poseidon2/config.rs @@ -0,0 +1,167 @@ +pub const WIDTH: usize = 12; +pub const D: u64 = 7; +pub const RATE: usize = 8; +pub const OUT: usize = 4; + +/// Generated by `poseidon2_round_numbers_128` +pub const ROUNDS_F: usize = 8; +pub const ROUNDS_F_HALF: usize = 4; +pub const ROUNDS_P: usize = 22; + +/// Generated randomly for ROUNDS_F +pub const EXTERNAL_CONSTANTS: [[u64; WIDTH]; ROUNDS_F] = [ + [ + 15492826721047263190, + 11728330187201910315, + 8836021247773420868, + 16777404051263952451, + 5510875212538051896, + 6173089941271892285, + 2927757366422211339, + 10340958981325008808, + 8541987352684552425, + 9739599543776434497, + 15073950188101532019, + 12084856431752384512, + ], + [ + 4584713381960671270, + 8807052963476652830, + 54136601502601741, + 4872702333905478703, + 5551030319979516287, + 12889366755535460989, + 16329242193178844328, + 412018088475211848, + 10505784623379650541, + 9758812378619434837, + 7421979329386275117, + 375240370024755551, + ], + [ + 3331431125640721931, + 15684937309956309981, + 578521833432107983, + 14379242000670861838, + 17922409828154900976, + 8153494278429192257, + 15904673920630731971, + 11217863998460634216, + 3301540195510742136, + 9937973023749922003, + 3059102938155026419, + 1895288289490976132, + ], + [ + 5580912693628927540, + 10064804080494788323, + 9582481583369602410, + 10186259561546797986, + 247426333829703916, + 13193193905461376067, + 6386232593701758044, + 17954717245501896472, + 1531720443376282699, + 2455761864255501970, + 11234429217864304495, + 4746959618548874102, + ], + [ + 13571697342473846203, + 17477857865056504753, + 15963032953523553760, + 16033593225279635898, + 14252634232868282405, + 8219748254835277737, + 7459165569491914711, + 15855939513193752003, + 16788866461340278896, + 7102224659693946577, + 3024718005636976471, + 13695468978618890430, + ], + [ + 8214202050877825436, + 2670727992739346204, + 16259532062589659211, + 11869922396257088411, + 3179482916972760137, + 13525476046633427808, + 3217337278042947412, + 14494689598654046340, + 15837379330312175383, + 8029037639801151344, + 2153456285263517937, + 8301106462311849241, + ], + [ + 13294194396455217955, + 17394768489610594315, + 12847609130464867455, + 14015739446356528640, + 5879251655839607853, + 9747000124977436185, + 8950393546890284269, + 10765765936405694368, + 14695323910334139959, + 16366254691123000864, + 15292774414889043182, + 10910394433429313384, + ], + [ + 17253424460214596184, + 3442854447664030446, + 3005570425335613727, + 10859158614900201063, + 9763230642109343539, + 6647722546511515039, + 909012944955815706, + 18101204076790399111, + 11588128829349125809, + 15863878496612806566, + 5201119062417750399, + 176665553780565743, + ], +]; + +/// Generated randomly for ROUNDS_P +pub const INTERNAL_CONSTANTS: [u64; ROUNDS_P] = [ + 11921381764981422944, + 10318423381711320787, + 8291411502347000766, + 229948027109387563, + 9152521390190983261, + 7129306032690285515, + 15395989607365232011, + 8641397269074305925, + 17256848792241043600, + 6046475228902245682, + 12041608676381094092, + 12785542378683951657, + 14546032085337914034, + 3304199118235116851, + 16499627707072547655, + 10386478025625759321, + 13475579315436919170, + 16042710511297532028, + 1411266850385657080, + 9024840976168649958, + 14047056970978379368, + 838728605080212101, +]; + +/// Taken from Plonk3 Poseidon2 implementation. https://github.com/Plonky3/Plonky3/blob/eeb4e37b20127c4daa871b2bad0df30a7c7380db/goldilocks/src/poseidon2.rs#L28 +pub const MATRIX_DIAG_12_U64: [u64; WIDTH] = [ + 0xc3b6c08e23ba9300, + 0xd84b5de94a324fb6, + 0x0d0c371c5b35b84f, + 0x7964f570e7188037, + 0x5daf18bbd996604b, + 0x6743bc47b9595257, + 0x5528b9362c59bb70, + 0xac45e25b7127b68b, + 0xa2077d7dfbb606b5, + 0xf3faac6faee378ae, + 0x0c6388b51545e883, + 0xd27dbb6944917b60, +]; diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs new file mode 100644 index 000000000..9e2958675 --- /dev/null +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -0,0 +1,532 @@ +//! Implementation of a Plonky2 gate for an entire Poseidon2 permutation over a +//! state of width 12 + +use core::marker::PhantomData; + +use anyhow::Result; + +use super::config::*; +use super::hash::Poseidon2; +use crate::field::extension::Extendable; +use crate::field::types::Field; +use crate::gates::gate::Gate; +use crate::gates::util::StridedConstraintConsumer; +use crate::hash::hash_types::RichField; +use crate::iop::ext_target::ExtensionTarget; +use crate::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef}; +use crate::iop::target::Target; +use crate::iop::wire::Wire; +use crate::iop::witness::{PartitionWitness, Witness, WitnessWrite}; +use crate::plonk::circuit_builder::CircuitBuilder; +use crate::plonk::circuit_data::CommonCircuitData; +use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase}; +use crate::util::serialization::{Buffer, IoResult, Read, Write}; + +/// Evaluates a full Poseidon2 permutation with 12 state elements. +/// +/// This also has some extra features to make it suitable for efficiently +/// verifying Merkle proofs. It has a flag which can be used to swap the first +/// four inputs with the next four, for ordering sibling digests. +#[derive(Debug, Default)] +pub struct Poseidon2Gate, const D: usize> { + _phantom: PhantomData, +} + +impl + Poseidon2, const D: usize> Poseidon2Gate { + pub fn new() -> Self { + Poseidon2Gate { + _phantom: PhantomData, + } + } + + /// The wire index for the `i`th input to the permutation. + pub fn wire_input(i: usize) -> usize { + i + } + + /// The wire index for the `i`th output to the permutation. + pub fn wire_output(i: usize) -> usize { + WIDTH + i + } + + /// If this is set to 1, the first four inputs will be swapped with the next + /// four inputs. This is useful for ordering hashes in Merkle proofs. + /// Otherwise, this should be set to 0. + pub const WIRE_SWAP: usize = 2 * WIDTH; + + const START_DELTA: usize = 2 * WIDTH + 1; + + /// A wire which stores `swap * (input[i + 4] - input[i])`; used to compute + /// the swapped inputs. + fn wire_delta(i: usize) -> usize { + assert!(i < 4); + Self::START_DELTA + i + } + + const START_ROUND_F_BEGIN: usize = Self::START_DELTA + 4; + + /// A wire which stores the input of the `i`-th S-box of the `round`-th + /// round of the first set of full rounds. + fn wire_full_sbox_0(round: usize, i: usize) -> usize { + debug_assert!( + round != 0, + "First round S-box inputs are not stored as wires" + ); + debug_assert!(round < ROUNDS_F_HALF); + debug_assert!(i < WIDTH); + Self::START_ROUND_F_BEGIN + WIDTH * (round - 1) + i + } + + const START_PARTIAL: usize = Self::START_ROUND_F_BEGIN + WIDTH * (ROUNDS_F_HALF - 1); + + /// A wire which stores the input of the S-box of the `round`-th round of + /// the partial rounds. + const fn wire_partial_sbox(round: usize) -> usize { + debug_assert!(round < ROUNDS_P); + Self::START_PARTIAL + round + } + + const START_ROUND_F_END: usize = Self::START_PARTIAL + ROUNDS_P; + + /// A wire which stores the input of the `i`-th S-box of the `round`-th + /// round of the second set of full rounds. + const fn wire_full_sbox_1(round: usize, i: usize) -> usize { + debug_assert!(round < ROUNDS_F_HALF); + debug_assert!(i < WIDTH); + Self::START_ROUND_F_END + WIDTH * round + i + } + + /// End of wire indices, exclusive. + const fn end() -> usize { + Self::START_ROUND_F_END + WIDTH * ROUNDS_F_HALF + } +} + +impl + Poseidon2, const D: usize> Gate for Poseidon2Gate { + fn id(&self) -> String { + format!("{:?}", self, WIDTH) + } + + fn serialize( + &self, + _dst: &mut Vec, + _common_data: &CommonCircuitData, + ) -> IoResult<()> { + Ok(()) + } + + fn deserialize(_src: &mut Buffer, _common_data: &CommonCircuitData) -> IoResult { + Ok(Poseidon2Gate::new()) + } + + fn eval_unfiltered(&self, vars: EvaluationVars) -> Vec { + let mut constraints = Vec::with_capacity(self.num_constraints()); + + // Assert that `swap` is binary. + let swap = vars.local_wires[Self::WIRE_SWAP]; + constraints.push(swap * (swap - F::Extension::ONE)); + + // Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`. + for i in 0..4 { + let input_lhs = vars.local_wires[Self::wire_input(i)]; + let input_rhs = vars.local_wires[Self::wire_input(i + 4)]; + let delta_i = vars.local_wires[Self::wire_delta(i)]; + constraints.push(swap * (input_rhs - input_lhs) - delta_i); + } + + // Compute the possibly-swapped input layer. + let mut state = [F::Extension::ZERO; WIDTH]; + for i in 0..4 { + let delta_i = vars.local_wires[Self::wire_delta(i)]; + let input_lhs = Self::wire_input(i); + let input_rhs = Self::wire_input(i + 4); + state[i] = vars.local_wires[input_lhs] + delta_i; + state[i + 4] = vars.local_wires[input_rhs] - delta_i; + } + for i in 8..WIDTH { + state[i] = vars.local_wires[Self::wire_input(i)]; + } + + // The initial linear layer. + ::external_linear_layer_extension(&mut state); + + // The first half of the external rounds. + for r in 0..ROUNDS_F_HALF { + ::add_rc_extension(&mut state, r); + if r != 0 { + for i in 0..WIDTH { + let sbox_in = vars.local_wires[Self::wire_full_sbox_0(r, i)]; + constraints.push(state[i] - sbox_in); + state[i] = sbox_in; + } + } + ::sbox_extension(&mut state); + ::external_linear_layer_extension(&mut state); + } + + // The internal rounds. + for r in 0..ROUNDS_P { + state[0] += F::Extension::from_canonical_u64(INTERNAL_CONSTANTS[r]); + let sbox_in = vars.local_wires[Self::wire_partial_sbox(r)]; + constraints.push(state[0] - sbox_in); + state[0] = sbox_in; + state[0] = ::sbox_p_extension(&state[0]); + ::internal_linear_layer_extension(&mut state); + } + + // The second half of the external rounds. + for r in ROUNDS_F_HALF..ROUNDS_F { + ::add_rc_extension(&mut state, r); + for i in 0..WIDTH { + let sbox_in = vars.local_wires[Self::wire_full_sbox_1(r - ROUNDS_F_HALF, i)]; + constraints.push(state[i] - sbox_in); + state[i] = sbox_in; + } + ::sbox_extension(&mut state); + ::external_linear_layer_extension(&mut state); + } + + for i in 0..WIDTH { + constraints.push(state[i] - vars.local_wires[Self::wire_output(i)]); + } + + constraints + } + + fn eval_unfiltered_base_one( + &self, + vars: EvaluationVarsBase, + mut yield_constr: StridedConstraintConsumer, + ) { + // Assert that `swap` is binary. + let swap = vars.local_wires[Self::WIRE_SWAP]; + yield_constr.one(swap * swap.sub_one()); + + // Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`. + for i in 0..4 { + let input_lhs = vars.local_wires[Self::wire_input(i)]; + let input_rhs = vars.local_wires[Self::wire_input(i + 4)]; + let delta_i = vars.local_wires[Self::wire_delta(i)]; + yield_constr.one(swap * (input_rhs - input_lhs) - delta_i); + } + + // Compute the possibly-swapped input layer. + let mut state = [F::ZERO; WIDTH]; + for i in 0..4 { + let delta_i = vars.local_wires[Self::wire_delta(i)]; + let input_lhs = Self::wire_input(i); + let input_rhs = Self::wire_input(i + 4); + state[i] = vars.local_wires[input_lhs] + delta_i; + state[i + 4] = vars.local_wires[input_rhs] - delta_i; + } + for i in 8..WIDTH { + state[i] = vars.local_wires[Self::wire_input(i)]; + } + + // The initial linear layer. + ::external_linear_layer(&mut state); + + // The first half of the external rounds. + for r in 0..ROUNDS_F_HALF { + ::add_rc(&mut state, r); + if r != 0 { + for i in 0..WIDTH { + let sbox_in = vars.local_wires[Self::wire_full_sbox_0(r, i)]; + yield_constr.one(state[i] - sbox_in); + state[i] = sbox_in; + } + } + ::sbox(&mut state); + ::external_linear_layer(&mut state); + } + + // The internal rounds. + for r in 0..ROUNDS_P { + state[0] += F::from_canonical_u64(INTERNAL_CONSTANTS[r]); + let sbox_in = vars.local_wires[Self::wire_partial_sbox(r)]; + yield_constr.one(state[0] - sbox_in); + state[0] = sbox_in; + state[0] = ::sbox_p(&state[0]); + ::internal_linear_layer(&mut state); + } + + // The second half of the external rounds. + for r in ROUNDS_F_HALF..ROUNDS_F { + ::add_rc(&mut state, r); + for i in 0..WIDTH { + let sbox_in = vars.local_wires[Self::wire_full_sbox_1(r - ROUNDS_F_HALF, i)]; + yield_constr.one(state[i] - sbox_in); + state[i] = sbox_in; + } + ::sbox(&mut state); + ::external_linear_layer(&mut state); + } + + for i in 0..WIDTH { + yield_constr.one(state[i] - vars.local_wires[Self::wire_output(i)]); + } + } + + fn eval_unfiltered_circuit( + &self, + builder: &mut CircuitBuilder, + vars: EvaluationTargets, + ) -> Vec> { + let mut constraints = Vec::with_capacity(self.num_constraints()); + + // Assert that `swap` is binary. + let swap = vars.local_wires[Self::WIRE_SWAP]; + constraints.push(builder.mul_sub_extension(swap, swap, swap)); + + // Assert that each delta wire is set properly: `delta_i = swap * (rhs - lhs)`. + for i in 0..4 { + let input_lhs = vars.local_wires[Self::wire_input(i)]; + let input_rhs = vars.local_wires[Self::wire_input(i + 4)]; + let delta_i = vars.local_wires[Self::wire_delta(i)]; + let diff = builder.sub_extension(input_rhs, input_lhs); + constraints.push(builder.mul_sub_extension(swap, diff, delta_i)); + } + + // Compute the possibly-swapped input layer. + let mut state = [builder.zero_extension(); WIDTH]; + for i in 0..4 { + let delta_i = vars.local_wires[Self::wire_delta(i)]; + let input_lhs = vars.local_wires[Self::wire_input(i)]; + let input_rhs = vars.local_wires[Self::wire_input(i + 4)]; + state[i] = builder.add_extension(input_lhs, delta_i); + state[i + 4] = builder.sub_extension(input_rhs, delta_i); + } + for i in 8..WIDTH { + state[i] = vars.local_wires[Self::wire_input(i)]; + } + + // The initial linear layer. + ::external_linear_layer_circuit(builder, &mut state); + + // The first half of the external rounds. + for r in 0..ROUNDS_F_HALF { + ::add_rc_circuit(builder, &mut state, r); + if r != 0 { + for i in 0..WIDTH { + let sbox_in = vars.local_wires[Self::wire_full_sbox_0(r, i)]; + constraints.push(builder.sub_extension(state[i], sbox_in)); + state[i] = sbox_in; + } + } + ::sbox_circuit(builder, &mut state); + ::external_linear_layer_circuit(builder, &mut state); + } + + // The internal rounds. + for r in 0..ROUNDS_P { + let round_constant = F::Extension::from_canonical_u64(INTERNAL_CONSTANTS[r]); + let round_constant = builder.constant_extension(round_constant); + state[0] = builder.add_extension(state[0], round_constant); + + let sbox_in = vars.local_wires[Self::wire_partial_sbox(r)]; + constraints.push(builder.sub_extension(state[0], sbox_in)); + state[0] = sbox_in; + state[0] = ::sbox_p_circuit(builder, state[0]); + ::internal_linear_layer_circuit(builder, &mut state); + } + + // The second half of the external rounds. + for r in ROUNDS_F_HALF..ROUNDS_F { + ::add_rc_circuit(builder, &mut state, r); + + for i in 0..WIDTH { + let sbox_in = vars.local_wires[Self::wire_full_sbox_1(r - ROUNDS_F_HALF, i)]; + constraints.push(builder.sub_extension(state[i], sbox_in)); + state[i] = sbox_in; + } + ::sbox_circuit(builder, &mut state); + ::external_linear_layer_circuit(builder, &mut state); + } + + for i in 0..WIDTH { + constraints + .push(builder.sub_extension(state[i], vars.local_wires[Self::wire_output(i)])); + } + + constraints + } + + fn generators(&self, row: usize, _local_constants: &[F]) -> Vec> { + let g = Poseidon2Generator:: { + row, + _phantom: PhantomData, + }; + vec![WitnessGeneratorRef::new(g.adapter())] + } + + fn num_wires(&self) -> usize { + Self::end() + } + + fn num_constants(&self) -> usize { + 0 + } + + fn degree(&self) -> usize { + 7 + } + + fn num_constraints(&self) -> usize { + WIDTH * (ROUNDS_F - 1) + ROUNDS_P + WIDTH + 1 + 4 + } +} + +#[derive(Debug, Default)] +pub struct Poseidon2Generator + Poseidon2, const D: usize> { + row: usize, + _phantom: PhantomData, +} + +impl + Poseidon2, const D: usize> SimpleGenerator + for Poseidon2Generator +{ + fn id(&self) -> String { + "Poseidon2Generator".to_string() + } + + fn dependencies(&self) -> Vec { + (0..WIDTH) + .map(|i| Poseidon2Gate::::wire_input(i)) + .chain(Some(Poseidon2Gate::::WIRE_SWAP)) + .map(|column| Target::wire(self.row, column)) + .collect() + } + + fn run_once( + &self, + witness: &PartitionWitness, + out_buffer: &mut GeneratedValues, + ) -> Result<()> { + let local_wire = |column| Wire { + row: self.row, + column, + }; + + let mut state = (0..WIDTH) + .map(|i| witness.get_wire(local_wire(Poseidon2Gate::::wire_input(i)))) + .collect::>(); + + let swap_value = witness.get_wire(local_wire(Poseidon2Gate::::WIRE_SWAP)); + debug_assert!(swap_value == F::ZERO || swap_value == F::ONE); + + for i in 0..4 { + let delta_i = swap_value * (state[i + 4] - state[i]); + out_buffer.set_wire(local_wire(Poseidon2Gate::::wire_delta(i)), delta_i)?; + } + + if swap_value == F::ONE { + for i in 0..4 { + state.swap(i, 4 + i); + } + } + + let mut state: [F; WIDTH] = state.try_into().unwrap(); + + ::external_linear_layer(&mut state); + + // The first half of the external rounds. + for r in 0..ROUNDS_F_HALF { + ::add_rc(&mut state, r); + if r != 0 { + for i in 0..WIDTH { + out_buffer.set_wire( + local_wire(Poseidon2Gate::::wire_full_sbox_0(r, i)), + state[i], + )?; + } + } + ::sbox(&mut state); + ::external_linear_layer(&mut state); + } + + // The internal rounds. + for r in 0..ROUNDS_P { + state[0] += F::from_canonical_u64(INTERNAL_CONSTANTS[r]); + out_buffer.set_wire( + local_wire(Poseidon2Gate::::wire_partial_sbox(r)), + state[0], + )?; + state[0] = ::sbox_p(&state[0]); + ::internal_linear_layer(&mut state); + } + + // The second half of the external rounds. + for r in ROUNDS_F_HALF..ROUNDS_F { + ::add_rc(&mut state, r); + for i in 0..WIDTH { + out_buffer.set_wire( + local_wire(Poseidon2Gate::::wire_full_sbox_1( + r - ROUNDS_F_HALF, + i, + )), + state[i], + )?; + } + ::sbox(&mut state); + ::external_linear_layer(&mut state); + } + + for i in 0..WIDTH { + out_buffer.set_wire(local_wire(Poseidon2Gate::::wire_output(i)), state[i])?; + } + + Ok(()) + } + + fn serialize(&self, dst: &mut Vec, _common_data: &CommonCircuitData) -> IoResult<()> { + dst.write_usize(self.row) + } + + fn deserialize(src: &mut Buffer, _common_data: &CommonCircuitData) -> IoResult { + let row = src.read_usize()?; + Ok(Self { + row, + _phantom: PhantomData, + }) + } +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + + use super::Poseidon2Gate; + use crate::field::goldilocks_field::GoldilocksField; + use crate::gates::gate_testing::{test_eval_fns, test_low_degree}; + use crate::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; + + #[test] + fn wire_indices() { + type F = GoldilocksField; + type Gate = Poseidon2Gate; + + assert_eq!(Gate::wire_input(0), 0); + assert_eq!(Gate::wire_input(11), 11); + assert_eq!(Gate::wire_output(0), 12); + assert_eq!(Gate::wire_output(11), 23); + assert_eq!(Gate::WIRE_SWAP, 24); + assert_eq!(Gate::wire_delta(0), 25); + assert_eq!(Gate::wire_delta(3), 28); + } + + #[test] + fn low_degree() { + type F = GoldilocksField; + let gate = Poseidon2Gate::::new(); + test_low_degree(gate) + } + + #[test] + fn eval_fns() -> Result<()> { + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + let gate = Poseidon2Gate::::new(); + test_eval_fns::(gate) + } +} diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs new file mode 100644 index 000000000..486e87897 --- /dev/null +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -0,0 +1,589 @@ +use core::fmt::Debug; + +use crate::field::extension::{Extendable, FieldExtension}; +use crate::field::goldilocks_field::GoldilocksField as F; +use crate::field::types::{Field, PrimeField64}; +use crate::hash::hash_types::{HashOut, RichField}; +use crate::hash::hashing::{compress, hash_n_to_hash_no_pad, PlonkyPermutation}; +use crate::iop::ext_target::ExtensionTarget; +use crate::iop::target::{BoolTarget, Target}; +use crate::plonk::circuit_builder::CircuitBuilder; +use crate::plonk::config::{AlgebraicHasher, Hasher}; + +use super::config::*; +use super::gate::Poseidon2Gate; + + + +pub trait Poseidon2: PrimeField64 { + fn poseidon2(input: [Self; WIDTH]) -> [Self; WIDTH] { + let mut state = input; + + Self::external_linear_layer(&mut state); + + Self::full_rounds(&mut state, 0); + Self::partial_rounds(&mut state); + Self::full_rounds(&mut state, ROUNDS_F_HALF); + + state + } + + #[inline] + fn full_rounds(state: &mut [Self; WIDTH], start: usize) { + for r in start..(start + ROUNDS_F_HALF) { + Self::add_rc(state, r); + Self::sbox(state); + Self::external_linear_layer(state); + } + } + + #[inline] + fn partial_rounds(state: &mut [Self; WIDTH]) { + for r in 0..ROUNDS_P { + state[0] += Self::from_canonical_u64(INTERNAL_CONSTANTS[r]); + state[0] = Self::sbox_p(&state[0]); + Self::internal_linear_layer(state); + } + } + + #[inline] + fn external_linear_layer(state: &mut [Self; 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) { + // Would be nice to find a better way to do this. + let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; + Self::apply_mat4_mut(&mut state_4); + state[i..i + 4].clone_from_slice(&state_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: [Self; 4] = + core::array::from_fn(|k| (0..WIDTH).step_by(4).map(|j| state[j + k]).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]; + } + } + + #[inline] + fn external_linear_layer_extension, const D: usize>( + state: &mut [F; 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) { + // Would be nice to find a better way to do this. + let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; + Self::apply_mat4_mut_extension(&mut state_4); + state[i..i + 4].clone_from_slice(&state_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: [F; 4] = + core::array::from_fn(|k| (0..WIDTH).step_by(4).map(|j| state[j + k]).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]; + } + } + + #[inline] + fn internal_linear_layer(state: &mut [Self; WIDTH]) { + let sum: Self = state.iter().cloned().sum(); + for i in 0..WIDTH { + state[i] *= Self::from_canonical_u64(MATRIX_DIAG_12_U64[i]); + state[i] += sum; + } + } + + #[inline] + fn internal_linear_layer_extension, const D: usize>( + state: &mut [F; WIDTH], + ) { + let mut sum = state[0]; + for i in 1..WIDTH { + sum += state[i]; + } + for i in 0..WIDTH { + state[i] *= F::from_canonical_u64(MATRIX_DIAG_12_U64[i]); + state[i] += sum; + } + } + + #[inline] + fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + + for i in 0..WIDTH { + unsafe { + state[i] = state[i].add_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); + } + } + } + + #[inline] + fn add_rc_extension, const D: usize>( + state: &mut [F; WIDTH], + external_round: usize, + ) { + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + + for i in 0..WIDTH { + state[i] += F::from_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); + } + } + + #[inline] + fn sbox(state: &mut [Self; WIDTH]) { + state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); + } + + #[inline] + fn sbox_extension, const D: usize>( + state: &mut [F; WIDTH], + ) { + state + .iter_mut() + .for_each(|a| *a = Self::sbox_p_extension(a)); + } + + #[inline] + fn sbox_p(a: &Self) -> Self { + a.exp_u64(D) + } + + #[inline] + fn sbox_p_extension, const D: usize>(a: &F) -> F { + a.exp_u64(super::config::D) + } + + // Multiply a 4-element vector x by: + // [ 2 3 1 1 ] + // [ 1 2 3 1 ] + // [ 1 1 2 3 ] + // [ 3 1 1 2 ]. + // This is more efficient than the previous matrix. + fn apply_mat4_mut(x: &mut [Self; 4]) { + let t01 = x[0] + x[1]; + let t23 = x[2] + x[3]; + let t0123 = t01 + t23; + let t01123 = t0123 + x[1]; + let t01233 = t0123 + x[3]; + // The order here is important. Need to overwrite x[0] and x[2] after x[1] and x[3]. + x[3] = t01233 + x[0].double(); // 3*x[0] + x[1] + x[2] + 2*x[3] + x[1] = t01123 + x[2].double(); // x[0] + 2*x[1] + 3*x[2] + x[3] + x[0] = t01123 + t01; // 2*x[0] + 3*x[1] + x[2] + x[3] + x[2] = t01233 + t23; // x[0] + x[1] + 2*x[2] + 3*x[3] + } + + fn apply_mat4_mut_extension, const D: usize>( + x: &mut [F; 4], + ) { + let t01 = x[0] + x[1]; + let t23 = x[2] + x[3]; + let t0123 = t01 + t23; + let t01123 = t0123 + x[1]; + let t01233 = t0123 + x[3]; + // The order here is important. Need to overwrite x[0] and x[2] after x[1] and x[3]. + x[3] = t01233 + x[0].double(); // 3*x[0] + x[1] + x[2] + 2*x[3] + x[1] = t01123 + x[2].double(); // x[0] + 2*x[1] + 3*x[2] + x[3] + x[0] = t01123 + t01; // 2*x[0] + 3*x[1] + x[2] + x[3] + x[2] = t01233 + t23; // x[0] + x[1] + 2*x[2] + 3*x[3] + } + + // In circuit functions + + fn external_linear_layer_circuit( + builder: &mut CircuitBuilder, + state: &mut [ExtensionTarget; WIDTH], + ) where + Self: RichField + Extendable, + { + // 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) { + // Would be nice to find a better way to do this. + let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; + Self::apply_mat4_mut_circuit(builder, &mut state_4); + state[i..i + 4].clone_from_slice(&state_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: [ExtensionTarget; 4] = core::array::from_fn(|k| { + (0..WIDTH) + .step_by(4) + .map(|j| state[j + k]) + .reduce(|acc, t| builder.add_extension(acc, t)) + .unwrap() + }); + + // 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] = builder.add_extension(state[i], sums[i % 4]); + } + } + + fn apply_mat4_mut_circuit( + builder: &mut CircuitBuilder, + x: &mut [ExtensionTarget; 4], + ) where + Self: RichField + Extendable, + { + let two = builder.constant_extension(Self::Extension::from_canonical_u64(2)); + + let t01 = builder.add_extension(x[0], x[1]); + let t23 = builder.add_extension(x[2], x[3]); + let t0123 = builder.add_extension(t01, t23); + let t01123 = builder.add_extension(t0123, x[1]); + let t01233 = builder.add_extension(t0123, x[3]); + // The order here is important. Need to overwrite x[0] and x[2] after x[1] and x[3]. + let dx0 = builder.mul_extension(x[0], two); + let dx2 = builder.mul_extension(x[2], two); + x[3] = builder.add_extension(t01233, dx0); // 3*x[0] + x[1] + x[2] + 2*x[3] + x[1] = builder.add_extension(t01123, dx2); // x[0] + 2*x[1] + 3*x[2] + x[3] + x[0] = builder.add_extension(t01123, t01); // 2*x[0] + 3*x[1] + x[2] + x[3] + x[2] = builder.add_extension(t01233, t23); // x[0] + x[1] + 2*x[2] + 3*x[3] + } + + fn matmul_m4_circuit( + builder: &mut CircuitBuilder, + input: &mut [ExtensionTarget; WIDTH], + ) where + Self: RichField + Extendable, + { + for i in 0..3 { + let t_0 = builder.mul_const_add_extension(Self::ONE, input[i * 4], input[i * 4 + 1]); + let t_1 = + builder.mul_const_add_extension(Self::ONE, input[i * 4 + 2], input[i * 4 + 3]); + let t_2 = builder.mul_const_add_extension(Self::TWO, input[i * 4 + 1], t_1); + let t_3 = builder.mul_const_add_extension(Self::TWO, input[i * 4 + 3], t_0); + + let four = Self::TWO + Self::TWO; + + let t_4 = builder.mul_const_add_extension(four, t_1, t_3); + let t_5 = builder.mul_const_add_extension(four, t_0, t_2); + let t_6 = builder.mul_const_add_extension(Self::ONE, t_3, t_5); + let t_7 = builder.mul_const_add_extension(Self::ONE, t_2, t_4); + + input[i * 4] = t_6; + input[i * 4 + 1] = t_5; + input[i * 4 + 2] = t_7; + input[i * 4 + 3] = t_4; + } + } + + fn add_rc_circuit( + builder: &mut CircuitBuilder, + input: &mut [ExtensionTarget; WIDTH], + rc_index: usize, + ) where + Self: RichField + Extendable, + { + for i in 0..WIDTH { + let round_constant = + Self::Extension::from_canonical_u64(EXTERNAL_CONSTANTS[rc_index][i]); + let round_constant = builder.constant_extension(round_constant); + input[i] = builder.add_extension(input[i], round_constant); + } + } + + fn sbox_circuit( + builder: &mut CircuitBuilder, + input: &mut [ExtensionTarget; WIDTH], + ) where + Self: RichField + Extendable, + { + for i in 0..WIDTH { + input[i] = Self::sbox_p_circuit(builder, input[i]); + } + } + + #[inline] + fn sbox_p_circuit( + builder: &mut CircuitBuilder, + input: ExtensionTarget, + ) -> ExtensionTarget + where + Self: RichField + Extendable, + { + builder.exp_u64_extension(input, super::config::D) + } + + fn internal_linear_layer_circuit( + builder: &mut CircuitBuilder, + input: &mut [ExtensionTarget; WIDTH], + ) where + Self: RichField + Extendable, + { + let sum = builder.add_many_extension([ + input[0], input[1], input[2], input[3], input[4], input[5], input[6], input[7], + input[8], input[9], input[10], input[11], + ]); + + for i in 0..WIDTH { + let round_constant = Self::Extension::from_canonical_u64(MATRIX_DIAG_12_U64[i]); + let round_constant = builder.constant_extension(round_constant); + + input[i] = builder.mul_add_extension(round_constant, input[i], sum); + } + } +} + +impl Poseidon2 for F {} + +#[derive(Copy, Clone, Default, Debug, PartialEq)] +pub struct Poseidon2Permutation { + state: [T; WIDTH], +} + +impl Eq for Poseidon2Permutation {} + +impl AsRef<[T]> for Poseidon2Permutation { + fn as_ref(&self) -> &[T] { + &self.state + } +} + +trait Permuter: Sized { + fn permute(input: [Self; WIDTH]) -> [Self; WIDTH]; +} + +impl Permuter for F { + fn permute(input: [Self; WIDTH]) -> [Self; WIDTH] { + ::poseidon2(input) + } +} + +impl Permuter for Target { + fn permute(_input: [Self; WIDTH]) -> [Self; WIDTH] { + panic!("Call `permute_swapped()` instead of `permute()`"); + } +} + +impl PlonkyPermutation + for Poseidon2Permutation +{ + const RATE: usize = RATE; + const WIDTH: usize = WIDTH; + + fn new>(elts: I) -> Self { + let mut perm = Self { + state: [T::default(); WIDTH], + }; + perm.set_from_iter(elts, 0); + perm + } + + fn set_elt(&mut self, elt: T, idx: usize) { + self.state[idx] = elt; + } + + fn set_from_slice(&mut self, elts: &[T], start_idx: usize) { + let begin = start_idx; + let end = start_idx + elts.len(); + self.state[begin..end].copy_from_slice(elts); + } + + fn set_from_iter>(&mut self, elts: I, start_idx: usize) { + for (s, e) in self.state[start_idx..].iter_mut().zip(elts) { + *s = e; + } + } + + fn permute(&mut self) { + self.state = T::permute(self.state); + } + + fn squeeze(&self) -> &[T] { + &self.state[..Self::RATE] + } +} + +/// Poseidon2 hash function. +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub struct Poseidon2Hash; +impl Hasher for Poseidon2Hash { + const HASH_SIZE: usize = 4 * 8; + type Hash = HashOut; + type Permutation = Poseidon2Permutation; + + fn hash_no_pad(input: &[F]) -> Self::Hash { + hash_n_to_hash_no_pad::(input) + } + + fn two_to_one(left: Self::Hash, right: Self::Hash) -> Self::Hash { + compress::(left, right) + } +} + +impl Poseidon2Hash { + pub fn hash_n_to_one( + input: &[>::Hash], + ) -> >::Hash { + assert!(!input.is_empty()); + + if input.len() == 1 { + return input[0]; + } + + let mut result = >::two_to_one(input[0], input[1]); + + for i in 2..input.len() { + result = >::two_to_one(result, input[i]); + } + + result + } +} + +impl AlgebraicHasher for Poseidon2Hash { + type AlgebraicPermutation = Poseidon2Permutation; + + fn permute_swapped( + inputs: Self::AlgebraicPermutation, + swap: BoolTarget, + builder: &mut CircuitBuilder, + ) -> Self::AlgebraicPermutation + where + F: RichField + Extendable, + { + let gate_type = Poseidon2Gate::::new(); + let gate = builder.add_gate(gate_type, vec![]); + + let swap_wire = Poseidon2Gate::::WIRE_SWAP; + let swap_wire = Target::wire(gate, swap_wire); + builder.connect(swap.target, swap_wire); + + // Route input wires. + let inputs = inputs.as_ref(); + for i in 0..WIDTH { + let in_wire = Poseidon2Gate::::wire_input(i); + let in_wire = Target::wire(gate, in_wire); + builder.connect(inputs[i], in_wire); + } + + // Collect output wires. + Self::AlgebraicPermutation::new( + (0..WIDTH).map(|i| Target::wire(gate, Poseidon2Gate::::wire_output(i))), + ) + } +} + +#[cfg(test)] +mod test { + use anyhow::Result; + use num::{BigUint, One}; + use p3_field::{AbstractField, PrimeField64 as _}; + use p3_goldilocks::Goldilocks; + use rand::{thread_rng, RngCore}; + + use super::*; + use crate::field::types::PrimeField64; + use crate::hash::hashing::hash_n_to_m_no_pad; + use crate::hash::poseidon2::p3::p3_poseidon2_hash_n_to_m_no_pad; + use crate::iop::witness::{PartialWitness, WitnessWrite}; + use crate::plonk::circuit_data::CircuitConfig; + use crate::plonk::config::PoseidonGoldilocksConfig; + + #[test] + fn test_poseidon2_with_plonky3() { + let mut rng = thread_rng(); + + let input: [u32; 12] = core::array::from_fn(|_| rng.next_u32()); + + let input_f = input + .iter() + .map(|&x| F::from_canonical_u64((x as u64) + 1073741824)) + .collect::>(); + let expected_output_f = hash_n_to_m_no_pad::>(&input_f, 12); + + let input_f3 = input + .iter() + .map(|&x| Goldilocks::from_canonical_u64((x as u64) + 1073741824)) + .collect::>(); + let expected_output_f3 = p3_poseidon2_hash_n_to_m_no_pad(&input_f3, 12); + + for i in 0..4 { + assert_eq!( + expected_output_f[i].to_canonical_u64(), + expected_output_f3[i].as_canonical_u64() + ); + } + } + + #[test] + fn test_poseidon2_gate() -> Result<()> { + let mut rng = thread_rng(); + + let input: [u32; 12] = core::array::from_fn(|_| rng.next_u32()); + let input_f = input + .iter() + .map(|&x| F::from_canonical_u64((x as u64) + 1073741824)) + .collect::>(); + + let expected_output = hash_n_to_m_no_pad::>(&input_f[0..8], 4); + + let mut builder = CircuitBuilder::::new(CircuitConfig::standard_recursion_config()); + + let input_target: [Target; 12] = input_f + .iter() + .map(|&x| builder.constant(x)) + .collect::>() + .try_into() + .unwrap(); + let output_target = + builder.hash_n_to_m_no_pad::(input_target[0..8].to_vec(), 4); + + let expected_output_target = builder.add_virtual_target_arr::<4>(); + for i in 0..4 { + builder.connect(expected_output_target[i], output_target[i]); + } + + let circuit = builder.build::(); + let mut pw = PartialWitness::new(); + pw.set_target_arr(&expected_output_target, &expected_output)?; + + let proof = circuit.prove(pw).unwrap(); + circuit.verify(proof.clone()) + } + + #[test] + fn test_poseidon2_gate_big() -> Result<()> { + let input_f: [F; 12] = + core::array::from_fn(|_| F::from_noncanonical_biguint(F::order() - BigUint::one())); + + let expected_output = hash_n_to_m_no_pad::>(&input_f[0..8], 4); + + let mut builder = CircuitBuilder::::new(CircuitConfig::standard_recursion_config()); + + let input_target: [Target; 12] = input_f + .iter() + .map(|&x| builder.constant(x)) + .collect::>() + .try_into() + .unwrap(); + let output_target = + builder.hash_n_to_m_no_pad::(input_target[0..8].to_vec(), 4); + + let expected_output_target = builder.add_virtual_target_arr::<4>(); + for i in 0..4 { + builder.connect(expected_output_target[i], output_target[i]); + } + + let circuit = builder.build::(); + let mut pw = PartialWitness::new(); + pw.set_target_arr(&expected_output_target, &expected_output)?; + + let proof = circuit.prove(pw).unwrap(); + circuit.verify(proof.clone()) + } +} diff --git a/plonky2/src/hash/poseidon2/mod.rs b/plonky2/src/hash/poseidon2/mod.rs new file mode 100644 index 000000000..3abdc563c --- /dev/null +++ b/plonky2/src/hash/poseidon2/mod.rs @@ -0,0 +1,7 @@ +pub mod config; +pub mod gate; +pub mod hash; + +#[cfg(test)] +pub mod p3; +pub mod pure; diff --git a/plonky2/src/hash/poseidon2/p3.rs b/plonky2/src/hash/poseidon2/p3.rs new file mode 100644 index 000000000..345bce362 --- /dev/null +++ b/plonky2/src/hash/poseidon2/p3.rs @@ -0,0 +1,67 @@ +use p3_field::AbstractField; +use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks}; +use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral}; +use p3_symmetric::Permutation; + +use super::config::*; + +// Poseidon2 from plonky3 +pub fn p3_poseidon2_hash_n_to_m_no_pad( + inputs: &[Goldilocks], + num_outputs: usize, +) -> Vec { + let external_linear_layer = Poseidon2ExternalMatrixGeneral; + let internal_linear_layer = DiffusionMatrixGoldilocks; + + let external_constants = EXTERNAL_CONSTANTS + .iter() + .map(|v| { + v.iter() + .map(|&x| Goldilocks::from_canonical_u64(x)) + .collect::>() + .try_into() + .unwrap() + }) + .collect::>(); + + let internal_constants = INTERNAL_CONSTANTS + .iter() + .map(|&x| Goldilocks::from_canonical_u64(x)) + .collect::>(); + + let poseidon = Poseidon2::< + Goldilocks, + Poseidon2ExternalMatrixGeneral, + DiffusionMatrixGoldilocks, + WIDTH, + D, + >::new( + ROUNDS_F, + external_constants, + external_linear_layer, + ROUNDS_P, + internal_constants, + internal_linear_layer, + ); + + let mut perm = [Goldilocks::zero(); WIDTH]; + + #[allow(clippy::manual_memcpy)] + for input_chunk in inputs.chunks(RATE) { + for i in 0..RATE.min(input_chunk.len()) { + perm[i] = input_chunk[i]; + } + poseidon.permute_mut(&mut perm); + } + + let mut outputs: Vec = Vec::new(); + loop { + for &item in perm[0..RATE].iter() { + outputs.push(item); + if outputs.len() == num_outputs { + return outputs; + } + } + poseidon.permute_mut(&mut perm); + } +} diff --git a/plonky2/src/hash/poseidon2/pure.rs b/plonky2/src/hash/poseidon2/pure.rs new file mode 100644 index 000000000..11ec03cec --- /dev/null +++ b/plonky2/src/hash/poseidon2/pure.rs @@ -0,0 +1,144 @@ +#![allow(clippy::all)] + +use p3_goldilocks::Goldilocks; +use plonky2_field::goldilocks_field::GoldilocksField; + +/// This is a test implementation of the Poseidon2 permutation without any gate or generator optimizations. +use crate::field::types::Field; +use crate::iop::target::Target; +use crate::plonk::circuit_builder::CircuitBuilder; + +use super::config::*; + +type F = GoldilocksField; +type Builder = CircuitBuilder; + +pub fn permute_swapped(builder: &mut Builder, inputs: &[Target; WIDTH]) -> [Target; WIDTH] { + let mut state: [Target; 12] = inputs.clone(); + external_permute_mut(builder, &mut state); + + // The first half of the external rounds. + for r in 0..ROUNDS_F_HALF { + let external_constants = EXTERNAL_CONSTANTS[r]; + let external_constants: [F; WIDTH] = external_constants + .iter() + .map(|&x| F::from_canonical_u64(x)) + .collect::>() + .try_into() + .unwrap(); + let external_constants: [Target; WIDTH] = + builder.constants(&external_constants).try_into().unwrap(); + add_rc(builder, &mut state, &external_constants); + sbox(builder, &mut state); + external_permute_mut(builder, &mut state); + } + + // The internal rounds. + for r in 0..ROUNDS_P { + let internal_constant = INTERNAL_CONSTANTS[r]; + let internal_constant = F::from_canonical_u64(internal_constant); + let internal_constant = builder.constant(internal_constant); + state[0] = builder.add(state[0], internal_constant); + state[0] = sbox_p(builder, state[0]); + internal_permute_mut(builder, &mut state); + } + + // The second half of the external rounds. + for r in ROUNDS_F_HALF..ROUNDS_F { + let external_constants = EXTERNAL_CONSTANTS[r]; + let external_constants: [F; WIDTH] = external_constants + .iter() + .map(|&x| F::from_canonical_u64(x)) + .collect::>() + .try_into() + .unwrap(); + let external_constants: [Target; WIDTH] = + builder.constants(&external_constants).try_into().unwrap(); + add_rc(builder, &mut state, &external_constants); + sbox(builder, &mut state); + external_permute_mut(builder, &mut state); + } + + state +} + +fn external_permute_mut(builder: &mut Builder, state: &mut [Target; 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) { + // Would be nice to find a better way to do this. + let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; + apply_mat4(builder, &mut state_4); + state[i..i + 4].clone_from_slice(&state_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: [Target; 4] = core::array::from_fn(|k| { + (0..WIDTH) + .step_by(4) + .map(|j| state[j + k]) + .reduce(|acc, t| builder.add(acc, t)) + .unwrap() + }); + + // 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] = builder.add(state[i], sums[i % 4]); + } +} + +fn add_rc(builder: &mut Builder, state: &mut [Target; WIDTH], external_constant: &[Target; WIDTH]) { + state + .iter_mut() + .zip(external_constant) + .for_each(|(a, b)| *a = builder.add(*a, *b)); +} + +fn sbox(builder: &mut Builder, state: &mut [Target; WIDTH]) { + state.iter_mut().for_each(|a| *a = sbox_p(builder, *a)); +} + +fn sbox_p(builder: &mut Builder, a: Target) -> Target { + let a2 = builder.mul(a, a); + let a3 = builder.mul(a2, a); + let a6 = builder.mul(a3, a3); + + builder.mul(a6, a) +} + +// Multiply a 4-element vector x by: +// [ 2 3 1 1 ] +// [ 1 2 3 1 ] +// [ 1 1 2 3 ] +// [ 3 1 1 2 ]. +// This is more efficient than the previous matrix. +fn apply_mat4(builder: &mut Builder, x: &mut [Target; 4]) { + let two = builder.constant(F::from_canonical_u64(2)); + + let t01 = builder.add(x[0], x[1]); + let t23 = builder.add(x[2], x[3]); + let t0123 = builder.add(t01, t23); + let t01123 = builder.add(t0123, x[1]); + let t01233 = builder.add(t0123, x[3]); + // The order here is important. Need to overwrite x[0] and x[2] after x[1] and x[3]. + let dx0 = builder.mul(x[0], two); + let dx2 = builder.mul(x[2], two); + x[3] = builder.add(t01233, dx0); // 3*x[0] + x[1] + x[2] + 2*x[3] + x[1] = builder.add(t01123, dx2); // x[0] + 2*x[1] + 3*x[2] + x[3] + x[0] = builder.add(t01123, t01); // 2*x[0] + 3*x[1] + x[2] + x[3] + x[2] = builder.add(t01233, t23); // x[0] + x[1] + 2*x[2] + 3*x[3] +} + +/// Given a vector v compute the matrix vector product (1 + diag(v))state with 1 denoting the constant matrix of ones. +pub fn internal_permute_mut(builder: &mut Builder, state: &mut [Target; WIDTH]) { + let sum = builder.add_many(state.iter().cloned()); + for i in 0..WIDTH { + let constant = MATRIX_DIAG_12_U64[i]; + let constant = F::from_canonical_u64(constant); + let constant = builder.constant(constant); + state[i] = builder.mul(state[i], constant); + state[i] = builder.add(state[i], sum); + } +} diff --git a/plonky2/src/plonk/config.rs b/plonky2/src/plonk/config.rs index 217c88976..9504ba3ca 100644 --- a/plonky2/src/plonk/config.rs +++ b/plonky2/src/plonk/config.rs @@ -20,6 +20,7 @@ use crate::hash::hash_types::{HashOut, RichField}; use crate::hash::hashing::PlonkyPermutation; use crate::hash::keccak::KeccakHash; use crate::hash::poseidon::PoseidonHash; +use crate::hash::poseidon2::hash::Poseidon2Hash; use crate::iop::target::{BoolTarget, Target}; use crate::plonk::circuit_builder::CircuitBuilder; @@ -111,8 +112,8 @@ pub struct PoseidonGoldilocksConfig; impl GenericConfig<2> for PoseidonGoldilocksConfig { type F = GoldilocksField; type FE = QuadraticExtension; - type Hasher = PoseidonHash; - type InnerHasher = PoseidonHash; + type Hasher = Poseidon2Hash; + type InnerHasher = Poseidon2Hash; } /// Configuration using truncated Keccak over the Goldilocks field. diff --git a/plonky2/src/recursion/cyclic_recursion.rs b/plonky2/src/recursion/cyclic_recursion.rs index df0fb95cd..83e711395 100644 --- a/plonky2/src/recursion/cyclic_recursion.rs +++ b/plonky2/src/recursion/cyclic_recursion.rs @@ -249,124 +249,124 @@ mod tests { builder.build::().common } - /// Uses cyclic recursion to build a hash chain. - /// The circuit has the following public input structure: - /// - Initial hash (4) - /// - Output for the tip of the hash chain (4) - /// - Chain length, i.e. the number of times the hash has been applied (1) - /// - VK for cyclic recursion (?) - #[test] - fn test_cyclic_recursion() -> Result<()> { - const D: usize = 2; - type C = PoseidonGoldilocksConfig; - type F = >::F; - - let config = CircuitConfig::standard_recursion_config(); - let mut builder = CircuitBuilder::::new(config); - let one = builder.one(); - - // Circuit that computes a repeated hash. - let initial_hash_target = builder.add_virtual_hash(); - builder.register_public_inputs(&initial_hash_target.elements); - let current_hash_in = builder.add_virtual_hash(); - let current_hash_out = - builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); - builder.register_public_inputs(¤t_hash_out.elements); - let counter = builder.add_virtual_public_input(); - - let mut common_data = common_data_for_recursion::(); - let verifier_data_target = builder.add_verifier_data_public_inputs(); - common_data.num_public_inputs = builder.num_public_inputs(); - - let condition = builder.add_virtual_bool_target_safe(); - - // Unpack inner proof's public inputs. - let inner_cyclic_proof_with_pis = builder.add_virtual_proof_with_pis(&common_data); - let inner_cyclic_pis = &inner_cyclic_proof_with_pis.public_inputs; - let inner_cyclic_initial_hash = HashOutTarget::try_from(&inner_cyclic_pis[0..4]).unwrap(); - let inner_cyclic_latest_hash = HashOutTarget::try_from(&inner_cyclic_pis[4..8]).unwrap(); - let inner_cyclic_counter = inner_cyclic_pis[8]; - - // Connect our initial hash to that of our inner proof. (If there is no inner proof, the - // initial hash will be unconstrained, which is intentional.) - builder.connect_hashes(initial_hash_target, inner_cyclic_initial_hash); - - // The input hash is the previous hash output if we have an inner proof, or the initial hash - // if this is the base case. - let actual_hash_in = - builder.select_hash(condition, inner_cyclic_latest_hash, initial_hash_target); - builder.connect_hashes(current_hash_in, actual_hash_in); - - // Our chain length will be inner_counter + 1 if we have an inner proof, or 1 if not. - let new_counter = builder.mul_add(condition.target, inner_cyclic_counter, one); - builder.connect(counter, new_counter); - - builder.conditionally_verify_cyclic_proof_or_dummy::( - condition, - &inner_cyclic_proof_with_pis, - &common_data, - )?; - - let cyclic_circuit_data = builder.build::(); - - let mut pw = PartialWitness::new(); - let initial_hash = [F::ZERO, F::ONE, F::TWO, F::from_canonical_usize(3)]; - let initial_hash_pis = initial_hash.into_iter().enumerate().collect(); - pw.set_bool_target(condition, false)?; - pw.set_proof_with_pis_target::( - &inner_cyclic_proof_with_pis, - &cyclic_base_proof( - &common_data, - &cyclic_circuit_data.verifier_only, - initial_hash_pis, - ), - )?; - pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; - let proof = cyclic_circuit_data.prove(pw)?; - check_cyclic_proof_verifier_data( - &proof, - &cyclic_circuit_data.verifier_only, - &cyclic_circuit_data.common, - )?; - cyclic_circuit_data.verify(proof.clone())?; - - // 1st recursive layer. - let mut pw = PartialWitness::new(); - pw.set_bool_target(condition, true)?; - pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; - pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; - let proof = cyclic_circuit_data.prove(pw)?; - check_cyclic_proof_verifier_data( - &proof, - &cyclic_circuit_data.verifier_only, - &cyclic_circuit_data.common, - )?; - cyclic_circuit_data.verify(proof.clone())?; - - // 2nd recursive layer. - let mut pw = PartialWitness::new(); - pw.set_bool_target(condition, true)?; - pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; - pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; - let proof = cyclic_circuit_data.prove(pw)?; - check_cyclic_proof_verifier_data( - &proof, - &cyclic_circuit_data.verifier_only, - &cyclic_circuit_data.common, - )?; - - // Verify that the proof correctly computes a repeated hash. - let initial_hash = &proof.public_inputs[..4]; - let hash = &proof.public_inputs[4..8]; - let counter = proof.public_inputs[8]; - let expected_hash: [F; 4] = iterate_poseidon( - initial_hash.try_into().unwrap(), - counter.to_canonical_u64() as usize, - ); - assert_eq!(hash, expected_hash); - - cyclic_circuit_data.verify(proof) - } + // /// Uses cyclic recursion to build a hash chain. + // /// The circuit has the following public input structure: + // /// - Initial hash (4) + // /// - Output for the tip of the hash chain (4) + // /// - Chain length, i.e. the number of times the hash has been applied (1) + // /// - VK for cyclic recursion (?) + // #[test] + // fn test_cyclic_recursion() -> Result<()> { + // const D: usize = 2; + // type C = PoseidonGoldilocksConfig; + // type F = >::F; + + // let config = CircuitConfig::standard_recursion_config(); + // let mut builder = CircuitBuilder::::new(config); + // let one = builder.one(); + + // // Circuit that computes a repeated hash. + // let initial_hash_target = builder.add_virtual_hash(); + // builder.register_public_inputs(&initial_hash_target.elements); + // let current_hash_in = builder.add_virtual_hash(); + // let current_hash_out = + // builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); + // builder.register_public_inputs(¤t_hash_out.elements); + // let counter = builder.add_virtual_public_input(); + + // let mut common_data = common_data_for_recursion::(); + // let verifier_data_target = builder.add_verifier_data_public_inputs(); + // common_data.num_public_inputs = builder.num_public_inputs(); + + // let condition = builder.add_virtual_bool_target_safe(); + + // // Unpack inner proof's public inputs. + // let inner_cyclic_proof_with_pis = builder.add_virtual_proof_with_pis(&common_data); + // let inner_cyclic_pis = &inner_cyclic_proof_with_pis.public_inputs; + // let inner_cyclic_initial_hash = HashOutTarget::try_from(&inner_cyclic_pis[0..4]).unwrap(); + // let inner_cyclic_latest_hash = HashOutTarget::try_from(&inner_cyclic_pis[4..8]).unwrap(); + // let inner_cyclic_counter = inner_cyclic_pis[8]; + + // // Connect our initial hash to that of our inner proof. (If there is no inner proof, the + // // initial hash will be unconstrained, which is intentional.) + // builder.connect_hashes(initial_hash_target, inner_cyclic_initial_hash); + + // // The input hash is the previous hash output if we have an inner proof, or the initial hash + // // if this is the base case. + // let actual_hash_in = + // builder.select_hash(condition, inner_cyclic_latest_hash, initial_hash_target); + // builder.connect_hashes(current_hash_in, actual_hash_in); + + // // Our chain length will be inner_counter + 1 if we have an inner proof, or 1 if not. + // let new_counter = builder.mul_add(condition.target, inner_cyclic_counter, one); + // builder.connect(counter, new_counter); + + // builder.conditionally_verify_cyclic_proof_or_dummy::( + // condition, + // &inner_cyclic_proof_with_pis, + // &common_data, + // )?; + + // let cyclic_circuit_data = builder.build::(); + + // let mut pw = PartialWitness::new(); + // let initial_hash = [F::ZERO, F::ONE, F::TWO, F::from_canonical_usize(3)]; + // let initial_hash_pis = initial_hash.into_iter().enumerate().collect(); + // pw.set_bool_target(condition, false)?; + // pw.set_proof_with_pis_target::( + // &inner_cyclic_proof_with_pis, + // &cyclic_base_proof( + // &common_data, + // &cyclic_circuit_data.verifier_only, + // initial_hash_pis, + // ), + // )?; + // pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; + // let proof = cyclic_circuit_data.prove(pw)?; + // check_cyclic_proof_verifier_data( + // &proof, + // &cyclic_circuit_data.verifier_only, + // &cyclic_circuit_data.common, + // )?; + // cyclic_circuit_data.verify(proof.clone())?; + + // // 1st recursive layer. + // let mut pw = PartialWitness::new(); + // pw.set_bool_target(condition, true)?; + // pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; + // pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; + // let proof = cyclic_circuit_data.prove(pw)?; + // check_cyclic_proof_verifier_data( + // &proof, + // &cyclic_circuit_data.verifier_only, + // &cyclic_circuit_data.common, + // )?; + // cyclic_circuit_data.verify(proof.clone())?; + + // // 2nd recursive layer. + // let mut pw = PartialWitness::new(); + // pw.set_bool_target(condition, true)?; + // pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; + // pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; + // let proof = cyclic_circuit_data.prove(pw)?; + // check_cyclic_proof_verifier_data( + // &proof, + // &cyclic_circuit_data.verifier_only, + // &cyclic_circuit_data.common, + // )?; + + // // Verify that the proof correctly computes a repeated hash. + // let initial_hash = &proof.public_inputs[..4]; + // let hash = &proof.public_inputs[4..8]; + // let counter = proof.public_inputs[8]; + // let expected_hash: [F; 4] = iterate_poseidon( + // initial_hash.try_into().unwrap(), + // counter.to_canonical_u64() as usize, + // ); + // assert_eq!(hash, expected_hash); + + // cyclic_circuit_data.verify(proof) + // } fn iterate_poseidon(initial_state: [F; 4], n: usize) -> [F; 4] { let mut current = initial_state; From 6e2b7df5172def309b5d208dd1ecf48c6789ec8b Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Mon, 10 Nov 2025 16:22:22 -0500 Subject: [PATCH 02/26] wip --- plonky2/src/hash/poseidon2/pure.rs | 1 - plonky2/src/recursion/cyclic_recursion.rs | 236 +++++++++++----------- plonky2/src/recursion/dummy_circuit.rs | 16 ++ 3 files changed, 134 insertions(+), 119 deletions(-) diff --git a/plonky2/src/hash/poseidon2/pure.rs b/plonky2/src/hash/poseidon2/pure.rs index 11ec03cec..2d24237e9 100644 --- a/plonky2/src/hash/poseidon2/pure.rs +++ b/plonky2/src/hash/poseidon2/pure.rs @@ -1,6 +1,5 @@ #![allow(clippy::all)] -use p3_goldilocks::Goldilocks; use plonky2_field::goldilocks_field::GoldilocksField; /// This is a test implementation of the Poseidon2 permutation without any gate or generator optimizations. diff --git a/plonky2/src/recursion/cyclic_recursion.rs b/plonky2/src/recursion/cyclic_recursion.rs index 83e711395..df0fb95cd 100644 --- a/plonky2/src/recursion/cyclic_recursion.rs +++ b/plonky2/src/recursion/cyclic_recursion.rs @@ -249,124 +249,124 @@ mod tests { builder.build::().common } - // /// Uses cyclic recursion to build a hash chain. - // /// The circuit has the following public input structure: - // /// - Initial hash (4) - // /// - Output for the tip of the hash chain (4) - // /// - Chain length, i.e. the number of times the hash has been applied (1) - // /// - VK for cyclic recursion (?) - // #[test] - // fn test_cyclic_recursion() -> Result<()> { - // const D: usize = 2; - // type C = PoseidonGoldilocksConfig; - // type F = >::F; - - // let config = CircuitConfig::standard_recursion_config(); - // let mut builder = CircuitBuilder::::new(config); - // let one = builder.one(); - - // // Circuit that computes a repeated hash. - // let initial_hash_target = builder.add_virtual_hash(); - // builder.register_public_inputs(&initial_hash_target.elements); - // let current_hash_in = builder.add_virtual_hash(); - // let current_hash_out = - // builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); - // builder.register_public_inputs(¤t_hash_out.elements); - // let counter = builder.add_virtual_public_input(); - - // let mut common_data = common_data_for_recursion::(); - // let verifier_data_target = builder.add_verifier_data_public_inputs(); - // common_data.num_public_inputs = builder.num_public_inputs(); - - // let condition = builder.add_virtual_bool_target_safe(); - - // // Unpack inner proof's public inputs. - // let inner_cyclic_proof_with_pis = builder.add_virtual_proof_with_pis(&common_data); - // let inner_cyclic_pis = &inner_cyclic_proof_with_pis.public_inputs; - // let inner_cyclic_initial_hash = HashOutTarget::try_from(&inner_cyclic_pis[0..4]).unwrap(); - // let inner_cyclic_latest_hash = HashOutTarget::try_from(&inner_cyclic_pis[4..8]).unwrap(); - // let inner_cyclic_counter = inner_cyclic_pis[8]; - - // // Connect our initial hash to that of our inner proof. (If there is no inner proof, the - // // initial hash will be unconstrained, which is intentional.) - // builder.connect_hashes(initial_hash_target, inner_cyclic_initial_hash); - - // // The input hash is the previous hash output if we have an inner proof, or the initial hash - // // if this is the base case. - // let actual_hash_in = - // builder.select_hash(condition, inner_cyclic_latest_hash, initial_hash_target); - // builder.connect_hashes(current_hash_in, actual_hash_in); - - // // Our chain length will be inner_counter + 1 if we have an inner proof, or 1 if not. - // let new_counter = builder.mul_add(condition.target, inner_cyclic_counter, one); - // builder.connect(counter, new_counter); - - // builder.conditionally_verify_cyclic_proof_or_dummy::( - // condition, - // &inner_cyclic_proof_with_pis, - // &common_data, - // )?; - - // let cyclic_circuit_data = builder.build::(); - - // let mut pw = PartialWitness::new(); - // let initial_hash = [F::ZERO, F::ONE, F::TWO, F::from_canonical_usize(3)]; - // let initial_hash_pis = initial_hash.into_iter().enumerate().collect(); - // pw.set_bool_target(condition, false)?; - // pw.set_proof_with_pis_target::( - // &inner_cyclic_proof_with_pis, - // &cyclic_base_proof( - // &common_data, - // &cyclic_circuit_data.verifier_only, - // initial_hash_pis, - // ), - // )?; - // pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; - // let proof = cyclic_circuit_data.prove(pw)?; - // check_cyclic_proof_verifier_data( - // &proof, - // &cyclic_circuit_data.verifier_only, - // &cyclic_circuit_data.common, - // )?; - // cyclic_circuit_data.verify(proof.clone())?; - - // // 1st recursive layer. - // let mut pw = PartialWitness::new(); - // pw.set_bool_target(condition, true)?; - // pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; - // pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; - // let proof = cyclic_circuit_data.prove(pw)?; - // check_cyclic_proof_verifier_data( - // &proof, - // &cyclic_circuit_data.verifier_only, - // &cyclic_circuit_data.common, - // )?; - // cyclic_circuit_data.verify(proof.clone())?; - - // // 2nd recursive layer. - // let mut pw = PartialWitness::new(); - // pw.set_bool_target(condition, true)?; - // pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; - // pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; - // let proof = cyclic_circuit_data.prove(pw)?; - // check_cyclic_proof_verifier_data( - // &proof, - // &cyclic_circuit_data.verifier_only, - // &cyclic_circuit_data.common, - // )?; - - // // Verify that the proof correctly computes a repeated hash. - // let initial_hash = &proof.public_inputs[..4]; - // let hash = &proof.public_inputs[4..8]; - // let counter = proof.public_inputs[8]; - // let expected_hash: [F; 4] = iterate_poseidon( - // initial_hash.try_into().unwrap(), - // counter.to_canonical_u64() as usize, - // ); - // assert_eq!(hash, expected_hash); - - // cyclic_circuit_data.verify(proof) - // } + /// Uses cyclic recursion to build a hash chain. + /// The circuit has the following public input structure: + /// - Initial hash (4) + /// - Output for the tip of the hash chain (4) + /// - Chain length, i.e. the number of times the hash has been applied (1) + /// - VK for cyclic recursion (?) + #[test] + fn test_cyclic_recursion() -> Result<()> { + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + let one = builder.one(); + + // Circuit that computes a repeated hash. + let initial_hash_target = builder.add_virtual_hash(); + builder.register_public_inputs(&initial_hash_target.elements); + let current_hash_in = builder.add_virtual_hash(); + let current_hash_out = + builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); + builder.register_public_inputs(¤t_hash_out.elements); + let counter = builder.add_virtual_public_input(); + + let mut common_data = common_data_for_recursion::(); + let verifier_data_target = builder.add_verifier_data_public_inputs(); + common_data.num_public_inputs = builder.num_public_inputs(); + + let condition = builder.add_virtual_bool_target_safe(); + + // Unpack inner proof's public inputs. + let inner_cyclic_proof_with_pis = builder.add_virtual_proof_with_pis(&common_data); + let inner_cyclic_pis = &inner_cyclic_proof_with_pis.public_inputs; + let inner_cyclic_initial_hash = HashOutTarget::try_from(&inner_cyclic_pis[0..4]).unwrap(); + let inner_cyclic_latest_hash = HashOutTarget::try_from(&inner_cyclic_pis[4..8]).unwrap(); + let inner_cyclic_counter = inner_cyclic_pis[8]; + + // Connect our initial hash to that of our inner proof. (If there is no inner proof, the + // initial hash will be unconstrained, which is intentional.) + builder.connect_hashes(initial_hash_target, inner_cyclic_initial_hash); + + // The input hash is the previous hash output if we have an inner proof, or the initial hash + // if this is the base case. + let actual_hash_in = + builder.select_hash(condition, inner_cyclic_latest_hash, initial_hash_target); + builder.connect_hashes(current_hash_in, actual_hash_in); + + // Our chain length will be inner_counter + 1 if we have an inner proof, or 1 if not. + let new_counter = builder.mul_add(condition.target, inner_cyclic_counter, one); + builder.connect(counter, new_counter); + + builder.conditionally_verify_cyclic_proof_or_dummy::( + condition, + &inner_cyclic_proof_with_pis, + &common_data, + )?; + + let cyclic_circuit_data = builder.build::(); + + let mut pw = PartialWitness::new(); + let initial_hash = [F::ZERO, F::ONE, F::TWO, F::from_canonical_usize(3)]; + let initial_hash_pis = initial_hash.into_iter().enumerate().collect(); + pw.set_bool_target(condition, false)?; + pw.set_proof_with_pis_target::( + &inner_cyclic_proof_with_pis, + &cyclic_base_proof( + &common_data, + &cyclic_circuit_data.verifier_only, + initial_hash_pis, + ), + )?; + pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; + let proof = cyclic_circuit_data.prove(pw)?; + check_cyclic_proof_verifier_data( + &proof, + &cyclic_circuit_data.verifier_only, + &cyclic_circuit_data.common, + )?; + cyclic_circuit_data.verify(proof.clone())?; + + // 1st recursive layer. + let mut pw = PartialWitness::new(); + pw.set_bool_target(condition, true)?; + pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; + pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; + let proof = cyclic_circuit_data.prove(pw)?; + check_cyclic_proof_verifier_data( + &proof, + &cyclic_circuit_data.verifier_only, + &cyclic_circuit_data.common, + )?; + cyclic_circuit_data.verify(proof.clone())?; + + // 2nd recursive layer. + let mut pw = PartialWitness::new(); + pw.set_bool_target(condition, true)?; + pw.set_proof_with_pis_target(&inner_cyclic_proof_with_pis, &proof)?; + pw.set_verifier_data_target(&verifier_data_target, &cyclic_circuit_data.verifier_only)?; + let proof = cyclic_circuit_data.prove(pw)?; + check_cyclic_proof_verifier_data( + &proof, + &cyclic_circuit_data.verifier_only, + &cyclic_circuit_data.common, + )?; + + // Verify that the proof correctly computes a repeated hash. + let initial_hash = &proof.public_inputs[..4]; + let hash = &proof.public_inputs[4..8]; + let counter = proof.public_inputs[8]; + let expected_hash: [F; 4] = iterate_poseidon( + initial_hash.try_into().unwrap(), + counter.to_canonical_u64() as usize, + ); + assert_eq!(hash, expected_hash); + + cyclic_circuit_data.verify(proof) + } fn iterate_poseidon(initial_state: [F; 4], n: usize) -> [F; 4] { let mut current = initial_state; diff --git a/plonky2/src/recursion/dummy_circuit.rs b/plonky2/src/recursion/dummy_circuit.rs index cd8559382..6b08b0ce8 100644 --- a/plonky2/src/recursion/dummy_circuit.rs +++ b/plonky2/src/recursion/dummy_circuit.rs @@ -90,6 +90,7 @@ where pub fn dummy_circuit, C: GenericConfig, const D: usize>( common_data: &CommonCircuitData, ) -> CircuitData { + println!("dummy_circuit received common_data.gates: {:?}", common_data.gates); let config = common_data.config.clone(); assert!( !common_data.config.zero_knowledge, @@ -113,6 +114,21 @@ pub fn dummy_circuit, C: GenericConfig, c } let circuit = builder.build::(); + + // Assert individual components to identify mismatches + assert_eq!(circuit.common.config, common_data.config, "config mismatch"); + assert_eq!(circuit.common.fri_params, common_data.fri_params, "fri_params mismatch"); + assert_eq!(circuit.common.gates, common_data.gates, "gates mismatch"); + assert_eq!(circuit.common.selectors_info, common_data.selectors_info, "selectors_info mismatch"); + assert_eq!(circuit.common.quotient_degree_factor, common_data.quotient_degree_factor, "quotient_degree_factor mismatch"); + assert_eq!(circuit.common.num_gate_constraints, common_data.num_gate_constraints, "num_gate_constraints mismatch"); + assert_eq!(circuit.common.num_constants, common_data.num_constants, "num_constants mismatch"); + assert_eq!(circuit.common.num_public_inputs, common_data.num_public_inputs, "num_public_inputs mismatch"); + assert_eq!(circuit.common.k_is, common_data.k_is, "k_is mismatch"); + assert_eq!(circuit.common.num_partial_products, common_data.num_partial_products, "num_partial_products mismatch"); + assert_eq!(circuit.common.num_lookup_polys, common_data.num_lookup_polys, "num_lookup_polys mismatch"); + assert_eq!(circuit.common.num_lookup_selectors, common_data.num_lookup_selectors, "num_lookup_selectors mismatch"); + assert_eq!(circuit.common.luts, common_data.luts, "luts mismatch"); assert_eq!(&circuit.common, common_data); circuit } From c099e7eafdf4668003ac16bd599bbdbfa15d7034 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 11:24:01 -0500 Subject: [PATCH 03/26] update benchmarks --- plonky2/examples/fibonacci.rs | 36 ++- plonky2/src/hash/poseidon2._rs | 342 ---------------------- plonky2/src/hash/poseidon2/gate.rs | 2 +- plonky2/src/hash/poseidon2/hash.rs | 7 +- plonky2/src/hash/poseidon2/pure.rs | 3 +- plonky2/src/plonk/config.rs | 12 + plonky2/src/recursion/cyclic_recursion.rs | 8 +- plonky2/src/recursion/dummy_circuit.rs | 52 +++- 8 files changed, 93 insertions(+), 369 deletions(-) delete mode 100644 plonky2/src/hash/poseidon2._rs diff --git a/plonky2/examples/fibonacci.rs b/plonky2/examples/fibonacci.rs index 578dc2424..7310ad164 100644 --- a/plonky2/examples/fibonacci.rs +++ b/plonky2/examples/fibonacci.rs @@ -1,14 +1,30 @@ + +use std::time::Instant; + +use env_logger; +use log::Level; use anyhow::Result; use plonky2::field::types::Field; use plonky2::iop::witness::{PartialWitness, WitnessWrite}; use plonky2::plonk::circuit_builder::CircuitBuilder; use plonky2::plonk::circuit_data::CircuitConfig; -use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; +use plonky2::plonk::config::{GenericConfig, Poseidon2GoldilocksConfig, PoseidonGoldilocksConfig}; +use plonky2::plonk::prover::prove; +use plonky2::util::timing::TimingTree; /// An example of using Plonky2 to prove a statement of the form /// "I know the 100th element of the Fibonacci sequence, starting with constants a and b." /// When a == 0 and b == 1, this is proving knowledge of the 100th (standard) Fibonacci number. fn main() -> Result<()> { + env_logger::Builder::from_default_env() + .filter_level(log::LevelFilter::Debug) + .init(); + work::()?; + work::() +} + + +fn work>() -> Result<()> { const D: usize = 2; type C = PoseidonGoldilocksConfig; type F = >::F; @@ -21,7 +37,7 @@ fn main() -> Result<()> { let initial_b = builder.add_virtual_target(); let mut prev_target = initial_a; let mut cur_target = initial_b; - for _ in 0..99 { + for _ in 0..9999 { let temp = builder.add(prev_target, cur_target); prev_target = cur_target; cur_target = temp; @@ -33,17 +49,29 @@ fn main() -> Result<()> { builder.register_public_input(cur_target); // Provide initial values. + let timer1 = Instant::now(); let mut pw = PartialWitness::new(); pw.set_target(initial_a, F::ZERO)?; pw.set_target(initial_b, F::ONE)?; let data = builder.build::(); - let proof = data.prove(pw)?; + let timer2 = Instant::now(); + + // Create a TimingTree to track detailed timing information + let mut timing = TimingTree::new("prove", Level::Debug); + let proof = prove::(&data.prover_only, &data.common, pw, &mut timing)?; + let timer3 = Instant::now(); + // Print the timing tree + timing.print(); + println!( "100th Fibonacci number mod |F| (starting with {}, {}) is: {}", proof.public_inputs[0], proof.public_inputs[1], proof.public_inputs[2] ); + println!("Build time: {:?}", timer2.duration_since(timer1)); + println!("Prove time: {:?}", timer3.duration_since(timer2)); + data.verify(proof) -} +} \ No newline at end of file diff --git a/plonky2/src/hash/poseidon2._rs b/plonky2/src/hash/poseidon2._rs deleted file mode 100644 index 9867a2f64..000000000 --- a/plonky2/src/hash/poseidon2._rs +++ /dev/null @@ -1,342 +0,0 @@ -//! Poseidon2 hash function implementation using Plonky3's Poseidon2. -//! -//! This module provides a wrapper around Plonky3's Poseidon2 implementation, -//! making it compatible with Plonky2's hash interface. - -#[cfg(not(feature = "std"))] -use alloc::vec::Vec; -use core::fmt::Debug; - -use p3_field::{PrimeCharacteristicRing, PrimeField64}; -use p3_goldilocks::{ - Goldilocks as GoldilocksP3, Poseidon2GoldilocksHL as Poseidon2P3Impl, - HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS, HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS, -}; -use p3_poseidon2::{ExternalLayerConstants, Poseidon2}; -use p3_symmetric::Permutation; - -use crate::field::goldilocks_field::GoldilocksField; -use crate::field::types::{Field, PrimeField64 as Plonky2PrimeField64}; -use crate::hash::hash_types::{HashOut, RichField, NUM_HASH_OUT_ELTS}; -use crate::hash::hashing::{compress, hash_n_to_hash_no_pad, PlonkyPermutation}; -use crate::iop::target::Target; -use crate::plonk::config::{AlgebraicHasher, Hasher}; - -/// Width of the Poseidon2 permutation (rate + capacity) -pub const POSEIDON2_WIDTH: usize = 12; -/// Rate of the sponge construction -pub const POSEIDON2_RATE: usize = 8; -/// Number of rounds for Poseidon2 (D parameter for x^D S-box) -pub const POSEIDON2_D: u64 = 7; - -/// Wrapper around Plonky3's Poseidon2 for Goldilocks field -pub struct Poseidon2Goldilocks { - /// The underlying Plonky3 Poseidon2 instance - permutation: Poseidon2P3Impl, -} - -// eraseamusegavelreign - -impl Debug for Poseidon2Goldilocks { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - f.debug_struct("Poseidon2Goldilocks").finish() - } -} - -impl Default for Poseidon2Goldilocks { - fn default() -> Self { - Self::new() - } -} - -impl Poseidon2Goldilocks { - /// Create a new Poseidon2 instance with default parameters - /// Uses Plonky3's Poseidon2 implementation for Goldilocks with precomputed constants - pub fn new() -> Self { - // Use Plonky3's Poseidon2 for Goldilocks with precomputed constants - // Convert the precomputed constants to the proper format - - // External round constants - convert 2D array of u64 to Vec<[Goldilocks; 8]> - let external_initial: Vec<[GoldilocksP3; 8]> = HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS[0] - .iter() - .map(|&round| round.map(GoldilocksP3::from_u64)) - .collect(); - let external_terminal: Vec<[GoldilocksP3; 8]> = HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS[1] - .iter() - .map(|&round| round.map(GoldilocksP3::from_u64)) - .collect(); - - // Internal round constants - convert Vec to Vec - let internal_constants: Vec = HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS - .iter() - .map(|&c| GoldilocksP3::from_u64(c)) - .collect(); - - let external_constants = ExternalLayerConstants::new(external_initial, external_terminal); - let permutation = Poseidon2::new(external_constants, internal_constants); - Self { permutation } - } - - /// Permute the state using Poseidon2 - pub fn permute(&self, state: [GoldilocksField; POSEIDON2_WIDTH]) -> [GoldilocksField; POSEIDON2_WIDTH] { - // Convert plonky2 fields to plonky3 fields - let p3_state: [GoldilocksP3; POSEIDON2_WIDTH] = state.map(|f| { - let val = Plonky2PrimeField64::to_canonical_u64(&f); - GoldilocksP3::from_u64(val) - }); - - // Apply permutation - let output = self.permutation.permute(p3_state); - - // Convert back to plonky2 fields - output.map(|f| GoldilocksField::from_canonical_u64(PrimeField64::as_canonical_u64(&f))) - } -} - -/// Permutation for Poseidon2 -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub struct Poseidon2Permutation([T; POSEIDON2_WIDTH]); - -impl Default for Poseidon2Permutation { - fn default() -> Self { - Self([T::default(); POSEIDON2_WIDTH]) - } -} - -impl AsRef<[T]> for Poseidon2Permutation { - fn as_ref(&self) -> &[T] { - &self.0 - } -} - -impl PlonkyPermutation for Poseidon2Permutation { - const RATE: usize = POSEIDON2_RATE; - const WIDTH: usize = POSEIDON2_WIDTH; - - fn new>(iter: I) -> Self { - let mut state = [GoldilocksField::ZERO; POSEIDON2_WIDTH]; - for (i, item) in iter.into_iter().take(POSEIDON2_WIDTH).enumerate() { - state[i] = item; - } - Self(state) - } - - fn set_elt(&mut self, elt: GoldilocksField, idx: usize) { - assert!(idx < POSEIDON2_WIDTH); - self.0[idx] = elt; - } - - fn set_from_iter>(&mut self, elts: I, start_idx: usize) { - assert!(start_idx <= POSEIDON2_WIDTH); - for (i, elt) in elts.into_iter().take(POSEIDON2_WIDTH - start_idx).enumerate() { - self.0[start_idx + i] = elt; - } - } - - fn set_from_slice(&mut self, elts: &[GoldilocksField], start_idx: usize) { - assert!(start_idx <= POSEIDON2_WIDTH); - let len = core::cmp::min(elts.len(), POSEIDON2_WIDTH - start_idx); - self.0[start_idx..start_idx + len].copy_from_slice(&elts[..len]); - } - - fn permute(&mut self) { - let poseidon2 = Poseidon2Goldilocks::new(); - self.0 = poseidon2.permute(self.0); - } - - fn squeeze(&self) -> &[GoldilocksField] { - &self.0[..POSEIDON2_RATE] - } -} - -impl PlonkyPermutation for Poseidon2Permutation { - const RATE: usize = POSEIDON2_RATE; - const WIDTH: usize = POSEIDON2_WIDTH; - - fn new>(iter: I) -> Self { - let mut state = [Target::default(); POSEIDON2_WIDTH]; - for (i, item) in iter.into_iter().take(POSEIDON2_WIDTH).enumerate() { - state[i] = item; - } - Self(state) - } - - fn set_elt(&mut self, elt: Target, idx: usize) { - assert!(idx < POSEIDON2_WIDTH); - self.0[idx] = elt; - } - - fn set_from_iter>(&mut self, elts: I, start_idx: usize) { - assert!(start_idx <= POSEIDON2_WIDTH); - for (i, elt) in elts.into_iter().take(POSEIDON2_WIDTH - start_idx).enumerate() { - self.0[start_idx + i] = elt; - } - } - - fn set_from_slice(&mut self, elts: &[Target], start_idx: usize) { - assert!(start_idx <= POSEIDON2_WIDTH); - let len = core::cmp::min(elts.len(), POSEIDON2_WIDTH - start_idx); - self.0[start_idx..start_idx + len].copy_from_slice(&elts[..len]); - } - - fn permute(&mut self) { - // For circuit targets, this is handled by the circuit builder - // The actual permutation is applied through gates - } - - fn squeeze(&self) -> &[Target] { - &self.0[..POSEIDON2_RATE] - } -} - -/// Poseidon2 hash function -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub struct Poseidon2Hash; - -impl Hasher for Poseidon2Hash { - const HASH_SIZE: usize = 4 * 8; // 4 field elements, 8 bytes each - type Hash = HashOut; - type Permutation = Poseidon2Permutation; - - fn hash_no_pad(input: &[GoldilocksField]) -> Self::Hash { - hash_n_to_hash_no_pad::(input) - } - - fn two_to_one(left: Self::Hash, right: Self::Hash) -> Self::Hash { - compress::(left, right) - } -} - -impl AlgebraicHasher for Poseidon2Hash { - type AlgebraicPermutation = Poseidon2Permutation; - - fn permute_swapped( - mut inputs: Self::AlgebraicPermutation, - swap: crate::iop::target::BoolTarget, - builder: &mut crate::plonk::circuit_builder::CircuitBuilder, - ) -> Self::AlgebraicPermutation - where - GoldilocksField: RichField + crate::field::extension::Extendable, - { - // Swap the first two chunks (first 4 elements with next 4 elements) if swap is true - // This is used for Merkle proof verification - let chunk_size = NUM_HASH_OUT_ELTS; - - for i in 0..chunk_size { - let left = inputs.0[i]; - let right = inputs.0[chunk_size + i]; - let (new_left, new_right) = ( - builder.select(swap, right, left), - builder.select(swap, left, right), - ); - inputs.0[i] = new_left; - inputs.0[chunk_size + i] = new_right; - } - - // Apply the permutation - // For now, we use witness generation since we don't have custom Poseidon2 gates yet - // In production, you would want to implement custom gates for efficiency - let output = builder.add_virtual_targets(POSEIDON2_WIDTH); - - // Add a constraint that will be checked during witness generation - // This is a placeholder - a full implementation would use custom gates - for _i in 0..POSEIDON2_WIDTH { - // The actual permutation will be computed during witness generation - // and verified by the constraints - let _ = builder.add_virtual_target(); - } - - Poseidon2Permutation(output.try_into().unwrap()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::field::types::{Field, Field64}; - - #[test] - fn test_poseidon2_permutation() { - let mut state = Poseidon2Permutation::::new( - (0..POSEIDON2_WIDTH).map(|i| GoldilocksField::from_canonical_u64(i as u64)) - ); - - let original_state = state.0; - state.permute(); - let permuted_state = state.0; - - // Check that permutation changes the state - assert_ne!(original_state, permuted_state); - - // Check that all outputs are valid field elements - for elem in &permuted_state { - assert!(elem.to_canonical_u64() < GoldilocksField::ORDER); - } - } - - #[test] - fn test_poseidon2_hash() { - let input: Vec = (0..8) - .map(GoldilocksField::from_canonical_u64) - .collect(); - - let hash = Poseidon2Hash::hash_no_pad(&input); - - // Verify hash has correct number of elements - assert_eq!(hash.elements.len(), NUM_HASH_OUT_ELTS); - - // Verify hash is deterministic - let hash2 = Poseidon2Hash::hash_no_pad(&input); - assert_eq!(hash, hash2); - } - - #[test] - fn test_poseidon2_two_to_one() { - let left = HashOut { - elements: [ - GoldilocksField::from_canonical_u64(1), - GoldilocksField::from_canonical_u64(2), - GoldilocksField::from_canonical_u64(3), - GoldilocksField::from_canonical_u64(4), - ], - }; - let right = HashOut { - elements: [ - GoldilocksField::from_canonical_u64(5), - GoldilocksField::from_canonical_u64(6), - GoldilocksField::from_canonical_u64(7), - GoldilocksField::from_canonical_u64(8), - ], - }; - - let hash = Poseidon2Hash::two_to_one(left, right); - - // Verify hash has correct number of elements - assert_eq!(hash.elements.len(), NUM_HASH_OUT_ELTS); - - // Verify hash is deterministic - let hash2 = Poseidon2Hash::two_to_one(left, right); - assert_eq!(hash, hash2); - - // Verify hash is different from inputs - assert_ne!(hash, left); - assert_ne!(hash, right); - } - - #[test] - fn test_sponge_construction() { - // Test that the sponge construction works correctly with multiple absorb/squeeze cycles - let input1: Vec = (0..16) - .map(GoldilocksField::from_canonical_u64) - .collect(); - let input2: Vec = (16..32) - .map(GoldilocksField::from_canonical_u64) - .collect(); - - let hash1 = Poseidon2Hash::hash_no_pad(&input1); - let hash2 = Poseidon2Hash::hash_no_pad(&input2); - - // Different inputs should produce different hashes - assert_ne!(hash1, hash2); - } -} diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index 9e2958675..acf16dc98 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -364,7 +364,7 @@ impl + Poseidon2, const D: usize> Gate for Po } fn num_constants(&self) -> usize { - 0 + 1 } fn degree(&self) -> usize { diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index 486e87897..1d4ec2148 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,5 +1,7 @@ use core::fmt::Debug; +use super::config::*; +use super::gate::Poseidon2Gate; use crate::field::extension::{Extendable, FieldExtension}; use crate::field::goldilocks_field::GoldilocksField as F; use crate::field::types::{Field, PrimeField64}; @@ -10,11 +12,6 @@ use crate::iop::target::{BoolTarget, Target}; use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::config::{AlgebraicHasher, Hasher}; -use super::config::*; -use super::gate::Poseidon2Gate; - - - pub trait Poseidon2: PrimeField64 { fn poseidon2(input: [Self; WIDTH]) -> [Self; WIDTH] { let mut state = input; diff --git a/plonky2/src/hash/poseidon2/pure.rs b/plonky2/src/hash/poseidon2/pure.rs index 2d24237e9..dda97d50c 100644 --- a/plonky2/src/hash/poseidon2/pure.rs +++ b/plonky2/src/hash/poseidon2/pure.rs @@ -2,13 +2,12 @@ use plonky2_field::goldilocks_field::GoldilocksField; +use super::config::*; /// This is a test implementation of the Poseidon2 permutation without any gate or generator optimizations. use crate::field::types::Field; use crate::iop::target::Target; use crate::plonk::circuit_builder::CircuitBuilder; -use super::config::*; - type F = GoldilocksField; type Builder = CircuitBuilder; diff --git a/plonky2/src/plonk/config.rs b/plonky2/src/plonk/config.rs index 9504ba3ca..64b67b2e6 100644 --- a/plonky2/src/plonk/config.rs +++ b/plonky2/src/plonk/config.rs @@ -106,10 +106,22 @@ pub trait GenericConfig: type InnerHasher: AlgebraicHasher; } + /// Configuration using Poseidon over the Goldilocks field. #[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Serialize)] pub struct PoseidonGoldilocksConfig; impl GenericConfig<2> for PoseidonGoldilocksConfig { + type F = GoldilocksField; + type FE = QuadraticExtension; + type Hasher = PoseidonHash; + type InnerHasher = PoseidonHash; +} + + +/// Configuration using Poseidon over the Goldilocks field. +#[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Serialize)] +pub struct Poseidon2GoldilocksConfig; +impl GenericConfig<2> for Poseidon2GoldilocksConfig { type F = GoldilocksField; type FE = QuadraticExtension; type Hasher = Poseidon2Hash; diff --git a/plonky2/src/recursion/cyclic_recursion.rs b/plonky2/src/recursion/cyclic_recursion.rs index df0fb95cd..983e7ff7c 100644 --- a/plonky2/src/recursion/cyclic_recursion.rs +++ b/plonky2/src/recursion/cyclic_recursion.rs @@ -209,7 +209,7 @@ mod tests { use crate::gates::noop::NoopGate; use crate::hash::hash_types::{HashOutTarget, RichField}; use crate::hash::hashing::hash_n_to_hash_no_pad; - use crate::hash::poseidon::{PoseidonHash, PoseidonPermutation}; + use crate::hash::poseidon2::hash::{Poseidon2, Poseidon2Hash, Poseidon2Permutation}; use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::circuit_data::{CircuitConfig, CommonCircuitData}; @@ -270,7 +270,7 @@ mod tests { builder.register_public_inputs(&initial_hash_target.elements); let current_hash_in = builder.add_virtual_hash(); let current_hash_out = - builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); + builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); builder.register_public_inputs(¤t_hash_out.elements); let counter = builder.add_virtual_public_input(); @@ -368,10 +368,10 @@ mod tests { cyclic_circuit_data.verify(proof) } - fn iterate_poseidon(initial_state: [F; 4], n: usize) -> [F; 4] { + fn iterate_poseidon(initial_state: [F; 4], n: usize) -> [F; 4] { let mut current = initial_state; for _ in 0..n { - current = hash_n_to_hash_no_pad::>(¤t).elements; + current = hash_n_to_hash_no_pad::>(¤t).elements; } current } diff --git a/plonky2/src/recursion/dummy_circuit.rs b/plonky2/src/recursion/dummy_circuit.rs index 6b08b0ce8..211047bc2 100644 --- a/plonky2/src/recursion/dummy_circuit.rs +++ b/plonky2/src/recursion/dummy_circuit.rs @@ -90,7 +90,10 @@ where pub fn dummy_circuit, C: GenericConfig, const D: usize>( common_data: &CommonCircuitData, ) -> CircuitData { - println!("dummy_circuit received common_data.gates: {:?}", common_data.gates); + println!( + "dummy_circuit received common_data.gates: {:?}", + common_data.gates + ); let config = common_data.config.clone(); assert!( !common_data.config.zero_knowledge, @@ -114,20 +117,47 @@ pub fn dummy_circuit, C: GenericConfig, c } let circuit = builder.build::(); - + // Assert individual components to identify mismatches assert_eq!(circuit.common.config, common_data.config, "config mismatch"); - assert_eq!(circuit.common.fri_params, common_data.fri_params, "fri_params mismatch"); + assert_eq!( + circuit.common.fri_params, common_data.fri_params, + "fri_params mismatch" + ); assert_eq!(circuit.common.gates, common_data.gates, "gates mismatch"); - assert_eq!(circuit.common.selectors_info, common_data.selectors_info, "selectors_info mismatch"); - assert_eq!(circuit.common.quotient_degree_factor, common_data.quotient_degree_factor, "quotient_degree_factor mismatch"); - assert_eq!(circuit.common.num_gate_constraints, common_data.num_gate_constraints, "num_gate_constraints mismatch"); - assert_eq!(circuit.common.num_constants, common_data.num_constants, "num_constants mismatch"); - assert_eq!(circuit.common.num_public_inputs, common_data.num_public_inputs, "num_public_inputs mismatch"); + assert_eq!( + circuit.common.selectors_info, common_data.selectors_info, + "selectors_info mismatch" + ); + assert_eq!( + circuit.common.quotient_degree_factor, common_data.quotient_degree_factor, + "quotient_degree_factor mismatch" + ); + assert_eq!( + circuit.common.num_gate_constraints, common_data.num_gate_constraints, + "num_gate_constraints mismatch" + ); + assert_eq!( + circuit.common.num_constants, common_data.num_constants, + "num_constants mismatch" + ); + assert_eq!( + circuit.common.num_public_inputs, common_data.num_public_inputs, + "num_public_inputs mismatch" + ); assert_eq!(circuit.common.k_is, common_data.k_is, "k_is mismatch"); - assert_eq!(circuit.common.num_partial_products, common_data.num_partial_products, "num_partial_products mismatch"); - assert_eq!(circuit.common.num_lookup_polys, common_data.num_lookup_polys, "num_lookup_polys mismatch"); - assert_eq!(circuit.common.num_lookup_selectors, common_data.num_lookup_selectors, "num_lookup_selectors mismatch"); + assert_eq!( + circuit.common.num_partial_products, common_data.num_partial_products, + "num_partial_products mismatch" + ); + assert_eq!( + circuit.common.num_lookup_polys, common_data.num_lookup_polys, + "num_lookup_polys mismatch" + ); + assert_eq!( + circuit.common.num_lookup_selectors, common_data.num_lookup_selectors, + "num_lookup_selectors mismatch" + ); assert_eq!(circuit.common.luts, common_data.luts, "luts mismatch"); assert_eq!(&circuit.common, common_data); circuit From 31e90748f3a01a6a607041c5e18c4fd1ad85ddb0 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 17:55:00 -0500 Subject: [PATCH 04/26] add more tests --- plonky2/examples/fibonacci.rs | 10 ++--- plonky2/src/hash/poseidon2/gate.rs | 70 ++++++++++++++++++++++++++++-- plonky2/src/plonk/config.rs | 2 - 3 files changed, 71 insertions(+), 11 deletions(-) diff --git a/plonky2/examples/fibonacci.rs b/plonky2/examples/fibonacci.rs index 7310ad164..434669aa7 100644 --- a/plonky2/examples/fibonacci.rs +++ b/plonky2/examples/fibonacci.rs @@ -1,9 +1,8 @@ - use std::time::Instant; +use anyhow::Result; use env_logger; use log::Level; -use anyhow::Result; use plonky2::field::types::Field; use plonky2::iop::witness::{PartialWitness, WitnessWrite}; use plonky2::plonk::circuit_builder::CircuitBuilder; @@ -23,7 +22,6 @@ fn main() -> Result<()> { work::() } - fn work>() -> Result<()> { const D: usize = 2; type C = PoseidonGoldilocksConfig; @@ -37,7 +35,7 @@ fn work>() -> Result<()> { let initial_b = builder.add_virtual_target(); let mut prev_target = initial_a; let mut cur_target = initial_b; - for _ in 0..9999 { + for _ in 0..99999 { let temp = builder.add(prev_target, cur_target); prev_target = cur_target; cur_target = temp; @@ -64,7 +62,7 @@ fn work>() -> Result<()> { // Print the timing tree timing.print(); - + println!( "100th Fibonacci number mod |F| (starting with {}, {}) is: {}", proof.public_inputs[0], proof.public_inputs[1], proof.public_inputs[2] @@ -74,4 +72,4 @@ fn work>() -> Result<()> { println!("Prove time: {:?}", timer3.duration_since(timer2)); data.verify(proof) -} \ No newline at end of file +} diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index acf16dc98..9f9413438 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -495,10 +495,13 @@ impl + Poseidon2, const D: usize> SimpleGenerator>::F; + + let config = CircuitConfig { + num_wires: 143, + ..CircuitConfig::standard_recursion_config() + }; + let mut builder = CircuitBuilder::new(config); + type Gate = Poseidon2Gate; + let gate = Gate::new(); + let row = builder.add_gate(gate, vec![]); + let circuit = builder.build_prover::(); + + let permutation_inputs = (0..WIDTH).map(F::from_canonical_usize).collect::>(); + + let mut inputs = PartialWitness::new(); + inputs + .set_wire( + Wire { + row, + column: Gate::WIRE_SWAP, + }, + F::ZERO, + ) + .unwrap(); + for i in 0..WIDTH { + inputs + .set_wire( + Wire { + row, + column: Gate::wire_input(i), + }, + permutation_inputs[i], + ) + .unwrap(); + } + + let witness = + generate_partial_witness(inputs, &circuit.prover_only, &circuit.common).unwrap(); + + let expected_outputs: [F; WIDTH] = F::poseidon2(permutation_inputs.try_into().unwrap()); + for i in 0..WIDTH { + let out = witness.get_wire(Wire { + row: 0, + column: Gate::wire_output(i), + }); + assert_eq!(out, expected_outputs[i]); + } } #[test] @@ -524,7 +588,7 @@ mod tests { #[test] fn eval_fns() -> Result<()> { const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let gate = Poseidon2Gate::::new(); test_eval_fns::(gate) diff --git a/plonky2/src/plonk/config.rs b/plonky2/src/plonk/config.rs index 64b67b2e6..545044148 100644 --- a/plonky2/src/plonk/config.rs +++ b/plonky2/src/plonk/config.rs @@ -106,7 +106,6 @@ pub trait GenericConfig: type InnerHasher: AlgebraicHasher; } - /// Configuration using Poseidon over the Goldilocks field. #[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Serialize)] pub struct PoseidonGoldilocksConfig; @@ -117,7 +116,6 @@ impl GenericConfig<2> for PoseidonGoldilocksConfig { type InnerHasher = PoseidonHash; } - /// Configuration using Poseidon over the Goldilocks field. #[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Serialize)] pub struct Poseidon2GoldilocksConfig; From 93251272368cfc86874f00a3a8e6aef3e68bf5bc Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 19:05:11 -0500 Subject: [PATCH 05/26] fix recursion tests --- .../src/recursion/conditional_recursive_verifier.rs | 4 ++-- plonky2/src/recursion/cyclic_recursion.rs | 13 +++++++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/plonky2/src/recursion/conditional_recursive_verifier.rs b/plonky2/src/recursion/conditional_recursive_verifier.rs index 0bca23f7f..ab79d3bd9 100644 --- a/plonky2/src/recursion/conditional_recursive_verifier.rs +++ b/plonky2/src/recursion/conditional_recursive_verifier.rs @@ -355,14 +355,14 @@ mod tests { use crate::gates::noop::NoopGate; use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_data::CircuitConfig; - use crate::plonk::config::PoseidonGoldilocksConfig; + use crate::plonk::config::Poseidon2GoldilocksConfig; use crate::recursion::dummy_circuit::{dummy_circuit, dummy_proof}; #[test] fn test_conditional_recursive_verifier() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_config(); diff --git a/plonky2/src/recursion/cyclic_recursion.rs b/plonky2/src/recursion/cyclic_recursion.rs index 983e7ff7c..4e8ce8e48 100644 --- a/plonky2/src/recursion/cyclic_recursion.rs +++ b/plonky2/src/recursion/cyclic_recursion.rs @@ -206,6 +206,7 @@ mod tests { use crate::field::extension::Extendable; use crate::field::types::{Field, PrimeField64}; + use crate::gates::constant::ConstantGate; use crate::gates::noop::NoopGate; use crate::hash::hash_types::{HashOutTarget, RichField}; use crate::hash::hashing::hash_n_to_hash_no_pad; @@ -213,7 +214,7 @@ mod tests { use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::circuit_data::{CircuitConfig, CommonCircuitData}; - use crate::plonk::config::{AlgebraicHasher, GenericConfig, PoseidonGoldilocksConfig}; + use crate::plonk::config::{AlgebraicHasher, GenericConfig, Poseidon2GoldilocksConfig}; use crate::recursion::cyclic_recursion::check_cyclic_proof_verifier_data; use crate::recursion::dummy_circuit::cyclic_base_proof; @@ -243,6 +244,14 @@ mod tests { let verifier_data = builder.add_virtual_verifier_data(data.common.config.fri_config.cap_height); builder.verify_proof::(&proof, &verifier_data, &data.common); + // Add a constant to ensure ConstantGate is included in the gate set + // This matches what the actual cyclic recursion circuit will do + builder.add_gate( + ConstantGate { + num_consts: builder.config.num_constants, + }, + vec![], + ); while builder.num_gates() < 1 << 12 { builder.add_gate(NoopGate, vec![]); } @@ -258,7 +267,7 @@ mod tests { #[test] fn test_cyclic_recursion() -> Result<()> { const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_config(); From 3b4c82dc972c2e5e4dd081cc0c3233272f48f232 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 19:16:02 -0500 Subject: [PATCH 06/26] fix bug in benchmarking --- plonky2/benches/hashing.rs | 53 ++++----------------- plonky2/examples/fibonacci.rs | 12 ++--- plonky2/src/recursion/recursive_verifier.rs | 16 +++---- 3 files changed, 21 insertions(+), 60 deletions(-) diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs index acbe5cb74..7d2270ee2 100644 --- a/plonky2/benches/hashing.rs +++ b/plonky2/benches/hashing.rs @@ -3,10 +3,10 @@ mod allocator; use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; use plonky2::field::goldilocks_field::GoldilocksField; use plonky2::field::types::Sample; -use plonky2::hash::hash_types::{BytesHash, HashOut, RichField}; +use plonky2::hash::hash_types::{BytesHash, RichField}; use plonky2::hash::keccak::KeccakHash; use plonky2::hash::poseidon::{Poseidon, SPONGE_WIDTH}; -use plonky2::hash::poseidon2::{Poseidon2Hash, POSEIDON2_WIDTH}; +use plonky2::hash::poseidon2::hash::Poseidon2; use plonky2::plonk::config::Hasher; use tynm::type_name; @@ -33,52 +33,16 @@ pub(crate) fn bench_poseidon(c: &mut Criterion) { ); } -pub(crate) fn bench_poseidon2(c: &mut Criterion) { +pub(crate) fn bench_poseidon2(c: &mut Criterion) { c.bench_function( &format!( - "poseidon2<{}, {POSEIDON2_WIDTH}>", - type_name::() + "poseidon2<{}, {SPONGE_WIDTH}>", + type_name::() ), |b| { b.iter_batched( - || { - let input: Vec = (0..POSEIDON2_WIDTH) - .map(|_| GoldilocksField::rand()) - .collect(); - input - }, - |input| Poseidon2Hash::hash_no_pad(&input), - BatchSize::SmallInput, - ) - }, - ); -} - -pub(crate) fn bench_poseidon2_two_to_one(c: &mut Criterion) { - c.bench_function( - &format!("poseidon2_two_to_one<{}>", type_name::()), - |b| { - b.iter_batched( - || { - let left = HashOut { - elements: [ - GoldilocksField::rand(), - GoldilocksField::rand(), - GoldilocksField::rand(), - GoldilocksField::rand(), - ], - }; - let right = HashOut { - elements: [ - GoldilocksField::rand(), - GoldilocksField::rand(), - GoldilocksField::rand(), - GoldilocksField::rand(), - ], - }; - (left, right) - }, - |(left, right)| Poseidon2Hash::two_to_one(left, right), + || F::rand_array::(), + |state| F::poseidon2(state), BatchSize::SmallInput, ) }, @@ -87,8 +51,7 @@ pub(crate) fn bench_poseidon2_two_to_one(c: &mut Criterion) { fn criterion_benchmark(c: &mut Criterion) { bench_poseidon::(c); - bench_poseidon2(c); - bench_poseidon2_two_to_one(c); + bench_poseidon2::(c); bench_keccak::(c); } diff --git a/plonky2/examples/fibonacci.rs b/plonky2/examples/fibonacci.rs index 434669aa7..ad751626a 100644 --- a/plonky2/examples/fibonacci.rs +++ b/plonky2/examples/fibonacci.rs @@ -24,18 +24,16 @@ fn main() -> Result<()> { fn work>() -> Result<()> { const D: usize = 2; - type C = PoseidonGoldilocksConfig; - type F = >::F; let config = CircuitConfig::standard_recursion_config(); - let mut builder = CircuitBuilder::::new(config); + let mut builder = CircuitBuilder::::new(config); // The arithmetic circuit. let initial_a = builder.add_virtual_target(); let initial_b = builder.add_virtual_target(); let mut prev_target = initial_a; let mut cur_target = initial_b; - for _ in 0..99999 { + for _ in 0..999999 { let temp = builder.add(prev_target, cur_target); prev_target = cur_target; cur_target = temp; @@ -49,15 +47,15 @@ fn work>() -> Result<()> { // Provide initial values. let timer1 = Instant::now(); let mut pw = PartialWitness::new(); - pw.set_target(initial_a, F::ZERO)?; - pw.set_target(initial_b, F::ONE)?; + pw.set_target(initial_a, C::F::ZERO)?; + pw.set_target(initial_b, C::F::ONE)?; let data = builder.build::(); let timer2 = Instant::now(); // Create a TimingTree to track detailed timing information let mut timing = TimingTree::new("prove", Level::Debug); - let proof = prove::(&data.prover_only, &data.common, pw, &mut timing)?; + let proof = prove::(&data.prover_only, &data.common, pw, &mut timing)?; let timer3 = Instant::now(); // Print the timing tree diff --git a/plonky2/src/recursion/recursive_verifier.rs b/plonky2/src/recursion/recursive_verifier.rs index 16a8ba85b..df52ea002 100644 --- a/plonky2/src/recursion/recursive_verifier.rs +++ b/plonky2/src/recursion/recursive_verifier.rs @@ -214,7 +214,7 @@ mod tests { use crate::gates::noop::NoopGate; use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_data::{CircuitConfig, VerifierOnlyCircuitData}; - use crate::plonk::config::{KeccakGoldilocksConfig, PoseidonGoldilocksConfig}; + use crate::plonk::config::{KeccakGoldilocksConfig, Poseidon2GoldilocksConfig}; use crate::plonk::proof::{CompressedProofWithPublicInputs, ProofWithPublicInputs}; use crate::plonk::prover::prove; use crate::util::timing::TimingTree; @@ -223,7 +223,7 @@ mod tests { fn test_recursive_verifier() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_zk_config(); @@ -239,7 +239,7 @@ mod tests { fn test_recursive_verifier_one_lookup() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_zk_config(); @@ -255,7 +255,7 @@ mod tests { fn test_recursive_verifier_two_luts() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_config(); @@ -271,7 +271,7 @@ mod tests { fn test_recursive_verifier_too_many_rows() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_config(); @@ -287,7 +287,7 @@ mod tests { fn test_recursive_recursive_verifier() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type F = >::F; let config = CircuitConfig::standard_recursion_config(); @@ -318,7 +318,7 @@ mod tests { fn test_size_optimized_recursion() -> Result<()> { init_logger(); const D: usize = 2; - type C = PoseidonGoldilocksConfig; + type C = Poseidon2GoldilocksConfig; type KC = KeccakGoldilocksConfig; type F = >::F; @@ -393,7 +393,7 @@ mod tests { fn test_recursive_verifier_multi_hash() -> Result<()> { init_logger(); const D: usize = 2; - type PC = PoseidonGoldilocksConfig; + type PC = Poseidon2GoldilocksConfig; type KC = KeccakGoldilocksConfig; type F = >::F; From e05586b5513e3046850053255c2507519bcd7ad7 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 20:12:09 -0500 Subject: [PATCH 07/26] wip --- field/src/goldilocks_field.rs | 17 ++++++++++++++ field/src/types.rs | 12 ++++++++++ plonky2/benches/hashing.rs | 5 +---- plonky2/src/hash/poseidon2/hash.rs | 36 +++++++++++++++++++----------- 4 files changed, 53 insertions(+), 17 deletions(-) diff --git a/field/src/goldilocks_field.rs b/field/src/goldilocks_field.rs index b0191ca59..37c53499b 100644 --- a/field/src/goldilocks_field.rs +++ b/field/src/goldilocks_field.rs @@ -156,10 +156,19 @@ impl Field for GoldilocksField { Self(n) } + #[inline(always)] + fn from_canonical_u64_unchecked(n: u64) -> Self { + Self(n) + } + fn from_noncanonical_u96((n_lo, n_hi): (u64, u32)) -> Self { reduce96((n_lo, n_hi)) } + fn from_noncanonical_u128_with_96_bits(n: u128) -> Self { + reduce128_with_96_bits(n) + } + fn from_noncanonical_u128(n: u128) -> Self { reduce128(n) } @@ -396,6 +405,14 @@ fn reduce96((x_lo, x_hi): (u64, u32)) -> GoldilocksField { GoldilocksField(t2) } +#[inline] +fn reduce128_with_96_bits(x: u128) -> GoldilocksField { + let (x_lo, x_hi) = split(x); // This is a no-op + let t1 = x_hi * EPSILON; + let t2 = unsafe { add_no_canonicalize_trashing_input(x_lo, t1) }; + GoldilocksField(t2) +} + /// Reduces to a 64-bit value. The result might not be in canonical form; it could be in between the /// field order and `2^64`. #[inline] diff --git a/field/src/types.rs b/field/src/types.rs index d714b7a84..0d1264283 100644 --- a/field/src/types.rs +++ b/field/src/types.rs @@ -310,6 +310,13 @@ pub trait Field: // TODO: Should probably be unsafe. fn from_canonical_u64(n: u64) -> Self; + /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()` + /// Does not perform any checks. + fn from_canonical_u64_unchecked(n: u64) -> Self { + // overload for non-basefield implementations + Self::from_canonical_u64(n) + } + /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()`. // TODO: Should probably be unsafe. fn from_canonical_u32(n: u32) -> Self { @@ -356,6 +363,11 @@ pub trait Field: Self::from_noncanonical_u128(n) } + fn from_noncanonical_u128_with_96_bits(n: u128) -> Self { + // Default implementation. + Self::from_noncanonical_u128(n) + } + fn exp_power_of_2(&self, power_log: usize) -> Self { let mut res = *self; for _ in 0..power_log { diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs index 7d2270ee2..0237e76dc 100644 --- a/plonky2/benches/hashing.rs +++ b/plonky2/benches/hashing.rs @@ -35,10 +35,7 @@ pub(crate) fn bench_poseidon(c: &mut Criterion) { pub(crate) fn bench_poseidon2(c: &mut Criterion) { c.bench_function( - &format!( - "poseidon2<{}, {SPONGE_WIDTH}>", - type_name::() - ), + &format!("poseidon2<{}, {SPONGE_WIDTH}>", type_name::()), |b| { b.iter_batched( || F::rand_array::(), diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index 1d4ec2148..5a3374200 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,4 +1,5 @@ use core::fmt::Debug; +use core::mem::transmute; use super::config::*; use super::gate::Poseidon2Gate; @@ -37,7 +38,7 @@ pub trait Poseidon2: PrimeField64 { #[inline] fn partial_rounds(state: &mut [Self; WIDTH]) { for r in 0..ROUNDS_P { - state[0] += Self::from_canonical_u64(INTERNAL_CONSTANTS[r]); + state[0] += Self::from_canonical_u64_unchecked(INTERNAL_CONSTANTS[r]); state[0] = Self::sbox_p(&state[0]); Self::internal_linear_layer(state); } @@ -95,8 +96,9 @@ pub trait Poseidon2: PrimeField64 { fn internal_linear_layer(state: &mut [Self; WIDTH]) { let sum: Self = state.iter().cloned().sum(); for i in 0..WIDTH { - state[i] *= Self::from_canonical_u64(MATRIX_DIAG_12_U64[i]); - state[i] += sum; + state[i] = sum.multiply_accumulate(state[i], Self::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i])); + // state[i] *= Self::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]); + // state[i] += sum; } } @@ -109,8 +111,13 @@ pub trait Poseidon2: PrimeField64 { sum += state[i]; } for i in 0..WIDTH { - state[i] *= F::from_canonical_u64(MATRIX_DIAG_12_U64[i]); - state[i] += sum; + state[i] = sum.multiply_accumulate( + state[i], + F::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]), + ); + + // state[i] *= F::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]); + // state[i] += sum; } } @@ -120,7 +127,7 @@ pub trait Poseidon2: PrimeField64 { for i in 0..WIDTH { unsafe { - state[i] = state[i].add_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); + state[i] += state[i].add_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); } } } @@ -133,7 +140,7 @@ pub trait Poseidon2: PrimeField64 { debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); for i in 0..WIDTH { - state[i] += F::from_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); + state[i] += F::from_canonical_u64_unchecked(EXTERNAL_CONSTANTS[external_round][i]); } } @@ -153,7 +160,10 @@ pub trait Poseidon2: PrimeField64 { #[inline] fn sbox_p(a: &Self) -> Self { - a.exp_u64(D) + let a2 = a.square(); + let a4 = a2.square(); + let a3 = *a * a2; + a3 * a4 } #[inline] @@ -235,7 +245,7 @@ pub trait Poseidon2: PrimeField64 { ) where Self: RichField + Extendable, { - let two = builder.constant_extension(Self::Extension::from_canonical_u64(2)); + let two = builder.constant_extension(Self::Extension::from_canonical_u64_unchecked(2)); let t01 = builder.add_extension(x[0], x[1]); let t23 = builder.add_extension(x[2], x[3]); @@ -287,7 +297,7 @@ pub trait Poseidon2: PrimeField64 { { for i in 0..WIDTH { let round_constant = - Self::Extension::from_canonical_u64(EXTERNAL_CONSTANTS[rc_index][i]); + Self::Extension::from_canonical_u64_unchecked(EXTERNAL_CONSTANTS[rc_index][i]); let round_constant = builder.constant_extension(round_constant); input[i] = builder.add_extension(input[i], round_constant); } @@ -327,7 +337,7 @@ pub trait Poseidon2: PrimeField64 { ]); for i in 0..WIDTH { - let round_constant = Self::Extension::from_canonical_u64(MATRIX_DIAG_12_U64[i]); + let round_constant = Self::Extension::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]); let round_constant = builder.constant_extension(round_constant); input[i] = builder.mul_add_extension(round_constant, input[i], sum); @@ -499,7 +509,7 @@ mod test { let input_f = input .iter() - .map(|&x| F::from_canonical_u64((x as u64) + 1073741824)) + .map(|&x| F::from_canonical_u64_unchecked((x as u64) + 1073741824)) .collect::>(); let expected_output_f = hash_n_to_m_no_pad::>(&input_f, 12); @@ -524,7 +534,7 @@ mod test { let input: [u32; 12] = core::array::from_fn(|_| rng.next_u32()); let input_f = input .iter() - .map(|&x| F::from_canonical_u64((x as u64) + 1073741824)) + .map(|&x| F::from_canonical_u64_unchecked((x as u64) + 1073741824)) .collect::>(); let expected_output = hash_n_to_m_no_pad::>(&input_f[0..8], 4); From 33dcaf939f91781daccd919a5258d07a2bf169b5 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 21:31:35 -0500 Subject: [PATCH 08/26] accelerating with neon --- field/src/goldilocks_field.rs | 5 - field/src/types.rs | 7 - .../arch/aarch64/poseidon_goldilocks_neon.rs | 6 +- plonky2/src/hash/poseidon2/hash.rs | 192 +++++++++++++----- 4 files changed, 145 insertions(+), 65 deletions(-) diff --git a/field/src/goldilocks_field.rs b/field/src/goldilocks_field.rs index 37c53499b..2c3d26955 100644 --- a/field/src/goldilocks_field.rs +++ b/field/src/goldilocks_field.rs @@ -156,11 +156,6 @@ impl Field for GoldilocksField { Self(n) } - #[inline(always)] - fn from_canonical_u64_unchecked(n: u64) -> Self { - Self(n) - } - fn from_noncanonical_u96((n_lo, n_hi): (u64, u32)) -> Self { reduce96((n_lo, n_hi)) } diff --git a/field/src/types.rs b/field/src/types.rs index 0d1264283..28a01f839 100644 --- a/field/src/types.rs +++ b/field/src/types.rs @@ -310,13 +310,6 @@ pub trait Field: // TODO: Should probably be unsafe. fn from_canonical_u64(n: u64) -> Self; - /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()` - /// Does not perform any checks. - fn from_canonical_u64_unchecked(n: u64) -> Self { - // overload for non-basefield implementations - Self::from_canonical_u64(n) - } - /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()`. // TODO: Should probably be unsafe. fn from_canonical_u32(n: u32) -> Self { diff --git a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs index a9328d069..69d8cd54b 100644 --- a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs +++ b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs @@ -58,7 +58,7 @@ const_assert!(check_mds_matrix()); /// Addition modulo ORDER accounting for wraparound. Correct only when a + b < 2**64 + ORDER. #[inline(always)] -unsafe fn add_with_wraparound(a: u64, b: u64) -> u64 { +pub unsafe fn add_with_wraparound(a: u64, b: u64) -> u64 { let res: u64; let adj: u64; asm!( @@ -872,12 +872,12 @@ unsafe fn partial_rounds( */ #[inline(always)] -fn unwrap_state(state: [GoldilocksField; 12]) -> [u64; 12] { +pub(crate) fn unwrap_state(state: [GoldilocksField; 12]) -> [u64; 12] { state.map(|s| s.0) } #[inline(always)] -fn wrap_state(state: [u64; 12]) -> [GoldilocksField; 12] { +pub(crate) fn wrap_state(state: [u64; 12]) -> [GoldilocksField; 12] { state.map(GoldilocksField) } diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index 5a3374200..cabce437f 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,5 +1,4 @@ use core::fmt::Debug; -use core::mem::transmute; use super::config::*; use super::gate::Poseidon2Gate; @@ -14,6 +13,7 @@ use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::config::{AlgebraicHasher, Hasher}; pub trait Poseidon2: PrimeField64 { + #[inline] fn poseidon2(input: [Self; WIDTH]) -> [Self; WIDTH] { let mut state = input; @@ -27,6 +27,7 @@ pub trait Poseidon2: PrimeField64 { } #[inline] + #[unroll::unroll_for_loops] fn full_rounds(state: &mut [Self; WIDTH], start: usize) { for r in start..(start + ROUNDS_F_HALF) { Self::add_rc(state, r); @@ -36,15 +37,17 @@ pub trait Poseidon2: PrimeField64 { } #[inline] + #[unroll::unroll_for_loops] fn partial_rounds(state: &mut [Self; WIDTH]) { for r in 0..ROUNDS_P { - state[0] += Self::from_canonical_u64_unchecked(INTERNAL_CONSTANTS[r]); + state[0] += Self::from_canonical_u64(INTERNAL_CONSTANTS[r]); state[0] = Self::sbox_p(&state[0]); Self::internal_linear_layer(state); } } #[inline] + #[unroll::unroll_for_loops] fn external_linear_layer(state: &mut [Self; 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'. @@ -62,12 +65,13 @@ pub trait Poseidon2: PrimeField64 { // 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]; - } + state.iter_mut().enumerate().for_each(|(i, x)| { + *x += sums[i % 4]; + }); } #[inline] + #[unroll::unroll_for_loops] fn external_linear_layer_extension, const D: usize>( state: &mut [F; WIDTH], ) { @@ -92,47 +96,25 @@ pub trait Poseidon2: PrimeField64 { } } - #[inline] - fn internal_linear_layer(state: &mut [Self; WIDTH]) { - let sum: Self = state.iter().cloned().sum(); - for i in 0..WIDTH { - state[i] = sum.multiply_accumulate(state[i], Self::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i])); - // state[i] *= Self::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]); - // state[i] += sum; - } - } + fn internal_linear_layer(state: &mut [Self; WIDTH]); #[inline] fn internal_linear_layer_extension, const D: usize>( state: &mut [F; WIDTH], ) { - let mut sum = state[0]; - for i in 1..WIDTH { - sum += state[i]; - } - for i in 0..WIDTH { - state[i] = sum.multiply_accumulate( - state[i], - F::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]), - ); - - // state[i] *= F::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]); - // state[i] += sum; - } + let sum: F = state.iter().cloned().sum(); + state + .iter_mut() + .zip(MATRIX_DIAG_12_U64.iter()) + .for_each(|(x, &m)| { + *x = sum.multiply_accumulate(*x, F::from_canonical_u64(m)); + }); } - #[inline] - fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { - debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); - - for i in 0..WIDTH { - unsafe { - state[i] += state[i].add_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); - } - } - } + fn add_rc(state: &mut [Self; WIDTH], external_round: usize); #[inline] + #[unroll::unroll_for_loops] fn add_rc_extension, const D: usize>( state: &mut [F; WIDTH], external_round: usize, @@ -140,14 +122,11 @@ pub trait Poseidon2: PrimeField64 { debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); for i in 0..WIDTH { - state[i] += F::from_canonical_u64_unchecked(EXTERNAL_CONSTANTS[external_round][i]); + state[i] += F::from_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); } } - #[inline] - fn sbox(state: &mut [Self; WIDTH]) { - state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); - } + fn sbox(state: &mut [Self; WIDTH]); #[inline] fn sbox_extension, const D: usize>( @@ -168,7 +147,11 @@ pub trait Poseidon2: PrimeField64 { #[inline] fn sbox_p_extension, const D: usize>(a: &F) -> F { - a.exp_u64(super::config::D) + debug_assert!(D == 7); + let a2 = a.square(); + let a4 = a2.square(); + let a3 = *a * a2; + a3 * a4 } // Multiply a 4-element vector x by: @@ -177,6 +160,7 @@ pub trait Poseidon2: PrimeField64 { // [ 1 1 2 3 ] // [ 3 1 1 2 ]. // This is more efficient than the previous matrix. + #[inline] fn apply_mat4_mut(x: &mut [Self; 4]) { let t01 = x[0] + x[1]; let t23 = x[2] + x[3]; @@ -190,6 +174,7 @@ pub trait Poseidon2: PrimeField64 { x[2] = t01233 + t23; // x[0] + x[1] + 2*x[2] + 3*x[3] } + #[inline] fn apply_mat4_mut_extension, const D: usize>( x: &mut [F; 4], ) { @@ -206,7 +191,8 @@ pub trait Poseidon2: PrimeField64 { } // In circuit functions - + #[inline] + #[unroll::unroll_for_loops] fn external_linear_layer_circuit( builder: &mut CircuitBuilder, state: &mut [ExtensionTarget; WIDTH], @@ -239,13 +225,15 @@ pub trait Poseidon2: PrimeField64 { } } + #[inline] + #[unroll::unroll_for_loops] fn apply_mat4_mut_circuit( builder: &mut CircuitBuilder, x: &mut [ExtensionTarget; 4], ) where Self: RichField + Extendable, { - let two = builder.constant_extension(Self::Extension::from_canonical_u64_unchecked(2)); + let two = builder.constant_extension(Self::Extension::from_canonical_u64(2)); let t01 = builder.add_extension(x[0], x[1]); let t23 = builder.add_extension(x[2], x[3]); @@ -261,6 +249,8 @@ pub trait Poseidon2: PrimeField64 { x[2] = builder.add_extension(t01233, t23); // x[0] + x[1] + 2*x[2] + 3*x[3] } + #[inline] + #[unroll::unroll_for_loops] fn matmul_m4_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -288,6 +278,8 @@ pub trait Poseidon2: PrimeField64 { } } + #[inline] + #[unroll::unroll_for_loops] fn add_rc_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -297,12 +289,14 @@ pub trait Poseidon2: PrimeField64 { { for i in 0..WIDTH { let round_constant = - Self::Extension::from_canonical_u64_unchecked(EXTERNAL_CONSTANTS[rc_index][i]); + Self::Extension::from_canonical_u64(EXTERNAL_CONSTANTS[rc_index][i]); let round_constant = builder.constant_extension(round_constant); input[i] = builder.add_extension(input[i], round_constant); } } + #[inline] + #[unroll::unroll_for_loops] fn sbox_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -325,6 +319,8 @@ pub trait Poseidon2: PrimeField64 { builder.exp_u64_extension(input, super::config::D) } + #[inline] + #[unroll::unroll_for_loops] fn internal_linear_layer_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -337,7 +333,7 @@ pub trait Poseidon2: PrimeField64 { ]); for i in 0..WIDTH { - let round_constant = Self::Extension::from_canonical_u64_unchecked(MATRIX_DIAG_12_U64[i]); + let round_constant = Self::Extension::from_canonical_u64(MATRIX_DIAG_12_U64[i]); let round_constant = builder.constant_extension(round_constant); input[i] = builder.mul_add_extension(round_constant, input[i], sum); @@ -345,7 +341,101 @@ pub trait Poseidon2: PrimeField64 { } } -impl Poseidon2 for F {} +impl Poseidon2 for F { + #[inline] + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + state + .iter_mut() + .zip(EXTERNAL_CONSTANTS[external_round].iter()) + .for_each(|(x, &m)| { + *x += Self::from_canonical_u64(m); + }); + } + + #[inline] + #[unroll::unroll_for_loops] + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + + unsafe { + use core::arch::aarch64::*; + use core::mem::transmute; + + let mut state_u64 = transmute::<[Self; WIDTH], [u64; WIDTH]>(*state); + let round_constants = &EXTERNAL_CONSTANTS[external_round]; + + // Process 2 elements at a time using NEON + for i in (0..WIDTH).step_by(2) { + let state_vec = vld1q_u64(state_u64[i..].as_ptr()); + let rc_vec = vld1q_u64(round_constants[i..].as_ptr()); + + // Add the round constants + let sum = vaddq_u64(state_vec, rc_vec); + + // Check for overflow (if sum < state_vec, we wrapped around) + let overflow_mask = vcltq_u64(sum, state_vec); + + // Add EPSILON (0xffffffff) where overflow occurred + let epsilon = vdupq_n_u64(0xffffffff); + let adjustment = vandq_u64(overflow_mask, epsilon); + let result = vaddq_u64(sum, adjustment); + + vst1q_u64(state_u64[i..].as_mut_ptr(), result); + } + *state = transmute::<[u64; WIDTH], [Self; WIDTH]>(state_u64); + } + } + + #[inline] + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + fn internal_linear_layer(state: &mut [Self; WIDTH]) { + let tmp = state + .iter() + .map(|&x| x.to_noncanonical_u64() as u128) + .sum::(); + let sum = Self::from_noncanonical_u128_with_96_bits(tmp); + state + .iter_mut() + .zip(MATRIX_DIAG_12_U64.iter()) + .for_each(|(x, &m)| { + *x = sum.multiply_accumulate(*x, Self::from_canonical_u64(m)); + }); + } + + #[inline] + #[unroll::unroll_for_loops] + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + fn internal_linear_layer(state: &mut [Self; WIDTH]) { + let tmp = state + .iter() + .map(|&x| x.to_noncanonical_u64() as u128) + .sum::(); + let sum = Self::from_noncanonical_u128_with_96_bits(tmp); + state + .iter_mut() + .zip(MATRIX_DIAG_12_U64.iter()) + .for_each(|(x, &m)| { + *x = sum.multiply_accumulate(*x, Self::from_canonical_u64(m)); + }); + } + + #[inline] + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + fn sbox(state: &mut [Self; WIDTH]) { + state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); + } + + #[inline(always)] + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + fn sbox(state: &mut [Self; WIDTH]) { + unsafe { + crate::hash::arch::aarch64::poseidon_goldilocks_neon::sbox_layer(state); + } + } +} #[derive(Copy, Clone, Default, Debug, PartialEq)] pub struct Poseidon2Permutation { @@ -433,6 +523,8 @@ impl Hasher for Poseidon2Hash { } impl Poseidon2Hash { + #[inline] + #[unroll::unroll_for_loops] pub fn hash_n_to_one( input: &[>::Hash], ) -> >::Hash { @@ -509,7 +601,7 @@ mod test { let input_f = input .iter() - .map(|&x| F::from_canonical_u64_unchecked((x as u64) + 1073741824)) + .map(|&x| F::from_canonical_u64((x as u64) + 1073741824)) .collect::>(); let expected_output_f = hash_n_to_m_no_pad::>(&input_f, 12); @@ -534,7 +626,7 @@ mod test { let input: [u32; 12] = core::array::from_fn(|_| rng.next_u32()); let input_f = input .iter() - .map(|&x| F::from_canonical_u64_unchecked((x as u64) + 1073741824)) + .map(|&x| F::from_canonical_u64((x as u64) + 1073741824)) .collect::>(); let expected_output = hash_n_to_m_no_pad::>(&input_f[0..8], 4); From c956fd830beafa10256d87695424c614b078b298 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 11 Nov 2025 22:32:25 -0500 Subject: [PATCH 09/26] clean up --- .../arch/aarch64/poseidon_goldilocks_neon.rs | 25 ++++++ plonky2/src/hash/poseidon2/hash.rs | 84 ++++++------------- 2 files changed, 49 insertions(+), 60 deletions(-) diff --git a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs index 69d8cd54b..17078ff99 100644 --- a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs +++ b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs @@ -915,3 +915,28 @@ pub unsafe fn mds_layer(state: &[GoldilocksField; WIDTH]) -> [GoldilocksField; W let state = mds_layer_full(state); wrap_state(state) } + +#[inline(always)] +#[unroll::unroll_for_loops] +pub unsafe fn vector_add(a: &[u64; WIDTH], b: &[u64; WIDTH]) -> [u64; WIDTH] { + let mut res = [0u64; WIDTH]; + // Process 2 elements at a time using NEON + for i in (0..WIDTH).step_by(2) { + let a_vec = vld1q_u64(a[i..].as_ptr()); + let b_vec = vld1q_u64(b[i..].as_ptr()); + + // Add the round constants + let sum = vaddq_u64(a_vec, b_vec); + + // Check for overflow (if sum < state_vec, we wrapped around) + let overflow_mask = vcltq_u64(sum, a_vec); + + // Add EPSILON (0xffffffff) where overflow occurred + let epsilon = vdupq_n_u64(0xffffffff); + let adjustment = vandq_u64(overflow_mask, epsilon); + let result = vaddq_u64(sum, adjustment); + + vst1q_u64(res[i..].as_mut_ptr(), result); + } + res +} diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index cabce437f..c12dac87d 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,4 +1,5 @@ use core::fmt::Debug; +use core::mem::transmute; use super::config::*; use super::gate::Poseidon2Gate; @@ -48,7 +49,7 @@ pub trait Poseidon2: PrimeField64 { #[inline] #[unroll::unroll_for_loops] - fn external_linear_layer(state: &mut [Self; WIDTH]) { + fn external_linear_layer(state: &mut [Self; 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) { @@ -63,11 +64,12 @@ pub trait Poseidon2: PrimeField64 { let sums: [Self; 4] = core::array::from_fn(|k| (0..WIDTH).step_by(4).map(|j| state[j + k]).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 - state.iter_mut().enumerate().for_each(|(i, x)| { - *x += sums[i % 4]; - }); + for i in 0..WIDTH { + state[i] += sums[i % 4]; + } } #[inline] @@ -96,7 +98,19 @@ pub trait Poseidon2: PrimeField64 { } } - fn internal_linear_layer(state: &mut [Self; WIDTH]); + #[inline] + #[unroll::unroll_for_loops] + fn internal_linear_layer(state: &mut [Self; WIDTH]) { + let tmp = state + .iter() + .map(|&x| x.to_noncanonical_u64() as u128) + .sum::(); + let sum = Self::from_noncanonical_u128_with_96_bits(tmp); + for i in 0..WIDTH { + state[i] = + sum.multiply_accumulate(state[i], Self::from_canonical_u64(MATRIX_DIAG_12_U64[i])); + } + } #[inline] fn internal_linear_layer_extension, const D: usize>( @@ -355,73 +369,23 @@ impl Poseidon2 for F { } #[inline] - #[unroll::unroll_for_loops] #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); unsafe { - use core::arch::aarch64::*; use core::mem::transmute; - let mut state_u64 = transmute::<[Self; WIDTH], [u64; WIDTH]>(*state); - let round_constants = &EXTERNAL_CONSTANTS[external_round]; - - // Process 2 elements at a time using NEON - for i in (0..WIDTH).step_by(2) { - let state_vec = vld1q_u64(state_u64[i..].as_ptr()); - let rc_vec = vld1q_u64(round_constants[i..].as_ptr()); - - // Add the round constants - let sum = vaddq_u64(state_vec, rc_vec); - - // Check for overflow (if sum < state_vec, we wrapped around) - let overflow_mask = vcltq_u64(sum, state_vec); + use crate::hash::arch::aarch64::poseidon_goldilocks_neon::vector_add; - // Add EPSILON (0xffffffff) where overflow occurred - let epsilon = vdupq_n_u64(0xffffffff); - let adjustment = vandq_u64(overflow_mask, epsilon); - let result = vaddq_u64(sum, adjustment); + let state_u64 = transmute::<[Self; WIDTH], [u64; WIDTH]>(*state); + let round_constants = &EXTERNAL_CONSTANTS[external_round]; - vst1q_u64(state_u64[i..].as_mut_ptr(), result); - } - *state = transmute::<[u64; WIDTH], [Self; WIDTH]>(state_u64); + let res = vector_add(&state_u64, round_constants); + *state = transmute::<[u64; WIDTH], [Self; WIDTH]>(res); } } - #[inline] - #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] - fn internal_linear_layer(state: &mut [Self; WIDTH]) { - let tmp = state - .iter() - .map(|&x| x.to_noncanonical_u64() as u128) - .sum::(); - let sum = Self::from_noncanonical_u128_with_96_bits(tmp); - state - .iter_mut() - .zip(MATRIX_DIAG_12_U64.iter()) - .for_each(|(x, &m)| { - *x = sum.multiply_accumulate(*x, Self::from_canonical_u64(m)); - }); - } - - #[inline] - #[unroll::unroll_for_loops] - #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] - fn internal_linear_layer(state: &mut [Self; WIDTH]) { - let tmp = state - .iter() - .map(|&x| x.to_noncanonical_u64() as u128) - .sum::(); - let sum = Self::from_noncanonical_u128_with_96_bits(tmp); - state - .iter_mut() - .zip(MATRIX_DIAG_12_U64.iter()) - .for_each(|(x, &m)| { - *x = sum.multiply_accumulate(*x, Self::from_canonical_u64(m)); - }); - } - #[inline] #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] fn sbox(state: &mut [Self; WIDTH]) { From 4be97e980f11c43c762cce495de2737dba0d75e1 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 10:17:28 -0500 Subject: [PATCH 10/26] update poseidon2's benchmark --- plonky2/Cargo.toml | 1 + plonky2/benches/hashing.rs | 92 +++++++++++++++++++++++++++++++++++++- 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/plonky2/Cargo.toml b/plonky2/Cargo.toml index 37bd6b25b..e267ca515 100644 --- a/plonky2/Cargo.toml +++ b/plonky2/Cargo.toml @@ -57,6 +57,7 @@ serde_cbor = { version = "0.11.2" } serde_json = { version = "1.0" } structopt = { version = "0.3.26", default-features = false } tynm = { version = "0.1.6", default-features = false } +circuit = { path = "../../zklighter-perps-circuits-plonky/circuit" } [target.'cfg(not(target_env = "msvc"))'.dev-dependencies] jemallocator = "0.5.0" diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs index 0237e76dc..4fd795bf7 100644 --- a/plonky2/benches/hashing.rs +++ b/plonky2/benches/hashing.rs @@ -2,7 +2,7 @@ mod allocator; use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; use plonky2::field::goldilocks_field::GoldilocksField; -use plonky2::field::types::Sample; +use plonky2::field::types::{Field, Sample}; use plonky2::hash::hash_types::{BytesHash, RichField}; use plonky2::hash::keccak::KeccakHash; use plonky2::hash::poseidon::{Poseidon, SPONGE_WIDTH}; @@ -10,6 +10,13 @@ use plonky2::hash::poseidon2::hash::Poseidon2; use plonky2::plonk::config::Hasher; use tynm::type_name; +use p3_field::AbstractField; +use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks}; +use p3_poseidon2::{Poseidon2 as P3Poseidon2, Poseidon2ExternalMatrixGeneral}; +use p3_symmetric::Permutation; + +use circuit::poseidon2::hash::Poseidon2 as LighterPoseidon2; + pub(crate) fn bench_keccak(c: &mut Criterion) { c.bench_function("keccak256", |b| { b.iter_batched( @@ -35,7 +42,7 @@ pub(crate) fn bench_poseidon(c: &mut Criterion) { pub(crate) fn bench_poseidon2(c: &mut Criterion) { c.bench_function( - &format!("poseidon2<{}, {SPONGE_WIDTH}>", type_name::()), + &format!("optimized poseidon2<{}, {SPONGE_WIDTH}>", type_name::()), |b| { b.iter_batched( || F::rand_array::(), @@ -46,9 +53,90 @@ pub(crate) fn bench_poseidon2(c: &mut Criterion) { ); } +pub(crate) fn bench_p3_poseidon2(c: &mut Criterion) { + const WIDTH: usize = 12; + const D: u64 = 7; + const ROUNDS_F: usize = 8; + const ROUNDS_P: usize = 22; + + // Create the Poseidon2 instance + let external_linear_layer = Poseidon2ExternalMatrixGeneral; + let internal_linear_layer = DiffusionMatrixGoldilocks; + + let external_constants = plonky2::hash::poseidon2::config::EXTERNAL_CONSTANTS + .iter() + .map(|v| { + v.iter() + .map(|&x| Goldilocks::from_canonical_u64(x)) + .collect::>() + .try_into() + .unwrap() + }) + .collect::>(); + + let internal_constants = plonky2::hash::poseidon2::config::INTERNAL_CONSTANTS + .iter() + .map(|&x| Goldilocks::from_canonical_u64(x)) + .collect::>(); + + let poseidon = P3Poseidon2::< + Goldilocks, + Poseidon2ExternalMatrixGeneral, + DiffusionMatrixGoldilocks, + WIDTH, + D, + >::new( + ROUNDS_F, + external_constants, + external_linear_layer, + ROUNDS_P, + internal_constants, + internal_linear_layer, + ); + + c.bench_function("plonky3's poseidon2", |b| { + b.iter_batched( + || { + let mut state = [Goldilocks::zero(); WIDTH]; + for i in 0..WIDTH { + state[i] = Goldilocks::from_canonical_u64(rand::random::()); + } + state + }, + |mut state| { + poseidon.permute_mut(&mut state); + state + }, + BatchSize::SmallInput, + ) + }); +} + +pub(crate) fn bench_circuit_poseidon2(c: &mut Criterion) { + const WIDTH: usize = 12; + + c.bench_function("lighter's poseidon2", |b| { + b.iter_batched( + || { + let mut state = [GoldilocksField::ZERO; WIDTH]; + for i in 0..WIDTH { + state[i] = GoldilocksField::from_canonical_u64(rand::random::()); + } + state + }, + |state| { + ::poseidon2(state) + }, + BatchSize::SmallInput, + ) + }); +} + fn criterion_benchmark(c: &mut Criterion) { bench_poseidon::(c); bench_poseidon2::(c); + bench_p3_poseidon2(c); + bench_circuit_poseidon2(c); bench_keccak::(c); } From ae3892679c7c2cb106fe687722e754607e1482a8 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 11:29:11 -0500 Subject: [PATCH 11/26] clean up --- plonky2/Cargo.toml | 1 - plonky2/benches/hashing.rs | 34 +++++------------------------- plonky2/src/fri/validate_shape.rs | 2 +- plonky2/src/hash/merkle_proofs.rs | 2 +- plonky2/src/hash/poseidon2/hash.rs | 4 +--- 5 files changed, 8 insertions(+), 35 deletions(-) diff --git a/plonky2/Cargo.toml b/plonky2/Cargo.toml index e267ca515..37bd6b25b 100644 --- a/plonky2/Cargo.toml +++ b/plonky2/Cargo.toml @@ -57,7 +57,6 @@ serde_cbor = { version = "0.11.2" } serde_json = { version = "1.0" } structopt = { version = "0.3.26", default-features = false } tynm = { version = "0.1.6", default-features = false } -circuit = { path = "../../zklighter-perps-circuits-plonky/circuit" } [target.'cfg(not(target_env = "msvc"))'.dev-dependencies] jemallocator = "0.5.0" diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs index 4fd795bf7..f19ad997b 100644 --- a/plonky2/benches/hashing.rs +++ b/plonky2/benches/hashing.rs @@ -1,8 +1,12 @@ mod allocator; use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; +use p3_field::AbstractField; +use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks}; +use p3_poseidon2::{Poseidon2 as P3Poseidon2, Poseidon2ExternalMatrixGeneral}; +use p3_symmetric::Permutation; use plonky2::field::goldilocks_field::GoldilocksField; -use plonky2::field::types::{Field, Sample}; +use plonky2::field::types::Sample; use plonky2::hash::hash_types::{BytesHash, RichField}; use plonky2::hash::keccak::KeccakHash; use plonky2::hash::poseidon::{Poseidon, SPONGE_WIDTH}; @@ -10,13 +14,6 @@ use plonky2::hash::poseidon2::hash::Poseidon2; use plonky2::plonk::config::Hasher; use tynm::type_name; -use p3_field::AbstractField; -use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks}; -use p3_poseidon2::{Poseidon2 as P3Poseidon2, Poseidon2ExternalMatrixGeneral}; -use p3_symmetric::Permutation; - -use circuit::poseidon2::hash::Poseidon2 as LighterPoseidon2; - pub(crate) fn bench_keccak(c: &mut Criterion) { c.bench_function("keccak256", |b| { b.iter_batched( @@ -112,31 +109,10 @@ pub(crate) fn bench_p3_poseidon2(c: &mut Criterion) { }); } -pub(crate) fn bench_circuit_poseidon2(c: &mut Criterion) { - const WIDTH: usize = 12; - - c.bench_function("lighter's poseidon2", |b| { - b.iter_batched( - || { - let mut state = [GoldilocksField::ZERO; WIDTH]; - for i in 0..WIDTH { - state[i] = GoldilocksField::from_canonical_u64(rand::random::()); - } - state - }, - |state| { - ::poseidon2(state) - }, - BatchSize::SmallInput, - ) - }); -} - fn criterion_benchmark(c: &mut Criterion) { bench_poseidon::(c); bench_poseidon2::(c); bench_p3_poseidon2(c); - bench_circuit_poseidon2(c); bench_keccak::(c); } diff --git a/plonky2/src/fri/validate_shape.rs b/plonky2/src/fri/validate_shape.rs index be675ed61..8805478f6 100644 --- a/plonky2/src/fri/validate_shape.rs +++ b/plonky2/src/fri/validate_shape.rs @@ -20,7 +20,7 @@ where F: RichField + Extendable, C: GenericConfig, { - validate_batch_fri_proof_shape::(proof, &[instance.clone()], params) + validate_batch_fri_proof_shape::(proof, std::slice::from_ref(instance), params) } pub(crate) fn validate_batch_fri_proof_shape( diff --git a/plonky2/src/hash/merkle_proofs.rs b/plonky2/src/hash/merkle_proofs.rs index 424e03ae6..74963a2a4 100644 --- a/plonky2/src/hash/merkle_proofs.rs +++ b/plonky2/src/hash/merkle_proofs.rs @@ -59,7 +59,7 @@ pub fn verify_merkle_proof_to_cap>( proof: &MerkleProof, ) -> Result<()> { verify_batch_merkle_proof_to_cap( - &[leaf_data.clone()], + &[leaf_data], &[proof.siblings.len()], leaf_index, merkle_cap, diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index c12dac87d..2264ced4e 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,5 +1,4 @@ use core::fmt::Debug; -use core::mem::transmute; use super::config::*; use super::gate::Poseidon2Gate; @@ -49,7 +48,7 @@ pub trait Poseidon2: PrimeField64 { #[inline] #[unroll::unroll_for_loops] - fn external_linear_layer(state: &mut [Self; WIDTH]){ + fn external_linear_layer(state: &mut [Self; 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) { @@ -64,7 +63,6 @@ pub trait Poseidon2: PrimeField64 { let sums: [Self; 4] = core::array::from_fn(|k| (0..WIDTH).step_by(4).map(|j| state[j + k]).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 { From 33aa4ed36c137ea936a0f5b4f1b217bdcf0cd47e Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 16:03:18 -0500 Subject: [PATCH 12/26] clean up benchmark --- plonky2/Cargo.toml | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/plonky2/Cargo.toml b/plonky2/Cargo.toml index 37bd6b25b..36f26fb79 100644 --- a/plonky2/Cargo.toml +++ b/plonky2/Cargo.toml @@ -65,29 +65,29 @@ jemallocator = "0.5.0" name = "generate_constants" required-features = ["rand_chacha"] -# [[bench]] -# name = "field_arithmetic" -# harness = false +[[bench]] +name = "field_arithmetic" +harness = false -# [[bench]] -# name = "ffts" -# harness = false +[[bench]] +name = "ffts" +harness = false [[bench]] name = "hashing" harness = false -# [[bench]] -# name = "merkle" -# harness = false +[[bench]] +name = "merkle" +harness = false -# [[bench]] -# name = "transpose" -# harness = false +[[bench]] +name = "transpose" +harness = false -# [[bench]] -# name = "reverse_index_bits" -# harness = false +[[bench]] +name = "reverse_index_bits" +harness = false # Display math equations properly in documentation [package.metadata.docs.rs] From ba5d35399ffad099655d3574e195454506e2d145 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 16:25:38 -0500 Subject: [PATCH 13/26] testing recursion for both poseidon and poseidon2 --- plonky2/src/recursion/cyclic_recursion.rs | 76 ++++++++++++++++------- 1 file changed, 54 insertions(+), 22 deletions(-) diff --git a/plonky2/src/recursion/cyclic_recursion.rs b/plonky2/src/recursion/cyclic_recursion.rs index 4e8ce8e48..956a8dba0 100644 --- a/plonky2/src/recursion/cyclic_recursion.rs +++ b/plonky2/src/recursion/cyclic_recursion.rs @@ -210,11 +210,14 @@ mod tests { use crate::gates::noop::NoopGate; use crate::hash::hash_types::{HashOutTarget, RichField}; use crate::hash::hashing::hash_n_to_hash_no_pad; - use crate::hash::poseidon2::hash::{Poseidon2, Poseidon2Hash, Poseidon2Permutation}; + use crate::hash::poseidon::PoseidonPermutation; + use crate::hash::poseidon2::hash::{Poseidon2, Poseidon2Permutation}; use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::circuit_data::{CircuitConfig, CommonCircuitData}; - use crate::plonk::config::{AlgebraicHasher, GenericConfig, Poseidon2GoldilocksConfig}; + use crate::plonk::config::{ + AlgebraicHasher, GenericConfig, Poseidon2GoldilocksConfig, PoseidonGoldilocksConfig, + }; use crate::recursion::cyclic_recursion::check_cyclic_proof_verifier_data; use crate::recursion::dummy_circuit::cyclic_base_proof; @@ -223,7 +226,9 @@ mod tests { F: RichField + Extendable, C: GenericConfig, const D: usize, - >() -> CommonCircuitData + >( + use_poseidon2: bool, + ) -> CommonCircuitData where C::Hasher: AlgebraicHasher, { @@ -244,14 +249,17 @@ mod tests { let verifier_data = builder.add_virtual_verifier_data(data.common.config.fri_config.cap_height); builder.verify_proof::(&proof, &verifier_data, &data.common); - // Add a constant to ensure ConstantGate is included in the gate set - // This matches what the actual cyclic recursion circuit will do - builder.add_gate( - ConstantGate { - num_consts: builder.config.num_constants, - }, - vec![], - ); + + if use_poseidon2 { + // Add a constant to ensure ConstantGate is included in the gate set + // This matches what the actual cyclic recursion circuit will do + builder.add_gate( + ConstantGate { + num_consts: builder.config.num_constants, + }, + vec![], + ); + } while builder.num_gates() < 1 << 12 { builder.add_gate(NoopGate, vec![]); } @@ -266,12 +274,19 @@ mod tests { /// - VK for cyclic recursion (?) #[test] fn test_cyclic_recursion() -> Result<()> { - const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; + cyclic_recursion::(false)?; + cyclic_recursion::(true) + } + fn cyclic_recursion + 'static, const D: usize>( + use_poseidon2: bool, + ) -> Result<()> + where + C::Hasher: AlgebraicHasher, + C::F: Poseidon2, + { let config = CircuitConfig::standard_recursion_config(); - let mut builder = CircuitBuilder::::new(config); + let mut builder = CircuitBuilder::::new(config); let one = builder.one(); // Circuit that computes a repeated hash. @@ -279,11 +294,11 @@ mod tests { builder.register_public_inputs(&initial_hash_target.elements); let current_hash_in = builder.add_virtual_hash(); let current_hash_out = - builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); + builder.hash_n_to_hash_no_pad::(current_hash_in.elements.to_vec()); builder.register_public_inputs(¤t_hash_out.elements); let counter = builder.add_virtual_public_input(); - let mut common_data = common_data_for_recursion::(); + let mut common_data = common_data_for_recursion::(use_poseidon2); let verifier_data_target = builder.add_verifier_data_public_inputs(); common_data.num_public_inputs = builder.num_public_inputs(); @@ -319,7 +334,12 @@ mod tests { let cyclic_circuit_data = builder.build::(); let mut pw = PartialWitness::new(); - let initial_hash = [F::ZERO, F::ONE, F::TWO, F::from_canonical_usize(3)]; + let initial_hash = [ + C::F::ZERO, + C::F::ONE, + C::F::TWO, + C::F::from_canonical_usize(3), + ]; let initial_hash_pis = initial_hash.into_iter().enumerate().collect(); pw.set_bool_target(condition, false)?; pw.set_proof_with_pis_target::( @@ -368,19 +388,31 @@ mod tests { let initial_hash = &proof.public_inputs[..4]; let hash = &proof.public_inputs[4..8]; let counter = proof.public_inputs[8]; - let expected_hash: [F; 4] = iterate_poseidon( + let expected_hash: [C::F; 4] = iterate_poseidon( initial_hash.try_into().unwrap(), counter.to_canonical_u64() as usize, + use_poseidon2, ); assert_eq!(hash, expected_hash); cyclic_circuit_data.verify(proof) } - fn iterate_poseidon(initial_state: [F; 4], n: usize) -> [F; 4] { + fn iterate_poseidon( + initial_state: [F; 4], + n: usize, + use_poseidon2: bool, + ) -> [F; 4] { let mut current = initial_state; - for _ in 0..n { - current = hash_n_to_hash_no_pad::>(¤t).elements; + + if use_poseidon2 { + for _ in 0..n { + current = hash_n_to_hash_no_pad::>(¤t).elements; + } + } else { + for _ in 0..n { + current = hash_n_to_hash_no_pad::>(¤t).elements; + } } current } From 45d75ece09997ce51b64e7c23ff79496157ee96d Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 16:32:45 -0500 Subject: [PATCH 14/26] clippy --- plonky2/src/plonk/config.rs | 2 +- plonky2/src/util/reducing.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/plonky2/src/plonk/config.rs b/plonky2/src/plonk/config.rs index 545044148..4bbcbc4bc 100644 --- a/plonky2/src/plonk/config.rs +++ b/plonky2/src/plonk/config.rs @@ -52,7 +52,7 @@ pub trait Hasher: Sized + Copy + Debug + Eq + PartialEq { fn hash_pad(input: &[F]) -> Self::Hash { let mut padded_input = input.to_vec(); padded_input.push(F::ONE); - while (padded_input.len() + 1) % Self::Permutation::RATE != 0 { + while !(padded_input.len() + 1).is_multiple_of(Self::Permutation::RATE) { padded_input.push(F::ZERO); } padded_input.push(F::ONE); diff --git a/plonky2/src/util/reducing.rs b/plonky2/src/util/reducing.rs index 89c8bd96c..1abe017b0 100644 --- a/plonky2/src/util/reducing.rs +++ b/plonky2/src/util/reducing.rs @@ -150,7 +150,7 @@ impl ReducingFactorTarget { let zero_ext = builder.zero_extension(); let mut acc = zero_ext; let mut reversed_terms = terms.to_vec(); - while reversed_terms.len() % max_coeffs_len != 0 { + while !reversed_terms.len().is_multiple_of(max_coeffs_len) { reversed_terms.push(zero); } reversed_terms.reverse(); @@ -200,7 +200,7 @@ impl ReducingFactorTarget { let zero_ext = builder.zero_extension(); let mut acc = zero_ext; let mut reversed_terms = terms.to_vec(); - while reversed_terms.len() % max_coeffs_len != 0 { + while !reversed_terms.len().is_multiple_of(max_coeffs_len) { reversed_terms.push(zero_ext); } reversed_terms.reverse(); From 842733ab031de0a53bea56fd033a2e7e295c6889 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 17:01:15 -0500 Subject: [PATCH 15/26] fix tests --- plonky2/src/hash/poseidon2/hash.rs | 35 +++++++++++++++++------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index 2264ced4e..a1d3c9c9c 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,5 +1,7 @@ use core::fmt::Debug; +use plonky2_field::ops::Square; + use super::config::*; use super::gate::Poseidon2Gate; use crate::field::extension::{Extendable, FieldExtension}; @@ -149,22 +151,9 @@ pub trait Poseidon2: PrimeField64 { .for_each(|a| *a = Self::sbox_p_extension(a)); } - #[inline] - fn sbox_p(a: &Self) -> Self { - let a2 = a.square(); - let a4 = a2.square(); - let a3 = *a * a2; - a3 * a4 - } + fn sbox_p(a: &Self) -> Self; - #[inline] - fn sbox_p_extension, const D: usize>(a: &F) -> F { - debug_assert!(D == 7); - let a2 = a.square(); - let a4 = a2.square(); - let a3 = *a * a2; - a3 * a4 - } + fn sbox_p_extension, const D: usize>(a: &F) -> F; // Multiply a 4-element vector x by: // [ 2 3 1 1 ] @@ -354,6 +343,22 @@ pub trait Poseidon2: PrimeField64 { } impl Poseidon2 for F { + #[inline] + fn sbox_p(a: &Self) -> Self { + let a2 = a.square(); + let a4 = a2.square(); + let a3 = *a * a2; + a3 * a4 + } + + #[inline] + fn sbox_p_extension, const D: usize>(a: &F) -> F { + let a2 = a.square(); + let a4 = a2.square(); + let a3 = *a * a2; + a3 * a4 + } + #[inline] #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { From 6c55225fe5ceda9e40e0ae7eae244cbb1c6646d9 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 17:35:42 -0500 Subject: [PATCH 16/26] make the tests more thorough --- .../conditional_recursive_verifier.rs | 25 +- plonky2/src/recursion/dummy_circuit.rs | 45 --- plonky2/src/recursion/recursive_verifier.rs | 285 +++++++++++++----- 3 files changed, 228 insertions(+), 127 deletions(-) diff --git a/plonky2/src/recursion/conditional_recursive_verifier.rs b/plonky2/src/recursion/conditional_recursive_verifier.rs index ab79d3bd9..6927866eb 100644 --- a/plonky2/src/recursion/conditional_recursive_verifier.rs +++ b/plonky2/src/recursion/conditional_recursive_verifier.rs @@ -349,28 +349,35 @@ mod tests { use anyhow::Result; use hashbrown::HashMap; + use plonky2_field::goldilocks_field::GoldilocksField; use super::*; use crate::field::types::Sample; use crate::gates::noop::NoopGate; use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_data::CircuitConfig; - use crate::plonk::config::Poseidon2GoldilocksConfig; + use crate::plonk::config::{Poseidon2GoldilocksConfig, PoseidonGoldilocksConfig}; use crate::recursion::dummy_circuit::{dummy_circuit, dummy_proof}; #[test] fn test_conditional_recursive_verifier() -> Result<()> { init_logger(); - const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; + conditional_recursive_verifier::()?; + conditional_recursive_verifier::() + } + + fn conditional_recursive_verifier + 'static>( + ) -> Result<()> + where + C::Hasher: AlgebraicHasher, + { let config = CircuitConfig::standard_recursion_config(); // Generate proof. - let mut builder = CircuitBuilder::::new(config.clone()); + let mut builder = CircuitBuilder::::new(config.clone()); let mut pw = PartialWitness::new(); let t = builder.add_virtual_target(); - pw.set_target(t, F::rand())?; + pw.set_target(t, C::F::rand())?; builder.register_public_input(t); let _t2 = builder.square(t); for _ in 0..64 { @@ -385,19 +392,19 @@ mod tests { let dummy_proof = dummy_proof(&dummy_data, HashMap::new())?; // Conditionally verify the two proofs. - let mut builder = CircuitBuilder::::new(config); + let mut builder = CircuitBuilder::::new(config); let mut pw = PartialWitness::new(); let pt = builder.add_virtual_proof_with_pis(&data.common); pw.set_proof_with_pis_target(&pt, &proof)?; let dummy_pt = builder.add_virtual_proof_with_pis(&data.common); - pw.set_proof_with_pis_target::(&dummy_pt, &dummy_proof)?; + pw.set_proof_with_pis_target::(&dummy_pt, &dummy_proof)?; let inner_data = builder.add_virtual_verifier_data(data.common.config.fri_config.cap_height); pw.set_verifier_data_target(&inner_data, &data.verifier_only)?; let dummy_inner_data = builder.add_virtual_verifier_data(data.common.config.fri_config.cap_height); pw.set_verifier_data_target(&dummy_inner_data, &dummy_data.verifier_only)?; - let b = builder.constant_bool(F::rand().0 % 2 == 0); + let b = builder.constant_bool(C::F::rand().0 % 2 == 0); builder.conditionally_verify_proof::( b, &pt, diff --git a/plonky2/src/recursion/dummy_circuit.rs b/plonky2/src/recursion/dummy_circuit.rs index 211047bc2..e9bdd4641 100644 --- a/plonky2/src/recursion/dummy_circuit.rs +++ b/plonky2/src/recursion/dummy_circuit.rs @@ -90,10 +90,6 @@ where pub fn dummy_circuit, C: GenericConfig, const D: usize>( common_data: &CommonCircuitData, ) -> CircuitData { - println!( - "dummy_circuit received common_data.gates: {:?}", - common_data.gates - ); let config = common_data.config.clone(); assert!( !common_data.config.zero_knowledge, @@ -118,47 +114,6 @@ pub fn dummy_circuit, C: GenericConfig, c let circuit = builder.build::(); - // Assert individual components to identify mismatches - assert_eq!(circuit.common.config, common_data.config, "config mismatch"); - assert_eq!( - circuit.common.fri_params, common_data.fri_params, - "fri_params mismatch" - ); - assert_eq!(circuit.common.gates, common_data.gates, "gates mismatch"); - assert_eq!( - circuit.common.selectors_info, common_data.selectors_info, - "selectors_info mismatch" - ); - assert_eq!( - circuit.common.quotient_degree_factor, common_data.quotient_degree_factor, - "quotient_degree_factor mismatch" - ); - assert_eq!( - circuit.common.num_gate_constraints, common_data.num_gate_constraints, - "num_gate_constraints mismatch" - ); - assert_eq!( - circuit.common.num_constants, common_data.num_constants, - "num_constants mismatch" - ); - assert_eq!( - circuit.common.num_public_inputs, common_data.num_public_inputs, - "num_public_inputs mismatch" - ); - assert_eq!(circuit.common.k_is, common_data.k_is, "k_is mismatch"); - assert_eq!( - circuit.common.num_partial_products, common_data.num_partial_products, - "num_partial_products mismatch" - ); - assert_eq!( - circuit.common.num_lookup_polys, common_data.num_lookup_polys, - "num_lookup_polys mismatch" - ); - assert_eq!( - circuit.common.num_lookup_selectors, common_data.num_lookup_selectors, - "num_lookup_selectors mismatch" - ); - assert_eq!(circuit.common.luts, common_data.luts, "luts mismatch"); assert_eq!(&circuit.common, common_data); circuit } diff --git a/plonky2/src/recursion/recursive_verifier.rs b/plonky2/src/recursion/recursive_verifier.rs index df52ea002..ebc0af9ca 100644 --- a/plonky2/src/recursion/recursive_verifier.rs +++ b/plonky2/src/recursion/recursive_verifier.rs @@ -205,6 +205,7 @@ mod tests { use anyhow::Result; use itertools::Itertools; use log::{info, Level}; + use plonky2_field::goldilocks_field::GoldilocksField; use super::*; use crate::fri::reduction_strategies::FriReductionStrategy; @@ -214,7 +215,9 @@ mod tests { use crate::gates::noop::NoopGate; use crate::iop::witness::{PartialWitness, WitnessWrite}; use crate::plonk::circuit_data::{CircuitConfig, VerifierOnlyCircuitData}; - use crate::plonk::config::{KeccakGoldilocksConfig, Poseidon2GoldilocksConfig}; + use crate::plonk::config::{ + KeccakGoldilocksConfig, Poseidon2GoldilocksConfig, PoseidonGoldilocksConfig, + }; use crate::plonk::proof::{CompressedProofWithPublicInputs, ProofWithPublicInputs}; use crate::plonk::prover::prove; use crate::util::timing::TimingTree; @@ -223,14 +226,27 @@ mod tests { fn test_recursive_verifier() -> Result<()> { init_logger(); const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; - let config = CircuitConfig::standard_recursion_zk_config(); + { + type C = PoseidonGoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_zk_config(); + + let (proof, vd, common_data) = dummy_proof::(&config, 4_000)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } - let (proof, vd, common_data) = dummy_proof::(&config, 4_000)?; - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, true, true)?; - test_serialization(&proof, &vd, &common_data)?; + { + type C = Poseidon2GoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_zk_config(); + + let (proof, vd, common_data) = dummy_proof::(&config, 4_000)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } Ok(()) } @@ -239,14 +255,26 @@ mod tests { fn test_recursive_verifier_one_lookup() -> Result<()> { init_logger(); const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; - let config = CircuitConfig::standard_recursion_zk_config(); - - let (proof, vd, common_data) = dummy_lookup_proof::(&config, 10)?; - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, true, true)?; - test_serialization(&proof, &vd, &common_data)?; + { + type C = PoseidonGoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_zk_config(); + + let (proof, vd, common_data) = dummy_lookup_proof::(&config, 10)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } + { + type C = Poseidon2GoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_zk_config(); + + let (proof, vd, common_data) = dummy_lookup_proof::(&config, 10)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } Ok(()) } @@ -255,15 +283,26 @@ mod tests { fn test_recursive_verifier_two_luts() -> Result<()> { init_logger(); const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; - let config = CircuitConfig::standard_recursion_config(); - - let (proof, vd, common_data) = dummy_two_luts_proof::(&config)?; - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, true, true)?; - test_serialization(&proof, &vd, &common_data)?; - + { + type C = PoseidonGoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_config(); + + let (proof, vd, common_data) = dummy_two_luts_proof::(&config)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } + { + type C = Poseidon2GoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_config(); + + let (proof, vd, common_data) = dummy_two_luts_proof::(&config)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } Ok(()) } @@ -271,15 +310,26 @@ mod tests { fn test_recursive_verifier_too_many_rows() -> Result<()> { init_logger(); const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; - let config = CircuitConfig::standard_recursion_config(); - - let (proof, vd, common_data) = dummy_too_many_rows_proof::(&config)?; - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, true, true)?; - test_serialization(&proof, &vd, &common_data)?; - + { + type C = PoseidonGoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_config(); + + let (proof, vd, common_data) = dummy_too_many_rows_proof::(&config)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } + { + type C = Poseidon2GoldilocksConfig; + type F = >::F; + let config = CircuitConfig::standard_recursion_config(); + + let (proof, vd, common_data) = dummy_too_many_rows_proof::(&config)?; + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + test_serialization(&proof, &vd, &common_data)?; + } Ok(()) } @@ -287,27 +337,65 @@ mod tests { fn test_recursive_recursive_verifier() -> Result<()> { init_logger(); const D: usize = 2; - type C = Poseidon2GoldilocksConfig; - type F = >::F; - - let config = CircuitConfig::standard_recursion_config(); - - // Start with a degree 2^14 proof - let (proof, vd, common_data) = dummy_proof::(&config, 16_000)?; - assert_eq!(common_data.degree_bits(), 14); - - // Shrink it to 2^13. - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, Some(13), false, false)?; - assert_eq!(common_data.degree_bits(), 13); - - // Shrink it to 2^12. - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, true, true)?; - assert_eq!(common_data.degree_bits(), 12); - - test_serialization(&proof, &vd, &common_data)?; + { + type C = PoseidonGoldilocksConfig; + type F = >::F; + + let config = CircuitConfig::standard_recursion_config(); + + // Start with a degree 2^14 proof + let (proof, vd, common_data) = dummy_proof::(&config, 16_000)?; + assert_eq!(common_data.degree_bits(), 14); + + // Shrink it to 2^13. + let (proof, vd, common_data) = recursive_proof::( + proof, + vd, + common_data, + &config, + Some(13), + false, + false, + )?; + assert_eq!(common_data.degree_bits(), 13); + + // Shrink it to 2^12. + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + assert_eq!(common_data.degree_bits(), 12); + + test_serialization(&proof, &vd, &common_data)?; + } + { + type C = Poseidon2GoldilocksConfig; + type F = >::F; + + let config = CircuitConfig::standard_recursion_config(); + + // Start with a degree 2^14 proof + let (proof, vd, common_data) = dummy_proof::(&config, 16_000)?; + assert_eq!(common_data.degree_bits(), 14); + + // Shrink it to 2^13. + let (proof, vd, common_data) = recursive_proof::( + proof, + vd, + common_data, + &config, + Some(13), + false, + false, + )?; + assert_eq!(common_data.degree_bits(), 13); + + // Shrink it to 2^12. + let (proof, vd, common_data) = + recursive_proof::(proof, vd, common_data, &config, None, true, true)?; + assert_eq!(common_data.degree_bits(), 12); + + test_serialization(&proof, &vd, &common_data)?; + } Ok(()) } @@ -316,20 +404,26 @@ mod tests { #[test] #[ignore] fn test_size_optimized_recursion() -> Result<()> { + size_optimized_recursion::()?; + size_optimized_recursion::() + } + + fn size_optimized_recursion>() -> Result<()> + where + C::Hasher: AlgebraicHasher, + { init_logger(); const D: usize = 2; - type C = Poseidon2GoldilocksConfig; type KC = KeccakGoldilocksConfig; - type F = >::F; let standard_config = CircuitConfig::standard_recursion_config(); // An initial dummy proof. - let (proof, vd, common_data) = dummy_proof::(&standard_config, 4_000)?; + let (proof, vd, common_data) = dummy_proof::(&standard_config, 4_000)?; assert_eq!(common_data.degree_bits(), 12); // A standard recursive proof. - let (proof, vd, common_data) = recursive_proof::( + let (proof, vd, common_data) = recursive_proof::( proof, vd, common_data, @@ -350,7 +444,7 @@ mod tests { }, ..standard_config }; - let (proof, vd, common_data) = recursive_proof::( + let (proof, vd, common_data) = recursive_proof::( proof, vd, common_data, @@ -373,7 +467,7 @@ mod tests { }, ..high_rate_config }; - let (proof, vd, common_data) = recursive_proof::( + let (proof, vd, common_data) = recursive_proof::( proof, vd, common_data, @@ -393,21 +487,66 @@ mod tests { fn test_recursive_verifier_multi_hash() -> Result<()> { init_logger(); const D: usize = 2; - type PC = Poseidon2GoldilocksConfig; type KC = KeccakGoldilocksConfig; - type F = >::F; - - let config = CircuitConfig::standard_recursion_config(); - let (proof, vd, common_data) = dummy_proof::(&config, 4_000)?; - - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, false, false)?; - test_serialization(&proof, &vd, &common_data)?; - - let (proof, vd, common_data) = - recursive_proof::(proof, vd, common_data, &config, None, false, false)?; - test_serialization(&proof, &vd, &common_data)?; + { + type PC = PoseidonGoldilocksConfig; + type F = >::F; + + let config = CircuitConfig::standard_recursion_config(); + let (proof, vd, common_data) = dummy_proof::(&config, 4_000)?; + + let (proof, vd, common_data) = recursive_proof::( + proof, + vd, + common_data, + &config, + None, + false, + false, + )?; + test_serialization(&proof, &vd, &common_data)?; + + let (proof, vd, common_data) = recursive_proof::( + proof, + vd, + common_data, + &config, + None, + false, + false, + )?; + test_serialization(&proof, &vd, &common_data)?; + } + { + type PC = Poseidon2GoldilocksConfig; + type F = >::F; + + let config = CircuitConfig::standard_recursion_config(); + let (proof, vd, common_data) = dummy_proof::(&config, 4_000)?; + + let (proof, vd, common_data) = recursive_proof::( + proof, + vd, + common_data, + &config, + None, + false, + false, + )?; + test_serialization(&proof, &vd, &common_data)?; + + let (proof, vd, common_data) = recursive_proof::( + proof, + vd, + common_data, + &config, + None, + false, + false, + )?; + test_serialization(&proof, &vd, &common_data)?; + } Ok(()) } From 52d6678832d722a3d5389473687c56e602b60a8a Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 17:45:04 -0500 Subject: [PATCH 17/26] try to fix CI --- .github/workflows/continuous-integration-workflow.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/continuous-integration-workflow.yml b/.github/workflows/continuous-integration-workflow.yml index f38bba611..4b51661d1 100644 --- a/.github/workflows/continuous-integration-workflow.yml +++ b/.github/workflows/continuous-integration-workflow.yml @@ -64,9 +64,8 @@ jobs: uses: actions/checkout@v4 - name: Install nightly toolchain - uses: dtolnay/rust-toolchain@master + uses: dtolnay/rust-toolchain@nightly with: - toolchain: nightly-2024-06-23 components: rustfmt, clippy - name: Rust Cache From e314e71adadab5c0260aab2a37e124e2cd1d38f5 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 17:47:25 -0500 Subject: [PATCH 18/26] more CI fix --- field/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/field/src/lib.rs b/field/src/lib.rs index c713db885..9a2ea4f9c 100644 --- a/field/src/lib.rs +++ b/field/src/lib.rs @@ -4,7 +4,6 @@ #![deny(rustdoc::broken_intra_doc_links)] #![deny(missing_debug_implementations)] #![feature(specialization)] -#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![cfg_attr(not(test), no_std)] extern crate alloc; From ce52ddd0891ebe7dabed59b67ccdacd3a4560fe0 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Wed, 12 Nov 2025 17:49:27 -0500 Subject: [PATCH 19/26] clippy --- plonky2/benches/hashing.rs | 6 +++--- plonky2/examples/fibonacci.rs | 1 - plonky2/src/batch_fri/oracle.rs | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/plonky2/benches/hashing.rs b/plonky2/benches/hashing.rs index f19ad997b..05e8731cb 100644 --- a/plonky2/benches/hashing.rs +++ b/plonky2/benches/hashing.rs @@ -95,9 +95,9 @@ pub(crate) fn bench_p3_poseidon2(c: &mut Criterion) { b.iter_batched( || { let mut state = [Goldilocks::zero(); WIDTH]; - for i in 0..WIDTH { - state[i] = Goldilocks::from_canonical_u64(rand::random::()); - } + state.iter_mut().for_each(|item| { + *item = Goldilocks::from_canonical_u64(rand::random::()); + }); state }, |mut state| { diff --git a/plonky2/examples/fibonacci.rs b/plonky2/examples/fibonacci.rs index ad751626a..be491aa21 100644 --- a/plonky2/examples/fibonacci.rs +++ b/plonky2/examples/fibonacci.rs @@ -1,7 +1,6 @@ use std::time::Instant; use anyhow::Result; -use env_logger; use log::Level; use plonky2::field::types::Field; use plonky2::iop::witness::{PartialWitness, WitnessWrite}; diff --git a/plonky2/src/batch_fri/oracle.rs b/plonky2/src/batch_fri/oracle.rs index 58deeaa3c..1f34cfef3 100644 --- a/plonky2/src/batch_fri/oracle.rs +++ b/plonky2/src/batch_fri/oracle.rs @@ -460,7 +460,7 @@ mod test { &fri_instances, &fri_openings, &fri_challenges, - &[merkle_cap.clone()], + std::slice::from_ref(&merkle_cap), &proof, &fri_params, )?; From c5362ac81d034e5de4c107709048fd01db4544a6 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Thu, 13 Nov 2025 08:10:08 -0500 Subject: [PATCH 20/26] roll back to lighter's version of poseidon for easy review --- .../arch/aarch64/poseidon_goldilocks_neon.rs | 25 ---- plonky2/src/hash/poseidon2/gate.rs | 68 +-------- plonky2/src/hash/poseidon2/hash.rs | 139 +++++------------- 3 files changed, 37 insertions(+), 195 deletions(-) diff --git a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs index 17078ff99..69d8cd54b 100644 --- a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs +++ b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs @@ -915,28 +915,3 @@ pub unsafe fn mds_layer(state: &[GoldilocksField; WIDTH]) -> [GoldilocksField; W let state = mds_layer_full(state); wrap_state(state) } - -#[inline(always)] -#[unroll::unroll_for_loops] -pub unsafe fn vector_add(a: &[u64; WIDTH], b: &[u64; WIDTH]) -> [u64; WIDTH] { - let mut res = [0u64; WIDTH]; - // Process 2 elements at a time using NEON - for i in (0..WIDTH).step_by(2) { - let a_vec = vld1q_u64(a[i..].as_ptr()); - let b_vec = vld1q_u64(b[i..].as_ptr()); - - // Add the round constants - let sum = vaddq_u64(a_vec, b_vec); - - // Check for overflow (if sum < state_vec, we wrapped around) - let overflow_mask = vcltq_u64(sum, a_vec); - - // Add EPSILON (0xffffffff) where overflow occurred - let epsilon = vdupq_n_u64(0xffffffff); - let adjustment = vandq_u64(overflow_mask, epsilon); - let result = vaddq_u64(sum, adjustment); - - vst1q_u64(res[i..].as_mut_ptr(), result); - } - res -} diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index 9f9413438..6227476fa 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -364,7 +364,7 @@ impl + Poseidon2, const D: usize> Gate for Po } fn num_constants(&self) -> usize { - 1 + 0 } fn degree(&self) -> usize { @@ -495,12 +495,9 @@ impl + Poseidon2, const D: usize> SimpleGenerator>::F; - - let config = CircuitConfig { - num_wires: 143, - ..CircuitConfig::standard_recursion_config() - }; - let mut builder = CircuitBuilder::new(config); - type Gate = Poseidon2Gate; - let gate = Gate::new(); - let row = builder.add_gate(gate, vec![]); - let circuit = builder.build_prover::(); - - let permutation_inputs = (0..WIDTH).map(F::from_canonical_usize).collect::>(); - - let mut inputs = PartialWitness::new(); - inputs - .set_wire( - Wire { - row, - column: Gate::WIRE_SWAP, - }, - F::ZERO, - ) - .unwrap(); - for i in 0..WIDTH { - inputs - .set_wire( - Wire { - row, - column: Gate::wire_input(i), - }, - permutation_inputs[i], - ) - .unwrap(); - } - - let witness = - generate_partial_witness(inputs, &circuit.prover_only, &circuit.common).unwrap(); - - let expected_outputs: [F; WIDTH] = F::poseidon2(permutation_inputs.try_into().unwrap()); - for i in 0..WIDTH { - let out = witness.get_wire(Wire { - row: 0, - column: Gate::wire_output(i), - }); - assert_eq!(out, expected_outputs[i]); - } } #[test] diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index a1d3c9c9c..1d4ec2148 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,7 +1,5 @@ use core::fmt::Debug; -use plonky2_field::ops::Square; - use super::config::*; use super::gate::Poseidon2Gate; use crate::field::extension::{Extendable, FieldExtension}; @@ -15,7 +13,6 @@ use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::config::{AlgebraicHasher, Hasher}; pub trait Poseidon2: PrimeField64 { - #[inline] fn poseidon2(input: [Self; WIDTH]) -> [Self; WIDTH] { let mut state = input; @@ -29,7 +26,6 @@ pub trait Poseidon2: PrimeField64 { } #[inline] - #[unroll::unroll_for_loops] fn full_rounds(state: &mut [Self; WIDTH], start: usize) { for r in start..(start + ROUNDS_F_HALF) { Self::add_rc(state, r); @@ -39,7 +35,6 @@ pub trait Poseidon2: PrimeField64 { } #[inline] - #[unroll::unroll_for_loops] fn partial_rounds(state: &mut [Self; WIDTH]) { for r in 0..ROUNDS_P { state[0] += Self::from_canonical_u64(INTERNAL_CONSTANTS[r]); @@ -49,7 +44,6 @@ pub trait Poseidon2: PrimeField64 { } #[inline] - #[unroll::unroll_for_loops] fn external_linear_layer(state: &mut [Self; 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'. @@ -73,7 +67,6 @@ pub trait Poseidon2: PrimeField64 { } #[inline] - #[unroll::unroll_for_loops] fn external_linear_layer_extension, const D: usize>( state: &mut [F; WIDTH], ) { @@ -99,16 +92,11 @@ pub trait Poseidon2: PrimeField64 { } #[inline] - #[unroll::unroll_for_loops] fn internal_linear_layer(state: &mut [Self; WIDTH]) { - let tmp = state - .iter() - .map(|&x| x.to_noncanonical_u64() as u128) - .sum::(); - let sum = Self::from_noncanonical_u128_with_96_bits(tmp); + let sum: Self = state.iter().cloned().sum(); for i in 0..WIDTH { - state[i] = - sum.multiply_accumulate(state[i], Self::from_canonical_u64(MATRIX_DIAG_12_U64[i])); + state[i] *= Self::from_canonical_u64(MATRIX_DIAG_12_U64[i]); + state[i] += sum; } } @@ -116,19 +104,28 @@ pub trait Poseidon2: PrimeField64 { fn internal_linear_layer_extension, const D: usize>( state: &mut [F; WIDTH], ) { - let sum: F = state.iter().cloned().sum(); - state - .iter_mut() - .zip(MATRIX_DIAG_12_U64.iter()) - .for_each(|(x, &m)| { - *x = sum.multiply_accumulate(*x, F::from_canonical_u64(m)); - }); + let mut sum = state[0]; + for i in 1..WIDTH { + sum += state[i]; + } + for i in 0..WIDTH { + state[i] *= F::from_canonical_u64(MATRIX_DIAG_12_U64[i]); + state[i] += sum; + } } - fn add_rc(state: &mut [Self; WIDTH], external_round: usize); + #[inline] + fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + + for i in 0..WIDTH { + unsafe { + state[i] = state[i].add_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); + } + } + } #[inline] - #[unroll::unroll_for_loops] fn add_rc_extension, const D: usize>( state: &mut [F; WIDTH], external_round: usize, @@ -140,7 +137,10 @@ pub trait Poseidon2: PrimeField64 { } } - fn sbox(state: &mut [Self; WIDTH]); + #[inline] + fn sbox(state: &mut [Self; WIDTH]) { + state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); + } #[inline] fn sbox_extension, const D: usize>( @@ -151,9 +151,15 @@ pub trait Poseidon2: PrimeField64 { .for_each(|a| *a = Self::sbox_p_extension(a)); } - fn sbox_p(a: &Self) -> Self; + #[inline] + fn sbox_p(a: &Self) -> Self { + a.exp_u64(D) + } - fn sbox_p_extension, const D: usize>(a: &F) -> F; + #[inline] + fn sbox_p_extension, const D: usize>(a: &F) -> F { + a.exp_u64(super::config::D) + } // Multiply a 4-element vector x by: // [ 2 3 1 1 ] @@ -161,7 +167,6 @@ pub trait Poseidon2: PrimeField64 { // [ 1 1 2 3 ] // [ 3 1 1 2 ]. // This is more efficient than the previous matrix. - #[inline] fn apply_mat4_mut(x: &mut [Self; 4]) { let t01 = x[0] + x[1]; let t23 = x[2] + x[3]; @@ -175,7 +180,6 @@ pub trait Poseidon2: PrimeField64 { x[2] = t01233 + t23; // x[0] + x[1] + 2*x[2] + 3*x[3] } - #[inline] fn apply_mat4_mut_extension, const D: usize>( x: &mut [F; 4], ) { @@ -192,8 +196,7 @@ pub trait Poseidon2: PrimeField64 { } // In circuit functions - #[inline] - #[unroll::unroll_for_loops] + fn external_linear_layer_circuit( builder: &mut CircuitBuilder, state: &mut [ExtensionTarget; WIDTH], @@ -226,8 +229,6 @@ pub trait Poseidon2: PrimeField64 { } } - #[inline] - #[unroll::unroll_for_loops] fn apply_mat4_mut_circuit( builder: &mut CircuitBuilder, x: &mut [ExtensionTarget; 4], @@ -250,8 +251,6 @@ pub trait Poseidon2: PrimeField64 { x[2] = builder.add_extension(t01233, t23); // x[0] + x[1] + 2*x[2] + 3*x[3] } - #[inline] - #[unroll::unroll_for_loops] fn matmul_m4_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -279,8 +278,6 @@ pub trait Poseidon2: PrimeField64 { } } - #[inline] - #[unroll::unroll_for_loops] fn add_rc_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -296,8 +293,6 @@ pub trait Poseidon2: PrimeField64 { } } - #[inline] - #[unroll::unroll_for_loops] fn sbox_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -320,8 +315,6 @@ pub trait Poseidon2: PrimeField64 { builder.exp_u64_extension(input, super::config::D) } - #[inline] - #[unroll::unroll_for_loops] fn internal_linear_layer_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -342,67 +335,7 @@ pub trait Poseidon2: PrimeField64 { } } -impl Poseidon2 for F { - #[inline] - fn sbox_p(a: &Self) -> Self { - let a2 = a.square(); - let a4 = a2.square(); - let a3 = *a * a2; - a3 * a4 - } - - #[inline] - fn sbox_p_extension, const D: usize>(a: &F) -> F { - let a2 = a.square(); - let a4 = a2.square(); - let a3 = *a * a2; - a3 * a4 - } - - #[inline] - #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] - fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { - debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); - state - .iter_mut() - .zip(EXTERNAL_CONSTANTS[external_round].iter()) - .for_each(|(x, &m)| { - *x += Self::from_canonical_u64(m); - }); - } - - #[inline] - #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] - fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { - debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); - - unsafe { - use core::mem::transmute; - - use crate::hash::arch::aarch64::poseidon_goldilocks_neon::vector_add; - - let state_u64 = transmute::<[Self; WIDTH], [u64; WIDTH]>(*state); - let round_constants = &EXTERNAL_CONSTANTS[external_round]; - - let res = vector_add(&state_u64, round_constants); - *state = transmute::<[u64; WIDTH], [Self; WIDTH]>(res); - } - } - - #[inline] - #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] - fn sbox(state: &mut [Self; WIDTH]) { - state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); - } - - #[inline(always)] - #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] - fn sbox(state: &mut [Self; WIDTH]) { - unsafe { - crate::hash::arch::aarch64::poseidon_goldilocks_neon::sbox_layer(state); - } - } -} +impl Poseidon2 for F {} #[derive(Copy, Clone, Default, Debug, PartialEq)] pub struct Poseidon2Permutation { @@ -490,8 +423,6 @@ impl Hasher for Poseidon2Hash { } impl Poseidon2Hash { - #[inline] - #[unroll::unroll_for_loops] pub fn hash_n_to_one( input: &[>::Hash], ) -> >::Hash { From 959ddbae761e013db223d72fc5c776d102d4f410 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Thu, 13 Nov 2025 08:21:24 -0500 Subject: [PATCH 21/26] Update gate.rs --- plonky2/src/hash/poseidon2/gate.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index 6227476fa..38b2b4a30 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -495,7 +495,7 @@ impl + Poseidon2, const D: usize> SimpleGenerator Date: Tue, 9 Dec 2025 11:12:33 -0500 Subject: [PATCH 22/26] address comments --- field/src/goldilocks_field.rs | 1 + field/src/lib.rs | 1 + plonky2/src/hash/poseidon2/gate.rs | 7 ++ plonky2/src/hash/poseidon2/mod.rs | 1 - plonky2/src/hash/poseidon2/pure.rs | 142 ----------------------------- 5 files changed, 9 insertions(+), 143 deletions(-) delete mode 100644 plonky2/src/hash/poseidon2/pure.rs diff --git a/field/src/goldilocks_field.rs b/field/src/goldilocks_field.rs index 2c3d26955..a4da1fee0 100644 --- a/field/src/goldilocks_field.rs +++ b/field/src/goldilocks_field.rs @@ -161,6 +161,7 @@ impl Field for GoldilocksField { } fn from_noncanonical_u128_with_96_bits(n: u128) -> Self { + debug_assert!(n < (1u128 << 96)); reduce128_with_96_bits(n) } diff --git a/field/src/lib.rs b/field/src/lib.rs index 9a2ea4f9c..c713db885 100644 --- a/field/src/lib.rs +++ b/field/src/lib.rs @@ -4,6 +4,7 @@ #![deny(rustdoc::broken_intra_doc_links)] #![deny(missing_debug_implementations)] #![feature(specialization)] +#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![cfg_attr(not(test), no_std)] extern crate alloc; diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index 38b2b4a30..b68b05067 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -498,6 +498,7 @@ mod tests { use super::Poseidon2Gate; use crate::field::goldilocks_field::GoldilocksField; use crate::gates::gate_testing::{test_eval_fns, test_low_degree}; + use crate::gates::poseidon::PoseidonGate; use crate::plonk::config::{GenericConfig, Poseidon2GoldilocksConfig}; #[test] @@ -518,6 +519,9 @@ mod tests { fn low_degree() { type F = GoldilocksField; let gate = Poseidon2Gate::::new(); + test_low_degree(gate); + + let gate = PoseidonGate::::new(); test_low_degree(gate) } @@ -526,6 +530,9 @@ mod tests { const D: usize = 2; type C = Poseidon2GoldilocksConfig; type F = >::F; + let gate = Poseidon2Gate::::new(); + test_eval_fns::(gate)?; + let gate = Poseidon2Gate::::new(); test_eval_fns::(gate) } diff --git a/plonky2/src/hash/poseidon2/mod.rs b/plonky2/src/hash/poseidon2/mod.rs index 3abdc563c..52ede9a4d 100644 --- a/plonky2/src/hash/poseidon2/mod.rs +++ b/plonky2/src/hash/poseidon2/mod.rs @@ -4,4 +4,3 @@ pub mod hash; #[cfg(test)] pub mod p3; -pub mod pure; diff --git a/plonky2/src/hash/poseidon2/pure.rs b/plonky2/src/hash/poseidon2/pure.rs deleted file mode 100644 index dda97d50c..000000000 --- a/plonky2/src/hash/poseidon2/pure.rs +++ /dev/null @@ -1,142 +0,0 @@ -#![allow(clippy::all)] - -use plonky2_field::goldilocks_field::GoldilocksField; - -use super::config::*; -/// This is a test implementation of the Poseidon2 permutation without any gate or generator optimizations. -use crate::field::types::Field; -use crate::iop::target::Target; -use crate::plonk::circuit_builder::CircuitBuilder; - -type F = GoldilocksField; -type Builder = CircuitBuilder; - -pub fn permute_swapped(builder: &mut Builder, inputs: &[Target; WIDTH]) -> [Target; WIDTH] { - let mut state: [Target; 12] = inputs.clone(); - external_permute_mut(builder, &mut state); - - // The first half of the external rounds. - for r in 0..ROUNDS_F_HALF { - let external_constants = EXTERNAL_CONSTANTS[r]; - let external_constants: [F; WIDTH] = external_constants - .iter() - .map(|&x| F::from_canonical_u64(x)) - .collect::>() - .try_into() - .unwrap(); - let external_constants: [Target; WIDTH] = - builder.constants(&external_constants).try_into().unwrap(); - add_rc(builder, &mut state, &external_constants); - sbox(builder, &mut state); - external_permute_mut(builder, &mut state); - } - - // The internal rounds. - for r in 0..ROUNDS_P { - let internal_constant = INTERNAL_CONSTANTS[r]; - let internal_constant = F::from_canonical_u64(internal_constant); - let internal_constant = builder.constant(internal_constant); - state[0] = builder.add(state[0], internal_constant); - state[0] = sbox_p(builder, state[0]); - internal_permute_mut(builder, &mut state); - } - - // The second half of the external rounds. - for r in ROUNDS_F_HALF..ROUNDS_F { - let external_constants = EXTERNAL_CONSTANTS[r]; - let external_constants: [F; WIDTH] = external_constants - .iter() - .map(|&x| F::from_canonical_u64(x)) - .collect::>() - .try_into() - .unwrap(); - let external_constants: [Target; WIDTH] = - builder.constants(&external_constants).try_into().unwrap(); - add_rc(builder, &mut state, &external_constants); - sbox(builder, &mut state); - external_permute_mut(builder, &mut state); - } - - state -} - -fn external_permute_mut(builder: &mut Builder, state: &mut [Target; 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) { - // Would be nice to find a better way to do this. - let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; - apply_mat4(builder, &mut state_4); - state[i..i + 4].clone_from_slice(&state_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: [Target; 4] = core::array::from_fn(|k| { - (0..WIDTH) - .step_by(4) - .map(|j| state[j + k]) - .reduce(|acc, t| builder.add(acc, t)) - .unwrap() - }); - - // 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] = builder.add(state[i], sums[i % 4]); - } -} - -fn add_rc(builder: &mut Builder, state: &mut [Target; WIDTH], external_constant: &[Target; WIDTH]) { - state - .iter_mut() - .zip(external_constant) - .for_each(|(a, b)| *a = builder.add(*a, *b)); -} - -fn sbox(builder: &mut Builder, state: &mut [Target; WIDTH]) { - state.iter_mut().for_each(|a| *a = sbox_p(builder, *a)); -} - -fn sbox_p(builder: &mut Builder, a: Target) -> Target { - let a2 = builder.mul(a, a); - let a3 = builder.mul(a2, a); - let a6 = builder.mul(a3, a3); - - builder.mul(a6, a) -} - -// Multiply a 4-element vector x by: -// [ 2 3 1 1 ] -// [ 1 2 3 1 ] -// [ 1 1 2 3 ] -// [ 3 1 1 2 ]. -// This is more efficient than the previous matrix. -fn apply_mat4(builder: &mut Builder, x: &mut [Target; 4]) { - let two = builder.constant(F::from_canonical_u64(2)); - - let t01 = builder.add(x[0], x[1]); - let t23 = builder.add(x[2], x[3]); - let t0123 = builder.add(t01, t23); - let t01123 = builder.add(t0123, x[1]); - let t01233 = builder.add(t0123, x[3]); - // The order here is important. Need to overwrite x[0] and x[2] after x[1] and x[3]. - let dx0 = builder.mul(x[0], two); - let dx2 = builder.mul(x[2], two); - x[3] = builder.add(t01233, dx0); // 3*x[0] + x[1] + x[2] + 2*x[3] - x[1] = builder.add(t01123, dx2); // x[0] + 2*x[1] + 3*x[2] + x[3] - x[0] = builder.add(t01123, t01); // 2*x[0] + 3*x[1] + x[2] + x[3] - x[2] = builder.add(t01233, t23); // x[0] + x[1] + 2*x[2] + 3*x[3] -} - -/// Given a vector v compute the matrix vector product (1 + diag(v))state with 1 denoting the constant matrix of ones. -pub fn internal_permute_mut(builder: &mut Builder, state: &mut [Target; WIDTH]) { - let sum = builder.add_many(state.iter().cloned()); - for i in 0..WIDTH { - let constant = MATRIX_DIAG_12_U64[i]; - let constant = F::from_canonical_u64(constant); - let constant = builder.constant(constant); - state[i] = builder.mul(state[i], constant); - state[i] = builder.add(state[i], sum); - } -} From ceb808eb1c51e6f05f14fa9a181e59e5c275f710 Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 9 Dec 2025 11:14:46 -0500 Subject: [PATCH 23/26] more debug assert --- field/src/types.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/field/src/types.rs b/field/src/types.rs index 28a01f839..153231c9a 100644 --- a/field/src/types.rs +++ b/field/src/types.rs @@ -358,6 +358,7 @@ pub trait Field: fn from_noncanonical_u128_with_96_bits(n: u128) -> Self { // Default implementation. + debug_assert!(n < (1u128 << 96)); Self::from_noncanonical_u128(n) } From 96e3fae6e1590e17b71c9792266151fe4a1ece9c Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 9 Dec 2025 11:15:46 -0500 Subject: [PATCH 24/26] fix typo --- plonky2/src/hash/poseidon2/gate.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index b68b05067..006b41be3 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -533,7 +533,7 @@ mod tests { let gate = Poseidon2Gate::::new(); test_eval_fns::(gate)?; - let gate = Poseidon2Gate::::new(); + let gate = PoseidonGate::::new(); test_eval_fns::(gate) } } From 7a900df066eb9ced49d01a1509a6648c52202b3b Mon Sep 17 00:00:00 2001 From: lighter-zz Date: Tue, 9 Dec 2025 11:32:24 -0500 Subject: [PATCH 25/26] remove stable feature --- field/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/field/src/lib.rs b/field/src/lib.rs index c713db885..9a2ea4f9c 100644 --- a/field/src/lib.rs +++ b/field/src/lib.rs @@ -4,7 +4,6 @@ #![deny(rustdoc::broken_intra_doc_links)] #![deny(missing_debug_implementations)] #![feature(specialization)] -#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))] #![cfg_attr(not(test), no_std)] extern crate alloc; From 6a056834ef49c6de8eda1888e26e2a70ac1aed35 Mon Sep 17 00:00:00 2001 From: "z.z." Date: Tue, 9 Dec 2025 12:49:38 -0500 Subject: [PATCH 26/26] optimize poseidon2 (#10) * optimize poseidon2 * Merge branch 'zz/impl_poseidon2_with_plonky3' into zz/optimize_poseidon2 * finally faster than poseidon * clean up * update num_constants for poseidon * reverting add_canonical_u64 * Update hash.rs * Update hash.rs --- .../arch/aarch64/poseidon_goldilocks_neon.rs | 25 ++ plonky2/src/hash/poseidon2/gate.rs | 67 ++++- plonky2/src/hash/poseidon2/hash.rs | 238 ++++++++++++------ 3 files changed, 254 insertions(+), 76 deletions(-) diff --git a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs index 69d8cd54b..17078ff99 100644 --- a/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs +++ b/plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs @@ -915,3 +915,28 @@ pub unsafe fn mds_layer(state: &[GoldilocksField; WIDTH]) -> [GoldilocksField; W let state = mds_layer_full(state); wrap_state(state) } + +#[inline(always)] +#[unroll::unroll_for_loops] +pub unsafe fn vector_add(a: &[u64; WIDTH], b: &[u64; WIDTH]) -> [u64; WIDTH] { + let mut res = [0u64; WIDTH]; + // Process 2 elements at a time using NEON + for i in (0..WIDTH).step_by(2) { + let a_vec = vld1q_u64(a[i..].as_ptr()); + let b_vec = vld1q_u64(b[i..].as_ptr()); + + // Add the round constants + let sum = vaddq_u64(a_vec, b_vec); + + // Check for overflow (if sum < state_vec, we wrapped around) + let overflow_mask = vcltq_u64(sum, a_vec); + + // Add EPSILON (0xffffffff) where overflow occurred + let epsilon = vdupq_n_u64(0xffffffff); + let adjustment = vandq_u64(overflow_mask, epsilon); + let result = vaddq_u64(sum, adjustment); + + vst1q_u64(res[i..].as_mut_ptr(), result); + } + res +} diff --git a/plonky2/src/hash/poseidon2/gate.rs b/plonky2/src/hash/poseidon2/gate.rs index 006b41be3..ff53335e2 100644 --- a/plonky2/src/hash/poseidon2/gate.rs +++ b/plonky2/src/hash/poseidon2/gate.rs @@ -1,6 +1,5 @@ //! Implementation of a Plonky2 gate for an entire Poseidon2 permutation over a //! state of width 12 - use core::marker::PhantomData; use anyhow::Result; @@ -495,10 +494,13 @@ impl + Poseidon2, const D: usize> SimpleGenerator>::F; + + let config = CircuitConfig { + num_wires: 143, + ..CircuitConfig::standard_recursion_config() + }; + let mut builder = CircuitBuilder::new(config); + type Gate = Poseidon2Gate; + let gate = Gate::new(); + let row = builder.add_gate(gate, vec![]); + let circuit = builder.build_prover::(); + + let permutation_inputs = (0..WIDTH).map(F::from_canonical_usize).collect::>(); + + let mut inputs = PartialWitness::new(); + inputs + .set_wire( + Wire { + row, + column: Gate::WIRE_SWAP, + }, + F::ZERO, + ) + .unwrap(); + for i in 0..WIDTH { + inputs + .set_wire( + Wire { + row, + column: Gate::wire_input(i), + }, + permutation_inputs[i], + ) + .unwrap(); + } + + let witness = + generate_partial_witness(inputs, &circuit.prover_only, &circuit.common).unwrap(); + + let expected_outputs: [F; WIDTH] = F::poseidon2(permutation_inputs.try_into().unwrap()); + for i in 0..WIDTH { + let out = witness.get_wire(Wire { + row: 0, + column: Gate::wire_output(i), + }); + assert_eq!(out, expected_outputs[i]); + } } #[test] diff --git a/plonky2/src/hash/poseidon2/hash.rs b/plonky2/src/hash/poseidon2/hash.rs index 1d4ec2148..5f9553fad 100644 --- a/plonky2/src/hash/poseidon2/hash.rs +++ b/plonky2/src/hash/poseidon2/hash.rs @@ -1,5 +1,7 @@ use core::fmt::Debug; +use plonky2_field::ops::Square; + use super::config::*; use super::gate::Poseidon2Gate; use crate::field::extension::{Extendable, FieldExtension}; @@ -13,6 +15,7 @@ use crate::plonk::circuit_builder::CircuitBuilder; use crate::plonk::config::{AlgebraicHasher, Hasher}; pub trait Poseidon2: PrimeField64 { + #[inline] fn poseidon2(input: [Self; WIDTH]) -> [Self; WIDTH] { let mut state = input; @@ -26,6 +29,7 @@ pub trait Poseidon2: PrimeField64 { } #[inline] + #[unroll::unroll_for_loops] fn full_rounds(state: &mut [Self; WIDTH], start: usize) { for r in start..(start + ROUNDS_F_HALF) { Self::add_rc(state, r); @@ -35,6 +39,7 @@ pub trait Poseidon2: PrimeField64 { } #[inline] + #[unroll::unroll_for_loops] fn partial_rounds(state: &mut [Self; WIDTH]) { for r in 0..ROUNDS_P { state[0] += Self::from_canonical_u64(INTERNAL_CONSTANTS[r]); @@ -44,29 +49,20 @@ pub trait Poseidon2: PrimeField64 { } #[inline] + #[unroll::unroll_for_loops] fn external_linear_layer(state: &mut [Self; 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) { - // Would be nice to find a better way to do this. - let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; - Self::apply_mat4_mut(&mut state_4); - state[i..i + 4].clone_from_slice(&state_4); + let mut state_u128: [u128; WIDTH] = [0u128; WIDTH]; + for i in 0..WIDTH { + state_u128[i] = state[i].to_noncanonical_u64() as u128; } - // 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: [Self; 4] = - core::array::from_fn(|k| (0..WIDTH).step_by(4).map(|j| state[j + k]).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 + external_linear_layer_u128(&mut state_u128); for i in 0..WIDTH { - state[i] += sums[i % 4]; + state[i] = Self::from_noncanonical_u128_with_96_bits(state_u128[i]); } } #[inline] + #[unroll::unroll_for_loops] fn external_linear_layer_extension, const D: usize>( state: &mut [F; WIDTH], ) { @@ -92,11 +88,12 @@ pub trait Poseidon2: PrimeField64 { } #[inline] + #[unroll::unroll_for_loops] fn internal_linear_layer(state: &mut [Self; WIDTH]) { - let sum: Self = state.iter().cloned().sum(); + let sum = sum_12(state); // hard coded for WIDTH = 12 for i in 0..WIDTH { - state[i] *= Self::from_canonical_u64(MATRIX_DIAG_12_U64[i]); - state[i] += sum; + state[i] = + sum.multiply_accumulate(state[i], Self::from_canonical_u64(MATRIX_DIAG_12_U64[i])); } } @@ -104,28 +101,19 @@ pub trait Poseidon2: PrimeField64 { fn internal_linear_layer_extension, const D: usize>( state: &mut [F; WIDTH], ) { - let mut sum = state[0]; - for i in 1..WIDTH { - sum += state[i]; - } - for i in 0..WIDTH { - state[i] *= F::from_canonical_u64(MATRIX_DIAG_12_U64[i]); - state[i] += sum; - } + let sum: F = state.iter().cloned().sum(); + state + .iter_mut() + .zip(MATRIX_DIAG_12_U64.iter()) + .for_each(|(x, &m)| { + *x = sum.multiply_accumulate(*x, F::from_canonical_u64(m)); + }); } - #[inline] - fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { - debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); - - for i in 0..WIDTH { - unsafe { - state[i] = state[i].add_canonical_u64(EXTERNAL_CONSTANTS[external_round][i]); - } - } - } + fn add_rc(state: &mut [Self; WIDTH], external_round: usize); #[inline] + #[unroll::unroll_for_loops] fn add_rc_extension, const D: usize>( state: &mut [F; WIDTH], external_round: usize, @@ -137,10 +125,7 @@ pub trait Poseidon2: PrimeField64 { } } - #[inline] - fn sbox(state: &mut [Self; WIDTH]) { - state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); - } + fn sbox(state: &mut [Self; WIDTH]); #[inline] fn sbox_extension, const D: usize>( @@ -151,35 +136,11 @@ pub trait Poseidon2: PrimeField64 { .for_each(|a| *a = Self::sbox_p_extension(a)); } - #[inline] - fn sbox_p(a: &Self) -> Self { - a.exp_u64(D) - } + fn sbox_p(a: &Self) -> Self; - #[inline] - fn sbox_p_extension, const D: usize>(a: &F) -> F { - a.exp_u64(super::config::D) - } - - // Multiply a 4-element vector x by: - // [ 2 3 1 1 ] - // [ 1 2 3 1 ] - // [ 1 1 2 3 ] - // [ 3 1 1 2 ]. - // This is more efficient than the previous matrix. - fn apply_mat4_mut(x: &mut [Self; 4]) { - let t01 = x[0] + x[1]; - let t23 = x[2] + x[3]; - let t0123 = t01 + t23; - let t01123 = t0123 + x[1]; - let t01233 = t0123 + x[3]; - // The order here is important. Need to overwrite x[0] and x[2] after x[1] and x[3]. - x[3] = t01233 + x[0].double(); // 3*x[0] + x[1] + x[2] + 2*x[3] - x[1] = t01123 + x[2].double(); // x[0] + 2*x[1] + 3*x[2] + x[3] - x[0] = t01123 + t01; // 2*x[0] + 3*x[1] + x[2] + x[3] - x[2] = t01233 + t23; // x[0] + x[1] + 2*x[2] + 3*x[3] - } + fn sbox_p_extension, const D: usize>(a: &F) -> F; + #[inline] fn apply_mat4_mut_extension, const D: usize>( x: &mut [F; 4], ) { @@ -196,7 +157,8 @@ pub trait Poseidon2: PrimeField64 { } // In circuit functions - + #[inline] + #[unroll::unroll_for_loops] fn external_linear_layer_circuit( builder: &mut CircuitBuilder, state: &mut [ExtensionTarget; WIDTH], @@ -206,10 +168,7 @@ pub trait Poseidon2: PrimeField64 { // 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) { - // Would be nice to find a better way to do this. - let mut state_4 = [state[i], state[i + 1], state[i + 2], state[i + 3]]; - Self::apply_mat4_mut_circuit(builder, &mut state_4); - state[i..i + 4].clone_from_slice(&state_4); + Self::apply_mat4_mut_circuit(builder, (&mut state[i..i + 4]).try_into().unwrap()); } // Now, we apply the outer circulant matrix (to compute the y_i values). @@ -229,6 +188,8 @@ pub trait Poseidon2: PrimeField64 { } } + #[inline] + #[unroll::unroll_for_loops] fn apply_mat4_mut_circuit( builder: &mut CircuitBuilder, x: &mut [ExtensionTarget; 4], @@ -251,6 +212,8 @@ pub trait Poseidon2: PrimeField64 { x[2] = builder.add_extension(t01233, t23); // x[0] + x[1] + 2*x[2] + 3*x[3] } + #[inline] + #[unroll::unroll_for_loops] fn matmul_m4_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -278,6 +241,8 @@ pub trait Poseidon2: PrimeField64 { } } + #[inline] + #[unroll::unroll_for_loops] fn add_rc_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -293,6 +258,8 @@ pub trait Poseidon2: PrimeField64 { } } + #[inline] + #[unroll::unroll_for_loops] fn sbox_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -315,6 +282,8 @@ pub trait Poseidon2: PrimeField64 { builder.exp_u64_extension(input, super::config::D) } + #[inline] + #[unroll::unroll_for_loops] fn internal_linear_layer_circuit( builder: &mut CircuitBuilder, input: &mut [ExtensionTarget; WIDTH], @@ -335,7 +304,106 @@ pub trait Poseidon2: PrimeField64 { } } -impl Poseidon2 for F {} +#[inline] +#[unroll::unroll_for_loops] +fn external_linear_layer_u128(state: &mut [u128; 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) { + // Multiply a 4-element vector x by: + // [ 2 3 1 1 ] + // [ 1 2 3 1 ] + // [ 1 1 2 3 ] + // [ 3 1 1 2 ]. + let t01 = state[i] + state[i + 1]; + let t23 = state[i + 2] + state[i + 3]; + let t0123 = t01 + t23; + + let x0 = state[i]; + let x2 = state[i + 2]; + + state[i] = t0123 + t01 + state[i + 1]; // 2*x[0] + 3*x[1] + x[2] + x[3] + state[i + 1] = t0123 + state[i + 1] + x2 + x2; // x[0] + 2*x[1] + 3*x[2] + x[3] + state[i + 2] = t0123 + t23 + state[i + 3]; // x[0] + x[1] + 2*x[2] + 3*x[3] + state[i + 3] = t0123 + state[i + 3] + x0 + x0; // 3*x[0] + x[1] + x[2] + 2*x[3] + } + // Now, we apply the outer circulant matrix (to compute the y_i values). + + // We first precompute the four sums of every four elements. + let mut sums = [0u128; 4]; + for i in 0..4 { + sums[i] = state[i] + state[i + 4] + state[i + 8]; + } + + // 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]; + } +} + +impl Poseidon2 for F { + #[inline] + fn sbox_p(a: &Self) -> Self { + let a2 = a.square(); + let a4 = a2.square(); + let a3 = *a * a2; + a3 * a4 + } + + #[inline] + fn sbox_p_extension, const D: usize>(a: &F) -> F { + let a2 = a.square(); + let a4 = a2.square(); + let a3 = *a * a2; + a3 * a4 + } + + #[inline] + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { + use plonky2_field::types::Field64; + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + state + .iter_mut() + .zip(EXTERNAL_CONSTANTS[external_round].iter()) + .for_each(|(x, &m)| { + *x = unsafe { x.add_canonical_u64(m) }; + }); + } + + #[inline] + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + fn add_rc(state: &mut [Self; WIDTH], external_round: usize) { + debug_assert!(external_round < EXTERNAL_CONSTANTS.len()); + + unsafe { + use core::mem::transmute; + + use crate::hash::arch::aarch64::poseidon_goldilocks_neon::vector_add; + + let state_u64 = transmute::<[Self; WIDTH], [u64; WIDTH]>(*state); + let round_constants = &EXTERNAL_CONSTANTS[external_round]; + + let res = vector_add(&state_u64, round_constants); + *state = transmute::<[u64; WIDTH], [Self; WIDTH]>(res); + } + } + + #[inline] + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + fn sbox(state: &mut [Self; WIDTH]) { + state.iter_mut().for_each(|a| *a = Self::sbox_p(a)); + } + + #[inline(always)] + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + fn sbox(state: &mut [Self; WIDTH]) { + unsafe { + crate::hash::arch::aarch64::poseidon_goldilocks_neon::sbox_layer(state); + } + } +} #[derive(Copy, Clone, Default, Debug, PartialEq)] pub struct Poseidon2Permutation { @@ -405,6 +473,26 @@ impl PlonkyPermutation< } } +#[inline] +/// Sum of 12 elements to u128; unrolled for performance. +fn sum_12(inputs: &[F]) -> F { + debug_assert!(inputs.len() == 12); + let tmp = inputs[0].to_noncanonical_u64() as u128 + + inputs[1].to_noncanonical_u64() as u128 + + inputs[2].to_noncanonical_u64() as u128 + + inputs[3].to_noncanonical_u64() as u128 + + inputs[4].to_noncanonical_u64() as u128 + + inputs[5].to_noncanonical_u64() as u128 + + inputs[6].to_noncanonical_u64() as u128 + + inputs[7].to_noncanonical_u64() as u128 + + inputs[8].to_noncanonical_u64() as u128 + + inputs[9].to_noncanonical_u64() as u128 + + inputs[10].to_noncanonical_u64() as u128 + + inputs[11].to_noncanonical_u64() as u128; + + F::from_noncanonical_u128_with_96_bits(tmp) +} + /// Poseidon2 hash function. #[derive(Copy, Clone, Debug, Eq, PartialEq)] pub struct Poseidon2Hash; @@ -423,6 +511,8 @@ impl Hasher for Poseidon2Hash { } impl Poseidon2Hash { + #[inline] + #[unroll::unroll_for_loops] pub fn hash_n_to_one( input: &[>::Hash], ) -> >::Hash {