halo2_poseidon: Refactor code so it compiles in its new crate

str4d · str4d · commit 0d883685d191 · 2024-12-16T13:24:41.000Z
diff --git 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<F>) -> Result<Self, Error> {
         chip.initial_state(&mut layouter).map(|state| Sponge {
             chip,
-            mode: Absorbing(
-                (0..RATE)
-                    .map(|_| None)
-                    .collect::<Vec<_>>()
-                    .try_into()
-                    .unwrap(),
-            ),
+            mode: Absorbing::init_empty(),
             state,
             _marker: PhantomData::default(),
         })
@@ -166,12 +159,10 @@ impl<
         mut layouter: impl Layouter<F>,
         value: PaddedWord<F>,
     ) -> 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<F>) -> Result<AssignedCell<F, F>, 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
@@ -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<PaddedWord<F>>)| {
+                    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<Vec<_>, Error> = (0..RATE).map(load_input_word).collect();
+                let input: Result<Vec<_>, 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<Self::Word, WIDTH>) -> Squeezing<Self::Word, RATE> {
-        Squeezing(
-            state[..RATE]
-                .iter()
-                .map(|word| Some(word.clone()))
-                .collect::<Vec<_>>()
-                .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
@@ -12,16 +12,16 @@ 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;
 
 use grain::SboxType;
 
 /// The type used to hold permutation state.
-pub(crate) type State<F, const T: usize> = [F; T];
+pub type State<F, const T: usize> = [F; T];
 
 /// The type used to hold sponge rate.
 pub(crate) type SpongeRate<F, const RATE: usize> = [Option<F>; 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<F: Field, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
+    state: &mut State<F, T>,
+    mds: &Mds<F, T>,
+    round_constants: &[[F; T]],
+) {
+    permute::<F, S, T, RATE>(state, mds, round_constants);
+}
+
 /// Runs the Poseidon permutation on the given state.
 pub(crate) fn permute<F: Field, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
     state: &mut State<F, T>,
@@ -176,16 +188,82 @@ impl<F, const RATE: usize> SpongeMode for Squeezing<F, RATE> {}
 
 impl<F: fmt::Debug, const RATE: usize> Absorbing<F, RATE> {
     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::<Vec<_>>()
                 .try_into()
                 .unwrap(),
         )
     }
 }
 
+impl<F, const RATE: usize> Absorbing<F, RATE> {
+    /// 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<F, RATE> {
+        &self.0
+    }
+}
+
+impl<F: fmt::Debug, const RATE: usize> Squeezing<F, RATE> {
+    /// 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::<Vec<_>>()
+                .try_into()
+                .unwrap(),
+        )
+    }
+}
+
+impl<F, const RATE: usize> Squeezing<F, RATE> {
+    /// Attempts to squeeze a value from the sponge state.
+    ///
+    /// Returns `None` if the sponge is empty.
+    pub fn squeeze(&mut self) -> Option<F> {
+        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
@@ -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
@@ -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<PermuteTestVector> {
+    pub fn permute() -> Vec<PermuteTestVector> {
         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<HashTestVector> {
+    pub fn hash() -> Vec<HashTestVector> {
         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<PermuteTestVector> {
+    pub fn permute() -> Vec<PermuteTestVector> {
         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<HashTestVector> {
+    pub fn hash() -> Vec<HashTestVector> {
         use HashTestVector as TestVector;
 
         // From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/hash/fq.py

Original file line number	Diff line number	Diff line change
`@@ -1,19 +1,19 @@`
`1`	`1`	//! Test vectors for [`OrchardNullifier`].
`2`	`2`
`3`		`-pub(crate) struct PermuteTestVector {`
`4`		`- pub(crate) initial_state: [[u8; 32]; 3],`
`5`		`- pub(crate) final_state: [[u8; 32]; 3],`
	`3`	`+pub struct PermuteTestVector {`
	`4`	`+ pub initial_state: [[u8; 32]; 3],`
	`5`	`+ pub final_state: [[u8; 32]; 3],`
`6`	`6`	`}`
`7`	`7`
`8`		`-pub(crate) struct HashTestVector {`
`9`		`- pub(crate) input: [[u8; 32]; 2],`
`10`		`- pub(crate) output: [u8; 32],`
	`8`	`+pub struct HashTestVector {`
	`9`	`+ pub input: [[u8; 32]; 2],`
	`10`	`+ pub output: [u8; 32],`
`11`	`11`	`}`
`12`	`12`
`13`		`-pub(crate) mod fp {`
	`13`	`+pub mod fp {`
`14`	`14`	`use super::*;`
`15`	`15`
`16`		`- pub(crate) fn permute() -> Vec<PermuteTestVector> {`
	`16`	`+ pub fn permute() -> Vec<PermuteTestVector> {`
`17`	`17`	`use PermuteTestVector as TestVector;`
`18`	`18`
`19`	`19`	`// From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/permute/fp.py`
`@@ -417,7 +417,7 @@ pub(crate) mod fp {`
`417`	`417`	`]`
`418`	`418`	`}`
`419`	`419`
`420`		`- pub(crate) fn hash() -> Vec<HashTestVector> {`
	`420`	`+ pub fn hash() -> Vec<HashTestVector> {`
`421`	`421`	`use HashTestVector as TestVector;`
`422`	`422`
`423`	`423`	`// From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/hash/fp.py`
`@@ -635,10 +635,10 @@ pub(crate) mod fp {`
`635`	`635`	`}`
`636`	`636`	`}`
`637`	`637`
`638`		`-pub(crate) mod fq {`
	`638`	`+pub mod fq {`
`639`	`639`	`use super::*;`
`640`	`640`
`641`		`- pub(crate) fn permute() -> Vec<PermuteTestVector> {`
	`641`	`+ pub fn permute() -> Vec<PermuteTestVector> {`
`642`	`642`	`use PermuteTestVector as TestVector;`
`643`	`643`
`644`	`644`	`// From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/permute/fq.py`
`@@ -1042,7 +1042,7 @@ pub(crate) mod fq {`
`1042`	`1042`	`]`
`1043`	`1043`	`}`
`1044`	`1044`
`1045`		`- pub(crate) fn hash() -> Vec<HashTestVector> {`
	`1045`	`+ pub fn hash() -> Vec<HashTestVector> {`
`1046`	`1046`	`use HashTestVector as TestVector;`
`1047`	`1047`
`1048`	`1048`	`// From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/orchard_poseidon/hash/fq.py`