light_merkle_tree_reference/
indexed.rs1use std::{fmt::Debug, marker::PhantomData};
2
3use light_hasher::{bigint::bigint_to_be_bytes_array, Hasher, HasherError};
4use light_indexed_array::{
5 array::{IndexedArray, IndexedElement},
6 errors::IndexedArrayError,
7 HIGHEST_ADDRESS_PLUS_ONE,
8};
9use num_bigint::BigUint;
10use num_traits::{CheckedAdd, CheckedSub, Num, ToBytes, Unsigned, Zero};
11use thiserror::Error;
12
13use crate::{MerkleTree, ReferenceMerkleTreeError};
14
15#[derive(Debug, Error, PartialEq)]
16pub enum IndexedReferenceMerkleTreeError {
17 #[error("NonInclusionProofFailedLowerBoundViolated")]
18 NonInclusionProofFailedLowerBoundViolated,
19 #[error("NonInclusionProofFailedHigherBoundViolated")]
20 NonInclusionProofFailedHigherBoundViolated,
21 #[error(transparent)]
22 Indexed(#[from] IndexedArrayError),
23 #[error(transparent)]
24 Reference(#[from] ReferenceMerkleTreeError),
25 #[error(transparent)]
26 Hasher(#[from] HasherError),
27}
28
29#[derive(Debug, Clone)]
30#[repr(C)]
31pub struct IndexedMerkleTree<H, I>
32where
33 H: Hasher,
34 I: CheckedAdd
35 + CheckedSub
36 + Copy
37 + Clone
38 + PartialOrd
39 + ToBytes
40 + TryFrom<usize>
41 + Unsigned
42 + Debug
43 + Into<usize>,
44{
45 pub merkle_tree: MerkleTree<H>,
46 pub indexed_array: IndexedArray<H, I>,
47 _index: PhantomData<I>,
48}
49
50impl<H, I> IndexedMerkleTree<H, I>
51where
52 H: Hasher,
53 I: CheckedAdd
54 + CheckedSub
55 + Copy
56 + Clone
57 + PartialOrd
58 + ToBytes
59 + TryFrom<usize>
60 + Unsigned
61 + Debug
62 + Into<usize>,
63{
64 pub fn new(
65 height: usize,
66 canopy_depth: usize,
67 ) -> Result<Self, IndexedReferenceMerkleTreeError> {
68 let mut merkle_tree = MerkleTree::new(height, canopy_depth);
69
70 let init_next_value = BigUint::from_str_radix(HIGHEST_ADDRESS_PLUS_ONE, 10).unwrap();
71 let indexed_array = IndexedArray::<H, I>::new(BigUint::zero(), init_next_value.clone());
72 let new_leaf = indexed_array
73 .get(0)
74 .unwrap()
75 .hash::<H>(&init_next_value)
76 .unwrap();
77 merkle_tree.append(&new_leaf)?;
78 assert_eq!(merkle_tree.leaf(0), new_leaf);
79 Ok(Self {
80 merkle_tree,
81 indexed_array,
82 _index: PhantomData,
83 })
84 }
85
86 pub fn get_path_of_leaf(
87 &self,
88 index: usize,
89 full: bool,
90 ) -> Result<Vec<[u8; 32]>, ReferenceMerkleTreeError> {
91 self.merkle_tree.get_path_of_leaf(index, full)
92 }
93
94 pub fn get_proof_of_leaf(
95 &self,
96 index: usize,
97 full: bool,
98 ) -> Result<Vec<[u8; 32]>, ReferenceMerkleTreeError> {
99 self.merkle_tree.get_proof_of_leaf(index, full)
100 }
101
102 pub fn root(&self) -> [u8; 32] {
103 self.merkle_tree.root()
104 }
105
106 pub fn update(
108 &mut self,
109 new_low_element: &IndexedElement<I>,
110 new_element: &IndexedElement<I>,
111 new_element_next_value: &BigUint,
112 ) -> Result<(), IndexedReferenceMerkleTreeError> {
113 let new_low_leaf = new_low_element.hash::<H>(&new_element.value)?;
115 println!("reference update new low leaf hash {:?}", new_low_leaf);
116 self.merkle_tree.update(
117 &new_low_leaf,
118 usize::try_from(new_low_element.index).unwrap(),
119 )?;
120 println!("reference updated root {:?}", self.merkle_tree.root());
121 let new_leaf = new_element.hash::<H>(new_element_next_value)?;
123 println!("reference update new leaf hash {:?}", new_leaf);
124 self.merkle_tree.append(&new_leaf)?;
125 println!("reference appended root {:?}", self.merkle_tree.root());
126
127 Ok(())
128 }
129
130 pub fn append(&mut self, value: &BigUint) -> Result<(), IndexedReferenceMerkleTreeError> {
132 println!("\n\nappending {:?}", value);
133 let nullifier_bundle = self.indexed_array.append(value)?;
134 println!("\n\nnullifier_bundle {:?}", nullifier_bundle);
135 self.update(
136 &nullifier_bundle.new_low_element,
137 &nullifier_bundle.new_element,
138 &nullifier_bundle.new_element_next_value,
139 )?;
140
141 Ok(())
142 }
143
144 pub fn get_non_inclusion_proof(
145 &self,
146 value: &BigUint,
147 ) -> Result<NonInclusionProof, IndexedReferenceMerkleTreeError> {
148 let (low_element, _next_value) =
149 self.indexed_array.find_low_element_for_nonexistent(value)?;
150 let merkle_proof =
151 self.get_proof_of_leaf(usize::try_from(low_element.index).unwrap(), true)?;
152 let higher_range_value = if low_element.next_index() == 0 {
153 self.indexed_array.highest_value.clone()
154 } else {
155 self.indexed_array
156 .get(low_element.next_index())
157 .unwrap()
158 .value
159 .clone()
160 };
161 let non_inclusion_proof = NonInclusionProof {
162 root: self.root(),
163 value: bigint_to_be_bytes_array::<32>(value).unwrap(),
164 leaf_lower_range_value: bigint_to_be_bytes_array::<32>(&low_element.value).unwrap(),
165 leaf_higher_range_value: bigint_to_be_bytes_array::<32>(&higher_range_value).unwrap(),
166 leaf_index: low_element.index.into(),
167 next_index: low_element.next_index(),
168 merkle_proof,
169 };
170 Ok(non_inclusion_proof)
171 }
172
173 pub fn verify_non_inclusion_proof(
174 &self,
175 proof: &NonInclusionProof,
176 ) -> Result<(), IndexedReferenceMerkleTreeError> {
177 let value_big_int = BigUint::from_bytes_be(&proof.value);
178 let lower_end_value = BigUint::from_bytes_be(&proof.leaf_lower_range_value);
179 if lower_end_value >= value_big_int {
180 return Err(IndexedReferenceMerkleTreeError::NonInclusionProofFailedLowerBoundViolated);
181 }
182 let higher_end_value = BigUint::from_bytes_be(&proof.leaf_higher_range_value);
183 if higher_end_value <= value_big_int {
184 return Err(
185 IndexedReferenceMerkleTreeError::NonInclusionProofFailedHigherBoundViolated,
186 );
187 }
188
189 let array_element = IndexedElement::<usize> {
190 value: lower_end_value,
191 index: proof.leaf_index,
192 next_index: proof.next_index,
193 };
194 let leaf_hash = array_element.hash::<H>(&higher_end_value)?;
195 self.merkle_tree
196 .verify(&leaf_hash, &proof.merkle_proof, proof.leaf_index)?;
197 Ok(())
198 }
199}
200
201#[derive(Debug)]
206pub struct NonInclusionProof {
207 pub root: [u8; 32],
208 pub value: [u8; 32],
209 pub leaf_lower_range_value: [u8; 32],
210 pub leaf_higher_range_value: [u8; 32],
211 pub leaf_index: usize,
212 pub next_index: usize,
213 pub merkle_proof: Vec<[u8; 32]>,
214}