privacy_core/commitment_tree/
poseidon.rs1use super::poseidon_primitives::{generate_constants, ConstantLength, Hash, Mds, Spec};
10use ff::Field;
11use halo2curves::bn256::Fr;
12
13pub const MERKLE_DEPTH_EVM: usize = 32;
15
16#[derive(Clone, Copy, Debug)]
18pub struct Bn254PoseidonMerkleSpec;
19
20impl Spec<Fr, 3, 2> for Bn254PoseidonMerkleSpec {
21 fn full_rounds() -> usize {
22 8
23 }
24 fn partial_rounds() -> usize {
25 56
26 }
27 fn sbox(val: Fr) -> Fr {
28 val.pow_vartime([5])
29 }
30 fn secure_mds() -> usize {
31 0
32 }
33 fn constants() -> (Vec<[Fr; 3]>, Mds<Fr, 3>, Mds<Fr, 3>) {
34 generate_constants::<Fr, Self, 3, 2>()
35 }
36}
37
38#[inline]
40pub fn merkle_compress(level: u8, left: Fr, right: Fr) -> Fr {
41 Hash::<Fr, Bn254PoseidonMerkleSpec, ConstantLength<3>, 3, 2>::init().hash([
42 Fr::from(level as u64),
43 left,
44 right,
45 ])
46}
47
48#[inline]
53pub fn poseidon_domain_pair(domain: u64, a: Fr, b: Fr) -> Fr {
54 Hash::<Fr, Bn254PoseidonMerkleSpec, ConstantLength<3>, 3, 2>::init().hash([
55 Fr::from(domain),
56 a,
57 b,
58 ])
59}
60
61pub fn merkle_root(position: u32, leaf: Fr, siblings: &[Fr; MERKLE_DEPTH_EVM]) -> Fr {
63 let mut node = leaf;
64 for (level, sibling) in siblings.iter().enumerate() {
65 let l = level as u8;
66 if (position >> level) & 1 == 0 {
67 node = merkle_compress(l, node, *sibling);
68 } else {
69 node = merkle_compress(l, *sibling, node);
70 }
71 }
72 node
73}
74
75#[derive(Debug)]
78pub struct Bn254IncrementalMerkleTree {
79 leaves: Vec<Fr>,
80 empty: [Fr; MERKLE_DEPTH_EVM + 1],
82}
83
84impl Bn254IncrementalMerkleTree {
85 pub fn new() -> Self {
86 let mut empty = [Fr::ZERO; MERKLE_DEPTH_EVM + 1];
87 for i in 1..=MERKLE_DEPTH_EVM {
88 empty[i] = merkle_compress((i - 1) as u8, empty[i - 1], empty[i - 1]);
89 }
90 Self { leaves: Vec::new(), empty }
91 }
92
93 pub fn append(&mut self, leaf: Fr) {
94 self.leaves.push(leaf);
95 }
96
97 pub fn len(&self) -> usize {
98 self.leaves.len()
99 }
100
101 pub fn is_empty(&self) -> bool {
102 self.leaves.is_empty()
103 }
104
105 pub fn root(&self) -> Fr {
106 self.subtree_hash(MERKLE_DEPTH_EVM, 0)
107 }
108
109 pub fn witness(&self, pos: u32) -> [Fr; MERKLE_DEPTH_EVM] {
111 assert!((pos as usize) < self.leaves.len(), "position out of tree");
112 let mut siblings = [Fr::ZERO; MERKLE_DEPTH_EVM];
113 for level in 0..MERKLE_DEPTH_EVM {
114 let sibling_node_idx = ((pos >> level) ^ 1) as usize;
115 siblings[level] = self.subtree_hash(level, sibling_node_idx);
116 }
117 siblings
118 }
119
120 fn subtree_hash(&self, level: usize, idx: usize) -> Fr {
121 let start = idx << level; if start >= self.leaves.len() {
123 return self.empty[level];
124 }
125 if level == 0 {
126 return self.leaves[start];
127 }
128 let left = self.subtree_hash(level - 1, idx * 2);
129 let right = self.subtree_hash(level - 1, idx * 2 + 1);
130 merkle_compress((level - 1) as u8, left, right)
131 }
132}
133
134impl Default for Bn254IncrementalMerkleTree {
135 fn default() -> Self {
136 Self::new()
137 }
138}