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