light_program_test/indexer/
address_tree.rs

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

light_program_test/indexer/
address_tree.rs
