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 std::marker::PhantomData;
|
||||||
|
|
||||||
use halo2::arithmetic::FieldExt;
|
use halo2::arithmetic::FieldExt;
|
||||||
|
@ -92,3 +93,178 @@ impl<F: FieldExt> Spec<F> for Generic<F> {
|
||||||
(round_constants, mds, mds_inv)
|
(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