Merge pull request #259 from zcash/222-rewrite-poseidon

Rewrite Poseidon implementation
This commit is contained in:
str4d 2021-12-15 13:28:26 +00:00 committed by GitHub
commit 8a4f4e347b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 357 additions and 300 deletions

View File

@ -101,10 +101,9 @@ where
},
)?;
let hasher = Hash::<_, _, S, _, WIDTH, RATE>::init(
let hasher = Hash::<_, _, S, ConstantLength<L>, WIDTH, RATE>::init(
chip,
layouter.namespace(|| "init"),
ConstantLength::<L>,
)?;
let output = hasher.hash(layouter.namespace(|| "hash"), message)?;
@ -210,7 +209,7 @@ fn bench_poseidon<S, const WIDTH: usize, const RATE: usize, const L: usize>(
.collect::<Vec<_>>()
.try_into()
.unwrap();
let output = poseidon::Hash::<_, S, _, WIDTH, RATE>::init(ConstantLength::<L>).hash(message);
let output = poseidon::Hash::<_, S, ConstantLength<L>, WIDTH, RATE>::init().hash(message);
let circuit = HashCircuit::<S, WIDTH, RATE, L> {
message: Some(message),

View File

@ -21,7 +21,9 @@ fn bench_primitives(c: &mut Criterion) {
let message = [pallas::Base::random(rng), pallas::Base::random(rng)];
group.bench_function("2-to-1", |b| {
b.iter(|| poseidon::Hash::<_, P128Pow5T3, _, 3, 2>::init(ConstantLength).hash(message))
b.iter(|| {
poseidon::Hash::<_, P128Pow5T3, ConstantLength<2>, 3, 2>::init().hash(message)
})
});
}

View File

@ -482,11 +482,11 @@ impl plonk::Circuit<pallas::Base> for Circuit {
// hash_old = poseidon_hash(nk, rho_old)
let hash_old = {
let poseidon_message = [nk.clone(), rho_old.clone()];
let poseidon_hasher = PoseidonHash::<_, _, poseidon::P128Pow5T3, _, 3, 2>::init(
config.poseidon_chip(),
layouter.namespace(|| "Poseidon init"),
ConstantLength::<2>,
)?;
let poseidon_hasher =
PoseidonHash::<_, _, poseidon::P128Pow5T3, ConstantLength<2>, 3, 2>::init(
config.poseidon_chip(),
layouter.namespace(|| "Poseidon init"),
)?;
poseidon_hasher.hash(
layouter.namespace(|| "Poseidon hash (nk, rho_old)"),
poseidon_message,

View File

@ -3,7 +3,9 @@
use std::array;
use std::convert::TryInto;
use std::fmt;
use std::marker::PhantomData;
use group::ff::Field;
use halo2::{
arithmetic::FieldExt,
circuit::{AssignedCell, Chip, Layouter},
@ -13,7 +15,18 @@ use halo2::{
mod pow5;
pub use pow5::{Pow5Chip, Pow5Config, StateWord};
use crate::primitives::poseidon::{ConstantLength, Domain, Spec, Sponge, SpongeState, State};
use crate::primitives::poseidon::{
Absorbing, ConstantLength, Domain, Spec, SpongeMode, Squeezing, State,
};
/// A word from the padded input to a Poseidon sponge.
#[derive(Clone, Debug)]
pub enum PaddedWord<F: Field> {
/// A message word provided by the prover.
Message(AssignedCell<F, F>),
/// A padding word, that will be fixed in the circuit parameters.
Padding(F),
}
/// The set of circuit instructions required to use the Poseidon permutation.
pub trait PoseidonInstructions<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>:
@ -30,34 +43,31 @@ pub trait PoseidonInstructions<F: FieldExt, S: Spec<F, T, RATE>, const T: usize,
) -> Result<State<Self::Word, T>, Error>;
}
/// The set of circuit instructions required to use the [`Duplex`] and [`Hash`] gadgets.
/// The set of circuit instructions required to use the [`Sponge`] and [`Hash`] gadgets.
///
/// [`Hash`]: self::Hash
pub trait PoseidonDuplexInstructions<
pub trait PoseidonSpongeInstructions<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
>: PoseidonInstructions<F, S, T, RATE>
{
/// Returns the initial empty state for the given domain.
fn initial_state(
&self,
layouter: &mut impl Layouter<F>,
domain: &impl Domain<F, T, RATE>,
) -> Result<State<Self::Word, T>, Error>;
fn initial_state(&self, layouter: &mut impl Layouter<F>)
-> Result<State<Self::Word, T>, Error>;
/// Pads the given input (according to the specified domain) and adds it to the state.
fn pad_and_add(
/// Adds the given input to the state.
fn add_input(
&self,
layouter: &mut impl Layouter<F>,
domain: &impl Domain<F, T, RATE>,
initial_state: &State<Self::Word, T>,
input: &SpongeState<Self::Word, RATE>,
input: &Absorbing<PaddedWord<F>, RATE>,
) -> Result<State<Self::Word, T>, Error>;
/// Extracts sponge output from the given state.
fn get_output(state: &State<Self::Word, T>) -> SpongeState<Self::Word, RATE>;
fn get_output(state: &State<Self::Word, T>) -> Squeezing<Self::Word, RATE>;
}
/// A word over which the Poseidon permutation operates.
@ -91,170 +101,176 @@ impl<
}
}
fn poseidon_duplex<
fn poseidon_sponge<
F: FieldExt,
PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
>(
chip: &PoseidonChip,
mut layouter: impl Layouter<F>,
domain: &D,
state: &mut State<PoseidonChip::Word, T>,
input: &SpongeState<PoseidonChip::Word, RATE>,
) -> Result<SpongeState<PoseidonChip::Word, RATE>, Error> {
*state = chip.pad_and_add(&mut layouter, domain, state, input)?;
input: Option<&Absorbing<PaddedWord<F>, RATE>>,
) -> Result<Squeezing<PoseidonChip::Word, RATE>, Error> {
if let Some(input) = input {
*state = chip.add_input(&mut layouter, state, input)?;
}
*state = chip.permute(&mut layouter, state)?;
Ok(PoseidonChip::get_output(state))
}
/// A Poseidon duplex sponge.
/// A Poseidon sponge.
#[derive(Debug)]
pub struct Duplex<
pub struct Sponge<
F: FieldExt,
PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
M: SpongeMode,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> {
chip: PoseidonChip,
sponge: Sponge<PoseidonChip::Word, RATE>,
mode: M,
state: State<PoseidonChip::Word, T>,
domain: D,
_marker: PhantomData<D>,
}
impl<
F: FieldExt,
PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> Duplex<F, PoseidonChip, S, D, T, RATE>
> Sponge<F, PoseidonChip, S, Absorbing<PaddedWord<F>, RATE>, D, T, RATE>
{
/// Constructs a new duplex sponge for the given Poseidon specification.
pub fn new(
chip: PoseidonChip,
mut layouter: impl Layouter<F>,
domain: D,
) -> Result<Self, Error> {
chip.initial_state(&mut layouter, &domain)
.map(|state| Duplex {
chip,
sponge: Sponge::Absorbing(
(0..RATE)
.map(|_| None)
.collect::<Vec<_>>()
.try_into()
.unwrap(),
),
state,
domain,
})
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(),
),
state,
_marker: PhantomData::default(),
})
}
/// Absorbs an element into the sponge.
pub fn absorb(
&mut self,
mut layouter: impl Layouter<F>,
value: AssignedCell<F, F>,
value: PaddedWord<F>,
) -> Result<(), Error> {
match self.sponge {
Sponge::Absorbing(ref mut input) => {
for entry in input.iter_mut() {
if entry.is_none() {
*entry = Some(value.into());
return Ok(());
}
}
// We've already absorbed as many elements as we can
let _ = poseidon_duplex(
&self.chip,
layouter.namespace(|| "PoseidonDuplex"),
&self.domain,
&mut self.state,
input,
)?;
self.sponge = Sponge::absorb(value.into());
}
Sponge::Squeezing(_) => {
// Drop the remaining output elements
self.sponge = Sponge::absorb(value.into());
for entry in self.mode.0.iter_mut() {
if entry.is_none() {
*entry = Some(value);
return Ok(());
}
}
// We've already absorbed as many elements as we can
let _ = poseidon_sponge(
&self.chip,
layouter.namespace(|| "PoseidonSponge"),
&mut self.state,
Some(&self.mode),
)?;
self.mode = Absorbing::init_with(value);
Ok(())
}
/// Squeezes an element from the sponge.
pub fn squeeze(&mut self, mut layouter: impl Layouter<F>) -> Result<AssignedCell<F, F>, Error> {
loop {
match self.sponge {
Sponge::Absorbing(ref input) => {
self.sponge = Sponge::Squeezing(poseidon_duplex(
&self.chip,
layouter.namespace(|| "PoseidonDuplex"),
&self.domain,
&mut self.state,
input,
)?);
}
Sponge::Squeezing(ref mut output) => {
for entry in output.iter_mut() {
if let Some(inner) = entry.take() {
return Ok(inner.into());
}
}
/// Transitions the sponge into its squeezing state.
#[allow(clippy::type_complexity)]
pub fn finish_absorbing(
mut self,
mut layouter: impl Layouter<F>,
) -> Result<Sponge<F, PoseidonChip, S, Squeezing<PoseidonChip::Word, RATE>, D, T, RATE>, Error>
{
let mode = poseidon_sponge(
&self.chip,
layouter.namespace(|| "PoseidonSponge"),
&mut self.state,
Some(&self.mode),
)?;
// We've already squeezed out all available elements
self.sponge = Sponge::Absorbing(
(0..RATE)
.map(|_| None)
.collect::<Vec<_>>()
.try_into()
.unwrap(),
);
}
}
}
Ok(Sponge {
chip: self.chip,
mode,
state: self.state,
_marker: PhantomData::default(),
})
}
}
/// A Poseidon hash function, built around a duplex sponge.
#[derive(Debug)]
pub struct Hash<
F: FieldExt,
PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
const T: usize,
const RATE: usize,
> {
duplex: Duplex<F, PoseidonChip, S, D, T, RATE>,
}
impl<
F: FieldExt,
PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> Sponge<F, PoseidonChip, S, Squeezing<PoseidonChip::Word, RATE>, D, T, RATE>
{
/// 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());
}
}
// We've already squeezed out all available elements
self.mode = poseidon_sponge(
&self.chip,
layouter.namespace(|| "PoseidonSponge"),
&mut self.state,
None,
)?;
}
}
}
/// A Poseidon hash function, built around a sponge.
#[derive(Debug)]
pub struct Hash<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> {
sponge: Sponge<F, PoseidonChip, S, Absorbing<PaddedWord<F>, RATE>, D, T, RATE>,
}
impl<
F: FieldExt,
PoseidonChip: PoseidonSpongeInstructions<F, S, D, T, RATE>,
S: Spec<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> Hash<F, PoseidonChip, S, D, T, RATE>
{
/// Initializes a new hasher.
pub fn init(chip: PoseidonChip, layouter: impl Layouter<F>, domain: D) -> Result<Self, Error> {
Duplex::new(chip, layouter, domain).map(|duplex| Hash { duplex })
pub fn init(chip: PoseidonChip, layouter: impl Layouter<F>) -> Result<Self, Error> {
Sponge::new(chip, layouter).map(|sponge| Hash { sponge })
}
}
impl<
F: FieldExt,
PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
PoseidonChip: PoseidonSpongeInstructions<F, S, ConstantLength<L>, T, RATE>,
S: Spec<F, T, RATE>,
const T: usize,
const RATE: usize,
@ -267,10 +283,16 @@ impl<
mut layouter: impl Layouter<F>,
message: [AssignedCell<F, F>; L],
) -> Result<AssignedCell<F, F>, Error> {
for (i, value) in array::IntoIter::new(message).enumerate() {
self.duplex
for (i, value) in array::IntoIter::new(message)
.map(PaddedWord::Message)
.chain(<ConstantLength<L> as Domain<F, RATE>>::padding(L).map(PaddedWord::Padding))
.enumerate()
{
self.sponge
.absorb(layouter.namespace(|| format!("absorb_{}", i)), value)?;
}
self.duplex.squeeze(layouter.namespace(|| "squeeze"))
self.sponge
.finish_absorbing(layouter.namespace(|| "finish absorbing"))?
.squeeze(layouter.namespace(|| "squeeze"))
}
}

View File

@ -8,9 +8,12 @@ use halo2::{
poly::Rotation,
};
use super::{PoseidonDuplexInstructions, PoseidonInstructions};
use crate::circuit::gadget::utilities::Var;
use crate::primitives::poseidon::{Domain, Mds, Spec, SpongeState, State};
use super::{PaddedWord, PoseidonInstructions, PoseidonSpongeInstructions};
use crate::primitives::poseidon::{Domain, Mds, Spec, State};
use crate::{
circuit::gadget::utilities::Var,
primitives::poseidon::{Absorbing, Squeezing},
};
/// Configuration for a [`Pow5Chip`].
#[derive(Clone, Debug)]
@ -268,17 +271,21 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
}
}
impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize>
PoseidonDuplexInstructions<F, S, WIDTH, RATE> for Pow5Chip<F, WIDTH, RATE>
impl<
F: FieldExt,
S: Spec<F, WIDTH, RATE>,
D: Domain<F, RATE>,
const WIDTH: usize,
const RATE: usize,
> PoseidonSpongeInstructions<F, S, D, WIDTH, RATE> for Pow5Chip<F, WIDTH, RATE>
{
fn initial_state(
&self,
layouter: &mut impl Layouter<F>,
domain: &impl Domain<F, WIDTH, RATE>,
) -> Result<State<Self::Word, WIDTH>, Error> {
let config = self.config();
let state = layouter.assign_region(
|| format!("initial state for domain {:?}", domain),
|| format!("initial state for domain {}", D::name()),
|mut region| {
let mut state = Vec::with_capacity(WIDTH);
let mut load_state_word = |i: usize, value: F| -> Result<_, Error> {
@ -296,7 +303,7 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
for i in 0..RATE {
load_state_word(i, F::zero())?;
}
load_state_word(RATE, domain.initial_capacity_element())?;
load_state_word(RATE, D::initial_capacity_element())?;
Ok(state)
},
@ -305,16 +312,15 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
Ok(state.try_into().unwrap())
}
fn pad_and_add(
fn add_input(
&self,
layouter: &mut impl Layouter<F>,
domain: &impl Domain<F, WIDTH, RATE>,
initial_state: &State<Self::Word, WIDTH>,
input: &SpongeState<Self::Word, RATE>,
input: &Absorbing<PaddedWord<F>, RATE>,
) -> Result<State<Self::Word, WIDTH>, Error> {
let config = self.config();
layouter.assign_region(
|| format!("pad-and-add for domain {:?}", domain),
|| format!("add input for domain {}", D::name()),
|mut region| {
config.s_pad_and_add.enable(&mut region, 1)?;
@ -334,19 +340,17 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
(0..WIDTH).map(load_state_word).collect();
let initial_state = initial_state?;
let padding_values = domain.padding();
// Load the input and padding into this region.
// Load the input into this region.
let load_input_word = |i: usize| {
let constraint_var = match (input[i].clone(), padding_values[i]) {
(Some(word), None) => word.0,
(None, Some(padding_value)) => region.assign_fixed(
let constraint_var = match input.0[i].clone() {
Some(PaddedWord::Message(word)) => word,
Some(PaddedWord::Padding(padding_value)) => region.assign_fixed(
|| format!("load pad_{}", i),
config.rc_b[i],
1,
|| Ok(padding_value),
)?,
_ => panic!("Input and padding don't match"),
_ => panic!("Input is not padded"),
};
constraint_var
.copy_advice(
@ -386,13 +390,15 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
)
}
fn get_output(state: &State<Self::Word, WIDTH>) -> SpongeState<Self::Word, RATE> {
state[..RATE]
.iter()
.map(|word| Some(word.clone()))
.collect::<Vec<_>>()
.try_into()
.unwrap()
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(),
)
}
}
@ -783,10 +789,9 @@ mod tests {
},
)?;
let hasher = Hash::<_, _, S, _, WIDTH, RATE>::init(
let hasher = Hash::<_, _, S, ConstantLength<L>, WIDTH, RATE>::init(
chip,
layouter.namespace(|| "init"),
ConstantLength::<L>,
)?;
let output = hasher.hash(layouter.namespace(|| "hash"), message)?;
@ -809,7 +814,7 @@ mod tests {
fn poseidon_hash() {
let message = [Fp::rand(), Fp::rand()];
let output =
poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength::<2>).hash(message);
poseidon::Hash::<_, OrchardNullifier, ConstantLength<2>, 3, 2>::init().hash(message);
let k = 6;
let circuit = HashCircuit::<OrchardNullifier, 3, 2, 2> {
@ -821,6 +826,22 @@ mod tests {
assert_eq!(prover.verify(), Ok(()))
}
#[test]
fn poseidon_hash_longer_input() {
let message = [Fp::rand(), Fp::rand(), Fp::rand()];
let output =
poseidon::Hash::<_, OrchardNullifier, ConstantLength<3>, 3, 2>::init().hash(message);
let k = 7;
let circuit = HashCircuit::<OrchardNullifier, 3, 2, 3> {
message: Some(message),
output: Some(output),
_spec: PhantomData,
};
let prover = MockProver::run(k, &circuit, vec![]).unwrap();
assert_eq!(prover.verify(), Ok(()))
}
#[test]
fn hash_test_vectors() {
for tv in crate::primitives::poseidon::test_vectors::fp::hash() {
@ -828,8 +849,8 @@ mod tests {
pallas::Base::from_repr(tv.input[0]).unwrap(),
pallas::Base::from_repr(tv.input[1]).unwrap(),
];
let output =
poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength).hash(message);
let output = poseidon::Hash::<_, OrchardNullifier, ConstantLength<2>, 3, 2>::init()
.hash(message);
let k = 6;
let circuit = HashCircuit::<OrchardNullifier, 3, 2, 2> {

View File

@ -24,8 +24,8 @@ use grain::SboxType;
/// The type used to hold permutation state.
pub(crate) type State<F, const T: usize> = [F; T];
/// The type used to hold duplex sponge state.
pub(crate) type SpongeState<F, const RATE: usize> = [Option<F>; RATE];
/// The type used to hold sponge rate.
pub(crate) type SpongeRate<F, const RATE: usize> = [Option<F>; RATE];
/// The type used to hold the MDS matrix and its inverse.
pub(crate) type Mds<F, const T: usize> = [[F; T]; T];
@ -124,14 +124,19 @@ pub(crate) fn permute<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RA
});
}
fn poseidon_duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
fn poseidon_sponge<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
state: &mut State<F, T>,
input: &SpongeState<F, RATE>,
pad_and_add: &dyn Fn(&mut State<F, T>, &SpongeState<F, RATE>),
input: Option<&Absorbing<F, RATE>>,
mds_matrix: &Mds<F, T>,
round_constants: &[[F; T]],
) -> SpongeState<F, RATE> {
pad_and_add(state, input);
) -> Squeezing<F, RATE> {
if let Some(Absorbing(input)) = input {
// `Iterator::zip` short-circuits when one iterator completes, so this will only
// mutate the rate portion of the state.
for (word, value) in state.iter_mut().zip(input.iter()) {
*word += value.expect("poseidon_sponge is called with a padded input");
}
}
permute::<F, S, T, RATE>(state, mds_matrix, round_constants);
@ -139,53 +144,70 @@ fn poseidon_duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE:
for (word, value) in output.iter_mut().zip(state.iter()) {
*word = Some(*value);
}
output
Squeezing(output)
}
mod private {
pub trait SealedSpongeMode {}
impl<F, const RATE: usize> SealedSpongeMode for super::Absorbing<F, RATE> {}
impl<F, const RATE: usize> SealedSpongeMode for super::Squeezing<F, RATE> {}
}
/// The state of the `Sponge`.
pub trait SpongeMode: private::SealedSpongeMode {}
/// The absorbing state of the `Sponge`.
#[derive(Debug)]
pub(crate) enum Sponge<F, const RATE: usize> {
Absorbing(SpongeState<F, RATE>),
Squeezing(SpongeState<F, RATE>),
}
pub struct Absorbing<F, const RATE: usize>(pub(crate) SpongeRate<F, RATE>);
impl<F: fmt::Debug, const RATE: usize> Sponge<F, RATE> {
pub(crate) fn absorb(val: F) -> Self {
let mut input: [Option<F>; RATE] = (0..RATE)
.map(|_| None)
.collect::<Vec<_>>()
.try_into()
.unwrap();
input[0] = Some(val);
Sponge::Absorbing(input)
/// The squeezing state of the `Sponge`.
#[derive(Debug)]
pub struct Squeezing<F, const RATE: usize>(pub(crate) SpongeRate<F, RATE>);
impl<F, const RATE: usize> SpongeMode for Absorbing<F, RATE> {}
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 {
Self(
iter::once(Some(val))
.chain((1..RATE).map(|_| None))
.collect::<Vec<_>>()
.try_into()
.unwrap(),
)
}
}
/// A Poseidon duplex sponge.
pub(crate) struct Duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> {
sponge: Sponge<F, RATE>,
/// A Poseidon sponge.
pub(crate) struct Sponge<
F: FieldExt,
S: Spec<F, T, RATE>,
M: SpongeMode,
const T: usize,
const RATE: usize,
> {
mode: M,
state: State<F, T>,
pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeState<F, RATE>)>,
mds_matrix: Mds<F, T>,
round_constants: Vec<[F; T]>,
_marker: PhantomData<S>,
}
impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> Duplex<F, S, T, RATE> {
/// Constructs a new duplex sponge for the given Poseidon specification.
pub(crate) fn new(
initial_capacity_element: F,
pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeState<F, RATE>)>,
) -> Self {
impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
Sponge<F, S, Absorbing<F, RATE>, T, RATE>
{
/// Constructs a new sponge for the given Poseidon specification.
pub(crate) fn new(initial_capacity_element: F) -> Self {
let (round_constants, mds_matrix, _) = S::constants();
let input = [None; RATE];
let mode = Absorbing([None; RATE]);
let mut state = [F::zero(); T];
state[RATE] = initial_capacity_element;
Duplex {
sponge: Sponge::Absorbing(input),
Sponge {
mode,
state,
pad_and_add,
mds_matrix,
round_constants,
_marker: PhantomData::default(),
@ -194,71 +216,78 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> Duplex
/// Absorbs an element into the sponge.
pub(crate) fn absorb(&mut self, value: F) {
match self.sponge {
Sponge::Absorbing(ref mut input) => {
for entry in input.iter_mut() {
if entry.is_none() {
*entry = Some(value);
return;
}
}
// We've already absorbed as many elements as we can
let _ = poseidon_duplex::<F, S, T, RATE>(
&mut self.state,
input,
&self.pad_and_add,
&self.mds_matrix,
&self.round_constants,
);
self.sponge = Sponge::absorb(value);
}
Sponge::Squeezing(_) => {
// Drop the remaining output elements
self.sponge = Sponge::absorb(value);
for entry in self.mode.0.iter_mut() {
if entry.is_none() {
*entry = Some(value);
return;
}
}
// We've already absorbed as many elements as we can
let _ = poseidon_sponge::<F, S, T, RATE>(
&mut self.state,
Some(&self.mode),
&self.mds_matrix,
&self.round_constants,
);
self.mode = Absorbing::init_with(value);
}
/// Transitions the sponge into its squeezing state.
pub(crate) fn finish_absorbing(mut self) -> Sponge<F, S, Squeezing<F, RATE>, T, RATE> {
let mode = poseidon_sponge::<F, S, T, RATE>(
&mut self.state,
Some(&self.mode),
&self.mds_matrix,
&self.round_constants,
);
Sponge {
mode,
state: self.state,
mds_matrix: self.mds_matrix,
round_constants: self.round_constants,
_marker: PhantomData::default(),
}
}
}
impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
Sponge<F, S, Squeezing<F, RATE>, T, RATE>
{
/// Squeezes an element from the sponge.
pub(crate) fn squeeze(&mut self) -> F {
loop {
match self.sponge {
Sponge::Absorbing(ref input) => {
self.sponge = Sponge::Squeezing(poseidon_duplex::<F, S, T, RATE>(
&mut self.state,
input,
&self.pad_and_add,
&self.mds_matrix,
&self.round_constants,
));
}
Sponge::Squeezing(ref mut output) => {
for entry in output.iter_mut() {
if let Some(e) = entry.take() {
return e;
}
}
// We've already squeezed out all available elements
self.sponge = Sponge::Absorbing([None; RATE]);
for entry in self.mode.0.iter_mut() {
if let Some(e) = entry.take() {
return e;
}
}
// We've already squeezed out all available elements
self.mode = poseidon_sponge::<F, S, T, RATE>(
&mut self.state,
None,
&self.mds_matrix,
&self.round_constants,
);
}
}
}
/// A domain in which a Poseidon hash function is being used.
pub trait Domain<F: FieldExt, const T: usize, const RATE: usize>: Copy + fmt::Debug {
pub trait Domain<F: FieldExt, const RATE: usize> {
/// Iterator that outputs padding field elements.
type Padding: IntoIterator<Item = F>;
/// The name of this domain, for debug formatting purposes.
fn name() -> String;
/// The initial capacity element, encoding this domain.
fn initial_capacity_element(&self) -> F;
fn initial_capacity_element() -> F;
/// The padding that will be added to each state word by [`Domain::pad_and_add`].
fn padding(&self) -> SpongeState<F, RATE>;
/// Returns a function that will update the given state with the given input to a
/// duplex permutation round, applying padding according to this domain specification.
fn pad_and_add(&self) -> Box<dyn Fn(&mut State<F, T>, &SpongeState<F, RATE>)>;
/// Returns the padding to be appended to the input.
fn padding(input_len: usize) -> Self::Padding;
}
/// A Poseidon hash function used with constant input length.
@ -267,59 +296,44 @@ pub trait Domain<F: FieldExt, const T: usize, const RATE: usize>: Copy + fmt::De
#[derive(Clone, Copy, Debug)]
pub struct ConstantLength<const L: usize>;
impl<F: FieldExt, const T: usize, const RATE: usize, const L: usize> Domain<F, T, RATE>
for ConstantLength<L>
{
fn initial_capacity_element(&self) -> F {
impl<F: FieldExt, const RATE: usize, const L: usize> Domain<F, RATE> for ConstantLength<L> {
type Padding = iter::Take<iter::Repeat<F>>;
fn name() -> String {
format!("ConstantLength<{}>", L)
}
fn initial_capacity_element() -> F {
// Capacity value is $length \cdot 2^64 + (o-1)$ where o is the output length.
// We hard-code an output length of 1.
F::from_u128((L as u128) << 64)
}
fn padding(&self) -> SpongeState<F, RATE> {
// For constant-input-length hashing, padding consists of the field elements being
// zero.
let mut padding = [None; RATE];
for word in padding.iter_mut().skip(L) {
*word = Some(F::zero());
}
padding
}
fn pad_and_add(&self) -> Box<dyn Fn(&mut State<F, T>, &SpongeState<F, RATE>)> {
Box::new(|state, input| {
// `Iterator::zip` short-circuits when one iterator completes, so this will only
// mutate the rate portion of the state.
for (word, value) in state.iter_mut().zip(input.iter()) {
// For constant-input-length hashing, padding consists of the field
// elements being zero, so we don't add anything to the state.
if let Some(value) = value {
*word += value;
}
}
})
fn padding(input_len: usize) -> Self::Padding {
assert_eq!(input_len, L);
// For constant-input-length hashing, we pad the input with zeroes to a multiple
// of RATE. On its own this would not be sponge-compliant padding, but the
// Poseidon authors encode the constant length into the capacity element, ensuring
// that inputs of different lengths do not share the same permutation.
let k = (L + RATE - 1) / RATE;
iter::repeat(F::zero()).take(k * RATE - L)
}
}
/// A Poseidon hash function, built around a duplex sponge.
/// A Poseidon hash function, built around a sponge.
pub struct Hash<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
D: Domain<F, RATE>,
const T: usize,
const RATE: usize,
> {
duplex: Duplex<F, S, T, RATE>,
domain: D,
sponge: Sponge<F, S, Absorbing<F, RATE>, T, RATE>,
_domain: PhantomData<D>,
}
impl<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
const T: usize,
const RATE: usize,
> fmt::Debug for Hash<F, S, D, T, RATE>
impl<F: FieldExt, S: Spec<F, T, RATE>, D: Domain<F, RATE>, const T: usize, const RATE: usize>
fmt::Debug for Hash<F, S, D, T, RATE>
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Hash")
@ -327,24 +341,19 @@ impl<
.field("rate", &RATE)
.field("R_F", &S::full_rounds())
.field("R_P", &S::partial_rounds())
.field("domain", &self.domain)
.field("domain", &D::name())
.finish()
}
}
impl<
F: FieldExt,
S: Spec<F, T, RATE>,
D: Domain<F, T, RATE>,
const T: usize,
const RATE: usize,
> Hash<F, S, D, T, RATE>
impl<F: FieldExt, S: Spec<F, T, RATE>, D: Domain<F, RATE>, const T: usize, const RATE: usize>
Hash<F, S, D, T, RATE>
{
/// Initializes a new hasher.
pub fn init(domain: D) -> Self {
pub fn init() -> Self {
Hash {
duplex: Duplex::new(domain.initial_capacity_element(), domain.pad_and_add()),
domain,
sponge: Sponge::new(D::initial_capacity_element()),
_domain: PhantomData::default(),
}
}
}
@ -354,10 +363,12 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize, const
{
/// Hashes the given input.
pub fn hash(mut self, message: [F; L]) -> F {
for value in array::IntoIter::new(message) {
self.duplex.absorb(value);
for value in
array::IntoIter::new(message).chain(<ConstantLength<L> as Domain<F, RATE>>::padding(L))
{
self.sponge.absorb(value);
}
self.duplex.squeeze()
self.sponge.finish_absorbing().squeeze()
}
}
@ -374,7 +385,7 @@ mod tests {
let (round_constants, mds, _) = OrchardNullifier::constants();
let hasher = Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength);
let hasher = Hash::<_, OrchardNullifier, ConstantLength<2>, 3, 2>::init();
let result = hasher.hash(message);
// The result should be equivalent to just directly applying the permutation and

View File

@ -292,7 +292,8 @@ mod tests {
Fp::from_repr(tv.input[1]).unwrap(),
];
let result = Hash::<_, super::P128Pow5T3, _, 3, 2>::init(ConstantLength).hash(message);
let result =
Hash::<_, super::P128Pow5T3, ConstantLength<2>, 3, 2>::init().hash(message);
assert_eq!(result.to_repr(), tv.output);
}
@ -303,7 +304,8 @@ mod tests {
Fq::from_repr(tv.input[1]).unwrap(),
];
let result = Hash::<_, super::P128Pow5T3, _, 3, 2>::init(ConstantLength).hash(message);
let result =
Hash::<_, super::P128Pow5T3, ConstantLength<2>, 3, 2>::init().hash(message);
assert_eq!(result.to_repr(), tv.output);
}

View File

@ -212,7 +212,7 @@ pub(crate) fn diversify_hash(d: &[u8; 11]) -> NonIdentityPallasPoint {
///
/// [concreteprfs]: https://zips.z.cash/protocol/nu5.pdf#concreteprfs
pub(crate) fn prf_nf(nk: pallas::Base, rho: pallas::Base) -> pallas::Base {
poseidon::Hash::<_, poseidon::P128Pow5T3, _, 3, 2>::init(poseidon::ConstantLength)
poseidon::Hash::<_, poseidon::P128Pow5T3, poseidon::ConstantLength<2>, 3, 2>::init()
.hash([nk, rho])
}