openvm_continuations/circuit/deferral/hook/decommit/
trace.rs1use std::borrow::BorrowMut;
2
3use openvm_circuit::arch::POSEIDON2_WIDTH;
4use openvm_cpu_backend::CpuBackend;
5use openvm_stark_backend::prover::AirProvingContext;
6use openvm_stark_sdk::config::baby_bear_poseidon2::{
7 poseidon2_compress_with_capacity, BabyBearPoseidon2Config, DIGEST_SIZE, F,
8};
9use p3_field::PrimeCharacteristicRing;
10use p3_matrix::dense::RowMajorMatrix;
11
12use crate::{
13 circuit::{
14 deferral::{hook::decommit::air::MerkleDecommitCols, DEF_INTERNAL_TAG, DEF_LEAF_TAG},
15 subair::{generate_cols_from_leaf_children, MerkleTreeCols},
16 },
17 utils::digests_to_poseidon2_input,
18};
19
20pub type IoCommit = ([F; DIGEST_SIZE], [F; DIGEST_SIZE]);
21
22pub struct MerkleDecommitTraceCtx {
23 pub proving_ctx: AirProvingContext<CpuBackend<BabyBearPoseidon2Config>>,
24 pub poseidon2_inputs: Vec<[F; POSEIDON2_WIDTH]>,
25 pub io_commits: Vec<IoCommit>,
26 pub merkle_root: [F; DIGEST_SIZE],
27}
28
29pub fn generate_proving_ctx(
30 leaf_children: Vec<IoCommit>,
31 num_real_leaves: usize,
32) -> MerkleDecommitTraceCtx {
33 assert!(
34 !leaf_children.is_empty(),
35 "deferral hook Merkle decommit requires at least one leaf"
36 );
37 assert!(
38 leaf_children.len().is_power_of_two(),
39 "deferral hook Merkle decommit requires a power-of-two number of leaves"
40 );
41 assert!(
42 (1..=leaf_children.len()).contains(&num_real_leaves),
43 "deferral hook Merkle decommit requires 1 <= num_real_leaves <= num_leaves"
44 );
45 let merkle_rows: Vec<MerkleTreeCols<F>> = generate_cols_from_leaf_children(leaf_children, true);
46 let width = MerkleDecommitCols::<u8>::width();
47 let height = merkle_rows.len();
48 let num_rows_f = F::from_usize(height);
49 let mut trace = vec![F::ZERO; height * width];
50 let mut poseidon2_inputs = Vec::with_capacity(2 * height.saturating_sub(1));
51 let mut io_commits = Vec::with_capacity(num_real_leaves);
52
53 for (row_idx, merkle_row) in merkle_rows.iter().copied().enumerate() {
54 let cols: &mut MerkleDecommitCols<F> =
55 trace[row_idx * width..(row_idx + 1) * width].borrow_mut();
56 cols.merkle_tree_cols = merkle_row;
57 cols.num_rows = num_rows_f;
58
59 let is_leaf = merkle_row.receive_type == F::ONE;
60 let should_send_commit = is_leaf && row_idx < num_real_leaves;
61 cols.send_commits = F::from_bool(should_send_commit);
62
63 if merkle_row.send_type != F::ZERO {
64 let tag = if is_leaf {
65 DEF_LEAF_TAG.map(F::from_u8)
66 } else {
67 DEF_INTERNAL_TAG.map(F::from_u8)
68 };
69 cols.tagged_left_child = poseidon2_compress_with_capacity(tag, merkle_row.left_child).0;
70 poseidon2_inputs.push(digests_to_poseidon2_input(tag, merkle_row.left_child));
71 poseidon2_inputs.push(digests_to_poseidon2_input(
72 cols.tagged_left_child,
73 merkle_row.right_child,
74 ));
75 }
76 if should_send_commit {
77 io_commits.push((merkle_row.left_child, merkle_row.right_child));
78 }
79 }
80
81 let merkle_root = merkle_rows[height - 2].parent;
82
83 MerkleDecommitTraceCtx {
84 proving_ctx: AirProvingContext::simple_no_pis(RowMajorMatrix::new(trace, width)),
85 poseidon2_inputs,
86 io_commits,
87 merkle_root,
88 }
89}