sp1_hypercube/operations/poseidon2/
permutation.rs

1//! AIR columns for the Poseidon2 hash function.
2
3use std::{
4    borrow::BorrowMut,
5    mem::{size_of, transmute},
6};
7
8use sp1_derive::AlignedBorrow;
9
10use crate::{
11    operations::poseidon2::{NUM_EXTERNAL_ROUNDS, NUM_INTERNAL_ROUNDS, WIDTH},
12    util::indices_arr,
13};
14
15/// A column map for a Poseidon2 AIR with degree 3 constraints.
16pub const POSEIDON2_DEGREE3_COL_MAP: Poseidon2Degree3Cols<usize> = make_col_map_degree3();
17
18/// The number of columns in a Poseidon2 AIR with degree 3 constraints.
19pub const NUM_POSEIDON2_DEGREE3_COLS: usize = size_of::<Poseidon2Degree3Cols<u8>>();
20
21/// Create a column map for [`Poseidon2Degree3`].
22const fn make_col_map_degree3() -> Poseidon2Degree3Cols<usize> {
23    let indices_arr = indices_arr::<NUM_POSEIDON2_DEGREE3_COLS>();
24    unsafe {
25        transmute::<[usize; NUM_POSEIDON2_DEGREE3_COLS], Poseidon2Degree3Cols<usize>>(indices_arr)
26    }
27}
28
29/// A column layout for a poseidon2 permutation with degree 3 constraints.
30#[derive(AlignedBorrow, Clone, Copy)]
31#[repr(C)]
32pub struct Poseidon2Degree3Cols<T: Copy> {
33    /// The Poseidon2 hasher state.
34    pub state: Poseidon2StateCols<T>,
35}
36
37/// The number of internal rounds minus one.
38///
39/// Making this constant explicit is useful for generating bindings.
40pub const NUM_INTERNAL_ROUNDS_M1: usize = NUM_INTERNAL_ROUNDS - 1;
41
42/// A column layout for the intermediate states of a Poseidon2 AIR across all rounds.
43#[derive(AlignedBorrow, Clone, Copy)]
44#[repr(C)]
45pub struct Poseidon2StateCols<T> {
46    /// State for all external Poseidon2 rounds.
47    pub external_rounds_state: [[T; WIDTH]; NUM_EXTERNAL_ROUNDS],
48    /// State for the first Poseidon2 internal round.
49    pub internal_rounds_state: [T; WIDTH],
50    /// State for the first elements of the remaining internal Poseidon2 rounds.
51    pub internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS_M1],
52    /// The final state of the Poseidon2 permutation.
53    pub output_state: [T; WIDTH],
54}
55
56/// Trait that describes getter functions for the permutation columns.
57pub trait Poseidon2Cols<T: Copy> {
58    /// State for all external Poseidon2 rounds.
59    fn external_rounds_state(&self) -> &[[T; WIDTH]];
60    /// State for the first Poseidon2 internal round.
61    fn internal_rounds_state(&self) -> &[T; WIDTH];
62    /// State for the first elements of the remaining internal Poseidon2 rounds.
63    fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS_M1];
64    /// The final state of the Poseidon2 permutation.
65    fn perm_output(&self) -> &[T; WIDTH];
66
67    /// A mutable view into the columns, as if `self` was [`Poseidon2StateCols`].
68    #[allow(clippy::type_complexity)]
69    fn get_cols_mut(
70        &mut self,
71    ) -> (&mut [[T; WIDTH]], &mut [T; WIDTH], &mut [T; NUM_INTERNAL_ROUNDS_M1], &mut [T; WIDTH]);
72}
73
74impl<T: Copy> Poseidon2Cols<T> for Poseidon2Degree3Cols<T> {
75    fn external_rounds_state(&self) -> &[[T; WIDTH]] {
76        &self.state.external_rounds_state
77    }
78
79    fn internal_rounds_state(&self) -> &[T; WIDTH] {
80        &self.state.internal_rounds_state
81    }
82
83    fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS_M1] {
84        &self.state.internal_rounds_s0
85    }
86
87    fn perm_output(&self) -> &[T; WIDTH] {
88        &self.state.output_state
89    }
90
91    fn get_cols_mut(
92        &mut self,
93    ) -> (&mut [[T; WIDTH]], &mut [T; WIDTH], &mut [T; NUM_INTERNAL_ROUNDS_M1], &mut [T; WIDTH])
94    {
95        (
96            &mut self.state.external_rounds_state,
97            &mut self.state.internal_rounds_state,
98            &mut self.state.internal_rounds_s0,
99            &mut self.state.output_state,
100        )
101    }
102}
103
104/// Convert a row to a mutable [`Poseidon2Cols`] instance.
105pub fn permutation_mut<'a, 'b: 'a, T, const DEGREE: usize>(
106    row: &'b mut [T],
107) -> Box<&'b mut (dyn Poseidon2Cols<T> + 'a)>
108where
109    T: Copy,
110{
111    if DEGREE == 3 {
112        let start = POSEIDON2_DEGREE3_COL_MAP.state.external_rounds_state[0][0];
113        let end = start + size_of::<Poseidon2Degree3Cols<u8>>();
114        let convert: &mut Poseidon2Degree3Cols<T> = row[start..end].borrow_mut();
115        Box::new(convert)
116    } else {
117        panic!("Unsupported degree");
118    }
119}
sp1_hypercube/operations/poseidon2/permutation.rs

sp1_hypercube/operations/poseidon2/
permutation.rs
