1
2use fawkes_crypto::{BorshSerialize, ff_uint::PrimeField};
3
4use crate::{constants,
5 fawkes_crypto::{
6 ff_uint::Num,
7 native::poseidon::{poseidon, MerkleProof},
8 rand::{self, Rng},
9 core::sizedvec::SizedVec
10 },
11 native::{
12 account::Account,
13 boundednum::BoundedNum,
14 note::Note,
15 params::{PoolParams},
16 tx::{make_delta, Tx, TransferPub, TransferSec, nullifier, tx_hash, tx_sign, out_commitment_hash},
17 key::{derive_key_a, derive_key_eta, derive_key_p_d},
18 tree::{TreePub, TreeSec},
19 delegated_deposit::{DelegatedDepositBatchPub, DelegatedDepositBatchSec, DelegatedDeposit}
20 }
21};
22
23use fawkes_crypto_keccak256::native::hash::keccak256;
24
25
26pub const N_ITEMS:usize = 1000;
27
28pub struct HashTreeState<P:PoolParams> {
29 pub hashes:Vec<Vec<Num<P::Fr>>>,
30 pub default_hashes: Vec<Num<P::Fr>>
31}
32
33impl<P:PoolParams> HashTreeState<P> {
34 pub fn new(params:&P) -> Self {
35 let default_hashes = {
36 std::iter::successors(Some(Num::ZERO), |t|
37 Some(poseidon([*t,*t].as_ref(), params.compress()))
38 ).skip(constants::OUTPLUSONELOG).take(constants::HEIGHT - constants::OUTPLUSONELOG+1).collect()
39 };
40
41 let hashes = (0..constants::HEIGHT - constants::OUTPLUSONELOG+1).map(|_| vec![]).collect();
42
43 Self {hashes, default_hashes}
44 }
45
46 pub fn push(&mut self, n:Num<P::Fr>, params:&P) {
47 let mut p = self.hashes[0].len();
48 self.hashes[0].push(n);
49
50 for i in 0..constants::HEIGHT - constants::OUTPLUSONELOG {
51 p >>= 1;
52 if self.hashes[i+1].len() <= p {
53 self.hashes[i+1].push(self.default_hashes[i+1]);
54 }
55 let left = self.cell(i, 2*p);
56 let right = self.cell(i, 2*p+1);
57 self.hashes[i+1][p] = poseidon([left, right].as_ref(), params.compress());
58 }
59 }
60
61 pub fn cell(&self, i:usize, j:usize) -> Num<P::Fr> {
62 if self.hashes[i].len() <= j {
63 self.default_hashes[i]
64 } else {
65 self.hashes[i][j]
66 }
67 }
68
69 pub fn merkle_proof(&self, id:usize) -> MerkleProof<P::Fr, { constants::HEIGHT - constants::OUTPLUSONELOG }> {
70 let sibling = (0..constants::HEIGHT - constants::OUTPLUSONELOG).map(|i| self.cell(i, (id>>i)^1)).collect();
71 let path = (0..constants::HEIGHT - constants::OUTPLUSONELOG).map(|i| (id>>i)&1==1).collect();
72 MerkleProof {sibling, path}
73 }
74
75 pub fn root(&self) -> Num<P::Fr> {
76 return self.cell(constants::HEIGHT - constants::OUTPLUSONELOG, 0)
77 }
78}
79
80pub struct State<P:PoolParams> {
81 pub hashes:Vec<Vec<Num<P::Fr>>>,
82 pub items:Vec<(Account<P::Fr>, Note<P::Fr>)>,
83 pub default_hashes:Vec<Num<P::Fr>>,
84 pub sigma:Num<P::Fs>,
85 pub account_id:usize,
86 pub note_id:Vec<usize>
87}
88
89impl<P:PoolParams> State<P> {
90 pub fn random_sample_state<R:Rng>(rng:&mut R, params:&P) -> Self {
91 let sigma = rng.gen();
92 let a = derive_key_a(sigma, params);
93 let eta = derive_key_eta(a.x, params);
94
95
96 let account_id = rng.gen_range(0, N_ITEMS);
97 let note_id = rand::seq::index::sample(rng, N_ITEMS, constants::IN).into_vec();
98
99
100 let mut items:Vec<(Account<_>, Note<_>)> = (0..N_ITEMS).map(|_| (Account::sample(rng, params), Note::sample(rng, params) )).collect();
101
102 for i in note_id.iter().cloned() {
103 items[i].1.p_d = derive_key_p_d(items[i].1.d.to_num(), eta, params).x;
104 }
105
106 items[account_id].0.p_d = derive_key_p_d(items[account_id].0.d.to_num(), eta, params).x;
107 items[account_id].0.i = BoundedNum::ZERO;
108
109 let mut default_hashes = vec![Num::ZERO;constants::HEIGHT+1];
110 let mut hashes = vec![];
111
112 for i in 0..constants::HEIGHT {
113 let t = default_hashes[i];
114 default_hashes[i+1] = poseidon([t,t].as_ref(), params.compress());
115 }
116
117 {
118 let mut t = vec![];
119 for j in 0..N_ITEMS {
120 let (a, n) = items[j].clone();
121 t.push(a.hash(params));
122 t.push(n.hash(params));
123 }
124 if t.len() & 1 == 1 {
125 t.push(default_hashes[0]);
126 }
127 hashes.push(t);
128 }
129
130 for i in 0..constants::HEIGHT {
131 let mut t = vec![];
132 for j in 0..hashes[i].len()>>1 {
133 t.push(poseidon([hashes[i][2*j],hashes[i][2*j+1]].as_ref(), params.compress()));
134 }
135 if t.len() & 1 == 1 {
136 t.push(default_hashes[i+1]);
137 }
138 hashes.push(t);
139 }
140
141 Self {
142 hashes,
143 items,
144 default_hashes,
145 sigma,
146 account_id,
147 note_id
148 }
149 }
150
151
152 pub fn random_sample_transfer<R:Rng>(&self, rng:&mut R, params:&P) -> (TransferPub<P::Fr>, TransferSec<P::Fr>) {
153
154 let zero_note = Note {
155 d: BoundedNum::ZERO,
156 p_d: Num::ZERO,
157 b: BoundedNum::ZERO,
158 t: BoundedNum::ZERO,
159 };
160
161 let root = self.root();
162 let index = N_ITEMS*2;
163 let a = derive_key_a(self.sigma, params);
164 let eta = derive_key_eta(a.x, params);
165 let nullifier = nullifier(self.hashes[0][self.account_id*2], eta, Num::from(self.account_id as u32 * 2), params);
166 let memo:Num<P::Fr> = rng.gen();
167
168
169 let mut input_value = self.items[self.account_id].0.b.to_num();
170 for &i in self.note_id.iter() {
171 input_value+=self.items[i].1.b.to_num();
172 }
173
174 let mut input_energy = self.items[self.account_id].0.e.to_num();
175 input_energy += self.items[self.account_id].0.b.to_num()*(Num::from((index-self.account_id*2) as u32)) ;
176
177
178 for &i in self.note_id.iter() {
179 input_energy+=self.items[i].1.b.to_num()*Num::from((index-(2*i+1)) as u32);
180 }
181
182 let mut out_account: Account<P::Fr> = Account::sample(rng, params);
183 out_account.b = BoundedNum::new(input_value);
184 out_account.e = BoundedNum::new(input_energy);
185 out_account.i = BoundedNum::new(Num::from(index as u32));
186 out_account.p_d = derive_key_p_d(out_account.d.to_num(), eta, params).x;
187
188
189 let mut out_note: Note<P::Fr> = Note::sample(rng, params);
190 out_note.b = BoundedNum::ZERO;
191
192 let mut input_hashes = vec![self.items[self.account_id].0.hash(params)];
193 for &i in self.note_id.iter() {
194 input_hashes.push(self.items[i].1.hash(params));
195 }
196
197 let out_notes:Vec<_> = std::iter::once(out_note).chain(core::iter::repeat(zero_note).take(constants::OUT-1)).collect();
198 let out_hashes:Vec<_> = std::iter::once(out_account.hash(params)).chain(out_notes.iter().map(|n| n.hash(params))).collect();
199 let out_commit = out_commitment_hash(&out_hashes, params);
200 let tx_hash = tx_hash(&input_hashes, out_commit, params);
201 let (eddsa_s,eddsa_r) = tx_sign(self.sigma, tx_hash, params);
202
203
204 let delta = make_delta::<P::Fr>(Num::ZERO, Num::ZERO, Num::from(index as u32), Num::ZERO);
205
206 let p = TransferPub::<P::Fr> {
207 root,
208 nullifier,
209 out_commit,
210 delta,
211 memo,
212 };
213
214
215
216
217
218 let tx = Tx {
219 input: (self.items[self.account_id].0.clone(), self.note_id.iter().map(|&i| self.items[i].1.clone()).collect()),
220 output: (out_account, out_notes.iter().cloned().collect() )
221 };
222
223
224
225 let s = TransferSec::<P::Fr> {
226 tx,
227 in_proof: (
228 self.merkle_proof(self.account_id*2),
229 self.note_id.iter().map(|&i| self.merkle_proof(i*2+1) ).collect()
230 ),
231 eddsa_s:eddsa_s.to_other().unwrap(),
232 eddsa_r,
233 eddsa_a:a.x
234 };
235
236 (p, s)
237 }
238
239 fn cell(&self, i:usize, j:usize) -> Num<P::Fr> {
240 if self.hashes[i].len() <= j {
241 self.default_hashes[i]
242 } else {
243 self.hashes[i][j]
244 }
245 }
246
247 fn merkle_proof(&self, id:usize) -> MerkleProof<P::Fr, { constants::HEIGHT }> {
248 let sibling = (0..constants::HEIGHT).map(|i| self.cell(i, (id>>i)^1)).collect();
249 let path = (0..constants::HEIGHT).map(|i| (id>>i)&1==1).collect();
250 MerkleProof {sibling, path}
251 }
252
253 fn root(&self) -> Num<P::Fr> {
254 return self.cell(constants::HEIGHT, 0)
255 }
256
257}
258
259
260pub fn random_sample_tree_update<P:PoolParams,R:Rng>(rng:&mut R, params:&P) -> (TreePub<P::Fr>, TreeSec<P::Fr>) {
261 use std::collections::HashMap;
262
263 let index_filled:usize = rng.gen_range(0, N_ITEMS);
264 let index_free = index_filled + 1;
265
266 const PATH_LENGTH:usize = constants::HEIGHT-constants::OUTPLUSONELOG;
267
268 let mut cell = HashMap::new();
269
270 let zero_leaf_value = {
271 let mut c = Num::ZERO;
272 for _ in 0..constants::OUTPLUSONELOG {
273 c = poseidon(&[c, c], params.compress());
274 }
275 c
276 };
277
278 let cell_defaults = {
279 let mut c = zero_leaf_value;
280 let mut res = vec![c;PATH_LENGTH+1];
281 for i in 1..PATH_LENGTH {
282 c = poseidon(&[c,c], params.compress());
283 res[i] = c;
284 }
285 res
286 };
287
288 macro_rules! cell_get {
289 ($h:expr, $i:expr) => { cell.get(&(($h),($i))).unwrap_or_else(||&cell_defaults[($h)]).clone() }
290 }
291
292 macro_rules! cell_set {
293 ($h:expr, $i:expr, $v:expr) => { cell.insert((($h),($i)), ($v)); }
294 }
295
296
297
298 let prev_leaf:Num<P::Fr> = rng.gen();
299 cell_set!(0, index_filled, prev_leaf);
300 for h in 0..PATH_LENGTH {
301 let index_level = index_filled>>h;
302 if index_level & 1 == 1 {
303 cell_set!(h, index_level^1, rng.gen());
304 }
305 }
306
307 for h in 1..PATH_LENGTH+1 {
308 let index = index_filled>>h;
309 let left = cell_get!(h-1, index*2);
310 let right = cell_get!(h-1, index*2+1);
311 let hash = poseidon(&[left,right], params.compress());
312 cell_set!(h, index, hash);
313 }
314
315
316
317
318 let path_filled = (0..PATH_LENGTH).map(|i| (index_filled>>i)&1==1).collect();
319 let sibling_filled:SizedVec<Num<P::Fr>, PATH_LENGTH> = (0..PATH_LENGTH).map(|h| cell_get!(h, (index_filled>>h)^1 )).collect();
320
321
322 let proof_filled = MerkleProof {
323 sibling: sibling_filled,
324 path: path_filled
325 };
326
327 let root_before = cell_get!(PATH_LENGTH, 0);
328
329 let path_free = (0..PATH_LENGTH).map(|i| (index_free>>i)&1==1).collect();
330 let sibling_free:SizedVec<Num<P::Fr>, PATH_LENGTH> = (0..PATH_LENGTH).map(|h| cell_get!(h, (index_free>>h)^1 )).collect();
331
332 let leaf = rng.gen();
333 cell_set!(0, index_free, leaf);
334
335 for h in 1..PATH_LENGTH+1 {
336 let index = index_free>>h;
337 let left = cell_get!(h-1, index*2);
338 let right = cell_get!(h-1, index*2+1);
339 let hash = poseidon(&[left,right], params.compress());
340 cell_set!(h, index, hash);
341 }
342
343 let root_after = cell_get!(PATH_LENGTH, 0);
344
345 let proof_free = MerkleProof {
346 sibling: sibling_free,
347 path: path_free
348 };
349
350 let p = TreePub {
351 root_before,
352 root_after,
353 leaf
354 };
355
356 let s = TreeSec {
357 proof_filled,
358 proof_free,
359 prev_leaf
360 };
361
362 (p,s)
363
364}
365
366pub fn serialize_scalars_and_delegated_deposits_be<Fr:PrimeField>(scalars:&[Num<Fr>], deposits:&[DelegatedDeposit<Fr>]) -> Vec<u8> {
367 deposits.iter().rev().flat_map(|d| {
368 let mut res = d.b.try_to_vec().unwrap();
369 res.extend(d.p_d.try_to_vec().unwrap());
370 res.extend(d.d.try_to_vec().unwrap());
371 res
372
373 })
374 .chain(scalars.iter().rev().flat_map(|s| s.try_to_vec().unwrap()))
375 .rev().collect::<Vec<_>>()
376}
377
378
379pub fn random_sample_delegated_deposit<P:PoolParams,R:Rng>(rng:&mut R, params:&P) -> (DelegatedDepositBatchPub<P::Fr>, DelegatedDepositBatchSec<P::Fr>) {
380
381 let deposits:SizedVec<_,{constants::DELEGATED_DEPOSITS_NUM}> = (0..constants::DELEGATED_DEPOSITS_NUM).map(|_| {
382 let n = Note::sample(rng, params);
383 DelegatedDeposit {
384 d:n.d,
385 p_d:n.p_d,
386 b:n.b,
387 }
388 }).collect();
389
390 let zero_note_hash = Note {
391 d:BoundedNum::ZERO,
392 p_d:Num::ZERO,
393 b:BoundedNum::ZERO,
394 t:BoundedNum::ZERO
395 }.hash(params);
396
397 let zero_account_hash = Account {
398 d: BoundedNum::ZERO,
399 p_d: Num::ZERO,
400 i: BoundedNum::ZERO,
401 b: BoundedNum::ZERO,
402 e: BoundedNum::ZERO,
403 }.hash(params);
404
405 let out_hash = std::iter::once(zero_account_hash)
406 .chain(deposits.iter().map(|d| d.to_note().hash(params)))
407 .chain(std::iter::repeat(zero_note_hash)).take(constants::OUT+1).collect::<Vec<_>>();
408
409 let _out_commitment_hash = out_commitment_hash(&out_hash, params);
410
411
412
413 let data = serialize_scalars_and_delegated_deposits_be(
414 &[_out_commitment_hash], deposits.as_slice());
415
416
417
418 let keccak_sum = {
419 let t = keccak256(&data);
420 let mut res = Num::ZERO;
421 for limb in t.iter() {
422 res = res * Num::from(256) + Num::from(*limb);
423 }
424 res
425 };
426
427 let p = DelegatedDepositBatchPub {keccak_sum};
428
429 let s = DelegatedDepositBatchSec {
430 deposits
431 };
432 (p,s)
433
434}