Merge pull request #831 from zcash/poseidon-primitive

Move Poseidon primitive into `halo2_poseidon`

halo2_gadgets/Cargo.toml

@@ -26,6 +26,7 @@ arrayvec = "0.7.0"
 bitvec = "1"
 ff = "0.13"
 group = "0.13"
+halo2_poseidon = { version = "0.0", path = "../halo2_poseidon", default-features = false }
 halo2_proofs = { version = "0.3", path = "../halo2_proofs", default-features = false }
 lazy_static = "1"
 pasta_curves = "0.5"
@@ -40,6 +41,7 @@ plotters = { version = "0.3.0", default-features = false, optional = true }
 
 [dev-dependencies]
 criterion = "0.3"
+halo2_poseidon = { version = "0.0", path = "../halo2_poseidon", default-features = false, features = ["test-dependencies"] }
 proptest = "1.0.0"
 sinsemilla = { version = "0.1", features = ["test-dependencies"] }
 

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;
 
@@ -13,7 +12,7 @@ use halo2_proofs::{
 mod pow5;
 pub use pow5::{Pow5Chip, Pow5Config, StateWord};
 
-pub mod primitives;
+pub use ::halo2_poseidon as primitives;
 use primitives::{Absorbing, ConstantLength, Domain, Spec, SpongeMode, Squeezing, State};
 
 /// A word from the padded input to a Poseidon sponge.
@@ -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

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,