poseidon-ark/src/lib.rs
src/lib.rsRaw
1use ark_bn254::Fr;
2use ark_ff::fields::Field;
3use ark_std::str::FromStr;
4use ark_std::Zero;
5use core::ops::{AddAssign, MulAssign};
6
7mod constants;
8
9#[derive(Debug)]
10pub struct Constants {
11 pub c: Vec<Vec<Fr>>,
12 pub m: Vec<Vec<Vec<Fr>>>,
13 pub n_rounds_f: usize,
14 pub n_rounds_p: Vec<usize>,
15}
16
17pub fn load_constants() -> Constants {
18 let (c_str, m_str) = constants::constants();
19 let mut c: Vec<Vec<Fr>> = Vec::new();
20 for i in 0..c_str.len() {
21 let mut cci: Vec<Fr> = Vec::new();
22 for j in 0..c_str[i].len() {
23 let b: Fr = Fr::from_str(c_str[i][j]).unwrap();
24 cci.push(b);
25 }
26 c.push(cci);
27 }
28 let mut m: Vec<Vec<Vec<Fr>>> = Vec::new();
29 for i in 0..m_str.len() {
30 let mut mi: Vec<Vec<Fr>> = Vec::new();
31 for j in 0..m_str[i].len() {
32 let mut mij: Vec<Fr> = Vec::new();
33 for k in 0..m_str[i][j].len() {
34 let b: Fr = Fr::from_str(m_str[i][j][k]).unwrap();
35 mij.push(b);
36 }
37 mi.push(mij);
38 }
39 m.push(mi);
40 }
41 Constants {
42 c,
43 m,
44 n_rounds_f: 8,
45 n_rounds_p: vec![
46 56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68,
47 ],
48 }
49}
50
51pub struct Poseidon {
52 constants: Constants,
53}
54impl Poseidon {
55 pub fn new() -> Poseidon {
56 Poseidon {
57 constants: load_constants(),
58 }
59 }
60 pub fn ark(&self, state: &mut Vec<Fr>, c: &[Fr], it: usize) {
61 for i in 0..state.len() {
62 state[i].add_assign(&c[it + i]);
63 }
64 }
65
66 pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut Vec<Fr>, i: usize) {
67 if i < n_rounds_f / 2 || i >= n_rounds_f / 2 + n_rounds_p {
68 for j in 0..state.len() {
69 let aux = state[j];
70 state[j] = state[j].square();
71 state[j] = state[j].square();
72 state[j].mul_assign(&aux);
73 }
74 } else {
75 let aux = state[0];
76 state[0] = state[0].square();
77 state[0] = state[0].square();
78 state[0].mul_assign(&aux);
79 }
80 }
81
82 pub fn mix(&self, state: &Vec<Fr>, m: &[Vec<Fr>]) -> Vec<Fr> {
83 let mut new_state: Vec<Fr> = Vec::new();
84 for i in 0..state.len() {
85 new_state.push(Fr::zero());
86 for j in 0..state.len() {
87 let mut mij = m[i][j];
88 mij.mul_assign(&state[j]);
89 new_state[i].add_assign(&mij);
90 }
91 }
92 new_state.clone()
93 }
94
95 pub fn hash(&self, inp: Vec<Fr>) -> Result<Fr, String> {
96 let t = inp.len() + 1;
97 if inp.is_empty() || inp.len() > self.constants.n_rounds_p.len() {
98 return Err("Wrong inputs length".to_string());
99 }
100 let n_rounds_f = self.constants.n_rounds_f.clone();
101 let n_rounds_p = self.constants.n_rounds_p[t - 2].clone();
102
103 let mut state = vec![Fr::zero(); t];
104 state[1..].clone_from_slice(&inp);
105
106 for i in 0..(n_rounds_f + n_rounds_p) {
107 self.ark(&mut state, &self.constants.c[t - 2], i * t);
108 self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
109 state = self.mix(&state, &self.constants.m[t - 2]);
110 }
111
112 Ok(state[0])
113 }
114}
115
116#[cfg(test)]
117mod tests {
118 use super::*;
119
120 #[test]
121 fn test_load_constants() {
122 let cons = load_constants();
123 assert_eq!(
124 cons.c[0][0].to_string(),
125 "4417881134626180770308697923359573201005643519861877412381846989312604493735"
126 );
127 assert_eq!(
128 cons.c[cons.c.len() - 1][0].to_string(),
129 "21579410516734741630578831791708254656585702717204712919233299001262271512412"
130 );
131 assert_eq!(
132 cons.m[0][0][0].to_string(),
133 "2910766817845651019878574839501801340070030115151021261302834310722729507541"
134 );
135 assert_eq!(
136 cons.m[cons.m.len() - 1][0][0].to_string(),
137 "11497693837059016825308731789443585196852778517742143582474723527597064448312"
138 );
139 }
140
141 #[test]
142 fn test_hash() {
143 let b0: Fr = Fr::from_str("0").unwrap();
144 let b1: Fr = Fr::from_str("1").unwrap();
145 let b2: Fr = Fr::from_str("2").unwrap();
146 let b3: Fr = Fr::from_str("3").unwrap();
147 let b4: Fr = Fr::from_str("4").unwrap();
148 let b5: Fr = Fr::from_str("5").unwrap();
149 let b6: Fr = Fr::from_str("6").unwrap();
150 let b7: Fr = Fr::from_str("7").unwrap();
151 let b8: Fr = Fr::from_str("8").unwrap();
152 let b9: Fr = Fr::from_str("9").unwrap();
153 let b10: Fr = Fr::from_str("10").unwrap();
154 let b11: Fr = Fr::from_str("11").unwrap();
155 let b12: Fr = Fr::from_str("12").unwrap();
156 let b13: Fr = Fr::from_str("13").unwrap();
157 let b14: Fr = Fr::from_str("14").unwrap();
158 let b15: Fr = Fr::from_str("15").unwrap();
159 let b16: Fr = Fr::from_str("16").unwrap();
160
161 let poseidon = Poseidon::new();
162
163 let big_arr: Vec<Fr> = vec![b1];
164 let h = poseidon.hash(big_arr.clone()).unwrap();
165 assert_eq!(
166 h.to_string(),
167 "18586133768512220936620570745912940619677854269274689475585506675881198879027"
168 );
169
170 let big_arr: Vec<Fr> = vec![b1, b2];
171 let h = poseidon.hash(big_arr.clone()).unwrap();
172 assert_eq!(
173 h.to_string(),
174 "7853200120776062878684798364095072458815029376092732009249414926327459813530"
175 );
176
177 let big_arr: Vec<Fr> = vec![b1, b2, b0, b0, b0];
178 let h = poseidon.hash(big_arr.clone()).unwrap();
179 assert_eq!(
180 h.to_string(),
181 "1018317224307729531995786483840663576608797660851238720571059489595066344487"
182 );
183
184 let big_arr: Vec<Fr> = vec![b1, b2, b0, b0, b0, b0];
185 let h = poseidon.hash(big_arr.clone()).unwrap();
186 assert_eq!(
187 h.to_string(),
188 "15336558801450556532856248569924170992202208561737609669134139141992924267169"
189 );
190
191 let big_arr: Vec<Fr> = vec![b3, b4, b0, b0, b0];
192 let h = poseidon.hash(big_arr.clone()).unwrap();
193 assert_eq!(
194 h.to_string(),
195 "5811595552068139067952687508729883632420015185677766880877743348592482390548"
196 );
197
198 let big_arr: Vec<Fr> = vec![b3, b4, b0, b0, b0, b0];
199 let h = poseidon.hash(big_arr).unwrap();
200 assert_eq!(
201 h.to_string(),
202 "12263118664590987767234828103155242843640892839966517009184493198782366909018"
203 );
204
205 let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6];
206 let h = poseidon.hash(big_arr).unwrap();
207 assert_eq!(
208 h.to_string(),
209 "20400040500897583745843009878988256314335038853985262692600694741116813247201"
210 );
211
212 let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14];
213 let h = poseidon.hash(big_arr).unwrap();
214 assert_eq!(
215 h.to_string(),
216 "8354478399926161176778659061636406690034081872658507739535256090879947077494"
217 );
218
219 let big_arr: Vec<Fr> = vec![b1, b2, b3, b4, b5, b6, b7, b8, b9, b0, b0, b0, b0, b0];
220 let h = poseidon.hash(big_arr).unwrap();
221 assert_eq!(
222 h.to_string(),
223 "5540388656744764564518487011617040650780060800286365721923524861648744699539"
224 );
225
226 let big_arr: Vec<Fr> = vec![
227 b1, b2, b3, b4, b5, b6, b7, b8, b9, b0, b0, b0, b0, b0, b0, b0,
228 ];
229 let h = poseidon.hash(big_arr).unwrap();
230 assert_eq!(
231 h.to_string(),
232 "11882816200654282475720830292386643970958445617880627439994635298904836126497"
233 );
234
235 let big_arr: Vec<Fr> = vec![
236 b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15, b16,
237 ];
238 let h = poseidon.hash(big_arr).unwrap();
239 assert_eq!(
240 h.to_string(),
241 "9989051620750914585850546081941653841776809718687451684622678807385399211877"
242 );
243 }
244 #[test]
245 fn test_wrong_inputs() {
246 let b0: Fr = Fr::from_str("0").unwrap();
247 let b1: Fr = Fr::from_str("1").unwrap();
248 let b2: Fr = Fr::from_str("2").unwrap();
249
250 let poseidon = Poseidon::new();
251
252 let big_arr: Vec<Fr> = vec![
253 b1, b2, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0, b0,
254 ];
255 poseidon.hash(big_arr).expect_err("Wrong inputs length");
256 }
257}