Skip to main content

light_client/indexer/types/
queue.rs

1use std::collections::HashMap;
2
3use super::super::IndexerError;
4
5#[derive(Debug, Clone, PartialEq, Default)]
6pub struct OutputQueueData {
7    pub leaf_indices: Vec<u64>,
8    pub account_hashes: Vec<[u8; 32]>,
9    pub old_leaves: Vec<[u8; 32]>,
10    pub first_queue_index: u64,
11    /// The tree's next_index - where new leaves will be appended
12    pub next_index: u64,
13    /// Pre-computed hash chains per ZKP batch (from on-chain)
14    pub leaves_hash_chains: Vec<[u8; 32]>,
15}
16
17/// V2 Input Queue Data
18#[derive(Debug, Clone, PartialEq, Default)]
19pub struct InputQueueData {
20    pub leaf_indices: Vec<u64>,
21    pub account_hashes: Vec<[u8; 32]>,
22    pub current_leaves: Vec<[u8; 32]>,
23    pub tx_hashes: Vec<[u8; 32]>,
24    /// Pre-computed nullifiers from indexer
25    pub nullifiers: Vec<[u8; 32]>,
26    pub first_queue_index: u64,
27    /// Pre-computed hash chains per ZKP batch (from on-chain)
28    pub leaves_hash_chains: Vec<[u8; 32]>,
29}
30
31/// State queue data with shared tree nodes for output and input queues
32#[derive(Debug, Clone, PartialEq, Default)]
33pub struct StateQueueData {
34    /// Shared deduplicated tree nodes for state queues (output + input)
35    /// node_index encoding: (level << 56) | position
36    pub nodes: Vec<u64>,
37    pub node_hashes: Vec<[u8; 32]>,
38    /// Initial root for the state tree (shared by output and input queues)
39    pub initial_root: [u8; 32],
40    /// Sequence number of the root
41    pub root_seq: u64,
42    /// Output queue data (if requested)
43    pub output_queue: Option<OutputQueueData>,
44    /// Input queue data (if requested)
45    pub input_queue: Option<InputQueueData>,
46}
47
48/// V2 Address Queue Data with deduplicated nodes
49/// Proofs are reconstructed from `nodes`/`node_hashes` using `low_element_indices`
50#[derive(Debug, Clone, PartialEq, Default)]
51pub struct AddressQueueData {
52    pub addresses: Vec<[u8; 32]>,
53    pub low_element_values: Vec<[u8; 32]>,
54    pub low_element_next_values: Vec<[u8; 32]>,
55    pub low_element_indices: Vec<u64>,
56    pub low_element_next_indices: Vec<u64>,
57    /// Deduplicated node indices - encoding: (level << 56) | position
58    pub nodes: Vec<u64>,
59    /// Hashes corresponding to each node index
60    pub node_hashes: Vec<[u8; 32]>,
61    pub initial_root: [u8; 32],
62    pub leaves_hash_chains: Vec<[u8; 32]>,
63    pub subtrees: Vec<[u8; 32]>,
64    /// Pagination offset for the returned queue slice.
65    pub start_index: u64,
66    /// Sparse tree insertion point / next index used to initialize staging trees.
67    pub tree_next_insertion_index: u64,
68    pub root_seq: u64,
69}
70
71impl AddressQueueData {
72    pub const ADDRESS_TREE_HEIGHT: usize =
73        light_prover_client::constants::DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize;
74
75    /// Reconstruct a single merkle proof for a given address index.
76    #[cfg(test)]
77    fn reconstruct_proof<const HEIGHT: usize>(
78        &self,
79        address_idx: usize,
80    ) -> Result<[[u8; 32]; HEIGHT], IndexerError> {
81        let node_lookup = self.build_node_lookup();
82        self.reconstruct_proof_with_lookup::<HEIGHT>(address_idx, &node_lookup)
83    }
84
85    /// Reconstruct a contiguous batch of proofs while reusing a single node lookup table.
86    pub fn reconstruct_proofs<const HEIGHT: usize>(
87        &self,
88        address_range: std::ops::Range<usize>,
89    ) -> Result<Vec<[[u8; 32]; HEIGHT]>, IndexerError> {
90        self.validate_proof_height::<HEIGHT>()?;
91        let available = self.proof_count();
92        if address_range.start > address_range.end {
93            return Err(IndexerError::InvalidParameters(format!(
94                "invalid address proof range {}..{}",
95                address_range.start, address_range.end
96            )));
97        }
98        if address_range.end > available {
99            return Err(IndexerError::InvalidParameters(format!(
100                "address proof range {}..{} exceeds available proofs {}",
101                address_range.start, address_range.end, available
102            )));
103        }
104        let node_lookup = self.build_node_lookup();
105        let mut proofs = Vec::with_capacity(address_range.len());
106
107        for address_idx in address_range {
108            proofs.push(self.reconstruct_proof_with_lookup::<HEIGHT>(address_idx, &node_lookup)?);
109        }
110
111        Ok(proofs)
112    }
113
114    /// Reconstruct all proofs for all addresses
115    pub fn reconstruct_all_proofs<const HEIGHT: usize>(
116        &self,
117    ) -> Result<Vec<[[u8; 32]; HEIGHT]>, IndexerError> {
118        self.reconstruct_proofs::<HEIGHT>(0..self.addresses.len())
119    }
120
121    fn build_node_lookup(&self) -> HashMap<u64, usize> {
122        let mut lookup = HashMap::with_capacity(self.nodes.len());
123        for (idx, node) in self.nodes.iter().copied().enumerate() {
124            lookup.entry(node).or_insert(idx);
125        }
126        lookup
127    }
128
129    fn proof_count(&self) -> usize {
130        self.addresses.len().min(self.low_element_indices.len())
131    }
132
133    fn reconstruct_proof_with_lookup<const HEIGHT: usize>(
134        &self,
135        address_idx: usize,
136        node_lookup: &HashMap<u64, usize>,
137    ) -> Result<[[u8; 32]; HEIGHT], IndexerError> {
138        let leaf_index = *self.low_element_indices.get(address_idx).ok_or_else(|| {
139            IndexerError::MissingResult {
140                context: "reconstruct_proof".to_string(),
141                message: format!(
142                    "address_idx {} out of bounds for low_element_indices (len {})",
143                    address_idx,
144                    self.low_element_indices.len(),
145                ),
146            }
147        })?;
148        let mut proof = [[0u8; 32]; HEIGHT];
149        let mut pos = leaf_index;
150
151        for (level, proof_element) in proof.iter_mut().enumerate() {
152            let sibling_pos = if pos.is_multiple_of(2) {
153                pos + 1
154            } else {
155                pos - 1
156            };
157            let sibling_idx = Self::encode_node_index(level, sibling_pos);
158            let hash_idx = node_lookup.get(&sibling_idx).copied().ok_or_else(|| {
159                IndexerError::MissingResult {
160                    context: "reconstruct_proof".to_string(),
161                    message: format!(
162                        "Missing proof node at level {} position {} (encoded: {})",
163                        level, sibling_pos, sibling_idx
164                    ),
165                }
166            })?;
167            let hash =
168                self.node_hashes
169                    .get(hash_idx)
170                    .ok_or_else(|| IndexerError::MissingResult {
171                        context: "reconstruct_proof".to_string(),
172                        message: format!(
173                            "node_hashes index {} out of bounds (len {})",
174                            hash_idx,
175                            self.node_hashes.len(),
176                        ),
177                    })?;
178            *proof_element = *hash;
179            pos /= 2;
180        }
181
182        Ok(proof)
183    }
184
185    /// Encode node index: (level << 56) | position
186    #[inline]
187    fn encode_node_index(level: usize, position: u64) -> u64 {
188        ((level as u64) << 56) | position
189    }
190
191    fn validate_proof_height<const HEIGHT: usize>(&self) -> Result<(), IndexerError> {
192        if HEIGHT == Self::ADDRESS_TREE_HEIGHT {
193            return Ok(());
194        }
195
196        Err(IndexerError::InvalidParameters(format!(
197            "address queue proofs require HEIGHT={} but got HEIGHT={}",
198            Self::ADDRESS_TREE_HEIGHT,
199            HEIGHT
200        )))
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use std::{collections::BTreeMap, hint::black_box, time::Instant};
207
208    use super::AddressQueueData;
209
210    fn hash_from_node(node_index: u64) -> [u8; 32] {
211        let mut hash = [0u8; 32];
212        hash[..8].copy_from_slice(&node_index.to_le_bytes());
213        hash[8..16].copy_from_slice(&node_index.rotate_left(17).to_le_bytes());
214        hash[16..24].copy_from_slice(&node_index.rotate_right(9).to_le_bytes());
215        hash[24..32].copy_from_slice(&(node_index ^ 0xA5A5_A5A5_A5A5_A5A5).to_le_bytes());
216        hash
217    }
218
219    fn build_queue_data<const HEIGHT: usize>(num_addresses: usize) -> AddressQueueData {
220        let low_element_indices = (0..num_addresses)
221            .map(|i| (i as u64).saturating_mul(2))
222            .collect::<Vec<_>>();
223        let mut nodes = BTreeMap::new();
224
225        for &leaf_index in &low_element_indices {
226            let mut pos = leaf_index;
227            for level in 0..HEIGHT {
228                let sibling_pos = if pos.is_multiple_of(2) {
229                    pos + 1
230                } else {
231                    pos - 1
232                };
233                let node_index = ((level as u64) << 56) | sibling_pos;
234                nodes
235                    .entry(node_index)
236                    .or_insert_with(|| hash_from_node(node_index));
237                pos /= 2;
238            }
239        }
240
241        let (nodes, node_hashes): (Vec<_>, Vec<_>) = nodes.into_iter().unzip();
242
243        AddressQueueData {
244            addresses: vec![[0u8; 32]; num_addresses],
245            low_element_values: vec![[1u8; 32]; num_addresses],
246            low_element_next_values: vec![[2u8; 32]; num_addresses],
247            low_element_indices,
248            low_element_next_indices: (0..num_addresses).map(|i| (i as u64) + 1).collect(),
249            nodes,
250            node_hashes,
251            initial_root: [9u8; 32],
252            leaves_hash_chains: vec![[3u8; 32]; num_addresses.max(1)],
253            subtrees: vec![[4u8; 32]; HEIGHT],
254            start_index: 0,
255            tree_next_insertion_index: 0,
256            root_seq: 0,
257        }
258    }
259
260    #[test]
261    fn batched_reconstruction_matches_individual_reconstruction() {
262        let queue = build_queue_data::<40>(128);
263
264        let expected = (0..queue.addresses.len())
265            .map(|i| queue.reconstruct_proof::<40>(i).unwrap())
266            .collect::<Vec<_>>();
267        let actual = queue
268            .reconstruct_proofs::<40>(0..queue.addresses.len())
269            .unwrap();
270
271        assert_eq!(actual, expected);
272    }
273
274    #[test]
275    #[ignore = "profiling helper"]
276    fn profile_reconstruct_proofs_batch() {
277        const HEIGHT: usize = 40;
278        const NUM_ADDRESSES: usize = 2_048;
279        const ITERS: usize = 25;
280
281        let queue = build_queue_data::<HEIGHT>(NUM_ADDRESSES);
282
283        let baseline_start = Instant::now();
284        for _ in 0..ITERS {
285            let proofs = (0..queue.addresses.len())
286                .map(|i| queue.reconstruct_proof::<HEIGHT>(i).unwrap())
287                .collect::<Vec<_>>();
288            black_box(proofs);
289        }
290        let baseline = baseline_start.elapsed();
291
292        let batched_start = Instant::now();
293        for _ in 0..ITERS {
294            black_box(
295                queue
296                    .reconstruct_proofs::<HEIGHT>(0..queue.addresses.len())
297                    .unwrap(),
298            );
299        }
300        let batched = batched_start.elapsed();
301
302        println!(
303            "queue reconstruction profile: addresses={}, height={}, iters={}, individual={:?}, batched={:?}, speedup={:.2}x",
304            NUM_ADDRESSES,
305            HEIGHT,
306            ITERS,
307            baseline,
308            batched,
309            baseline.as_secs_f64() / batched.as_secs_f64(),
310        );
311    }
312}
313
314/// V2 Queue Elements Result with deduplicated node data
315#[derive(Debug, Clone, PartialEq, Default)]
316pub struct QueueElementsResult {
317    pub state_queue: Option<StateQueueData>,
318    pub address_queue: Option<AddressQueueData>,
319}
320
321/// Queue leaf index item returned by getQueueLeafIndices
322#[derive(Debug, Clone, PartialEq, Default)]
323pub struct QueueLeafIndex {
324    pub hash: [u8; 32],
325    pub leaf_index: u64,
326    pub queue_index: u64,
327}