use core::array::from_fn;
use std::iter::once;
use ff::Field;
use midnight_proofs::{
circuit::{Chip, Layouter, Region, Value},
plonk::{Advice, Column, ConstraintSystem, Constraints, Error, Expression, Fixed, Selector},
poly::Rotation,
};
#[cfg(any(test, feature = "testing"))]
use {crate::testing_utils::FromScratch, midnight_proofs::plonk::Instance};
use super::{
constants::{PoseidonField, NB_FULL_ROUNDS, NB_PARTIAL_ROUNDS, RATE, WIDTH},
full_round_cpu, partial_round_cpu_for_circuits,
round_skips::PreComputedRoundCircuit,
NB_POSEIDON_ADVICE_COLS, NB_POSEIDON_FIXED_COLS,
};
#[cfg(any(test, feature = "testing"))]
use crate::field::{
native::{NB_ARITH_COLS, NB_ARITH_FIXED_COLS},
NativeConfig,
};
use crate::{
field::NativeChip,
instructions::{ArithInstructions, AssignmentInstructions, SpongeInstructions},
types::AssignedNative,
utils::ComposableChip,
CircuitField,
};
pub(crate) const NB_SKIPS_CIRCUIT: usize = 5;
pub(super) type AssignedRegister<F> = [AssignedNative<F>; WIDTH];
#[derive(Clone, Debug)]
pub struct AssignedPoseidonState<F: CircuitField> {
pub(super) register: AssignedRegister<F>,
pub(super) queue: Vec<AssignedNative<F>>,
pub(super) squeeze_position: usize,
input_len: Option<usize>,
}
#[derive(Clone, Debug)]
pub struct PoseidonConfig<F: PoseidonField> {
q_full_round: Selector,
q_partial_round: Selector,
register_cols: [Column<Advice>; NB_POSEIDON_ADVICE_COLS],
constant_cols: [Column<Fixed>; NB_POSEIDON_FIXED_COLS],
pre_computed: PreComputedRoundCircuit<F>,
}
#[derive(Clone, Debug)]
pub struct PoseidonChip<F: PoseidonField> {
config: PoseidonConfig<F>,
native_chip: NativeChip<F>,
}
impl<F: PoseidonField> Chip<F> for PoseidonChip<F> {
type Config = PoseidonConfig<F>;
type Loaded = ();
fn config(&self) -> &Self::Config {
&self.config
}
fn loaded(&self) -> &Self::Loaded {
&()
}
}
impl<F: PoseidonField> PoseidonChip<F> {
pub const fn register_size() -> usize {
WIDTH
}
pub const fn rate() -> usize {
RATE
}
pub const fn nb_full_rounds() -> usize {
NB_FULL_ROUNDS
}
pub const fn nb_partial_rounds() -> usize {
NB_PARTIAL_ROUNDS
}
}
fn linear_layer<F: PoseidonField>(
inputs: [Expression<F>; WIDTH],
outputs: [Expression<F>; WIDTH],
constants: [Expression<F>; WIDTH],
) -> [Expression<F>; WIDTH] {
let mut ids = constants;
for (i, output) in outputs.into_iter().enumerate() {
ids[i] = &ids[i] - output;
#[allow(clippy::needless_range_loop)]
for j in 0..WIDTH {
ids[i] = &ids[i] + Expression::Constant(F::MDS[i][j]) * &inputs[j];
}
}
ids
}
pub(crate) fn sbox<F: Field>(x: Expression<F>) -> Expression<F> {
x.clone() * x.square().square()
}
impl<F: PoseidonField> ComposableChip<F> for PoseidonChip<F> {
type SharedResources = (
[Column<Advice>; NB_POSEIDON_ADVICE_COLS],
[Column<Fixed>; NB_POSEIDON_FIXED_COLS],
);
type InstructionDeps = NativeChip<F>;
fn new(config: &PoseidonConfig<F>, native_chip: &Self::InstructionDeps) -> Self {
Self {
config: config.clone(),
native_chip: native_chip.clone(),
}
}
fn configure(
meta: &mut ConstraintSystem<F>,
shared_res: &Self::SharedResources,
) -> PoseidonConfig<F> {
let register_cols = shared_res.0;
let constant_cols = shared_res.1;
let q_full_round = meta.selector();
let q_partial_round = meta.complex_selector();
register_cols[..WIDTH].iter().for_each(|col| meta.enable_equality(*col));
meta.create_gate("full_round_gate", |meta| {
let inputs_and_hints: [(Expression<F>, Expression<F>); WIDTH] = from_fn(|i| {
(
meta.query_advice(register_cols[i], Rotation::cur()),
meta.query_advice(register_cols[i + WIDTH], Rotation::cur()),
)
});
let outputs = from_fn(|i| meta.query_advice(register_cols[i], Rotation::next()));
let constants = from_fn(|i| meta.query_fixed(constant_cols[i], Rotation::cur()));
let sboxed_inputs = inputs_and_hints.clone().map(|(x, x3)| x.square() * x3);
Constraints::with_selector(
q_full_round,
[
inputs_and_hints.map(|(x, x3)| x.clone() * x.square() - x3),
linear_layer(sboxed_inputs, outputs, constants),
]
.concat(),
)
});
let pre_computed = PreComputedRoundCircuit::<F>::init();
let ids = pre_computed.partial_round_id;
meta.create_gate("partial_round_gate", |meta| {
let inputs = register_cols
.iter()
.map(|col| meta.query_advice(*col, Rotation::cur()))
.collect::<Vec<_>>();
let round_constants = constant_cols
.iter()
.map(|col| meta.query_fixed(*col, Rotation::cur()))
.collect::<Vec<_>>();
let outputs = register_cols[0..WIDTH]
.iter()
.map(|col| meta.query_advice(*col, Rotation::next()))
.collect::<Vec<_>>();
let constraints = ids.to_expression(&inputs);
let output_lin_constraints =
(0..WIDTH - 1).map(|i| &round_constants[i] - &outputs[i] + &constraints[i]);
let input_pow_constraints = (WIDTH - 1..WIDTH + NB_SKIPS_CIRCUIT - 1)
.map(|i| &round_constants[i] - &inputs[i + 1] + &constraints[i]);
let output_pow_constraint: Expression<F> =
&round_constants[WIDTH + NB_SKIPS_CIRCUIT - 1] - &outputs[WIDTH - 1]
+ &constraints[WIDTH + NB_SKIPS_CIRCUIT - 1];
let constraints = output_lin_constraints
.chain(input_pow_constraints)
.chain(once(output_pow_constraint))
.collect::<Vec<_>>();
Constraints::with_additive_selector(q_partial_round, constraints)
});
PoseidonConfig {
q_full_round,
q_partial_round,
register_cols,
constant_cols,
pre_computed,
}
}
fn load(&self, _layouter: &mut impl Layouter<F>) -> Result<(), Error> {
Ok(())
}
}
impl<F: PoseidonField> PoseidonChip<F> {
fn assign_constants_full(
&self,
region: &mut Region<'_, F>,
round_index: usize,
offset: usize,
) -> Result<(), Error> {
let round_constants = if round_index == NB_FULL_ROUNDS + NB_PARTIAL_ROUNDS - 1 {
[F::ZERO; WIDTH]
} else {
F::ROUND_CONSTANTS[round_index + 1]
};
for (col, constant) in self.config.constant_cols[0..WIDTH].iter().zip(round_constants) {
region.assign_fixed(
|| "load constant for a full round",
*col,
offset,
|| Value::known(constant),
)?;
}
Ok(())
}
fn assign_constants_partial(
&self,
region: &mut Region<'_, F>,
round_batch_index: usize,
offset: usize,
) -> Result<(), Error> {
for (col, constant) in self
.config
.constant_cols
.iter()
.zip(&self.config.pre_computed.round_constants[round_batch_index])
{
region.assign_fixed(
|| format!("load constant for partial a round with {NB_SKIPS_CIRCUIT} skips"),
*col,
offset,
|| Value::known(*constant),
)?;
}
Ok(())
}
fn full_round(
&self,
region: &mut Region<'_, F>,
inputs: &mut [AssignedNative<F>; WIDTH],
round_index: usize,
offset: &mut usize,
) -> Result<(), Error> {
self.config.q_full_round.enable(region, *offset)?;
self.assign_constants_full(region, round_index, *offset)?;
for (x, col) in inputs.iter().zip(self.config.register_cols[WIDTH..(2 * WIDTH)].iter()) {
region.assign_advice(
|| "full round hint",
*col,
*offset,
|| x.value().map(|x| *x * x.square()),
)?;
}
*offset += 1;
let outputs = Value::from_iter(inputs.iter().map(|x| x.value().copied()))
.map(|inputs: Vec<F>| {
let mut inputs = inputs;
full_round_cpu(round_index, &mut inputs);
inputs
})
.transpose_vec(WIDTH);
outputs
.iter()
.zip(self.config.register_cols[0..WIDTH].iter())
.zip(inputs)
.try_for_each(|((output, column), input)| {
region
.assign_advice(|| "full round output", *column, *offset, || *output)
.map(|assigned| *input = assigned)
})
}
fn partial_round(
&self,
region: &mut Region<'_, F>,
inputs: &mut [AssignedNative<F>], round_batch_index: usize,
offset: &mut usize,
) -> Result<(), Error> {
self.config.q_partial_round.enable(region, *offset)?;
self.assign_constants_partial(region, round_batch_index, *offset)?;
let outputs = Value::from_iter(inputs.iter().map(|x| x.value().copied()))
.map(|inputs: Vec<F>| {
let mut state = inputs;
let skip_advice_vals = partial_round_cpu_for_circuits(
&self.config.pre_computed,
round_batch_index,
&mut state,
);
for (col, skip_advice) in self.config.register_cols[WIDTH..WIDTH + NB_SKIPS_CIRCUIT]
.iter()
.zip(skip_advice_vals)
{
let _ = region.assign_advice(
|| "partial round intermediary inputs",
*col,
*offset,
|| Value::known(skip_advice),
);
}
state
})
.transpose_vec(WIDTH);
*offset += 1;
outputs
.iter()
.zip(self.config.register_cols[0..WIDTH].iter())
.zip(inputs)
.try_for_each(|((output, column), input)| {
region
.assign_advice(|| "partial round outputs", *column, *offset, || *output)
.map(|assigned| *input = assigned)
})
}
pub(super) fn permutation(
&self,
layouter: &mut impl Layouter<F>,
inputs: &AssignedRegister<F>,
) -> Result<AssignedRegister<F>, Error> {
layouter.assign_region(
|| "permutation layout",
|mut region| {
let mut offset: usize = 0;
let mut state: AssignedRegister<F> = self
.native_chip
.add_constants_in_region(
&mut region,
inputs,
&F::ROUND_CONSTANTS[0],
&mut offset,
)?
.try_into()
.unwrap();
assert!(
state.iter().zip(self.config.register_cols).all(|(acell, col)| {
let col1: Column<Advice> = acell.cell().column.try_into().unwrap();
col1 == col
})
);
for round_index in 0..NB_FULL_ROUNDS / 2 {
self.full_round(&mut region, &mut state, round_index, &mut offset)?;
}
for round_batch_index in 0..NB_PARTIAL_ROUNDS / (1 + NB_SKIPS_CIRCUIT) {
self.partial_round(&mut region, &mut state, round_batch_index, &mut offset)?;
}
for round_index in
(NB_FULL_ROUNDS / 2 + NB_PARTIAL_ROUNDS..).take(NB_FULL_ROUNDS / 2)
{
self.full_round(&mut region, &mut state, round_index, &mut offset)?;
}
Ok(state)
},
)
}
}
impl<F: PoseidonField> SpongeInstructions<F, AssignedNative<F>, AssignedNative<F>>
for PoseidonChip<F>
{
type State = AssignedPoseidonState<F>;
fn init(
&self,
layouter: &mut impl Layouter<F>,
input_len: Option<usize>,
) -> Result<Self::State, Error> {
let zero = self.native_chip.assign_fixed(layouter, F::ZERO)?;
let mut register: AssignedRegister<F> = vec![zero; WIDTH].try_into().unwrap();
register[RATE] = self.native_chip.assign_fixed(
layouter,
F::from_u128(input_len.map(|len| len as u128).unwrap_or(1 << 64)),
)?;
Ok(AssignedPoseidonState {
register,
queue: Vec::new(),
squeeze_position: 0,
input_len,
})
}
fn absorb(
&self,
_layouter: &mut impl Layouter<F>,
state: &mut Self::State,
inputs: &[AssignedNative<F>],
) -> Result<(), Error> {
state.queue.extend(inputs.to_vec());
state.squeeze_position = 0;
Ok(())
}
fn squeeze(
&self,
layouter: &mut impl Layouter<F>,
state: &mut Self::State,
) -> Result<AssignedNative<F>, Error> {
if state.squeeze_position > 0 {
if state.input_len.is_some() {
panic!(
"Attempting to squeeze multiple times a fixed-size Poseidon sponge (Circuit)."
)
};
debug_assert!(state.queue.is_empty());
let output = state.register[state.squeeze_position % RATE].clone();
state.squeeze_position = (state.squeeze_position + 1) % RATE;
return Ok(output);
}
match state.input_len {
None => {
let padding =
self.native_chip.assign_fixed(layouter, F::from(state.queue.len() as u64))?;
state.queue.push(padding);
}
Some(len) => {
if state.queue.len() != len {
panic!("Inconsistent lengths in fixed-size Poseidon sponge (Circuit). Expected: {}, found: {}.", len, state.queue.len())
};
}
}
for chunk in state.queue.chunks(RATE) {
for (entry, value) in state.register.iter_mut().zip(chunk.iter()) {
*entry = self.native_chip.add(layouter, entry, value)?;
}
state.register = self.permutation(layouter, &state.register)?;
}
state.queue = Vec::new();
state.squeeze_position = 1 % RATE;
Ok(state.register[0].clone())
}
}
#[cfg(any(test, feature = "testing"))]
impl<F: PoseidonField> FromScratch<F> for PoseidonChip<F> {
type Config = (NativeConfig, PoseidonConfig<F>);
fn new_from_scratch(config: &Self::Config) -> Self {
let native_chip = NativeChip::new(&config.0, &());
PoseidonChip::new(&config.1, &native_chip)
}
fn configure_from_scratch(
meta: &mut ConstraintSystem<F>,
advice_columns: &mut Vec<Column<Advice>>,
fixed_columns: &mut Vec<Column<Fixed>>,
instance_columns: &[Column<Instance>; 2],
) -> Self::Config {
let nb_advice_needed = std::cmp::max(NB_POSEIDON_ADVICE_COLS, NB_ARITH_COLS);
let nb_fixed_needed = std::cmp::max(NB_POSEIDON_FIXED_COLS, NB_ARITH_FIXED_COLS);
while advice_columns.len() < nb_advice_needed {
advice_columns.push(meta.advice_column());
}
while fixed_columns.len() < nb_fixed_needed {
fixed_columns.push(meta.fixed_column());
}
let native_config = NativeChip::configure(
meta,
&(
advice_columns[..NB_ARITH_COLS].try_into().unwrap(),
fixed_columns[..NB_ARITH_FIXED_COLS].try_into().unwrap(),
*instance_columns,
),
);
let poseidon_config = PoseidonChip::configure(
meta,
&(
advice_columns[..NB_POSEIDON_ADVICE_COLS].try_into().unwrap(),
fixed_columns[..NB_POSEIDON_FIXED_COLS].try_into().unwrap(),
),
);
(native_config, poseidon_config)
}
fn load_from_scratch(&self, layouter: &mut impl Layouter<F>) -> Result<(), Error> {
self.native_chip.load_from_scratch(layouter)
}
}
#[cfg(test)]
mod tests {
use midnight_proofs::{circuit::SimpleFloorPlanner, dev::MockProver, plonk::Circuit};
use super::*;
use crate::{
field::NativeGadget,
hash::poseidon::{permutation_cpu, round_skips::PreComputedRoundCPU},
instructions::{sponge::tests::test_sponge, AssertionInstructions},
utils::circuit_modeling::{circuit_to_json, cost_measure_end, cost_measure_start},
};
#[derive(Clone, Debug, Default)]
struct PermCircuit<F> {
inputs: [Value<F>; WIDTH],
expected: [F; WIDTH],
}
impl<F: PoseidonField> Circuit<F> for PermCircuit<F> {
type Config = <PoseidonChip<F> as FromScratch<F>>::Config;
type FloorPlanner = SimpleFloorPlanner;
type Params = ();
fn without_witnesses(&self) -> Self {
unreachable!()
}
fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
let committed_instance_column = meta.instance_column();
let instance_column = meta.instance_column();
PoseidonChip::configure_from_scratch(
meta,
&mut vec![],
&mut vec![],
&[committed_instance_column, instance_column],
)
}
fn synthesize(
&self,
config: Self::Config,
mut layouter: impl Layouter<F>,
) -> Result<(), Error> {
let poseidon_chip = PoseidonChip::new_from_scratch(&config);
let inputs: AssignedRegister<F> = poseidon_chip
.native_chip
.assign_many(&mut layouter, &self.inputs)?
.try_into()
.unwrap();
cost_measure_start(&mut layouter);
let outputs = poseidon_chip.permutation(&mut layouter, &inputs)?;
cost_measure_end(&mut layouter);
for (out, expected) in outputs.iter().zip(self.expected.iter()) {
poseidon_chip.native_chip.assert_equal_to_fixed(&mut layouter, out, *expected)?;
}
poseidon_chip.load_from_scratch(&mut layouter)
}
}
pub fn run_permutation_test<F>(inputs: [F; WIDTH], cost_model: bool)
where
F: PoseidonField + ff::FromUniformBytes<64> + Ord,
{
let pre_computed = PreComputedRoundCPU::init();
let mut expected = inputs;
permutation_cpu(&pre_computed, &mut expected);
let circuit = PermCircuit {
inputs: inputs.map(Value::known),
expected,
};
MockProver::run(&circuit, vec![vec![], vec![]]).unwrap().assert_satisfied();
if cost_model {
circuit_to_json("Poseidon", "one_permutation", circuit);
}
}
fn run_sponge_test<F>(field: &str, cost_model: bool)
where
F: PoseidonField + ff::FromUniformBytes<64> + Ord,
{
println!(
">> Testing Poseidon Sponge (field {field}, {} partial-round skips)",
NB_SKIPS_CIRCUIT
);
test_sponge::<
F,
AssignedNative<F>,
AssignedNative<F>,
PoseidonChip<F>,
NativeGadget<F, _, _>,
>(cost_model, "Poseidon");
println!("=> Done.\n")
}
#[test]
fn permutation_test() {
let inputs = [midnight_curves::Fq::from(0); WIDTH];
run_permutation_test(inputs, true);
}
#[test]
fn sponge_test() {
run_sponge_test::<midnight_curves::Fq>("blstrs", true);
}
}