openvm_circuit/system/poseidon2/
chip.rs

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