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 @@
 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,15 +147,9 @@
     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 @@
         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 @@
             &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(())
     }
@@ -203,7 +195,7 @@
             chip: self.chip,
             mode,
             state: self.state,
             _marker: PhantomData::default(),
         })
     }
 }
@@ -220,10 +212,8 @@
     /// 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

@@ -238,31 +238,31 @@
                 // Load the initial state into this region.
                 let state = Pow5State::load(&mut region, config, initial_state)?;
 
                 let state = (0..config.half_full_rounds).fold(Ok(state), |res, r| {
                     res.and_then(|state| state.full_round(&mut region, config, r, r))
                 })?;
 
                 let state = (0..config.half_partial_rounds).fold(Ok(state), |res, r| {
                     res.and_then(|state| {
                         state.partial_round(
                             &mut region,
                             config,
                             config.half_full_rounds + 2 * r,
                             config.half_full_rounds + r,
                         )
                     })
                 })?;
 
                 let state = (0..config.half_full_rounds).fold(Ok(state), |res, r| {
                     res.and_then(|state| {
                         state.full_round(
                             &mut region,
                             config,
                             config.half_full_rounds + 2 * config.half_partial_rounds + r,
                             config.half_full_rounds + config.half_partial_rounds + r,
                         )
                     })
                 })?;
 
                 Ok(state.0)
             },
@@ -340,19 +340,20 @@
                 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 @@
 
                     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 @@
     }
 
     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"))
     }
 }
 
@@ -448,7 +448,7 @@
                     .value()
                     .map(|v| *v + config.round_constants[round][idx])
             });
             let r: Value<Vec<F>> = q.map(|q| q.map(|q| q.pow(&config.alpha))).collect();
             let m = &config.m_reg;
             let state = m.iter().map(|m_i| {
                 r.as_ref().map(|r| {
@@ -474,7 +474,7 @@
             let p: Value<Vec<_>> = self.0.iter().map(|word| word.0.value().cloned()).collect();
 
             let r: Value<Vec<_>> = p.map(|p| {
                 let r_0 = (p[0] + config.round_constants[round][0]).pow(&config.alpha);
                 let r_i = p[1..]
                     .iter()
                     .enumerate()
@@ -514,7 +514,7 @@
             }
 
             let r_mid: Value<Vec<_>> = p_mid.map(|p| {
                 let r_0 = (p[0] + config.round_constants[round + 1][0]).pow(&config.alpha);
                 let r_i = p[1..]
                     .iter()
                     .enumerate()
@@ -687,7 +687,7 @@
                 .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,