1#![no_std]
2
3use ark_bn254::Fr;
4use ark_ff::{fields::Field, Zero};
5use ark_std::{string::String, string::ToString, vec, vec::Vec};
6use core::ops::{AddAssign, MulAssign};
7
8mod static_constants;
9use static_constants::*;
10
11pub struct Poseidon;
12impl Poseidon {
13 pub fn new() -> Poseidon {
14 Poseidon
15 }
16 pub fn ark(&self, state: &mut Vec<Fr>, c: &[Fr], it: usize) {
17 for i in 0..state.len() {
18 state[i].add_assign(&c[it + i]);
19 }
20 }
21
22 pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut Vec<Fr>, i: usize) {
23 if i < n_rounds_f / 2 || i >= n_rounds_f / 2 + n_rounds_p {
24 for j in 0..state.len() {
25 let aux = state[j];
26 state[j] = state[j].square();
27 state[j] = state[j].square();
28 state[j].mul_assign(&aux);
29 }
30 } else {
31 let aux = state[0];
32 state[0] = state[0].square();
33 state[0] = state[0].square();
34 state[0].mul_assign(&aux);
35 }
36 }
37
38 pub fn mix(&self, state: &Vec<Fr>, m: &[&[Fr]]) -> Vec<Fr> {
39 let mut new_state: Vec<Fr> = Vec::new();
40 for i in 0..state.len() {
41 new_state.push(Fr::zero());
42 for j in 0..state.len() {
43 let mut mij = m[i][j];
44 mij.mul_assign(&state[j]);
45 new_state[i].add_assign(&mij);
46 }
47 }
48 new_state.clone()
49 }
50
51 pub fn hash(&self, inp: Vec<Fr>) -> Result<Fr, String> {
52 let t = inp.len() + 1;
53 if inp.is_empty() || inp.len() > N_ROUNDS_P.len() {
54 return Err("Wrong inputs length".to_string());
55 }
56 let n_rounds_f = N_ROUNDS_F;
57 let n_rounds_p = N_ROUNDS_P[t - 2];
58
59 let mut state = vec![Fr::zero(); t];
60 state[1..].clone_from_slice(&inp);
61
62 for i in 0..(n_rounds_f + n_rounds_p) {
63 self.ark(&mut state, C_CONSTANTS[t - 2], i * t);
64 self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
65 state = self.mix(&state, M_CONSTANTS[t - 2]);
66 }
67
68 Ok(state[0])
69 }
70}
71
72#[cfg(test)]
73mod tests {
74 use super::*;
75 use ark_std::str::FromStr;
76
77 #[test]
78 fn test_load_constants() {
79 assert_eq!(
80 C_CONSTANTS[0][0].to_string(),
81 "4417881134626180770308697923359573201005643519861877412381846989312604493735"
82 );
83 assert_eq!(
84 C_CONSTANTS[C_CONSTANTS.len() - 1][0].to_string(),
85 "21579410516734741630578831791708254656585702717204712919233299001262271512412"
86 );
87 assert_eq!(
88 M_CONSTANTS[0][0][0].to_string(),
89 "2910766817845651019878574839501801340070030115151021261302834310722729507541"
90 );
91 assert_eq!(
92 M_CONSTANTS[M_CONSTANTS.len() - 1][0][0].to_string(),
93 "11497693837059016825308731789443585196852778517742143582474723527597064448312"
94 );
95 }
96
97 #[test]
98 fn test_poseidon_reference_vectors() {
99 let poseidon = Poseidon::new();
102
103 let input = vec![Fr::from_str("1").unwrap()];
106 let result = poseidon.hash(input).unwrap();
107 let expected = Fr::from_str("18586133768512220936620570745912940619677854269274689475585506675881198879027").unwrap();
108 assert_eq!(result, expected, "hash([1]) should match known good value");
109
110 let input = vec![Fr::from_str("1").unwrap(), Fr::from_str("2").unwrap()];
112 let result = poseidon.hash(input).unwrap();
113 let expected = Fr::from_str("7853200120776062878684798364095072458815029376092732009249414926327459813530").unwrap();
114 assert_eq!(result, expected, "hash([1, 2]) should match known good value");
115
116 let input = vec![
119 Fr::from_str("1").unwrap(),
120 Fr::from_str("2").unwrap(),
121 Fr::from_str("0").unwrap(),
122 Fr::from_str("0").unwrap(),
123 Fr::from_str("0").unwrap()
124 ];
125 let result = poseidon.hash(input).unwrap();
126 let expected = Fr::from_str("1018317224307729531995786483840663576608797660851238720571059489595066344487").unwrap();
127 assert_eq!(result, expected, "hash([1, 2, 0, 0, 0]) should match known good value");
128
129 let input = vec![
131 Fr::from_str("1").unwrap(),
132 Fr::from_str("2").unwrap(),
133 Fr::from_str("3").unwrap(),
134 Fr::from_str("4").unwrap(),
135 Fr::from_str("5").unwrap(),
136 Fr::from_str("6").unwrap()
137 ];
138 let result = poseidon.hash(input).unwrap();
139 let expected = Fr::from_str("20400040500897583745843009878988256314335038853985262692600694741116813247201").unwrap();
140 assert_eq!(result, expected, "hash([1, 2, 3, 4, 5, 6]) should match known good value");
141
142 let input = vec![
144 Fr::from_str("1").unwrap(), Fr::from_str("2").unwrap(), Fr::from_str("3").unwrap(),
145 Fr::from_str("4").unwrap(), Fr::from_str("5").unwrap(), Fr::from_str("6").unwrap(),
146 Fr::from_str("7").unwrap(), Fr::from_str("8").unwrap(), Fr::from_str("9").unwrap(),
147 Fr::from_str("10").unwrap(), Fr::from_str("11").unwrap(), Fr::from_str("12").unwrap(),
148 Fr::from_str("13").unwrap(), Fr::from_str("14").unwrap()
149 ];
150 let result = poseidon.hash(input).unwrap();
151 let expected = Fr::from_str("8354478399926161176778659061636406690034081872658507739535256090879947077494").unwrap();
152 assert_eq!(result, expected, "Large input hash should match known good value");
153 }
154
155 #[test]
156 fn test_poseidon_properties() {
157 let poseidon = Poseidon::new();
158
159 let input = vec![Fr::from_str("12345").unwrap(), Fr::from_str("67890").unwrap()];
161 let result1 = poseidon.hash(input.clone()).unwrap();
162 let result2 = poseidon.hash(input.clone()).unwrap();
163 assert_eq!(result1, result2, "Hash should be deterministic");
164
165 let input1 = vec![Fr::from_str("1").unwrap(), Fr::from_str("2").unwrap()];
167 let input2 = vec![Fr::from_str("2").unwrap(), Fr::from_str("1").unwrap()];
168 let result1 = poseidon.hash(input1).unwrap();
169 let result2 = poseidon.hash(input2).unwrap();
170 assert_ne!(result1, result2, "Different inputs should give different outputs");
171
172 let input1 = vec![Fr::from_str("1").unwrap()];
174 let input2 = vec![Fr::from_str("2").unwrap()];
175 let result1 = poseidon.hash(input1).unwrap();
176 let result2 = poseidon.hash(input2).unwrap();
177 assert_ne!(result1, result2, "Small input change should change output significantly");
178 }
179
180 #[test]
181 fn test_hash() {
182 let b0: Fr = Fr::from_str("0").unwrap();
183 let b1: Fr = Fr::from_str("1").unwrap();
184 let b2: Fr = Fr::from_str("2").unwrap();
185 let b3: Fr = Fr::from_str("3").unwrap();
186 let b4: Fr = Fr::from_str("4").unwrap();
187 let b5: Fr = Fr::from_str("5").unwrap();
188 let b6: Fr = Fr::from_str("6").unwrap();
189 let b7: Fr = Fr::from_str("7").unwrap();
190 let b8: Fr = Fr::from_str("8").unwrap();
191 let b9: Fr = Fr::from_str("9").unwrap();
192 let b10: Fr = Fr::from_str("10").unwrap();
193 let b11: Fr = Fr::from_str("11").unwrap();
194 let b12: Fr = Fr::from_str("12").unwrap();
195 let b13: Fr = Fr::from_str("13").unwrap();
196 let b14: Fr = Fr::from_str("14").unwrap();
197 let b15: Fr = Fr::from_str("15").unwrap();
198 let b16: Fr = Fr::from_str("16").unwrap();
199
200 let poseidon = Poseidon::new();
201
202 let big_arr: Vec<Fr> = vec![b1];
203 let h = poseidon.hash(big_arr.clone()).unwrap();
204 assert_eq!(
205 h.to_string(),
206 "18586133768512220936620570745912940619677854269274689475585506675881198879027"
207 );
208
209 let big_arr: Vec<Fr> = vec![b1, b2];
210 let h = poseidon.hash(big_arr.clone()).unwrap();
211 assert_eq!(
212 h.to_string(),
213 "7853200120776062878684798364095072458815029376092732009249414926327459813530"
214 );
215
216 let big_arr: Vec<Fr> = vec![b1, b2, b0, b0, b0];
217 let h = poseidon.hash(big_arr.clone()).unwrap();
218 assert_eq!(
219 h.to_string(),
220 "1018317224307729531995786483840663576608797660851238720571059489595066344487"
221 );
222
223 let big_arr: Vec<Fr> = vec![b1, b2, b0, b0, b0, b0];
224 let h = poseidon.hash(big_arr.clone()).unwrap();
225 assert_eq!(
226 h.to_string(),
227 "15336558801450556532856248569924170992202208561737609669134139141992924267169"
228 );
229
230 let big_arr: Vec<Fr> = vec![b3, b4, b0, b0, b0];
231 let h = poseidon.hash(big_arr.clone()).unwrap();
232 assert_eq!(
233 h.to_string(),
234 "5811595552068139067952687508729883632420015185677766880877743348592482390548"
235 );
236
237 let big_arr: Vec<Fr> = vec![b3, b4, b0, b0, b0, b0];
238 let h = poseidon.hash(big_arr).unwrap();
239 assert_eq!(
240 h.to_string(),
241 "12263118664590987767234828103155242843640892839966517009184493198782366909018"
242 );
243
244 let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6];
245 let h = poseidon.hash(big_arr).unwrap();
246 assert_eq!(
247 h.to_string(),
248 "20400040500897583745843009878988256314335038853985262692600694741116813247201"
249 );
250
251 let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14];
252 let h = poseidon.hash(big_arr).unwrap();
253 assert_eq!(
254 h.to_string(),
255 "8354478399926161176778659061636406690034081872658507739535256090879947077494"
256 );
257
258 let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6, b7, b8, b9, b0, b0, b0, b0, b0];
259 let h = poseidon.hash(big_arr).unwrap();
260 assert_eq!(
261 h.to_string(),
262 "5540388656744764564518487011617040650780060800286365721923524861648744699539"
263 );
264
265 let big_arr: Vec<Fr> = vec![
266 b1, b2, b3, b4, b5, b6, b7, b8, b9, b0, b0, b0, b0, b0, b0, b0,
267 ];
268 let h = poseidon.hash(big_arr).unwrap();
269 assert_eq!(
270 h.to_string(),
271 "11882816200654282475720830292386643970958445617880627439994635298904836126497"
272 );
273
274 let big_arr: Vec<Fr> = vec![
275 b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15, b16,
276 ];
277 let h = poseidon.hash(big_arr).unwrap();
278 assert_eq!(
279 h.to_string(),
280 "9989051620750914585850546081941653841776809718687451684622678807385399211877"
281 );
282 }
283 #[test]
284 fn test_wrong_inputs() {
285 let b0: Fr = Fr::from_str("0").unwrap();
286 let b1: Fr = Fr::from_str("1").unwrap();
287 let b2: Fr = Fr::from_str("2").unwrap();
288
289 let poseidon = Poseidon::new();
290
291 let big_arr: Vec<Fr> = vec![
292 b1, b2, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0,
293 ];
294 poseidon.hash(big_arr).expect_err("Wrong inputs length");
295 }
296}