privacy_core/commitment_tree/
frozen.rs1use ff::{Field, PrimeField};
15use halo2curves::bn256::Fr;
16
17use super::poseidon::poseidon_domain_pair;
18
19pub const FROZEN_IMT_DEPTH: usize = 20;
22pub const FROZEN_IMT_LEAF_DOMAIN: u64 = 240;
24pub const FROZEN_IMT_NODE_D0: u64 = 256;
26
27#[inline]
29pub fn frozen_imt_leaf(val: Fr, next_val: Fr) -> Fr {
30 poseidon_domain_pair(FROZEN_IMT_LEAF_DOMAIN, val, next_val)
31}
32
33fn fr_lt(a: &Fr, b: &Fr) -> bool {
35 let (ar, br) = (a.to_repr(), b.to_repr());
36 let (a, b) = (ar.as_ref(), br.as_ref());
37 for i in (0..a.len()).rev() {
39 if a[i] != b[i] {
40 return a[i] < b[i];
41 }
42 }
43 false
44}
45
46pub fn fr_from_be_bytes(be: &[u8; 32]) -> Option<Fr> {
48 let mut le = *be;
49 le.reverse();
50 Option::from(Fr::from_repr(le.into()))
51}
52
53pub fn fr_to_le_bytes(fr: Fr) -> [u8; 32] {
55 fr.to_repr().into()
56}
57
58pub fn fr_to_be_bytes(fr: Fr) -> [u8; 32] {
60 let mut be = fr_to_le_bytes(fr);
61 be.reverse();
62 be
63}
64
65pub fn fr_to_le_hex(fr: Fr) -> String {
67 format!("0x{}", hex::encode(fr_to_le_bytes(fr)))
68}
69
70#[derive(Clone, Copy, Debug)]
71struct Leaf {
72 val: Fr,
73 next_val: Fr,
74}
75
76#[derive(Clone, Debug)]
78pub struct FrozenNonMembershipWitness {
79 pub low_val: Fr,
80 pub low_next_val: Fr,
81 pub siblings: [Fr; FROZEN_IMT_DEPTH],
82 pub path_bits: [Fr; FROZEN_IMT_DEPTH],
83}
84
85#[derive(Clone, Debug)]
87pub struct FrozenImt {
88 leaves: Vec<Leaf>,
91}
92
93impl FrozenImt {
94 pub fn new() -> Self {
96 Self { leaves: vec![Leaf { val: Fr::ZERO, next_val: Fr::ZERO }] }
97 }
98
99 pub fn from_frozen_values(values: &[Fr]) -> Self {
102 let mut t = Self::new();
103 for &v in values {
104 t.insert(v);
105 }
106 t
107 }
108
109 pub fn len(&self) -> usize {
111 self.leaves.len()
112 }
113
114 pub fn frozen_values(&self) -> Vec<Fr> {
116 self.leaves[1..].iter().map(|l| l.val).collect()
117 }
118
119 pub fn contains(&self, v: Fr) -> bool {
121 self.leaves.iter().any(|l| l.val == v)
122 }
123
124 pub fn insert(&mut self, v: Fr) -> bool {
127 if v == Fr::ZERO || self.contains(v) {
128 return false;
129 }
130 let pred = self.bracketing_index(v);
131 let new_leaf = Leaf { val: v, next_val: self.leaves[pred].next_val };
132 self.leaves[pred].next_val = v;
133 self.leaves.push(new_leaf);
134 true
135 }
136
137 fn bracketing_index(&self, v: Fr) -> usize {
141 for (i, l) in self.leaves.iter().enumerate() {
142 let above_low = fr_lt(&l.val, &v);
143 let below_high = l.next_val == Fr::ZERO || fr_lt(&v, &l.next_val);
144 if above_low && below_high {
145 return i;
146 }
147 }
148 0 }
150
151 #[inline]
152 fn leaf_hash(&self, i: usize) -> Fr {
153 frozen_imt_leaf(self.leaves[i].val, self.leaves[i].next_val)
154 }
155
156 fn empty_at(&self, level: usize) -> Fr {
158 let mut e = Fr::ZERO;
159 for i in 0..level {
160 e = poseidon_domain_pair(FROZEN_IMT_NODE_D0 + i as u64, e, e);
161 }
162 e
163 }
164
165 fn subtree_hash(&self, level: usize, idx: usize) -> Fr {
167 let start = idx << level;
168 if start >= self.leaves.len() {
169 return self.empty_at(level);
170 }
171 if level == 0 {
172 return self.leaf_hash(start);
173 }
174 let left = self.subtree_hash(level - 1, idx * 2);
175 let right = self.subtree_hash(level - 1, idx * 2 + 1);
176 poseidon_domain_pair(FROZEN_IMT_NODE_D0 + (level - 1) as u64, left, right)
177 }
178
179 pub fn root(&self) -> Fr {
181 self.subtree_hash(FROZEN_IMT_DEPTH, 0)
182 }
183
184 fn witness_at(&self, pos: usize) -> ([Fr; FROZEN_IMT_DEPTH], [Fr; FROZEN_IMT_DEPTH]) {
186 let mut siblings = [Fr::ZERO; FROZEN_IMT_DEPTH];
187 let mut path_bits = [Fr::ZERO; FROZEN_IMT_DEPTH];
188 for level in 0..FROZEN_IMT_DEPTH {
189 path_bits[level] = if (pos >> level) & 1 == 1 { Fr::ONE } else { Fr::ZERO };
190 siblings[level] = self.subtree_hash(level, (pos >> level) ^ 1);
191 }
192 (siblings, path_bits)
193 }
194
195 pub fn non_membership_witness(&self, cmx: Fr) -> Option<FrozenNonMembershipWitness> {
197 if self.contains(cmx) {
198 return None;
199 }
200 let pos = self.bracketing_index(cmx);
201 let low = self.leaves[pos];
202 let (siblings, path_bits) = self.witness_at(pos);
203 Some(FrozenNonMembershipWitness {
204 low_val: low.val,
205 low_next_val: low.next_val,
206 siblings,
207 path_bits,
208 })
209 }
210}
211
212impl Default for FrozenImt {
213 fn default() -> Self {
214 Self::new()
215 }
216}
217
218#[cfg(test)]
219mod tests {
220 use super::*;
221
222 #[test]
225 fn empty_root_matches_perc20_constant() {
226 const DEC: &str =
227 "9079151408671112139333676443195611613776084922747126087146403043120709007371";
228 let expected = Fr::from_str_vartime(DEC).unwrap();
229 assert_eq!(FrozenImt::new().root(), expected);
230 }
231
232 #[test]
235 fn witness_reproduces_root_and_brackets() {
236 let mut t = FrozenImt::new();
237 for v in [50u64, 10, 99, 7] {
238 assert!(t.insert(Fr::from(v)));
239 }
240 let root = t.root();
241 let cmx = Fr::from(42u64); let w = t.non_membership_witness(cmx).expect("non-member");
243 assert!(fr_lt(&w.low_val, &cmx));
244 assert!(w.low_next_val == Fr::ZERO || fr_lt(&cmx, &w.low_next_val));
245 assert_eq!(recompute_root(&w), root, "witness must reproduce rt_frozen");
246 }
247
248 #[test]
250 fn frozen_value_has_no_witness() {
251 let mut t = FrozenImt::new();
252 t.insert(Fr::from(123u64));
253 assert!(t.non_membership_witness(Fr::from(123u64)).is_none());
254 }
255
256 #[test]
261 fn rebuild_from_values_reproduces_root() {
262 let mut t = FrozenImt::new();
263 for v in [3u64, 1, 4, 1, 5, 9, 2, 6] {
264 t.insert(Fr::from(v));
265 }
266 let rebuilt = FrozenImt::from_frozen_values(&t.frozen_values());
267 assert_eq!(rebuilt.frozen_values(), t.frozen_values());
268 assert_eq!(rebuilt.root(), t.root());
269 }
270
271 fn recompute_root(w: &FrozenNonMembershipWitness) -> Fr {
273 let mut level = frozen_imt_leaf(w.low_val, w.low_next_val);
274 for i in 0..FROZEN_IMT_DEPTH {
275 let bit = w.path_bits[i];
276 let diff = level - w.siblings[i];
277 let left = level - bit * diff;
278 let right = w.siblings[i] + bit * diff;
279 level = poseidon_domain_pair(FROZEN_IMT_NODE_D0 + i as u64, left, right);
280 }
281 level
282 }
283}