diff --git a/halo2_gadgets/src/poseidon.rs b/halo2_gadgets/src/poseidon.rs index 07ba0df0a..34c1daa9e 100644 --- a/halo2_gadgets/src/poseidon.rs +++ b/halo2_gadgets/src/poseidon.rs @@ -1,6 +1,5 @@ //! The Poseidon algebraic hash function. -use std::convert::TryInto; use std::fmt; use std::marker::PhantomData; @@ -148,13 +147,7 @@ impl< pub fn new(chip: PoseidonChip, mut layouter: impl Layouter) -> Result { chip.initial_state(&mut layouter).map(|state| Sponge { chip, - mode: Absorbing( - (0..RATE) - .map(|_| None) - .collect::>() - .try_into() - .unwrap(), - ), + mode: Absorbing::init_empty(), state, _marker: PhantomData::default(), }) @@ -166,12 +159,10 @@ impl< mut layouter: impl Layouter, value: PaddedWord, ) -> Result<(), Error> { - for entry in self.mode.0.iter_mut() { - if entry.is_none() { - *entry = Some(value); - return Ok(()); - } - } + let value = match self.mode.absorb(value) { + Ok(()) => return Ok(()), + Err(value) => value, + }; // We've already absorbed as many elements as we can let _ = poseidon_sponge( @@ -180,7 +171,8 @@ impl< &mut self.state, Some(&self.mode), )?; - self.mode = Absorbing::init_with(value); + self.mode = Absorbing::init_empty(); + self.mode.absorb(value).expect("state is not full"); Ok(()) } @@ -220,10 +212,8 @@ impl< /// Squeezes an element from the sponge. pub fn squeeze(&mut self, mut layouter: impl Layouter) -> Result, Error> { loop { - for entry in self.mode.0.iter_mut() { - if let Some(inner) = entry.take() { - return Ok(inner.into()); - } + if let Some(value) = self.mode.squeeze() { + return Ok(value.into()); } // We've already squeezed out all available elements diff --git a/halo2_gadgets/src/poseidon/pow5.rs b/halo2_gadgets/src/poseidon/pow5.rs index 704bda1da..5c68a2c9c 100644 --- a/halo2_gadgets/src/poseidon/pow5.rs +++ b/halo2_gadgets/src/poseidon/pow5.rs @@ -340,19 +340,20 @@ impl< let initial_state = initial_state?; // Load the input into this region. - let load_input_word = |i: usize| { - let (cell, value) = match input.0[i].clone() { + let load_input_word = |(i, input_word): (usize, &Option>)| { + let (cell, value) = match input_word { Some(PaddedWord::Message(word)) => (word.cell(), word.value().copied()), Some(PaddedWord::Padding(padding_value)) => { + let value = Value::known(*padding_value); let cell = region .assign_fixed( || format!("load pad_{}", i), config.rc_b[i], 1, - || Value::known(padding_value), + || value, )? .cell(); - (cell, Value::known(padding_value)) + (cell, value) } _ => panic!("Input is not padded"), }; @@ -366,7 +367,12 @@ impl< Ok(StateWord(var)) }; - let input: Result, Error> = (0..RATE).map(load_input_word).collect(); + let input: Result, Error> = input + .expose_inner() + .iter() + .enumerate() + .map(load_input_word) + .collect(); let input = input?; // Constrain the output. @@ -394,14 +400,8 @@ impl< } fn get_output(state: &State) -> Squeezing { - Squeezing( - state[..RATE] - .iter() - .map(|word| Some(word.clone())) - .collect::>() - .try_into() - .unwrap(), - ) + let vals = state[..RATE].to_vec(); + Squeezing::init_full(vals.try_into().expect("correct length")) } } @@ -687,7 +687,7 @@ mod tests { .try_into() .unwrap(); let (round_constants, mds, _) = S::constants(); - poseidon::permute::<_, S, WIDTH, RATE>( + poseidon::test_only_permute::<_, S, WIDTH, RATE>( &mut expected_final_state, &mds, &round_constants, diff --git a/halo2_poseidon/src/lib.rs b/halo2_poseidon/src/lib.rs index 6c8d133dc..7745db9f3 100644 --- a/halo2_poseidon/src/lib.rs +++ b/halo2_poseidon/src/lib.rs @@ -12,8 +12,8 @@ pub(crate) mod fq; pub(crate) mod grain; pub(crate) mod mds; -#[cfg(test)] -pub(crate) mod test_vectors; +#[cfg(any(test, feature = "test-dependencies"))] +pub mod test_vectors; mod p128pow5t3; pub use p128pow5t3::P128Pow5T3; @@ -21,7 +21,7 @@ pub use p128pow5t3::P128Pow5T3; use grain::SboxType; /// The type used to hold permutation state. -pub(crate) type State = [F; T]; +pub type State = [F; T]; /// The type used to hold sponge rate. pub(crate) type SpongeRate = [Option; RATE]; @@ -83,6 +83,18 @@ pub fn generate_constants< (round_constants, mds, mds_inv) } +/// Runs the Poseidon permutation on the given state. +/// +/// Exposed for testing purposes only. +#[cfg(feature = "test-dependencies")] +pub fn test_only_permute, const T: usize, const RATE: usize>( + state: &mut State, + mds: &Mds, + round_constants: &[[F; T]], +) { + permute::(state, mds, round_constants); +} + /// Runs the Poseidon permutation on the given state. pub(crate) fn permute, const T: usize, const RATE: usize>( state: &mut State, @@ -176,9 +188,16 @@ impl SpongeMode for Squeezing {} impl Absorbing { pub(crate) fn init_with(val: F) -> Self { + let mut state = Self::init_empty(); + state.absorb(val).expect("state is not full"); + state + } + + /// Initializes an empty sponge in the absorbing state. + pub fn init_empty() -> Self { Self( - iter::once(Some(val)) - .chain((1..RATE).map(|_| None)) + (0..RATE) + .map(|_| None) .collect::>() .try_into() .unwrap(), @@ -186,6 +205,65 @@ impl Absorbing { } } +impl Absorbing { + /// Attempts to absorb a value into the sponge state. + /// + /// Returns the value if it was not absorbed because the sponge is full. + pub fn absorb(&mut self, value: F) -> Result<(), F> { + for entry in self.0.iter_mut() { + if entry.is_none() { + *entry = Some(value); + return Ok(()); + } + } + // Sponge is full. + Err(value) + } + + /// Exposes the inner state of the sponge. + /// + /// This is a low-level API, requiring a detailed understanding of this specific + /// Poseidon implementation to use correctly and securely. It is exposed for use by + /// the circuit implementation in `halo2_gadgets`, and may be removed from the public + /// API if refactoring enables the circuit implementation to move into this crate. + pub fn expose_inner(&self) -> &SpongeRate { + &self.0 + } +} + +impl Squeezing { + /// Initializes a full sponge in the squeezing state. + /// + /// This is a low-level API, requiring a detailed understanding of this specific + /// Poseidon implementation to use correctly and securely. It is exposed for use by + /// the circuit implementation in `halo2_gadgets`, and may be removed from the public + /// API if refactoring enables the circuit implementation to move into this crate. + pub fn init_full(vals: [F; RATE]) -> Self { + Self( + vals.into_iter() + .map(Some) + .collect::>() + .try_into() + .unwrap(), + ) + } +} + +impl Squeezing { + /// Attempts to squeeze a value from the sponge state. + /// + /// Returns `None` if the sponge is empty. + pub fn squeeze(&mut self) -> Option { + for entry in self.0.iter_mut() { + if let Some(inner) = entry.take() { + return Some(inner); + } + } + // Sponge is empty. + None + } +} + /// A Poseidon sponge. pub(crate) struct Sponge< F: Field, diff --git a/halo2_poseidon/src/p128pow5t3.rs b/halo2_poseidon/src/p128pow5t3.rs index ff5317406..c7640cd7a 100644 --- a/halo2_poseidon/src/p128pow5t3.rs +++ b/halo2_poseidon/src/p128pow5t3.rs @@ -1,4 +1,4 @@ -use halo2_proofs::arithmetic::Field; +use ff::Field; use pasta_curves::{pallas::Base as Fp, vesta::Base as Fq}; use super::{Mds, Spec}; @@ -73,9 +73,7 @@ mod tests { super::{fp, fq}, Fp, Fq, }; - use crate::poseidon::primitives::{ - generate_constants, permute, ConstantLength, Hash, Mds, Spec, - }; + use crate::{generate_constants, permute, ConstantLength, Hash, Mds, Spec}; /// The same Poseidon specification as poseidon::P128Pow5T3, but constructed /// such that its constants will be generated at runtime. @@ -257,7 +255,7 @@ mod tests { { let (round_constants, mds, _) = super::P128Pow5T3::constants(); - for tv in crate::poseidon::primitives::test_vectors::fp::permute() { + for tv in crate::test_vectors::fp::permute() { let mut state = [ Fp::from_repr(tv.initial_state[0]).unwrap(), Fp::from_repr(tv.initial_state[1]).unwrap(), @@ -275,7 +273,7 @@ mod tests { { let (round_constants, mds, _) = super::P128Pow5T3::constants(); - for tv in crate::poseidon::primitives::test_vectors::fq::permute() { + for tv in crate::test_vectors::fq::permute() { let mut state = [ Fq::from_repr(tv.initial_state[0]).unwrap(), Fq::from_repr(tv.initial_state[1]).unwrap(), @@ -293,7 +291,7 @@ mod tests { #[test] fn hash_test_vectors() { - for tv in crate::poseidon::primitives::test_vectors::fp::hash() { + for tv in crate::test_vectors::fp::hash() { let message = [ Fp::from_repr(tv.input[0]).unwrap(), Fp::from_repr(tv.input[1]).unwrap(), @@ -305,7 +303,7 @@ mod tests { assert_eq!(result.to_repr(), tv.output); } - for tv in crate::poseidon::primitives::test_vectors::fq::hash() { + for tv in crate::test_vectors::fq::hash() { let message = [ Fq::from_repr(tv.input[0]).unwrap(), Fq::from_repr(tv.input[1]).unwrap(), diff --git a/halo2_poseidon/src/test_vectors.rs b/halo2_poseidon/src/test_vectors.rs index 1c345a485..d7ec29893 100644 --- a/halo2_poseidon/src/test_vectors.rs +++ b/halo2_poseidon/src/test_vectors.rs @@ -1,19 +1,19 @@ //! Test vectors for [`OrchardNullifier`]. -pub(crate) struct PermuteTestVector { - pub(crate) initial_state: [[u8; 32]; 3], - pub(crate) final_state: [[u8; 32]; 3], +pub struct PermuteTestVector { + pub initial_state: [[u8; 32]; 3], + pub final_state: [[u8; 32]; 3], } -pub(crate) struct HashTestVector { - pub(crate) input: [[u8; 32]; 2], - pub(crate) output: [u8; 32], +pub struct HashTestVector { + pub input: [[u8; 32]; 2], + pub output: [u8; 32], } -pub(crate) mod fp { +pub mod fp { use super::*; - pub(crate) fn permute() -> Vec { + pub fn permute() -> Vec { use PermuteTestVector as TestVector; // From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/permute/fp.py @@ -417,7 +417,7 @@ pub(crate) mod fp { ] } - pub(crate) fn hash() -> Vec { + pub fn hash() -> Vec { use HashTestVector as TestVector; // From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/hash/fp.py @@ -635,10 +635,10 @@ pub(crate) mod fp { } } -pub(crate) mod fq { +pub mod fq { use super::*; - pub(crate) fn permute() -> Vec { + pub fn permute() -> Vec { use PermuteTestVector as TestVector; // From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/permute/fq.py @@ -1042,7 +1042,7 @@ pub(crate) mod fq { ] } - pub(crate) fn hash() -> Vec { + pub fn hash() -> Vec { use HashTestVector as TestVector; // From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/hash/fq.py