openvm_circuit/system/poseidon2/
chip.rs

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

openvm_circuit/system/poseidon2/
chip.rs
