mirror of https://github.com/zcash/halo2.git
Poseidon duplex sponge and hash function
This commit is contained in:
parent
9a2c1b0217
commit
266705166f
|
@ -1,3 +1,4 @@
|
|||
use std::iter;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use halo2::arithmetic::FieldExt;
|
||||
|
@ -92,3 +93,178 @@ impl<F: FieldExt> Spec<F> for Generic<F> {
|
|||
(round_constants, mds, mds_inv)
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs the Poseidon permutation on the given state.
|
||||
fn permute<F: FieldExt, S: Spec<F>>(
|
||||
state: &mut [F],
|
||||
spec: &S,
|
||||
mds: &[Vec<F>],
|
||||
round_constants: &[Vec<F>],
|
||||
) {
|
||||
// TODO: Remove this when we can use const generics.
|
||||
assert!(state.len() == spec.arity());
|
||||
|
||||
let r_f = spec.full_rounds() / 2;
|
||||
let r_p = spec.partial_rounds();
|
||||
|
||||
let apply_mds = |state: &mut [F]| {
|
||||
let new_state: Vec<_> = mds
|
||||
.iter()
|
||||
.map(|mds_row| {
|
||||
mds_row
|
||||
.iter()
|
||||
.zip(state.iter())
|
||||
.fold(F::zero(), |acc, (mds, word)| acc + *mds * *word)
|
||||
})
|
||||
.collect();
|
||||
for (word, new_word) in state.iter_mut().zip(new_state.into_iter()) {
|
||||
*word = new_word;
|
||||
}
|
||||
};
|
||||
|
||||
let full_round = |state: &mut [F], rcs: &[F]| {
|
||||
for (word, rc) in state.iter_mut().zip(rcs.iter()) {
|
||||
*word = spec.sbox(*word + rc);
|
||||
}
|
||||
apply_mds(state);
|
||||
};
|
||||
|
||||
let part_round = |state: &mut [F], rcs: &[F]| {
|
||||
for (word, rc) in state.iter_mut().zip(rcs.iter()) {
|
||||
*word += rc;
|
||||
}
|
||||
state[0] = spec.sbox(state[0]);
|
||||
apply_mds(state);
|
||||
};
|
||||
|
||||
iter::empty()
|
||||
.chain(iter::repeat(&full_round as &dyn Fn(&mut [F], &[F])).take(r_f))
|
||||
.chain(iter::repeat(&part_round as &dyn Fn(&mut [F], &[F])).take(r_p))
|
||||
.chain(iter::repeat(&full_round as &dyn Fn(&mut [F], &[F])).take(r_f))
|
||||
.zip(round_constants.iter())
|
||||
.fold(state, |state, (round, rcs)| {
|
||||
round(state, rcs);
|
||||
state
|
||||
});
|
||||
}
|
||||
|
||||
fn pad_and_add<F: FieldExt>(state: &mut [F], input: &[F]) {
|
||||
let padding = state.len() - input.len();
|
||||
// TODO: Decide on a padding strategy (currently padding with all-ones)
|
||||
for (word, val) in state
|
||||
.iter_mut()
|
||||
.zip(input.iter().chain(iter::repeat(&F::one()).take(padding)))
|
||||
{
|
||||
*word += val;
|
||||
}
|
||||
}
|
||||
|
||||
enum SpongeState<F: FieldExt> {
|
||||
Absorbing(Vec<F>),
|
||||
Squeezing(Vec<F>),
|
||||
}
|
||||
|
||||
/// A Poseidon duplex sponge.
|
||||
pub struct Duplex<F: FieldExt, S: Spec<F>> {
|
||||
spec: S,
|
||||
sponge: Option<SpongeState<F>>,
|
||||
state: Vec<F>,
|
||||
rate: usize,
|
||||
mds_matrix: Vec<Vec<F>>,
|
||||
round_constants: Vec<Vec<F>>,
|
||||
}
|
||||
|
||||
impl<F: FieldExt, S: Spec<F>> Duplex<F, S> {
|
||||
/// Constructs a new duplex sponge with the given rate.
|
||||
pub fn new(spec: S, rate: usize) -> Self {
|
||||
assert!(rate < spec.arity());
|
||||
|
||||
let state = vec![F::zero(); spec.arity()];
|
||||
let (round_constants, mds_matrix, _) = spec.constants();
|
||||
|
||||
Duplex {
|
||||
spec,
|
||||
sponge: Some(SpongeState::Absorbing(vec![])),
|
||||
state,
|
||||
rate,
|
||||
mds_matrix,
|
||||
round_constants,
|
||||
}
|
||||
}
|
||||
|
||||
fn process(&mut self, input: &[F]) -> Vec<F> {
|
||||
pad_and_add(&mut self.state[..self.rate], input);
|
||||
|
||||
permute(
|
||||
&mut self.state,
|
||||
&self.spec,
|
||||
&self.mds_matrix,
|
||||
&self.round_constants,
|
||||
);
|
||||
|
||||
self.state[..self.rate].to_vec()
|
||||
}
|
||||
|
||||
/// Absorbs an element into the sponge.
|
||||
pub fn absorb(&mut self, value: F) {
|
||||
match self.sponge.take().unwrap() {
|
||||
SpongeState::Absorbing(mut input) => {
|
||||
if input.len() < self.rate {
|
||||
input.push(value);
|
||||
self.sponge = Some(SpongeState::Absorbing(input));
|
||||
return;
|
||||
}
|
||||
|
||||
// We've already absorbed as many elements as we can
|
||||
let _ = self.process(&input);
|
||||
self.sponge = Some(SpongeState::Absorbing(vec![value]));
|
||||
}
|
||||
SpongeState::Squeezing(_) => {
|
||||
// Drop the remaining output elements
|
||||
self.sponge = Some(SpongeState::Absorbing(vec![value]));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Squeezes an element from the sponge.
|
||||
pub fn squeeze(&mut self) -> F {
|
||||
loop {
|
||||
match self.sponge.take().unwrap() {
|
||||
SpongeState::Absorbing(input) => {
|
||||
self.sponge = Some(SpongeState::Squeezing(self.process(&input)));
|
||||
}
|
||||
SpongeState::Squeezing(mut output) => {
|
||||
if !output.is_empty() {
|
||||
let ret = output.remove(0);
|
||||
self.sponge = Some(SpongeState::Squeezing(output));
|
||||
return ret;
|
||||
}
|
||||
|
||||
// We've already squeezed out all available elements
|
||||
self.sponge = Some(SpongeState::Absorbing(vec![]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A Poseidon hash function, built around a duplex sponge.
|
||||
pub struct Hash<F: FieldExt, S: Spec<F>>(Duplex<F, S>);
|
||||
|
||||
impl<F: FieldExt, S: Spec<F>> Hash<F, S> {
|
||||
/// Initializes a new hasher.
|
||||
pub fn init(spec: S, rate: usize) -> Self {
|
||||
Hash(Duplex::new(spec, rate))
|
||||
}
|
||||
|
||||
/// Updates the hasher with the given value.
|
||||
pub fn update(&mut self, value: F) {
|
||||
self.0.absorb(value);
|
||||
}
|
||||
|
||||
/// Finalizes the hasher, returning its output.
|
||||
pub fn finalize(mut self) -> F {
|
||||
// TODO: Check which state element other implementations use.
|
||||
self.0.squeeze()
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue