Skip to main content

darklake_poseidon_ark/
lib.rs

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        // Test vectors to ensure our implementation produces consistent, deterministic results
100        // These are regression tests based on our current working implementation
101        let poseidon = Poseidon::new();
102
103        // Original working test vectors from the existing test_hash function
104        // hash([1]) - confirmed working value
105        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        // hash([1, 2]) - confirmed working value  
111        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        // Test additional cases for comprehensive coverage
117        // hash([1, 2, 0, 0, 0]) - 5-element case
118        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        // hash([1, 2, 3, 4, 5, 6]) - 6-element case
130        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        // Larger input vector test
143        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        // Test that hash is deterministic
160        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        // Test that different inputs give different outputs
166        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        // Test avalanche effect - small change in input should drastically change output
173        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}