poseidon: Make gadget tests generic over WIDTH, RATE

This commit is contained in:
therealyingtong 2021-11-03 14:37:02 +01:00
parent b63c868591
commit 9b76556503
2 changed files with 69 additions and 61 deletions

View File

@ -615,40 +615,37 @@ mod tests {
}, },
primitives::poseidon::{self, ConstantLength, P128Pow5T3 as OrchardNullifier, Spec}, primitives::poseidon::{self, ConstantLength, P128Pow5T3 as OrchardNullifier, Spec},
}; };
use std::convert::TryInto;
use std::marker::PhantomData;
const WIDTH: usize = 3; struct PermuteCircuit<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize>(
const RATE: usize = 2; PhantomData<S>,
);
struct PermuteCircuit {} impl<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize> Circuit<Fp>
for PermuteCircuit<S, WIDTH, RATE>
impl Circuit<Fp> for PermuteCircuit { {
type Config = Pow5Config<Fp, WIDTH, RATE>; type Config = Pow5Config<Fp, WIDTH, RATE>;
type FloorPlanner = SimpleFloorPlanner; type FloorPlanner = SimpleFloorPlanner;
fn without_witnesses(&self) -> Self { fn without_witnesses(&self) -> Self {
PermuteCircuit {} PermuteCircuit::<S, WIDTH, RATE>(PhantomData)
} }
fn configure(meta: &mut ConstraintSystem<Fp>) -> Pow5Config<Fp, WIDTH, RATE> { fn configure(meta: &mut ConstraintSystem<Fp>) -> Pow5Config<Fp, WIDTH, RATE> {
let state = [ let state = (0..WIDTH).map(|_| meta.advice_column()).collect::<Vec<_>>();
meta.advice_column(),
meta.advice_column(),
meta.advice_column(),
];
let partial_sbox = meta.advice_column(); let partial_sbox = meta.advice_column();
let rc_a = [ let rc_a = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
meta.fixed_column(), let rc_b = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
meta.fixed_column(),
meta.fixed_column(),
];
let rc_b = [
meta.fixed_column(),
meta.fixed_column(),
meta.fixed_column(),
];
Pow5Chip::configure::<OrchardNullifier>(meta, state, partial_sbox, rc_a, rc_b) Pow5Chip::configure::<S>(
meta,
state.try_into().unwrap(),
partial_sbox,
rc_a.try_into().unwrap(),
rc_b.try_into().unwrap(),
)
} }
fn synthesize( fn synthesize(
@ -659,7 +656,7 @@ mod tests {
let initial_state = layouter.assign_region( let initial_state = layouter.assign_region(
|| "prepare initial state", || "prepare initial state",
|mut region| { |mut region| {
let mut state_word = |i: usize| -> Result<_, Error> { let state_word = |i: usize| {
let value = Some(Fp::from(i as u64)); let value = Some(Fp::from(i as u64));
let var = region.assign_advice( let var = region.assign_advice(
|| format!("load state_{}", i), || format!("load state_{}", i),
@ -670,22 +667,27 @@ mod tests {
Ok(StateWord { var, value }) Ok(StateWord { var, value })
}; };
Ok([state_word(0)?, state_word(1)?, state_word(2)?]) let state: Result<Vec<_>, Error> = (0..WIDTH).map(state_word).collect();
Ok(state?.try_into().unwrap())
}, },
)?; )?;
let chip = Pow5Chip::construct(config.clone()); let chip = Pow5Chip::construct(config.clone());
let final_state = <Pow5Chip<_, WIDTH, RATE> as PoseidonInstructions< let final_state = <Pow5Chip<_, WIDTH, RATE> as PoseidonInstructions<
Fp, Fp,
OrchardNullifier, S,
WIDTH, WIDTH,
2, RATE,
>>::permute(&chip, &mut layouter, &initial_state)?; >>::permute(&chip, &mut layouter, &initial_state)?;
// For the purpose of this test, compute the real final state inline. // For the purpose of this test, compute the real final state inline.
let mut expected_final_state = [Fp::zero(), Fp::one(), Fp::from_u64(2)]; let mut expected_final_state = (0..WIDTH)
let (round_constants, mds, _) = OrchardNullifier::constants(); .map(|idx| Fp::from_u64(idx as u64))
poseidon::permute::<_, OrchardNullifier, WIDTH, RATE>( .collect::<Vec<_>>()
.try_into()
.unwrap();
let (round_constants, mds, _) = S::constants();
poseidon::permute::<_, S, WIDTH, RATE>(
&mut expected_final_state, &mut expected_final_state,
&mds, &mds,
&round_constants, &round_constants,
@ -704,9 +706,11 @@ mod tests {
region.constrain_equal(final_state[i].var, var) region.constrain_equal(final_state[i].var, var)
}; };
final_state_word(0)?; for i in 0..(WIDTH - 1) {
final_state_word(1)?; final_state_word(i)?;
final_state_word(2) }
final_state_word(WIDTH - 1)
}, },
) )
} }
@ -715,49 +719,49 @@ mod tests {
#[test] #[test]
fn poseidon_permute() { fn poseidon_permute() {
let k = 6; let k = 6;
let circuit = PermuteCircuit {}; let circuit = PermuteCircuit::<OrchardNullifier, 3, 2>(PhantomData);
let prover = MockProver::run(k, &circuit, vec![]).unwrap(); let prover = MockProver::run(k, &circuit, vec![]).unwrap();
assert_eq!(prover.verify(), Ok(())) assert_eq!(prover.verify(), Ok(()))
} }
#[derive(Default)] struct HashCircuit<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize> {
struct HashCircuit {
message: Option<[Fp; 2]>, message: Option<[Fp; 2]>,
// For the purpose of this test, witness the result. // For the purpose of this test, witness the result.
// TODO: Move this into an instance column. // TODO: Move this into an instance column.
output: Option<Fp>, output: Option<Fp>,
_spec: PhantomData<S>,
} }
impl Circuit<Fp> for HashCircuit { impl<S: Spec<Fp, WIDTH, RATE>, const WIDTH: usize, const RATE: usize> Circuit<Fp>
for HashCircuit<S, WIDTH, RATE>
{
type Config = Pow5Config<Fp, WIDTH, RATE>; type Config = Pow5Config<Fp, WIDTH, RATE>;
type FloorPlanner = SimpleFloorPlanner; type FloorPlanner = SimpleFloorPlanner;
fn without_witnesses(&self) -> Self { fn without_witnesses(&self) -> Self {
Self::default() Self {
message: None,
output: None,
_spec: PhantomData,
}
} }
fn configure(meta: &mut ConstraintSystem<Fp>) -> Pow5Config<Fp, WIDTH, RATE> { fn configure(meta: &mut ConstraintSystem<Fp>) -> Pow5Config<Fp, WIDTH, RATE> {
let state = [ let state = (0..WIDTH).map(|_| meta.advice_column()).collect::<Vec<_>>();
meta.advice_column(),
meta.advice_column(),
meta.advice_column(),
];
let partial_sbox = meta.advice_column(); let partial_sbox = meta.advice_column();
let rc_a = [ let rc_a = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
meta.fixed_column(), let rc_b = (0..WIDTH).map(|_| meta.fixed_column()).collect::<Vec<_>>();
meta.fixed_column(),
meta.fixed_column(),
];
let rc_b = [
meta.fixed_column(),
meta.fixed_column(),
meta.fixed_column(),
];
meta.enable_constant(rc_b[0]); meta.enable_constant(rc_b[0]);
Pow5Chip::configure::<OrchardNullifier>(meta, state, partial_sbox, rc_a, rc_b) Pow5Chip::configure::<S>(
meta,
state.try_into().unwrap(),
partial_sbox,
rc_a.try_into().unwrap(),
rc_b.try_into().unwrap(),
)
} }
fn synthesize( fn synthesize(
@ -770,7 +774,7 @@ mod tests {
let message = layouter.assign_region( let message = layouter.assign_region(
|| "load message", || "load message",
|mut region| { |mut region| {
let mut message_word = |i: usize| -> Result<_, Error> { let message_word = |i: usize| {
let value = self.message.map(|message_vals| message_vals[i]); let value = self.message.map(|message_vals| message_vals[i]);
let cell = region.assign_advice( let cell = region.assign_advice(
|| format!("load message_{}", i), || format!("load message_{}", i),
@ -781,14 +785,15 @@ mod tests {
Ok(CellValue::new(cell, value)) Ok(CellValue::new(cell, value))
}; };
Ok([message_word(0)?, message_word(1)?]) let message: Result<Vec<_>, Error> = (0..RATE).map(message_word).collect();
Ok(message?.try_into().unwrap())
}, },
)?; )?;
let hasher = Hash::<_, _, OrchardNullifier, _, WIDTH, 2>::init( let hasher = Hash::<_, _, S, _, WIDTH, RATE>::init(
chip, chip,
layouter.namespace(|| "init"), layouter.namespace(|| "init"),
ConstantLength::<2>, ConstantLength::<RATE>,
)?; )?;
let output = hasher.hash(layouter.namespace(|| "hash"), message)?; let output = hasher.hash(layouter.namespace(|| "hash"), message)?;
@ -814,9 +819,10 @@ mod tests {
poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength::<2>).hash(message); poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength::<2>).hash(message);
let k = 6; let k = 6;
let circuit = HashCircuit { let circuit = HashCircuit::<OrchardNullifier, 3, 2> {
message: Some(message), message: Some(message),
output: Some(output), output: Some(output),
_spec: PhantomData,
}; };
let prover = MockProver::run(k, &circuit, vec![]).unwrap(); let prover = MockProver::run(k, &circuit, vec![]).unwrap();
assert_eq!(prover.verify(), Ok(())) assert_eq!(prover.verify(), Ok(()))
@ -833,9 +839,10 @@ mod tests {
poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength).hash(message); poseidon::Hash::<_, OrchardNullifier, _, 3, 2>::init(ConstantLength).hash(message);
let k = 6; let k = 6;
let circuit = HashCircuit { let circuit = HashCircuit::<OrchardNullifier, 3, 2> {
message: Some(message), message: Some(message),
output: Some(output), output: Some(output),
_spec: PhantomData,
}; };
let prover = MockProver::run(k, &circuit, vec![]).unwrap(); let prover = MockProver::run(k, &circuit, vec![]).unwrap();
assert_eq!(prover.verify(), Ok(())); assert_eq!(prover.verify(), Ok(()));
@ -853,9 +860,10 @@ mod tests {
.titled("Poseidon Chip Layout", ("sans-serif", 60)) .titled("Poseidon Chip Layout", ("sans-serif", 60))
.unwrap(); .unwrap();
let circuit = HashCircuit { let circuit = HashCircuit::<OrchardNullifier, 3, 2> {
message: None, message: None,
output: None, output: None,
_spec: PhantomData,
}; };
halo2::dev::CircuitLayout::default() halo2::dev::CircuitLayout::default()
.render(6, &circuit, &root) .render(6, &circuit, &root)

View File

@ -30,7 +30,7 @@ pub(crate) type SpongeState<F, const RATE: usize> = [Option<F>; RATE];
pub(crate) type Mds<F, const T: usize> = [[F; T]; T]; pub(crate) type Mds<F, const T: usize> = [[F; T]; T];
/// A specification for a Poseidon permutation. /// A specification for a Poseidon permutation.
pub trait Spec<F: FieldExt, const T: usize, const RATE: usize> { pub trait Spec<F: FieldExt, const T: usize, const RATE: usize>: fmt::Debug {
/// The number of full rounds for this specification. /// The number of full rounds for this specification.
/// ///
/// This must be an even number. /// This must be an even number.