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![56, 57, 56, 60, 60, 63, 64, 63],
46 }
47}
48
49pub struct Poseidon {
50 constants: Constants,
51}
52impl Poseidon {
53 pub fn new() -> Poseidon {
54 Poseidon {
55 constants: load_constants(),
56 }
57 }
58 pub fn ark(&self, state: &mut Vec<Fr>, c: &[Fr], it: usize) {
59 for i in 0..state.len() {
60 state[i].add_assign(&c[it + i]);
61 }
62 }
63
64 pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut Vec<Fr>, i: usize) {
65 if i < n_rounds_f / 2 || i >= n_rounds_f / 2 + n_rounds_p {
66 for j in 0..state.len() {
67 let aux = state[j];
68 state[j] = state[j].square();
69 state[j] = state[j].square();
70 state[j].mul_assign(&aux);
71 }
72 } else {
73 let aux = state[0];
74 state[0] = state[0].square();
75 state[0] = state[0].square();
76 state[0].mul_assign(&aux);
77 }
78 }
79
80 pub fn mix(&self, state: &Vec<Fr>, m: &[Vec<Fr>]) -> Vec<Fr> {
81 let mut new_state: Vec<Fr> = Vec::new();
82 for i in 0..state.len() {
83 new_state.push(Fr::zero());
84 for j in 0..state.len() {
85 let mut mij = m[i][j];
86 mij.mul_assign(&state[j]);
87 new_state[i].add_assign(&mij);
88 }
89 }
90 new_state.clone()
91 }
92
93 pub fn hash(&self, inp: Vec<Fr>) -> Result<Fr, String> {
94 let t = inp.len() + 1;
95 if inp.is_empty() || inp.len() >= self.constants.n_rounds_p.len() - 1 {
96 return Err("Wrong inputs length".to_string());
97 }
98 let n_rounds_f = self.constants.n_rounds_f.clone();
99 let n_rounds_p = self.constants.n_rounds_p[t - 2].clone();
100
101 let mut state = vec![Fr::zero(); t];
102 state[1..].clone_from_slice(&inp);
103
104 for i in 0..(n_rounds_f + n_rounds_p) {
105 self.ark(&mut state, &self.constants.c[t - 2], i * t);
106 self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
107 state = self.mix(&state, &self.constants.m[t - 2]);
108 }
109
110 Ok(state[0])
111 }
112}
113
114#[cfg(test)]
115mod tests {
116 use super::*;
117
118 #[test]
119 fn test_load_constants() {
120 let cons = load_constants();
121 assert_eq!(
122 cons.c[0][0].to_string(),
123 "4417881134626180770308697923359573201005643519861877412381846989312604493735"
124 );
125 assert_eq!(
126 cons.c[cons.c.len() - 1][0].to_string(),
127 "14715728137766105031387583973733149375806784983272780095398485311648630967927"
128 );
129 assert_eq!(
130 cons.m[0][0][0].to_string(),
131 "2910766817845651019878574839501801340070030115151021261302834310722729507541"
132 );
133 assert_eq!(
134 cons.m[cons.m.len() - 1][0][0].to_string(),
135 "708458300293891745856425423607721463509413916954480913172999113933455141974"
136 );
137 }
138
139 #[test]
140 fn test_hash() {
141 let b0: Fr = Fr::from_str("0").unwrap();
142 let b1: Fr = Fr::from_str("1").unwrap();
143 let b2: Fr = Fr::from_str("2").unwrap();
144 let b3: Fr = Fr::from_str("3").unwrap();
145 let b4: Fr = Fr::from_str("4").unwrap();
146 let b5: Fr = Fr::from_str("5").unwrap();
147 let b6: Fr = Fr::from_str("6").unwrap();
148
149 let mut big_arr: Vec<Fr> = Vec::new();
150 big_arr.push(b1.clone());
151 let poseidon = Poseidon::new();
152 let h = poseidon.hash(big_arr.clone()).unwrap();
153 assert_eq!(
154 h.to_string(),
155 "18586133768512220936620570745912940619677854269274689475585506675881198879027"
156 );
157
158 let mut big_arr: Vec<Fr> = Vec::new();
159 big_arr.push(b1.clone());
160 big_arr.push(b2.clone());
161 let poseidon = Poseidon::new();
162 let h = poseidon.hash(big_arr.clone()).unwrap();
163 assert_eq!(
164 h.to_string(),
165 "7853200120776062878684798364095072458815029376092732009249414926327459813530"
166 );
167
168 let mut big_arr: Vec<Fr> = Vec::new();
169 big_arr.push(b1.clone());
170 big_arr.push(b2.clone());
171 big_arr.push(b0.clone());
172 big_arr.push(b0.clone());
173 big_arr.push(b0.clone());
174 let poseidon = Poseidon::new();
175 let h = poseidon.hash(big_arr.clone()).unwrap();
176 assert_eq!(
177 h.to_string(),
178 "1018317224307729531995786483840663576608797660851238720571059489595066344487"
179 );
180
181 let mut big_arr: Vec<Fr> = Vec::new();
182 big_arr.push(b1.clone());
183 big_arr.push(b2.clone());
184 big_arr.push(b0.clone());
185 big_arr.push(b0.clone());
186 big_arr.push(b0.clone());
187 big_arr.push(b0.clone());
188 let poseidon = Poseidon::new();
189 let h = poseidon.hash(big_arr.clone()).unwrap();
190 assert_eq!(
191 h.to_string(),
192 "15336558801450556532856248569924170992202208561737609669134139141992924267169"
193 );
194
195 let mut big_arr: Vec<Fr> = Vec::new();
196 big_arr.push(b3.clone());
197 big_arr.push(b4.clone());
198 big_arr.push(b0.clone());
199 big_arr.push(b0.clone());
200 big_arr.push(b0.clone());
201 let poseidon = Poseidon::new();
202 let h = poseidon.hash(big_arr.clone()).unwrap();
203 assert_eq!(
204 h.to_string(),
205 "5811595552068139067952687508729883632420015185677766880877743348592482390548"
206 );
207
208 let mut big_arr: Vec<Fr> = Vec::new();
209 big_arr.push(b3.clone());
210 big_arr.push(b4.clone());
211 big_arr.push(b0.clone());
212 big_arr.push(b0.clone());
213 big_arr.push(b0.clone());
214 big_arr.push(b0.clone());
215 let h = poseidon.hash(big_arr.clone()).unwrap();
216 assert_eq!(
217 h.to_string(),
218 "12263118664590987767234828103155242843640892839966517009184493198782366909018"
219 );
220
221 let mut big_arr: Vec<Fr> = Vec::new();
222 big_arr.push(b1.clone());
223 big_arr.push(b2.clone());
224 big_arr.push(b3.clone());
225 big_arr.push(b4.clone());
226 big_arr.push(b5.clone());
227 big_arr.push(b6.clone());
228 let h = poseidon.hash(big_arr.clone()).unwrap();
229 assert_eq!(
230 h.to_string(),
231 "20400040500897583745843009878988256314335038853985262692600694741116813247201"
232 );
233 }
234}