openvm_verify_stark_circuit/commit/
air.rs

1use std::{array::from_fn, borrow::Borrow};
2
3use openvm_circuit_primitives::{ColumnsAir, SubAir};
4use openvm_continuations::circuit::subair::{
5    MerkleRootBus, MerkleTreeCols, MerkleTreeInternalBus, MerkleTreeSubAir,
6};
7use openvm_recursion_circuit::{bus::Poseidon2CompressBus, utils::assert_zeros};
8use openvm_stark_backend::{
9    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
10};
11use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
12use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
13use p3_field::PrimeCharacteristicRing;
14use p3_matrix::Matrix;
15
16use crate::{
17    bus::{OutputValBus, OutputValMessage},
18    output::VALS_IN_DIGEST,
19};
20
21/**
22 * Builds a binary Merkle tree to decommit and expose or emit the raw user public values.
23 * Constrains that:
24 * - leaf nodes read single digests, compress with zeros, and compute leaf hashes
25 * - internal nodes receive children from an internal permutation bus
26 * - root commitment is sent to `MerkleRootBus`
27 * - leaf payload is sent on `OutputValBus` starting at OUTPUT_USER_PVS_START_IDX
28 */
29#[derive(ColumnsAir)]
30#[columns_via(MerkleTreeCols<u8>)]
31pub struct UserPvsCommitValuesAir {
32    pub subair: MerkleTreeSubAir,
33    pub output_val_bus: OutputValBus,
34    num_user_pvs: usize,
35}
36
37impl UserPvsCommitValuesAir {
38    pub fn new(
39        poseidon2_compress_bus: Poseidon2CompressBus,
40        merkle_root_bus: MerkleRootBus,
41        merkle_tree_internal_bus: MerkleTreeInternalBus,
42        output_val_bus: OutputValBus,
43        num_user_pvs: usize,
44    ) -> Self {
45        // Each leaf consumes `DIGEST_SIZE` public values, which are compressed with zeros
46        // to compute the leaf hash. We require at least one leaf, and a full binary tree.
47        assert!(num_user_pvs >= DIGEST_SIZE);
48        assert!(num_user_pvs.is_multiple_of(DIGEST_SIZE));
49        assert!((num_user_pvs / DIGEST_SIZE).is_power_of_two());
50
51        UserPvsCommitValuesAir {
52            subair: MerkleTreeSubAir::new(
53                poseidon2_compress_bus,
54                merkle_root_bus,
55                merkle_tree_internal_bus,
56                0,
57                false,
58            ),
59            output_val_bus,
60            num_user_pvs,
61        }
62    }
63}
64
65impl<F> BaseAir<F> for UserPvsCommitValuesAir {
66    fn width(&self) -> usize {
67        MerkleTreeCols::<u8>::width()
68    }
69}
70impl<F> BaseAirWithPublicValues<F> for UserPvsCommitValuesAir {}
71impl<F> PartitionedBaseAir<F> for UserPvsCommitValuesAir {}
72
73impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB>
74    for UserPvsCommitValuesAir
75{
76    fn eval(&self, builder: &mut AB) {
77        let main = builder.main();
78        let (local, next) = (
79            main.row_slice(0).expect("window should have two elements"),
80            main.row_slice(1).expect("window should have two elements"),
81        );
82
83        let local: &MerkleTreeCols<AB::Var> = (*local).borrow();
84        let next: &MerkleTreeCols<AB::Var> = (*next).borrow();
85
86        let num_rows = AB::F::from_usize(2 * self.num_user_pvs / DIGEST_SIZE);
87        self.subair
88            .eval(builder, (local, next, num_rows.into(), None));
89
90        /*
91         * Send the left_child of each leaf node to output_values to be processed
92         * elsewhere. Note that in this case, this AIR has no public values. Also,
93         * output_val_bus expects to receive the app_exe_commit and app_vm_commit
94         * at indices 0..OUTPUT_USER_PVS_START_IDX. Note a row is a leaf node if
95         * its receive_type == 1.
96         */
97        let is_leaf = local.receive_type * (AB::Expr::TWO - local.receive_type);
98        assert_zeros(&mut builder.when(is_leaf.clone()), local.right_child);
99
100        const OUTPUT_USER_PVS_START_IDX: usize = (2 * DIGEST_SIZE) / VALS_IN_DIGEST;
101        const OUTPUT_VAL_MSGS_PER_ROW: usize = DIGEST_SIZE / VALS_IN_DIGEST;
102
103        for (i, output_values) in local.left_child.chunks_exact(VALS_IN_DIGEST).enumerate() {
104            self.output_val_bus.send(
105                builder,
106                OutputValMessage {
107                    values: from_fn(|i| output_values[i].into()),
108                    idx: AB::Expr::from_usize(OUTPUT_USER_PVS_START_IDX + i)
109                        + local.row_idx * AB::Expr::from_usize(OUTPUT_VAL_MSGS_PER_ROW),
110                },
111                is_leaf.clone(),
112            );
113        }
114    }
115}