light_indexed_merkle_tree/
reference.rs

1use std::marker::PhantomData;
2
3use light_bounded_vec::{BoundedVec, BoundedVecError};
4use light_concurrent_merkle_tree::light_hasher::{errors::HasherError, Hasher};
5use light_hasher::bigint::bigint_to_be_bytes_array;
6use light_merkle_tree_reference::{MerkleTree, ReferenceMerkleTreeError};
7use num_bigint::BigUint;
8use num_traits::{CheckedAdd, CheckedSub, Num, ToBytes, Unsigned};
9use thiserror::Error;
10
11use crate::{
12    array::{IndexedArray, IndexedElement},
13    errors::IndexedMerkleTreeError,
14    HIGHEST_ADDRESS_PLUS_ONE,
15};
16
17#[derive(Debug, Error, PartialEq)]
18pub enum IndexedReferenceMerkleTreeError {
19    #[error("NonInclusionProofFailedLowerBoundViolated")]
20    NonInclusionProofFailedLowerBoundViolated,
21    #[error("NonInclusionProofFailedHigherBoundViolated")]
22    NonInclusionProofFailedHigherBoundViolated,
23    #[error(transparent)]
24    Indexed(#[from] IndexedMerkleTreeError),
25    #[error(transparent)]
26    Reference(#[from] ReferenceMerkleTreeError),
27    #[error(transparent)]
28    Hasher(#[from] HasherError),
29    #[error(transparent)]
30    BoundedVec(#[from] BoundedVecError),
31}
32
33#[derive(Debug, Clone)]
34#[repr(C)]
35pub struct IndexedMerkleTree<H, I>
36where
37    H: Hasher,
38    I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
39{
40    pub merkle_tree: MerkleTree<H>,
41    _index: PhantomData<I>,
42}
43
44impl<H, I> IndexedMerkleTree<H, I>
45where
46    H: Hasher,
47    I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
48    usize: From<I>,
49{
50    pub fn new(
51        height: usize,
52        canopy_depth: usize,
53    ) -> Result<Self, IndexedReferenceMerkleTreeError> {
54        let mut merkle_tree = MerkleTree::new(height, canopy_depth);
55
56        // Append the first low leaf, which has value 0 and does not point
57        // to any other leaf yet.
58        // This low leaf is going to be updated during the first `update`
59        // operation.
60        merkle_tree.append(&H::zero_indexed_leaf())?;
61
62        Ok(Self {
63            merkle_tree,
64            _index: PhantomData,
65        })
66    }
67
68    /// Initializes the reference indexed merkle tree on par with the
69    /// on-chain indexed concurrent merkle tree.
70    /// Inserts the ranges 0 - BN254 Field Size - 1 into the tree.
71    pub fn init(&mut self) -> Result<(), IndexedReferenceMerkleTreeError> {
72        let mut indexed_array = IndexedArray::<H, I>::default();
73        let init_value = BigUint::from_str_radix(HIGHEST_ADDRESS_PLUS_ONE, 10).unwrap();
74        let nullifier_bundle = indexed_array.append(&init_value)?;
75        let new_low_leaf = nullifier_bundle
76            .new_low_element
77            .hash::<H>(&nullifier_bundle.new_element.value)?;
78
79        self.merkle_tree.update(&new_low_leaf, 0)?;
80        let new_leaf = nullifier_bundle
81            .new_element
82            .hash::<H>(&nullifier_bundle.new_element_next_value)?;
83        self.merkle_tree.append(&new_leaf)?;
84        Ok(())
85    }
86
87    /// Initializes the reference indexed merkle tree on par with the
88    /// on-chain indexed concurrent merkle tree.
89    /// Inserts the ranges 0 - BN254 Field Size - 1 into the tree.
90    pub fn init_opt(&mut self) -> Result<(), IndexedReferenceMerkleTreeError> {
91        let mut indexed_array = IndexedArray::<H, I>::default();
92        let init_value = BigUint::from_str_radix(HIGHEST_ADDRESS_PLUS_ONE, 10).unwrap();
93        let nullifier_bundle = indexed_array.append(&init_value)?;
94        let new_low_leaf = nullifier_bundle
95            .new_low_element
96            .hash::<H>(&nullifier_bundle.new_element.value)?;
97
98        self.merkle_tree.update(&new_low_leaf, 0)?;
99        let new_leaf = nullifier_bundle
100            .new_element
101            .hash::<H>(&nullifier_bundle.new_element_next_value)?;
102        self.merkle_tree.append(&new_leaf)?;
103        Ok(())
104    }
105
106    pub fn get_path_of_leaf(
107        &self,
108        index: usize,
109        full: bool,
110    ) -> Result<BoundedVec<[u8; 32]>, IndexedReferenceMerkleTreeError> {
111        let path = self.merkle_tree.get_path_of_leaf(index, full)?;
112        Ok(BoundedVec::from_slice(path.as_slice()))
113    }
114
115    pub fn get_proof_of_leaf(
116        &self,
117        index: usize,
118        full: bool,
119    ) -> Result<BoundedVec<[u8; 32]>, IndexedReferenceMerkleTreeError> {
120        Ok(BoundedVec::from_slice(
121            self.merkle_tree.get_proof_of_leaf(index, full)?.as_slice(),
122        ))
123    }
124
125    pub fn root(&self) -> [u8; 32] {
126        self.merkle_tree.root()
127    }
128
129    // TODO: rename input values
130    pub fn update(
131        &mut self,
132        new_low_element: &IndexedElement<I>,
133        new_element: &IndexedElement<I>,
134        new_element_next_value: &BigUint,
135    ) -> Result<(), IndexedReferenceMerkleTreeError> {
136        // Update the low element.
137        let new_low_leaf = new_low_element.hash::<H>(&new_element.value)?;
138        println!("reference update new low leaf hash {:?}", new_low_leaf);
139        self.merkle_tree
140            .update(&new_low_leaf, usize::from(new_low_element.index))?;
141        println!("reference updated root {:?}", self.merkle_tree.root());
142        // Append the new element.
143        let new_leaf = new_element.hash::<H>(new_element_next_value)?;
144        println!("reference update new leaf hash {:?}", new_leaf);
145        self.merkle_tree.append(&new_leaf)?;
146        println!("reference appended root {:?}", self.merkle_tree.root());
147
148        Ok(())
149    }
150
151    // TODO: add append with new value, so that we don't need to compute the lowlevel values manually
152    pub fn append(
153        &mut self,
154        value: &BigUint,
155        indexed_array: &mut IndexedArray<H, I>,
156    ) -> Result<(), IndexedReferenceMerkleTreeError> {
157        println!("appending {:?}", value);
158        let nullifier_bundle = indexed_array.append(value).unwrap();
159        self.update(
160            &nullifier_bundle.new_low_element,
161            &nullifier_bundle.new_element,
162            &nullifier_bundle.new_element_next_value,
163        )?;
164
165        Ok(())
166    }
167
168    pub fn get_non_inclusion_proof(
169        &self,
170        value: &BigUint,
171        indexed_array: &IndexedArray<H, I>,
172    ) -> Result<NonInclusionProof, IndexedReferenceMerkleTreeError> {
173        let (low_element, _next_value) = indexed_array.find_low_element_for_nonexistent(value)?;
174        let merkle_proof = self
175            .get_proof_of_leaf(usize::from(low_element.index), true)
176            .unwrap();
177        let higher_range_value = indexed_array
178            .get(low_element.next_index())
179            .unwrap()
180            .value
181            .clone();
182        Ok(NonInclusionProof {
183            root: self.root(),
184            value: bigint_to_be_bytes_array::<32>(value).unwrap(),
185            leaf_lower_range_value: bigint_to_be_bytes_array::<32>(&low_element.value).unwrap(),
186            leaf_higher_range_value: bigint_to_be_bytes_array::<32>(&higher_range_value).unwrap(),
187            leaf_index: low_element.index.into(),
188            next_index: low_element.next_index(),
189            merkle_proof,
190        })
191    }
192
193    pub fn verify_non_inclusion_proof(
194        &self,
195        proof: &NonInclusionProof,
196    ) -> Result<(), IndexedReferenceMerkleTreeError> {
197        let value_big_int = BigUint::from_bytes_be(&proof.value);
198        let lower_end_value = BigUint::from_bytes_be(&proof.leaf_lower_range_value);
199        if lower_end_value >= value_big_int {
200            return Err(IndexedReferenceMerkleTreeError::NonInclusionProofFailedLowerBoundViolated);
201        }
202        let higher_end_value = BigUint::from_bytes_be(&proof.leaf_higher_range_value);
203        if higher_end_value <= value_big_int {
204            return Err(
205                IndexedReferenceMerkleTreeError::NonInclusionProofFailedHigherBoundViolated,
206            );
207        }
208
209        let array_element = IndexedElement::<usize> {
210            value: lower_end_value,
211            index: proof.leaf_index,
212            next_index: proof.next_index,
213        };
214        let leaf_hash = array_element.hash::<H>(&higher_end_value)?;
215        self.merkle_tree
216            .verify(&leaf_hash, proof.merkle_proof.as_slice(), proof.leaf_index)?;
217        Ok(())
218    }
219}
220
221// TODO: check why next_index is usize while index is I
222/// We prove non-inclusion by:
223/// 1. Showing that value is greater than leaf_lower_range_value and less than leaf_higher_range_value
224/// 2. Showing that the leaf_hash H(leaf_lower_range_value, leaf_next_index, leaf_higher_value) is included in the root (Merkle tree)
225#[derive(Debug)]
226pub struct NonInclusionProof {
227    pub root: [u8; 32],
228    pub value: [u8; 32],
229    pub leaf_lower_range_value: [u8; 32],
230    pub leaf_higher_range_value: [u8; 32],
231    pub leaf_index: usize,
232    pub next_index: usize,
233    pub merkle_proof: BoundedVec<[u8; 32]>,
234}