openvm_circuit/system/poseidon2/
chip.rs

Help
1use std::{
2    array,
3    sync::{atomic::AtomicU32, Arc},
4};
5
6use openvm_poseidon2_air::{Poseidon2Config, Poseidon2SubChip};
7use openvm_stark_backend::{
8    interaction::{BusIndex, LookupBus},
9    p3_field::PrimeField32,
10};
11use rustc_hash::FxHashMap;
12
13use super::{
14    air::Poseidon2PeripheryAir, PERIPHERY_POSEIDON2_CHUNK_SIZE, PERIPHERY_POSEIDON2_WIDTH,
15};
16use crate::arch::hasher::{Hasher, HasherChip};
17
18#[derive(Debug)]
19pub struct Poseidon2PeripheryBaseChip<F: PrimeField32, const SBOX_REGISTERS: usize> {
20    pub air: Arc<Poseidon2PeripheryAir<F, SBOX_REGISTERS>>,
21    pub subchip: Poseidon2SubChip<F, SBOX_REGISTERS>,
22    pub records: FxHashMap<[F; PERIPHERY_POSEIDON2_WIDTH], AtomicU32>,
23}
24
25impl<F: PrimeField32, const SBOX_REGISTERS: usize> Poseidon2PeripheryBaseChip<F, SBOX_REGISTERS> {
26    pub fn new(poseidon2_config: Poseidon2Config<F>, bus_idx: BusIndex) -> Self {
27        let subchip = Poseidon2SubChip::new(poseidon2_config.constants);
28        Self {
29            air: Arc::new(Poseidon2PeripheryAir::new(
30                subchip.air.clone(),
31                LookupBus::new(bus_idx),
32            )),
33            subchip,
34            records: FxHashMap::default(),
35        }
36    }
37}
38
39impl<F: PrimeField32, const SBOX_REGISTERS: usize> Hasher<PERIPHERY_POSEIDON2_CHUNK_SIZE, F>
40    for Poseidon2PeripheryBaseChip<F, SBOX_REGISTERS>
41{
42    fn compress(
43        &self,
44        lhs: &[F; PERIPHERY_POSEIDON2_CHUNK_SIZE],
45        rhs: &[F; PERIPHERY_POSEIDON2_CHUNK_SIZE],
46    ) -> [F; PERIPHERY_POSEIDON2_CHUNK_SIZE] {
47        let mut input_state = [F::ZERO; PERIPHERY_POSEIDON2_WIDTH];
48        input_state[..PERIPHERY_POSEIDON2_CHUNK_SIZE].copy_from_slice(lhs);
49        input_state[PERIPHERY_POSEIDON2_CHUNK_SIZE..].copy_from_slice(rhs);
50
51        let output = self.subchip.permute(input_state);
52        array::from_fn(|i| output[i])
53    }
54}
55
56impl<F: PrimeField32, const SBOX_REGISTERS: usize> HasherChip<PERIPHERY_POSEIDON2_CHUNK_SIZE, F>
57    for Poseidon2PeripheryBaseChip<F, SBOX_REGISTERS>
58{
59    /// Key method for Hasher trait.
60    ///
61    /// Takes two chunks, hashes them, and returns the result. Total width 3 * CHUNK, exposed in `direct_interaction_width()`.
62    ///
63    /// No interactions with other chips.
64    fn compress_and_record(
65        &mut self,
66        lhs: &[F; PERIPHERY_POSEIDON2_CHUNK_SIZE],
67        rhs: &[F; PERIPHERY_POSEIDON2_CHUNK_SIZE],
68    ) -> [F; PERIPHERY_POSEIDON2_CHUNK_SIZE] {
69        let mut input = [F::ZERO; PERIPHERY_POSEIDON2_WIDTH];
70        input[..PERIPHERY_POSEIDON2_CHUNK_SIZE].copy_from_slice(lhs);
71        input[PERIPHERY_POSEIDON2_CHUNK_SIZE..].copy_from_slice(rhs);
72
73        let count = self.records.entry(input).or_insert(AtomicU32::new(0));
74        count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
75
76        let output = self.subchip.permute(input);
77        array::from_fn(|i| output[i])
78    }
79}
openvm_circuit/system/poseidon2/chip.rs

openvm_circuit/system/poseidon2/
chip.rs
