light_program_test/indexer/
address_tree.rs

1use std::fmt::Debug;
2
3#[cfg(feature = "devenv")]
4use light_batched_merkle_tree::constants::DEFAULT_BATCH_ROOT_HISTORY_LEN;
5use light_client::{
6    fee::FeeConfig,
7    indexer::{AddressMerkleTreeAccounts, IndexerError},
8};
9use light_hasher::Poseidon;
10use light_indexed_merkle_tree::{
11    array::{IndexedArray, IndexedElement, IndexedElementBundle},
12    reference::IndexedMerkleTree,
13};
14use light_prover_client::proof_types::non_inclusion::v2::NonInclusionMerkleProofInputs;
15use light_sdk::constants::STATE_MERKLE_TREE_ROOTS;
16use num_bigint::{BigInt, BigUint};
17use num_traits::ops::bytes::FromBytes;
18
19#[cfg(not(feature = "devenv"))]
20use super::test_indexer::DEFAULT_BATCH_ROOT_HISTORY_LEN;
21
22#[derive(Debug, Clone)]
23pub enum IndexedMerkleTreeVersion {
24    V1(Box<IndexedMerkleTree<Poseidon, usize>>),
25    V2(Box<light_merkle_tree_reference::indexed::IndexedMerkleTree<Poseidon, usize>>),
26}
27
28#[derive(Debug, Clone)]
29pub struct AddressMerkleTreeBundle {
30    pub rollover_fee: i64,
31    pub merkle_tree: IndexedMerkleTreeVersion,
32    indexed_array: Box<IndexedArray<Poseidon, usize>>,
33    pub accounts: AddressMerkleTreeAccounts,
34    pub queue_elements: Vec<[u8; 32]>,
35}
36
37impl AddressMerkleTreeBundle {
38    pub fn new_v1(accounts: AddressMerkleTreeAccounts) -> Result<Self, IndexerError> {
39        let height = 26;
40        let canopy = 10;
41        let mut merkle_tree = IndexedMerkleTree::<Poseidon, usize>::new(height, canopy)
42            .map_err(|_| IndexerError::InvalidResponseData)?;
43        merkle_tree.merkle_tree.root_history_array_len = Some(STATE_MERKLE_TREE_ROOTS);
44        let mut merkle_tree = Box::new(merkle_tree);
45        merkle_tree.init()?;
46        let mut indexed_array = Box::<IndexedArray<Poseidon, usize>>::default();
47        indexed_array.init()?;
48        Ok(AddressMerkleTreeBundle {
49            merkle_tree: IndexedMerkleTreeVersion::V1(merkle_tree),
50            indexed_array,
51            accounts,
52            rollover_fee: FeeConfig::default().address_queue_rollover as i64,
53            queue_elements: vec![],
54        })
55    }
56
57    pub fn new_v2(accounts: AddressMerkleTreeAccounts) -> Result<Self, IndexerError> {
58        let height = 40;
59        let canopy = 0;
60        let mut merkle_tree = light_merkle_tree_reference::indexed::IndexedMerkleTree::<
61            Poseidon,
62            usize,
63        >::new(height, canopy)
64        .map_err(|_| IndexerError::InvalidResponseData)?;
65        merkle_tree.merkle_tree.root_history_array_len =
66            Some(DEFAULT_BATCH_ROOT_HISTORY_LEN as usize);
67        let merkle_tree = IndexedMerkleTreeVersion::V2(Box::new(merkle_tree));
68
69        Ok(AddressMerkleTreeBundle {
70            merkle_tree,
71            indexed_array: Box::default(),
72            accounts,
73            rollover_fee: FeeConfig::default().address_queue_rollover as i64,
74            queue_elements: vec![],
75        })
76    }
77
78    pub fn get_v1_indexed_merkle_tree(&self) -> Option<&IndexedMerkleTree<Poseidon, usize>> {
79        match &self.merkle_tree {
80            IndexedMerkleTreeVersion::V1(tree) => Some(tree),
81            _ => None,
82        }
83    }
84
85    pub fn get_v1_indexed_merkle_tree_mut(
86        &mut self,
87    ) -> Option<&mut IndexedMerkleTree<Poseidon, usize>> {
88        match &mut self.merkle_tree {
89            IndexedMerkleTreeVersion::V1(tree) => Some(tree),
90            _ => None,
91        }
92    }
93
94    pub fn get_v2_indexed_merkle_tree(
95        &self,
96    ) -> Option<&light_merkle_tree_reference::indexed::IndexedMerkleTree<Poseidon, usize>> {
97        match &self.merkle_tree {
98            IndexedMerkleTreeVersion::V2(tree) => Some(tree),
99            _ => None,
100        }
101    }
102
103    pub fn get_v2_indexed_merkle_tree_mut(
104        &mut self,
105    ) -> Option<&mut light_merkle_tree_reference::indexed::IndexedMerkleTree<Poseidon, usize>> {
106        match &mut self.merkle_tree {
107            IndexedMerkleTreeVersion::V2(tree) => Some(tree),
108            _ => None,
109        }
110    }
111
112    pub fn get_subtrees(&self) -> Vec<[u8; 32]> {
113        match &self.merkle_tree {
114            IndexedMerkleTreeVersion::V1(tree) => tree.merkle_tree.get_subtrees(),
115            IndexedMerkleTreeVersion::V2(tree) => tree.merkle_tree.get_subtrees(),
116        }
117    }
118
119    pub fn root(&self) -> [u8; 32] {
120        match &self.merkle_tree {
121            IndexedMerkleTreeVersion::V1(tree) => tree.merkle_tree.root(),
122            IndexedMerkleTreeVersion::V2(tree) => tree.merkle_tree.root(),
123        }
124    }
125
126    pub fn find_low_element_for_nonexistent(
127        &self,
128        value: &BigUint,
129    ) -> Result<(IndexedElement<usize>, BigUint), IndexerError> {
130        match &self.merkle_tree {
131            IndexedMerkleTreeVersion::V1(_) => Ok(self
132                .indexed_array
133                .find_low_element_for_nonexistent(value)
134                .map_err(|_| IndexerError::InvalidResponseData)?),
135            IndexedMerkleTreeVersion::V2(tree) => {
136                let (indexed_element, next_value) = tree
137                    .indexed_array
138                    .find_low_element_for_nonexistent(value)
139                    .map_err(|_| IndexerError::InvalidResponseData)?;
140                Ok((
141                    IndexedElement {
142                        index: indexed_element.index,
143                        value: indexed_element.value.clone(),
144                        next_index: indexed_element.next_index,
145                    },
146                    next_value,
147                ))
148            }
149        }
150    }
151
152    pub fn new_element_with_low_element_index(
153        &self,
154        index: usize,
155        value: &BigUint,
156    ) -> Result<IndexedElementBundle<usize>, IndexerError> {
157        match &self.merkle_tree {
158            IndexedMerkleTreeVersion::V1(_) => Ok(self
159                .indexed_array
160                .new_element_with_low_element_index(index, value)
161                .map_err(|_| IndexerError::InvalidResponseData)?),
162            IndexedMerkleTreeVersion::V2(tree) => {
163                let res = tree
164                    .indexed_array
165                    .new_element_with_low_element_index(index, value)
166                    .map_err(|_| IndexerError::InvalidResponseData)?;
167                Ok(IndexedElementBundle {
168                    new_element: IndexedElement {
169                        index: res.new_element.index,
170                        value: res.new_element.value.clone(),
171                        next_index: res.new_element.next_index,
172                    },
173                    new_low_element: IndexedElement {
174                        index: res.new_low_element.index,
175                        value: res.new_low_element.value.clone(),
176                        next_index: res.new_low_element.next_index,
177                    },
178                    new_element_next_value: res.new_element_next_value.clone(),
179                })
180            }
181        }
182    }
183
184    pub fn get_proof_of_leaf(
185        &self,
186        index: usize,
187        full: bool,
188    ) -> Result<Vec<[u8; 32]>, IndexerError> {
189        match &self.merkle_tree {
190            IndexedMerkleTreeVersion::V1(tree) => Ok(tree
191                .get_proof_of_leaf(index, full)
192                .map_err(|_| IndexerError::InvalidResponseData)?
193                .to_vec()),
194            IndexedMerkleTreeVersion::V2(tree) => Ok(tree
195                .get_proof_of_leaf(index, full)
196                .map_err(|_| IndexerError::InvalidResponseData)?),
197        }
198    }
199
200    pub fn append(&mut self, value: &BigUint) -> Result<(), IndexerError> {
201        match &mut self.merkle_tree {
202            IndexedMerkleTreeVersion::V1(tree) => {
203                tree.append(value, &mut self.indexed_array)
204                    .map_err(|_| IndexerError::InvalidResponseData)?;
205                Ok(())
206            }
207            IndexedMerkleTreeVersion::V2(tree) => {
208                tree.append(value)
209                    .map_err(|_| IndexerError::InvalidResponseData)?;
210                Ok(())
211            }
212        }
213    }
214
215    pub fn get_non_inclusion_proof_inputs(
216        &self,
217        value: &[u8; 32],
218    ) -> Result<NonInclusionMerkleProofInputs, IndexerError> {
219        match &self.merkle_tree {
220            IndexedMerkleTreeVersion::V1(tree) => Ok(get_non_inclusion_proof_inputs(
221                value,
222                tree,
223                &self.indexed_array,
224            )),
225            IndexedMerkleTreeVersion::V2(merkle_tree) => {
226                let non_inclusion_proof = merkle_tree
227                    .get_non_inclusion_proof(&BigUint::from_be_bytes(value))
228                    .map_err(|_| IndexerError::InvalidResponseData)?;
229                let proof = non_inclusion_proof
230                    .merkle_proof
231                    .iter()
232                    .map(|x| BigInt::from_be_bytes(x))
233                    .collect();
234                Ok(NonInclusionMerkleProofInputs {
235                    root: BigInt::from_be_bytes(merkle_tree.root().as_slice()),
236                    value: BigInt::from_be_bytes(value),
237                    leaf_lower_range_value: BigInt::from_be_bytes(
238                        &non_inclusion_proof.leaf_lower_range_value,
239                    ),
240                    leaf_higher_range_value: BigInt::from_be_bytes(
241                        &non_inclusion_proof.leaf_higher_range_value,
242                    ),
243                    merkle_proof_hashed_indexed_element_leaf: proof,
244                    index_hashed_indexed_element_leaf: BigInt::from(non_inclusion_proof.leaf_index),
245                    next_index: BigInt::from(non_inclusion_proof.next_index),
246                })
247            }
248        }
249    }
250
251    pub fn right_most_index(&self) -> usize {
252        match &self.merkle_tree {
253            IndexedMerkleTreeVersion::V1(tree) => tree.merkle_tree.rightmost_index,
254            IndexedMerkleTreeVersion::V2(tree) => tree.merkle_tree.rightmost_index,
255        }
256    }
257
258    pub fn append_with_low_element_index(
259        &mut self,
260        index: usize,
261        value: &BigUint,
262    ) -> Result<IndexedElementBundle<usize>, IndexerError> {
263        match &mut self.merkle_tree {
264            IndexedMerkleTreeVersion::V1(_) => Ok(self
265                .indexed_array
266                .append_with_low_element_index(index, value)
267                .map_err(|_| IndexerError::InvalidResponseData)?),
268            IndexedMerkleTreeVersion::V2(_) => {
269                unimplemented!("append_with_low_element_index")
270            }
271        }
272    }
273
274    pub fn sequence_number(&self) -> u64 {
275        match &self.merkle_tree {
276            IndexedMerkleTreeVersion::V1(tree) => tree.merkle_tree.sequence_number as u64,
277            IndexedMerkleTreeVersion::V2(tree) => tree.merkle_tree.sequence_number as u64,
278        }
279    }
280
281    pub fn height(&self) -> usize {
282        match &self.merkle_tree {
283            IndexedMerkleTreeVersion::V1(tree) => tree.merkle_tree.height,
284            IndexedMerkleTreeVersion::V2(tree) => tree.merkle_tree.height,
285        }
286    }
287
288    pub fn get_path_of_leaf(
289        &self,
290        index: usize,
291        full: bool,
292    ) -> Result<Vec<[u8; 32]>, IndexerError> {
293        match &self.merkle_tree {
294            IndexedMerkleTreeVersion::V1(tree) => Ok(tree
295                .get_path_of_leaf(index, full)
296                .map_err(|_| IndexerError::InvalidResponseData)?
297                .to_vec()),
298            IndexedMerkleTreeVersion::V2(tree) => Ok(tree
299                .get_path_of_leaf(index, full)
300                .map_err(|_| IndexerError::InvalidResponseData)?),
301        }
302    }
303
304    pub fn indexed_array_v1(&self) -> Option<&IndexedArray<Poseidon, usize>> {
305        println!(
306            "indexed_array_v2: merkle_tree pubkey: {:?}",
307            self.accounts.merkle_tree
308        );
309        match &self.merkle_tree {
310            IndexedMerkleTreeVersion::V1(_) => Some(&self.indexed_array),
311            _ => None,
312        }
313    }
314
315    pub fn indexed_array_v2(
316        &self,
317    ) -> Option<&light_indexed_array::array::IndexedArray<Poseidon, usize>> {
318        println!(
319            "indexed_array_v2: merkle_tree pubkey: {:?}",
320            self.accounts.merkle_tree
321        );
322        match &self.merkle_tree {
323            IndexedMerkleTreeVersion::V2(tree) => Some(&tree.indexed_array),
324            _ => None,
325        }
326    }
327
328    pub fn update(
329        &mut self,
330        new_low_element: &IndexedElement<usize>,
331        new_element: &IndexedElement<usize>,
332        new_element_next_value: &BigUint,
333    ) -> Result<(), IndexerError> {
334        match &mut self.merkle_tree {
335            IndexedMerkleTreeVersion::V1(tree) => {
336                Ok(tree.update(new_low_element, new_element, new_element_next_value)?)
337            }
338            IndexedMerkleTreeVersion::V2(tree) => {
339                let new_low_element = light_indexed_array::array::IndexedElement::<usize> {
340                    index: new_low_element.index,
341                    value: new_low_element.value.clone(),
342                    next_index: new_low_element.next_index,
343                };
344                let new_element = light_indexed_array::array::IndexedElement::<usize> {
345                    index: new_element.index,
346                    value: new_element.value.clone(),
347                    next_index: new_element.next_index,
348                };
349                tree.update(&new_low_element, &new_element, new_element_next_value)
350                    .unwrap();
351                Ok(())
352            }
353        }
354    }
355}
356
357// TODO: eliminate use of BigInt in favor of BigUint
358pub fn get_non_inclusion_proof_inputs(
359    value: &[u8; 32],
360    merkle_tree: &light_indexed_merkle_tree::reference::IndexedMerkleTree<
361        light_hasher::Poseidon,
362        usize,
363    >,
364    indexed_array: &IndexedArray<light_hasher::Poseidon, usize>,
365) -> NonInclusionMerkleProofInputs {
366    let non_inclusion_proof = merkle_tree
367        .get_non_inclusion_proof(&BigUint::from_be_bytes(value), indexed_array)
368        .unwrap();
369    let proof = non_inclusion_proof
370        .merkle_proof
371        .iter()
372        .map(|x| BigInt::from_be_bytes(x))
373        .collect();
374    NonInclusionMerkleProofInputs {
375        root: BigInt::from_be_bytes(merkle_tree.root().as_slice()),
376        value: BigInt::from_be_bytes(value),
377        leaf_lower_range_value: BigInt::from_be_bytes(&non_inclusion_proof.leaf_lower_range_value),
378        leaf_higher_range_value: BigInt::from_be_bytes(
379            &non_inclusion_proof.leaf_higher_range_value,
380        ),
381        merkle_proof_hashed_indexed_element_leaf: proof,
382        index_hashed_indexed_element_leaf: BigInt::from(non_inclusion_proof.leaf_index),
383        next_index: BigInt::from(non_inclusion_proof.next_index),
384    }
385}
light_program_test/indexer/address_tree.rs

light_program_test/indexer/
address_tree.rs
