Skip to main content

privacy_core/commitment_tree/
poseidon.rs

1//! Poseidon-based Merkle hashing on the BN254 scalar field (`Fr`).
2//!
3//! Lifted verbatim from `orchard-bn254`'s `poseidon_merkle_bn254` (the merkle-tree parts
4//! only), re-sourced to the **standalone** crates `halo2_poseidon` (Poseidon primitives)
5//! and `halo2curves` (BN254 `Fr`) so it carries **no `halo2_proofs` / `halo2_gadgets`
6//! dependency**. The Poseidon spec/constants are byte-identical to the prover + on-chain
7//! `PoseidonT3` (verified against the live commitment-tree root).
8
9use super::poseidon_primitives::{generate_constants, ConstantLength, Hash, Mds, Spec};
10use ff::Field;
11use halo2curves::bn256::Fr;
12
13/// Orchard-compatible depth for the incremental note commitment tree.
14pub const MERKLE_DEPTH_EVM: usize = 32;
15
16/// Poseidon-128, width 3, rate 2, x^5 S-box; Grain-generated constants for BN254 `Fr`.
17#[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/// One Merkle layer: `H(level || left || right)` with domain separation on `level`.
39#[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/// `Poseidon(domain, a, b)` over `ConstantLength<3>` — the width-3 two-input hash with
49/// a `u64` domain tag. Byte-identical to the prover's `poseidon_merkle_bn254::poseidon2`
50/// and to `merkle_compress` when `domain == level`. Used for the frozen Indexed MT,
51/// whose node domains exceed the `u8` range of `merkle_compress`.
52#[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
61/// Full Merkle root (depth [`MERKLE_DEPTH_EVM`]) over `Fr` leaves from a sibling path.
62pub 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/// Append-only incremental Merkle tree of fixed depth [`MERKLE_DEPTH_EVM`] (32) over BN254
76/// scalar leaves. Empty leaves default to [`Fr::ZERO`].
77#[derive(Debug)]
78pub struct Bn254IncrementalMerkleTree {
79    leaves: Vec<Fr>,
80    /// `empty[l]` is the root of a depth-`l` all-zero subtree (`empty[0]` = `Fr::ZERO`).
81    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    /// Authentication path (siblings) for the leaf at `pos`. Panics if `pos >= len()`.
110    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; // idx * 2^level
122        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}