openvm_recursion_circuit/transcript/merkle_verify/
trace.rs

1use core::borrow::BorrowMut;
2
3use itertools::Itertools;
4pub use openvm_poseidon2_air::POSEIDON2_WIDTH;
5use openvm_stark_backend::{keygen::types::MultiStarkVerifyingKey, proof::Proof, SystemParams};
6use openvm_stark_sdk::config::baby_bear_poseidon2::{
7    poseidon2_compress_with_capacity, BabyBearPoseidon2Config, CHUNK, DIGEST_SIZE, F,
8};
9use p3_field::{PrimeCharacteristicRing, PrimeField32};
10
11use crate::{
12    system::Preflight,
13    transcript::merkle_verify::air::{compute_cum_sum, MerkleVerifyCols, MerkleVerifyLog},
14};
15
16#[tracing::instrument(level = "trace", skip_all)]
17pub fn generate_trace(
18    mvk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
19    proofs: &[Proof<BabyBearPoseidon2Config>],
20    preflights: &[Preflight],
21    params: &SystemParams,
22    required_height: Option<usize>,
23) -> Option<(Vec<F>, Vec<[F; POSEIDON2_WIDTH]>)> {
24    let k = params.k_whir();
25    let width = MerkleVerifyCols::<F>::width();
26    let num_leaves: usize = 1 << k;
27    let mut poseidon2_compress_inputs = vec![];
28    let logs_per_proof = preflights
29        .iter()
30        .map(|preflight| build_merkle_logs(preflight, params))
31        .collect_vec();
32    // vec of vec, index by proof and then merkle_proof
33    let num_rows_per_proof_per_merkle = preflights
34        .iter()
35        .enumerate()
36        .map(|(proof_idx, _)| {
37            logs_per_proof[proof_idx]
38                .iter()
39                .map(|log| log.depth + num_leaves) // (d + 1) + (num_leaves - 1)
40                .collect_vec()
41        })
42        .collect_vec();
43    // cum sum of proof
44    let num_rows_cum_sums: Vec<usize> =
45        compute_cum_sum(&num_rows_per_proof_per_merkle, |x| x.iter().sum::<usize>());
46    // for each proof, cum sum of merkle proofs
47    let num_rows_cum_sums_within_proof: Vec<Vec<usize>> = num_rows_per_proof_per_merkle
48        .iter()
49        .map(|x| compute_cum_sum(x, |y| *y))
50        .collect();
51    let num_valid_rows = *num_rows_cum_sums.last().unwrap();
52    let height = if let Some(height) = required_height {
53        if num_valid_rows > height {
54            return None;
55        }
56        height
57    } else {
58        num_valid_rows.next_power_of_two()
59    };
60    let mut trace = vec![F::ZERO; height * width];
61    let mut cur_hash = [F::ZERO; DIGEST_SIZE];
62    let mut cur_idx = 0;
63
64    // layer 0: num_leaves = 2^k hashes
65    // layer 1: half the hashes
66    // ... layer k - 1: 2 hashes
67    // layer k: final hash --> into merkle proof
68    let mut leaf_tree = vec![vec![[F::ZERO; DIGEST_SIZE]; num_leaves]; k + 1];
69
70    for (row_idx, row) in trace.chunks_mut(width).take(num_valid_rows).enumerate() {
71        let proof_idx = num_rows_cum_sums.partition_point(|&x| x <= row_idx);
72        let preflight = &preflights[proof_idx];
73        let proof = &proofs[proof_idx];
74        let idx_in_proof = if proof_idx == 0 {
75            row_idx
76        } else {
77            row_idx - num_rows_cum_sums[proof_idx - 1]
78        };
79        let merkle_proof_idx =
80            num_rows_cum_sums_within_proof[proof_idx].partition_point(|&x| x <= idx_in_proof);
81        // i: the index [0, total_depth) within a merkle proof
82        let i = if merkle_proof_idx == 0 {
83            idx_in_proof
84        } else {
85            idx_in_proof - num_rows_cum_sums_within_proof[proof_idx][merkle_proof_idx - 1]
86        };
87
88        let &MerkleVerifyLog {
89            merkle_idx,
90            depth,
91            query_idx,
92            commit_major,
93            commit_minor,
94        } = &logs_per_proof[proof_idx][merkle_proof_idx];
95
96        if i == 0 {
97            if commit_major == 0 {
98                for (coset_idx, states) in preflight.initial_row_states[commit_minor][query_idx]
99                    .iter()
100                    .enumerate()
101                {
102                    let post_state = states.last().unwrap();
103                    leaf_tree[0][coset_idx].copy_from_slice(&post_state[..CHUNK]);
104                }
105            } else {
106                for (coset_idx, state) in preflight.codeword_states[commit_major - 1][query_idx]
107                    .iter()
108                    .enumerate()
109                {
110                    leaf_tree[0][coset_idx].copy_from_slice(&state[..CHUNK]);
111                }
112            }
113        }
114
115        let mut stacking_commits = vec![proof.common_main_commit];
116        for (air_id, data) in &preflight.proof_shape.sorted_trace_vdata {
117            stacking_commits.extend(
118                mvk.inner.per_air[*air_id]
119                    .preprocessed_data
120                    .as_ref()
121                    .into_iter()
122                    .map(|pdata| pdata.commit)
123                    .chain(data.cached_commitments.iter().cloned()),
124            );
125        }
126
127        let cols: &mut MerkleVerifyCols<F> = row.borrow_mut();
128        // determine the layer and offset in the leaf_tree
129        // 0th layer: [0, num_leaves / 2)
130        // 1st layer: [num_leaves / 2, num_leaves / 2 + num_leaves / 4)
131        // kth layer: 1 final hash (that goes into merkle proof)
132        let combination_indices = compute_combination_indices(k, i);
133
134        cols.is_valid = F::ONE;
135        cols.proof_idx = F::from_usize(proof_idx);
136        cols.commit_major = F::from_usize(commit_major);
137        cols.commit_minor = F::from_usize(commit_minor);
138        cols.total_depth = F::from_usize(depth + k + 1);
139
140        if i == depth + num_leaves - 1 {
141            cols.is_last_merkle = F::ONE;
142        }
143
144        if let Some(combination_indices) = combination_indices {
145            // combining leaves part
146            cols.left =
147                leaf_tree[combination_indices.source_layer][combination_indices.left_source_index];
148            cols.right =
149                leaf_tree[combination_indices.source_layer][combination_indices.right_source_index];
150            let output = poseidon2_compress_with_capacity(cols.left, cols.right).0;
151            leaf_tree[combination_indices.result_layer][combination_indices.result_index] = output;
152            cols.compression_output = output;
153
154            cols.merkle_idx_bit_src = F::from_usize(merkle_idx); // const idx for leaves part
155            cols.current_idx_bit_src = F::from_usize(merkle_idx);
156            cols.height = F::from_usize(combination_indices.source_layer);
157            cols.is_last_leaf = F::from_bool(combination_indices.source_layer + 1 == k);
158            cols.recv_flag = F::TWO;
159
160            let mut input_state = [F::ZERO; POSEIDON2_WIDTH];
161            input_state[..DIGEST_SIZE].copy_from_slice(&cols.left);
162            input_state[DIGEST_SIZE..].copy_from_slice(&cols.right);
163            poseidon2_compress_inputs.push(input_state);
164
165            cols.is_combining_leaves = F::ONE;
166            cols.leaf_sub_idx = F::from_usize(combination_indices.result_index);
167        } else {
168            // merkle proof part
169            debug_assert!(i >= num_leaves - 1);
170            if i == num_leaves - 1 {
171                // The first row of the merkle proof part, initialize cur_hash and cur_idx
172                cur_hash = leaf_tree[k][0];
173                cur_idx = merkle_idx;
174            }
175            let pos = i + 1 - num_leaves;
176            let is_last = pos == depth;
177            let whir_proof = &proof.whir_proof;
178            let sibling = match (commit_major, is_last) {
179                (0, true) => stacking_commits[commit_minor],
180                (0, false) => whir_proof.initial_round_merkle_proofs[commit_minor][query_idx][pos],
181                (idx, true) => whir_proof.codeword_commits[idx - 1],
182                (idx, false) => whir_proof.codeword_merkle_proofs[idx - 1][query_idx][pos],
183            };
184
185            if cur_idx % 2 == 0 {
186                cols.left = cur_hash;
187                cols.right = sibling;
188                cols.recv_flag = F::ZERO;
189            } else {
190                cols.left = sibling;
191                cols.right = cur_hash;
192                cols.recv_flag = F::ONE;
193            }
194
195            let output = poseidon2_compress_with_capacity(cols.left, cols.right).0;
196            cols.compression_output = output;
197            let mut input_state = [F::ZERO; POSEIDON2_WIDTH];
198            input_state[..DIGEST_SIZE].copy_from_slice(&cols.left);
199            input_state[DIGEST_SIZE..].copy_from_slice(&cols.right);
200            poseidon2_compress_inputs.push(input_state);
201
202            cur_hash = output;
203            cur_idx /= 2;
204
205            cols.merkle_idx_bit_src = F::from_usize(merkle_idx);
206            cols.current_idx_bit_src = F::from_usize(cur_idx);
207            cols.height = F::from_usize(i + 1 - num_leaves + k);
208            cols.is_combining_leaves = F::ZERO;
209        }
210    }
211
212    Some((trace, poseidon2_compress_inputs))
213}
214
215fn build_merkle_logs(preflight: &Preflight, params: &SystemParams) -> Vec<MerkleVerifyLog> {
216    let num_whir_rounds = params.num_whir_rounds();
217    let mut logs = Vec::new();
218    let mut log_rs_domain_size = params.l_skip + params.n_stack + params.log_blowup;
219    let mut query_offset = 0;
220    for round_idx in 0..num_whir_rounds {
221        let num_queries = params.whir.rounds[round_idx].num_queries;
222        for query_idx in 0..num_queries {
223            let sample = preflight.whir.queries[query_offset];
224            let merkle_idx = sample.as_canonical_u32() as usize;
225            let depth = log_rs_domain_size - params.k_whir();
226
227            if round_idx == 0 {
228                for commit_minor in 0..preflight.initial_row_states.len() {
229                    logs.push(MerkleVerifyLog {
230                        merkle_idx,
231                        depth,
232                        query_idx,
233                        commit_major: 0,
234                        commit_minor,
235                    });
236                }
237            } else {
238                logs.push(MerkleVerifyLog {
239                    merkle_idx,
240                    depth,
241                    query_idx,
242                    commit_major: round_idx,
243                    commit_minor: 0,
244                });
245            }
246
247            query_offset += 1;
248        }
249        log_rs_domain_size -= 1;
250    }
251
252    logs
253}
254
255// Represents the necessary indices for a single combining operation:
256// combining arr[source_layer][left_source_index] and arr[source_layer][right_source_index]
257// to produce arr[result_layer][result_index].
258#[derive(Debug, PartialEq)]
259pub struct CombinationIndices {
260    /// The index of the array (vector) containing the two elements to be combined (`arr[j]`).
261    pub source_layer: usize,
262    /// The index of the left element in the source layer (`arr[j][2*c]`).
263    pub left_source_index: usize,
264    /// The index of the right element in the source layer (`arr[j][2*c + 1]`).
265    pub right_source_index: usize,
266    /// The index of the array (vector) where the result is stored (`arr[j+1]`).
267    pub result_layer: usize,
268    /// The index of the result element in the result layer (`arr[j+1][c]`).
269    pub result_index: usize,
270}
271
272/// Calculates the layer and indices for the i-th combining operation in a complete binary tree
273/// with 2^k leaves.
274///
275/// # Arguments
276/// * `k` - The power defining the number of leaves (2^k). Assumed constant for the tree structure.
277/// * `i` - The zero-based, overall index of the combining operation (0 <= i < 2^k - 1).
278///
279/// # Returns
280/// An `Option<CombinationIndices>` containing the location of the operation, or `None` if `i` is
281/// out of bounds.
282pub fn compute_combination_indices(k: usize, i: usize) -> Option<CombinationIndices> {
283    if k == 0 {
284        // A tree with 2^0 = 1 leaf has no combining operations.
285        return None;
286    }
287
288    // The total number of combining operations in a complete binary tree is 2^k - 1.
289    // 1 << k is equivalent to 2^k.
290    let total_operations = (1 << k) - 1;
291
292    if i >= total_operations {
293        return None; // Index is out of bounds
294    }
295
296    let mut current_index = i;
297    let mut source_layer = 0;
298
299    // The number of combining operations at source_layer `j` (which combines elements from
300    // arr[j] into arr[j+1]) is 2^(k - (j + 1)).
301
302    // We iterate through layers, subtracting the number of operations in that layer
303    // until the `current_index` falls within the range of the current layer.
304    while source_layer < k {
305        // Calculate C_j = 2^(k - (j + 1))
306        let exponent = k - (source_layer + 1);
307        let combinations_in_layer = 1 << exponent; // 2^exponent
308
309        if current_index < combinations_in_layer {
310            // Found the correct layer!
311            let index_within_layer = current_index;
312
313            return Some(CombinationIndices {
314                source_layer,
315                // The two source elements are always at 2*c and 2*c + 1
316                left_source_index: 2 * index_within_layer,
317                right_source_index: 2 * index_within_layer + 1,
318                // The result is stored in the next layer, at index c
319                result_layer: source_layer + 1,
320                result_index: index_within_layer,
321            });
322        }
323
324        // Subtract the count for the current layer and move up to the next layer
325        current_index -= combinations_in_layer;
326        source_layer += 1;
327    }
328
329    // This line should technically be unreachable due to the initial i < total_operations check.
330    None
331}
332
333#[cfg(feature = "cuda")]
334pub mod cuda {
335    use itertools::Itertools;
336    use openvm_cuda_backend::{base::DeviceMatrix, prelude::F};
337    use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
338
339    use super::*;
340    use crate::{
341        cuda::{
342            preflight::PreflightGpu, proof::ProofGpu, types::MerkleVerifyRecord,
343            vk::VerifyingKeyGpu,
344        },
345        transcript::{cuda_abi, cuda_tracegen::TranscriptBlob},
346    };
347
348    const MAX_SUPPORTED_K: usize = 4;
349
350    pub(crate) struct MerkleVerifyBlob {
351        pub records: Vec<MerkleVerifyRecord>,
352        pub leaf_hashes: Vec<F>,
353        pub sibling_hashes: Vec<F>,
354        pub proof_row_starts: Vec<usize>,
355        pub total_rows: usize,
356        pub num_leaves: usize,
357        pub k: usize,
358        pub num_proofs: usize,
359        pub poseidon2_buffer_offset: usize,
360    }
361
362    impl MerkleVerifyBlob {
363        pub fn new(
364            child_vk: &VerifyingKeyGpu,
365            proofs: &[ProofGpu],
366            preflights: &[PreflightGpu],
367            poseidon2_buffer_offset: usize,
368        ) -> Self {
369            assert_eq!(proofs.len(), preflights.len());
370            let k = child_vk.system_params.k_whir();
371            assert!(
372                k <= MAX_SUPPORTED_K,
373                "unsupported k={} for CUDA merkle verify (max {})",
374                k,
375                MAX_SUPPORTED_K
376            );
377            let num_leaves = 1 << k;
378
379            let mut records = Vec::new();
380            let mut leaf_hashes = Vec::new();
381            let mut sibling_hashes = Vec::new();
382            let mut proof_row_starts = Vec::with_capacity(proofs.len());
383
384            let stacking_commits_per_proof = proofs
385                .iter()
386                .zip(preflights)
387                .map(|(proof, preflight)| {
388                    build_stacking_commits(&child_vk.cpu, &proof.cpu, &preflight.cpu)
389                })
390                .collect_vec();
391
392            let logs_per_proof: Vec<_> = preflights
393                .iter()
394                .map(|preflight| build_merkle_logs(&preflight.cpu, &child_vk.system_params))
395                .collect();
396
397            let mut total_rows = 0usize;
398            for (proof_idx, proof) in proofs.iter().enumerate() {
399                let preflight = &preflights[proof_idx].cpu;
400                proof_row_starts.push(total_rows);
401                for (merkle_proof_idx, log) in logs_per_proof[proof_idx].iter().enumerate() {
402                    let num_rows = log.depth + num_leaves;
403                    let leaf_offset = leaf_hashes.len();
404                    for coset_idx in 0..num_leaves {
405                        if log.commit_major == 0 {
406                            let states = &preflight.initial_row_states[log.commit_minor]
407                                [log.query_idx][coset_idx];
408                            let post_state = states.last().unwrap();
409                            leaf_hashes.extend_from_slice(&post_state[..CHUNK]);
410                        } else {
411                            let state = &preflight.codeword_states[log.commit_major - 1]
412                                [log.query_idx][coset_idx];
413                            leaf_hashes.extend_from_slice(&state[..CHUNK]);
414                        }
415                    }
416                    let path =
417                        build_merkle_path(log, &proof.cpu, &stacking_commits_per_proof[proof_idx]);
418                    let sibling_offset = sibling_hashes.len();
419                    for sibling in path {
420                        sibling_hashes.extend_from_slice(&sibling);
421                    }
422                    records.push(MerkleVerifyRecord {
423                        proof_idx: proof_idx as u16,
424                        merkle_proof_idx: merkle_proof_idx as u16,
425                        start_row: total_rows as u32,
426                        num_rows: num_rows as u32,
427                        depth: log.depth as u16,
428                        merkle_idx: log.merkle_idx as u32,
429                        commit_major: log.commit_major as u16,
430                        commit_minor: log.commit_minor as u16,
431                        leaf_hash_offset: leaf_offset as u32,
432                        siblings_offset: sibling_offset as u32,
433                    });
434                    total_rows += num_rows;
435                }
436            }
437
438            Self {
439                records,
440                leaf_hashes,
441                sibling_hashes,
442                proof_row_starts,
443                total_rows,
444                num_leaves,
445                k,
446                num_proofs: proofs.len(),
447                poseidon2_buffer_offset,
448            }
449        }
450    }
451
452    #[tracing::instrument(level = "trace", skip_all)]
453    pub(crate) fn generate_trace(
454        blob: &TranscriptBlob,
455        device_ctx: &GpuDeviceCtx,
456        required_height: Option<usize>,
457    ) -> Option<DeviceMatrix<F>> {
458        let merkle_blob = &blob.merkle_verify_blob;
459        let trace_width = MerkleVerifyCols::<F>::width();
460        let trace_height = if let Some(height) = required_height {
461            if height == 0 || merkle_blob.total_rows > height {
462                return None;
463            }
464            height
465        } else {
466            merkle_blob.total_rows.next_power_of_two().max(1)
467        };
468        let mut trace = DeviceMatrix::with_capacity_on(trace_height, trace_width, device_ctx);
469        trace.buffer().fill_zero_on(device_ctx).unwrap();
470
471        if merkle_blob.total_rows == 0 || merkle_blob.records.is_empty() {
472            return Some(trace);
473        }
474
475        let d_records = merkle_blob.records.to_device_on(device_ctx).unwrap();
476        let d_leaf_hashes = merkle_blob.leaf_hashes.to_device_on(device_ctx).unwrap();
477        let d_siblings = merkle_blob.sibling_hashes.to_device_on(device_ctx).unwrap();
478        let d_proof_row_starts: DeviceBuffer<usize> = merkle_blob
479            .proof_row_starts
480            .to_device_on(device_ctx)
481            .unwrap();
482        let scratch_len = merkle_blob.records.len() * merkle_blob.num_leaves * DIGEST_SIZE;
483        let d_leaf_scratch = DeviceBuffer::<F>::with_capacity_on(scratch_len, device_ctx);
484        d_leaf_scratch.fill_zero_on(device_ctx).unwrap();
485
486        unsafe {
487            cuda_abi::merkle_verify_tracegen(
488                &mut trace,
489                &d_records,
490                &d_leaf_hashes,
491                &d_siblings,
492                merkle_blob.num_leaves,
493                merkle_blob.k,
494                &blob.poseidon2_buffer,
495                merkle_blob.poseidon2_buffer_offset,
496                merkle_blob.total_rows,
497                &d_proof_row_starts,
498                merkle_blob.num_proofs,
499                &d_leaf_scratch,
500                device_ctx.stream.as_raw(),
501            )
502            .expect("failed to launch merkle verify tracegen");
503        }
504
505        Some(trace)
506    }
507
508    fn build_stacking_commits(
509        mvk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
510        proof: &Proof<BabyBearPoseidon2Config>,
511        preflight: &Preflight,
512    ) -> Vec<[F; DIGEST_SIZE]> {
513        let mut commits = vec![proof.common_main_commit];
514        for (air_id, data) in &preflight.proof_shape.sorted_trace_vdata {
515            if let Some(pdata) = mvk.inner.per_air[*air_id].preprocessed_data.as_ref() {
516                commits.push(pdata.commit);
517            }
518            commits.extend(data.cached_commitments.iter());
519        }
520        commits
521    }
522
523    fn build_merkle_path(
524        log: &MerkleVerifyLog,
525        proof: &Proof<BabyBearPoseidon2Config>,
526        stacking_commits: &[[F; DIGEST_SIZE]],
527    ) -> Vec<[F; DIGEST_SIZE]> {
528        let whir_proof = &proof.whir_proof;
529        (0..=log.depth)
530            .map(|pos| {
531                let is_last = pos == log.depth;
532                match (log.commit_major, is_last) {
533                    (0, true) => stacking_commits[log.commit_minor],
534                    (0, false) => {
535                        whir_proof.initial_round_merkle_proofs[log.commit_minor][log.query_idx][pos]
536                    }
537                    (idx, true) => whir_proof.codeword_commits[idx - 1],
538                    (idx, false) => whir_proof.codeword_merkle_proofs[idx - 1][log.query_idx][pos],
539                }
540            })
541            .collect()
542    }
543}
openvm_recursion_circuit/transcript/merkle_verify/trace.rs

openvm_recursion_circuit/transcript/merkle_verify/
trace.rs
