poseidon_permutation/
r1cs.rs

1#![allow(non_snake_case)]
2use ark_std::vec::Vec;
3
4use ark_r1cs_std::{fields::fp::FpVar, prelude::*};
5use ark_relations::r1cs::ConstraintSystemRef;
6use decaf377::Fq;
7use poseidon_parameters::v1::{Alpha, MatrixOperations, PoseidonParameters};
8
9/// Represents a Poseidon permutation instance.
10pub struct InstanceVar<
11    const STATE_SIZE: usize,
12    const STATE_SIZE_MINUS_1: usize,
13    const NUM_MDS_ELEMENTS: usize,
14    const NUM_STATE_SIZE_MINUS_1_ELEMENTS: usize,
15    const NUM_ROUND_ROWS: usize,
16    const NUM_ROUND_COLS: usize,
17    const NUM_ROUND_ELEMENTS: usize,
18    const NUM_PARTIAL_ROUNDS: usize,
19> {
20    /// Parameters for this instance of Poseidon.
21    pub parameters: PoseidonParameters<
22        STATE_SIZE,
23        STATE_SIZE_MINUS_1,
24        NUM_MDS_ELEMENTS,
25        NUM_STATE_SIZE_MINUS_1_ELEMENTS,
26        NUM_ROUND_ROWS,
27        NUM_ROUND_COLS,
28        NUM_ROUND_ELEMENTS,
29        NUM_PARTIAL_ROUNDS,
30    >,
31
32    /// Constraint system
33    pub cs: ConstraintSystemRef<Fq>,
34
35    /// Current state
36    pub state_words: Vec<FpVar<Fq>>,
37}
38
39impl<
40        const STATE_SIZE: usize,
41        const STATE_SIZE_MINUS_1: usize,
42        const NUM_MDS_ELEMENTS: usize,
43        const NUM_STATE_SIZE_MINUS_1_ELEMENTS: usize,
44        const NUM_ROUND_ROWS: usize,
45        const NUM_ROUND_COLS: usize,
46        const NUM_ROUND_ELEMENTS: usize,
47        const NUM_PARTIAL_ROUNDS: usize,
48    >
49    InstanceVar<
50        STATE_SIZE,
51        STATE_SIZE_MINUS_1,
52        NUM_MDS_ELEMENTS,
53        NUM_STATE_SIZE_MINUS_1_ELEMENTS,
54        NUM_ROUND_ROWS,
55        NUM_ROUND_COLS,
56        NUM_ROUND_ELEMENTS,
57        NUM_PARTIAL_ROUNDS,
58    >
59{
60    /// Fixed width hash from n:1. Outputs a Fq given `t` input words.
61    pub fn n_to_1_fixed_hash(
62        parameters: PoseidonParameters<
63            STATE_SIZE,
64            STATE_SIZE_MINUS_1,
65            NUM_MDS_ELEMENTS,
66            NUM_STATE_SIZE_MINUS_1_ELEMENTS,
67            NUM_ROUND_ROWS,
68            NUM_ROUND_COLS,
69            NUM_ROUND_ELEMENTS,
70            NUM_PARTIAL_ROUNDS,
71        >,
72        cs: ConstraintSystemRef<Fq>,
73        input_words: [FpVar<Fq>; STATE_SIZE],
74    ) -> FpVar<Fq> {
75        // t = rate + capacity
76
77        let mut instance = InstanceVar {
78            parameters,
79            cs,
80            state_words: input_words.to_vec(),
81        };
82
83        // Apply Poseidon permutation.
84        instance.permute();
85
86        // Emit a single element since this is a n:1 hash.
87        instance.state_words[1].clone()
88    }
89
90    /// Poseidon permutation.
91    pub fn permute(&mut self) {
92        let R_f = self.parameters.rounds.full() / 2;
93        let R_P = self.parameters.rounds.partial();
94        let mut round_constants_counter = 0;
95        let round_constants: [Fq; NUM_ROUND_ELEMENTS] = self.parameters.arc.inner_elements();
96
97        // First full rounds
98        for _ in 0..R_f {
99            // Apply `AddRoundConstants` layer
100            for i in 0..STATE_SIZE {
101                self.state_words[i] += round_constants[round_constants_counter];
102                round_constants_counter += 1;
103            }
104            self.full_sub_words();
105            self.mix_layer_mds();
106        }
107
108        // Partial rounds
109        for _ in 0..R_P {
110            // Apply `AddRoundConstants` layer
111            for i in 0..STATE_SIZE {
112                self.state_words[i] += round_constants[round_constants_counter];
113                round_constants_counter += 1;
114            }
115            self.partial_sub_words();
116            self.mix_layer_mds();
117        }
118
119        // Final full rounds
120        for _ in 0..R_f {
121            // Apply `AddRoundConstants` layer
122            for i in 0..STATE_SIZE {
123                self.state_words[i] += round_constants[round_constants_counter];
124                round_constants_counter += 1;
125            }
126            self.full_sub_words();
127            self.mix_layer_mds();
128        }
129    }
130
131    /// Applies the partial `SubWords` layer.
132    fn partial_sub_words(&mut self) {
133        match self.parameters.alpha {
134            Alpha::Exponent(exp) => {
135                self.state_words[0] = (self.state_words[0])
136                    .pow_by_constant([exp as u64])
137                    .expect("can compute pow")
138            }
139            Alpha::Inverse => unimplemented!("err: inverse alpha not implemented"),
140        }
141    }
142
143    /// Applies the full `SubWords` layer.
144    fn full_sub_words(&mut self) {
145        match self.parameters.alpha {
146            Alpha::Exponent(exp) => {
147                for i in 0..STATE_SIZE {
148                    self.state_words[i] = (self.state_words[i])
149                        .pow_by_constant([exp as u64])
150                        .expect("can compute pow");
151                }
152            }
153            Alpha::Inverse => {
154                unimplemented!("err: inverse alpha not implemented")
155            }
156        }
157    }
158
159    /// Applies the `MixLayer` using the MDS matrix.
160    fn mix_layer_mds(&mut self) {
161        self.state_words = self
162            .parameters
163            .mds
164            .0
165             .0
166            .iter_rows()
167            .map(|row| {
168                let temp_vec: Vec<FpVar<Fq>> = row
169                    .iter()
170                    .zip(&self.state_words)
171                    .map(|(x, y)| {
172                        FpVar::<Fq>::new_constant(self.cs.clone(), x).expect("can create constant")
173                            * y
174                    })
175                    .collect();
176                let result = temp_vec.iter().sum();
177                result
178            })
179            .collect();
180    }
181}