openvm_static_verifier/hash/
poseidon2.rs

1//! Halo2 implementation of poseidon2 perm for Bn254
2//! sbox degree 5
3
4use halo2_base::{
5    gates::{range::RangeChip, GateInstructions, RangeInstructions},
6    utils::ScalarField,
7    AssignedValue, Context,
8    QuantumCell::{self, Constant},
9};
10pub(crate) use openvm_stark_sdk::config::baby_bear_bn254_poseidon2::{
11    BABY_BEAR_RATE as MULTI_FIELD32_RATE, BN254_RATE as POSEIDON2_RATE, DIGEST_WIDTH,
12};
13
14use crate::{field::baby_bear::ReducedBabyBearWire, Fr};
15
16const MULTI_FIELD32_NUM_F_ELMS: usize = MULTI_FIELD32_RATE / POSEIDON2_RATE;
17
18#[derive(Clone, Debug)]
19pub struct Poseidon2State<F: ScalarField, const T: usize> {
20    pub s: [AssignedValue<F>; T],
21}
22
23#[derive(Debug, Clone)]
24pub struct Poseidon2Params<F: ScalarField, const T: usize> {
25    /// Number of full rounds
26    pub rounds_f: usize,
27    pub rounds_p: usize,
28    pub mat_internal_diag_m_1: [F; T],
29    pub external_rc: Vec<[F; T]>,
30    pub internal_rc: Vec<F>,
31}
32
33impl<F: ScalarField, const T: usize> Poseidon2Params<F, T> {
34    pub fn new(
35        rounds_f: usize,
36        rounds_p: usize,
37        mat_internal_diag_m_1: [F; T],
38        external_rc: Vec<[F; T]>,
39        internal_rc: Vec<F>,
40    ) -> Self {
41        Self {
42            rounds_f,
43            rounds_p,
44            mat_internal_diag_m_1,
45            external_rc,
46            internal_rc,
47        }
48    }
49}
50
51impl<F: ScalarField, const T: usize> Poseidon2State<F, T> {
52    pub fn new(state: [AssignedValue<F>; T]) -> Self {
53        Self { s: state }
54    }
55    /// Perform permutation on this state.
56    ///
57    /// ATTENTION: inputs.len() needs to be fixed at compile time.
58    pub fn permutation(
59        &mut self,
60        ctx: &mut Context<F>,
61        gate: &impl GateInstructions<F>,
62        params: &Poseidon2Params<F, T>,
63    ) {
64        let rounds_f_beginning = params.rounds_f / 2;
65
66        // First half of the full round
67        self.matmul_external(ctx, gate);
68        for r in 0..rounds_f_beginning {
69            self.add_rc(ctx, gate, params.external_rc[r]);
70            self.sbox(ctx, gate);
71            self.matmul_external(ctx, gate);
72        }
73
74        for r in 0..params.rounds_p {
75            self.s[0] = gate.add(ctx, self.s[0], Constant(params.internal_rc[r]));
76            self.s[0] = Self::x_power5(ctx, gate, self.s[0]);
77            self.matmul_internal(ctx, gate, params.mat_internal_diag_m_1);
78        }
79
80        for r in rounds_f_beginning..params.rounds_f {
81            self.add_rc(ctx, gate, params.external_rc[r]);
82            self.sbox(ctx, gate);
83            self.matmul_external(ctx, gate);
84        }
85    }
86
87    fn x_power5(
88        ctx: &mut Context<F>,
89        gate: &impl GateInstructions<F>,
90        x: AssignedValue<F>,
91    ) -> AssignedValue<F> {
92        let x2 = gate.mul(ctx, x, x);
93        let x4 = gate.mul(ctx, x2, x2);
94        gate.mul(ctx, x, x4)
95    }
96
97    fn sbox(&mut self, ctx: &mut Context<F>, gate: &impl GateInstructions<F>) {
98        for x in self.s.iter_mut() {
99            *x = Self::x_power5(ctx, gate, *x);
100        }
101    }
102
103    fn matmul_external(&mut self, ctx: &mut Context<F>, gate: &impl GateInstructions<F>) {
104        // M_E for T=2: circ(2, 1) = [[2,1],[1,2]]
105        // M_E for T=3: circ(2, 1, 1)
106        assert!(T == 2 || T == 3);
107
108        let sum = gate.sum(ctx, self.s.iter().copied());
109        for (i, x) in self.s.iter_mut().enumerate() {
110            // This is the same as `*x = gate.add(ctx, *x, sum)` but we save a cell by reusing
111            // `sum`:
112            if i % 2 == 0 {
113                ctx.assign_region(
114                    [
115                        QuantumCell::Witness(*x.value() + sum.value()),
116                        QuantumCell::Existing(*x),
117                        QuantumCell::Constant(-F::ONE),
118                        QuantumCell::Existing(sum),
119                    ],
120                    [0],
121                );
122                *x = ctx.get(-4);
123            } else {
124                ctx.assign_region(
125                    [
126                        QuantumCell::Existing(*x),
127                        QuantumCell::Constant(F::ONE),
128                        QuantumCell::Witness(*x.value() + sum.value()),
129                    ],
130                    [-1],
131                );
132                *x = ctx.get(-1);
133            }
134        }
135    }
136
137    fn add_rc(
138        &mut self,
139        ctx: &mut Context<F>,
140        gate: &impl GateInstructions<F>,
141        round_constants: [F; T],
142    ) {
143        for (x, rc) in self.s.iter_mut().zip(round_constants.iter()) {
144            *x = gate.add(ctx, *x, Constant(*rc));
145        }
146    }
147
148    fn matmul_internal(
149        &mut self,
150        ctx: &mut Context<F>,
151        gate: &impl GateInstructions<F>,
152        mat_internal_diag_m_1: [F; T],
153    ) {
154        assert!(T == 2 || T == 3);
155        let sum = gate.sum(ctx, self.s.iter().copied());
156        #[allow(clippy::needless_range_loop)]
157        for i in 0..T {
158            // This is the same as `self.s[i] = gate.mul_add(ctx, self.s[i],
159            // Constant(mat_internal_diag_m_1[i]), sum)` but we save a cell by reusing `sum`.
160            if i % 2 == 0 {
161                ctx.assign_region(
162                    [
163                        QuantumCell::Witness(
164                            *self.s[i].value() * mat_internal_diag_m_1[i] + sum.value(),
165                        ),
166                        QuantumCell::Existing(self.s[i]),
167                        QuantumCell::Constant(-mat_internal_diag_m_1[i]),
168                        QuantumCell::Existing(sum),
169                    ],
170                    [0],
171                );
172                self.s[i] = ctx.get(-4);
173            } else {
174                ctx.assign_region(
175                    [
176                        QuantumCell::Existing(self.s[i]),
177                        QuantumCell::Constant(mat_internal_diag_m_1[i]),
178                        QuantumCell::Witness(
179                            *self.s[i].value() * mat_internal_diag_m_1[i] + sum.value(),
180                        ),
181                    ],
182                    [-1],
183                );
184                self.s[i] = ctx.get(-1);
185            }
186        }
187    }
188}
189
190pub(crate) fn pack_base_2_31_cells(
191    ctx: &mut Context<Fr>,
192    gate: &impl GateInstructions<Fr>,
193    values: &[ReducedBabyBearWire],
194) -> AssignedValue<Fr> {
195    assert!(
196        values.len() <= MULTI_FIELD32_NUM_F_ELMS,
197        "base-2^31 packing supports at most {MULTI_FIELD32_NUM_F_ELMS} limbs"
198    );
199    let base = Fr::from(1u64 << 31);
200    gate.inner_product(
201        ctx,
202        values.iter().map(|value| value.value()),
203        core::iter::successors(Some(Fr::from(1u64)), |power| Some(*power * base))
204            .take(values.len())
205            .map(Constant),
206    )
207}
208
209pub(crate) fn hash_babybear_slice_to_digest(
210    ctx: &mut Context<Fr>,
211    range: &RangeChip<Fr>,
212    values: &[ReducedBabyBearWire],
213) -> AssignedValue<Fr> {
214    let gate = range.gate();
215    let params = &*super::POSEIDON2_PARAMS;
216    let zero = ctx.load_zero();
217    let mut state = Poseidon2State::new(core::array::from_fn(|_| zero));
218    for block_chunk in values.chunks(MULTI_FIELD32_RATE) {
219        for (chunk_id, chunk) in block_chunk.chunks(MULTI_FIELD32_NUM_F_ELMS).enumerate() {
220            state.s[chunk_id] = pack_base_2_31_cells(ctx, gate, chunk);
221        }
222        state.permutation(ctx, gate, params);
223    }
224    state.s[0]
225}
226
227pub(crate) fn compress_bn254_digests(
228    ctx: &mut Context<Fr>,
229    range: &RangeChip<Fr>,
230    left: AssignedValue<Fr>,
231    right: AssignedValue<Fr>,
232) -> AssignedValue<Fr> {
233    let gate = range.gate();
234    let params = &*super::POSEIDON2_COMPRESS_PARAMS;
235    let mut state = Poseidon2State::new([left, right]);
236    state.permutation(ctx, gate, params);
237    state.s[0]
238}
openvm_static_verifier/hash/poseidon2.rs

openvm_static_verifier/hash/
poseidon2.rs
