Skip to main content

privacy_core/commitment_tree/
frozen.rs

1//! Frozen-set (compliance blacklist) **Indexed Merkle Tree** over BN254 `Fr`.
2//!
3//! Mirrors PERC20 `circuits/frozen_cmx_nonmember.circom` and the prover's
4//! `poseidon_merkle_bn254` IMT helpers — the root produced here is the on-chain
5//! `cmxFrozenRoot()` / circuit `rt_frozen` and **must** match byte-for-byte.
6//!
7//! Non-membership of `cmx` is proven by a single Merkle inclusion of a sorted
8//! "low-leaf" `(val, next_val)` that strictly brackets it (`val < cmx < next_val`,
9//! `next_val == 0` being the +inf sentinel). The tree is kept as a sorted linked
10//! list: inserting `v` splices it after its predecessor (update the predecessor's
11//! `next`, append the new leaf). Depth bounds the blacklist *capacity* (`2^DEPTH`),
12//! not the key space, so distinct `cmx` never collide regardless of depth.
13
14use ff::{Field, PrimeField};
15use halo2curves::bn256::Fr;
16
17use super::poseidon::poseidon_domain_pair;
18
19/// Indexed-MT depth = log2(capacity). Must equal circom `FROZEN_IMT_DEPTH()` and
20/// `poseidon_merkle_bn254::FROZEN_IMT_DEPTH`.
21pub const FROZEN_IMT_DEPTH: usize = 20;
22/// Leaf domain: `leaf = Poseidon(LEAF_DOMAIN, val, next_val)`.
23pub const FROZEN_IMT_LEAF_DOMAIN: u64 = 240;
24/// Internal-node domain base: node at level `i` = `Poseidon(NODE_D0 + i, left, right)`.
25pub const FROZEN_IMT_NODE_D0: u64 = 256;
26
27/// Low-leaf hash, matching `FrozenCmxNonMember`'s `leaf` gate.
28#[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
33/// Canonical unsigned compare `a < b` over `Fr` (via little-endian `to_repr`).
34fn 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    // `to_repr` is little-endian; compare most-significant byte first.
38    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
46/// Big-endian 32 bytes (EVM `uint256`) → `Fr`. `None` if not a canonical field element.
47pub 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
53/// `Fr` → little-endian 32 bytes (the indexer's on-the-wire convention).
54pub fn fr_to_le_bytes(fr: Fr) -> [u8; 32] {
55    fr.to_repr().into()
56}
57
58/// `Fr` → big-endian 32 bytes (EVM `uint256` order; inverse of [`fr_from_be_bytes`]).
59pub 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
65/// `Fr` → `0x`-prefixed little-endian 32-byte hex (matches `/merkle_path` siblings).
66pub 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/// Non-membership witness for one `cmx`, matching `FrozenCmxNonMember`'s inputs.
77#[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/// Sorted Indexed Merkle Tree of frozen `cmx` values.
86#[derive(Clone, Debug)]
87pub struct FrozenImt {
88    /// Leaves in insertion order; `leaves[0]` is always the `{0, 0}` sentinel
89    /// (val 0, next +inf) so that an empty blacklist still brackets every `cmx > 0`.
90    leaves: Vec<Leaf>,
91}
92
93impl FrozenImt {
94    /// Empty blacklist: just the `{0, 0}` sentinel at index 0.
95    pub fn new() -> Self {
96        Self { leaves: vec![Leaf { val: Fr::ZERO, next_val: Fr::ZERO }] }
97    }
98
99    /// Rebuild a tree from its frozen values (e.g. a persisted checkpoint), in
100    /// the order they were inserted. Positions are reproduced exactly.
101    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    /// Number of leaves (including the sentinel).
110    pub fn len(&self) -> usize {
111        self.leaves.len()
112    }
113
114    /// The frozen values in insertion order, excluding the `{0,0}` sentinel.
115    pub fn frozen_values(&self) -> Vec<Fr> {
116        self.leaves[1..].iter().map(|l| l.val).collect()
117    }
118
119    /// True if `v` is currently frozen.
120    pub fn contains(&self, v: Fr) -> bool {
121        self.leaves.iter().any(|l| l.val == v)
122    }
123
124    /// Freeze `v` by splicing it after its predecessor. Returns `false` (no-op) if
125    /// `v` is already frozen or `v == 0` (the reserved sentinel value).
126    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    /// Index of the leaf whose interval `(val, next_val)` contains `v`
138    /// (`next_val == 0` = +inf). For a well-formed tree exactly one such leaf
139    /// exists when `v` is not itself a leaf value.
140    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 // unreachable for a well-formed tree; fall back to the sentinel.
149    }
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    /// Digest of an all-empty subtree of height `level` (empty leaf slot = 0).
157    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    /// Hash of the subtree rooted at `(level, idx)` (recursive, empty-padded).
166    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    /// Compliance root `rt_frozen`.
180    pub fn root(&self) -> Fr {
181        self.subtree_hash(FROZEN_IMT_DEPTH, 0)
182    }
183
184    /// Authentication path (siblings + path bits) for the leaf at `pos`.
185    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    /// Non-membership witness for `cmx`, or `None` if `cmx` is currently frozen.
196    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    /// The empty-blacklist root must match the PERC20 prover / circuit constant
223    /// (`poseidon_merkle_bn254::frozen_empty_tree_root` → MAINNET_RT_FROZEN_DEC).
224    #[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    /// A non-membership witness recomputes the live root (mirrors the circuit's
233    /// inclusion check) and brackets the queried cmx.
234    #[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); // not frozen; sits between 10 and 50
242        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    /// A frozen value has no non-membership witness.
249    #[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    /// Rebuilding from `frozen_values()` reproduces the exact same tree (positions
257    /// and therefore root) — the property persistence relies on. (The IMT root is
258    /// position-dependent, so insertion order is part of the committed state and is
259    /// preserved by replaying values in order.)
260    #[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    /// Recompute the IMT root from a witness exactly as `FrozenCmxNonMember` does.
272    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}