openvm_continuations/circuit/deferral/hook/decommit/
trace.rs

1use 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}