openvm_continuations/circuit/root/commit/
trace.rs1use 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 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 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}