openvm_continuations/circuit/root/commit/
trace.rs

1use std::borrow::BorrowMut;
2
3use openvm_circuit::arch::POSEIDON2_WIDTH;
4use openvm_circuit_primitives::encoder::Encoder;
5use openvm_cpu_backend::CpuBackend;
6use openvm_stark_backend::{prover::AirProvingContext, StarkProtocolConfig};
7use openvm_stark_sdk::config::baby_bear_poseidon2::{DIGEST_SIZE, F};
8use p3_field::PrimeCharacteristicRing;
9use p3_matrix::dense::RowMajorMatrix;
10
11use crate::{
12    circuit::{
13        root::commit::{MerkleTreeCols, MAX_ENCODER_DEGREE},
14        subair::generate_cols_from_leaf_children,
15    },
16    utils::digests_to_poseidon2_input,
17};
18
19pub fn generate_proving_ctx<SC: StarkProtocolConfig<F = F>>(
20    user_pvs: Vec<F>,
21) -> (AirProvingContext<CpuBackend<SC>>, Vec<[F; POSEIDON2_WIDTH]>) {
22    let num_user_pvs = user_pvs.len();
23
24    // Each leaf consumes `DIGEST_SIZE` public values, which is padded and hashed before
25    // being inserted into the Merkle tree. We require at least one leaf (so at least one
26    // Poseidon2 hash), and a full binary tree.
27    debug_assert!(num_user_pvs >= DIGEST_SIZE);
28    debug_assert!(num_user_pvs.is_multiple_of(DIGEST_SIZE));
29    debug_assert!((num_user_pvs / DIGEST_SIZE).is_power_of_two());
30
31    // One selector per leaf PV chunk (each leaf consumes 1 digest).
32    let encoder = Encoder::new(num_user_pvs / DIGEST_SIZE, MAX_ENCODER_DEGREE, true);
33
34    let num_leaves = num_user_pvs / DIGEST_SIZE;
35    let leaf_children = user_pvs
36        .chunks_exact(DIGEST_SIZE)
37        .map(|digest| {
38            (
39                digest.try_into().expect("digest sized chunk"),
40                [F::ZERO; DIGEST_SIZE],
41            )
42        })
43        .collect();
44    let rows = generate_cols_from_leaf_children(leaf_children, false);
45    let poseidon2_compress_inputs = rows
46        .iter()
47        .take(rows.len() - 1)
48        .map(|row| digests_to_poseidon2_input(row.left_child, row.right_child))
49        .collect();
50    debug_assert_eq!(rows.len(), 2 * num_leaves);
51
52    let const_width = MerkleTreeCols::<u8>::width();
53    let width = const_width + encoder.width();
54    let mut trace = vec![F::ZERO; rows.len() * width];
55
56    for (row_idx, (chunk, row)) in trace.chunks_mut(width).zip(rows).enumerate() {
57        let cols: &mut MerkleTreeCols<F> = chunk[..const_width].borrow_mut();
58        *cols = row;
59
60        if row_idx < num_leaves {
61            chunk[const_width..].copy_from_slice(
62                encoder
63                    .get_flag_pt(row_idx)
64                    .into_iter()
65                    .map(F::from_u32)
66                    .collect::<Vec<_>>()
67                    .as_slice(),
68            );
69        }
70    }
71
72    let common_main = RowMajorMatrix::new(trace, width);
73    let ctx = AirProvingContext::simple(common_main, user_pvs);
74    (ctx, poseidon2_compress_inputs)
75}