openvm_continuations/circuit/deferral/inner/input/
air.rs

1use std::{array::from_fn, borrow::Borrow};
2
3use openvm_circuit_primitives::{
4    utils::not, ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
5};
6use openvm_recursion_circuit::{
7    bus::{
8        CachedCommitBus, CachedCommitBusMessage, Poseidon2PermuteBus, Poseidon2PermuteMessage,
9        PublicValuesBus, PublicValuesBusMessage,
10    },
11    prelude::DIGEST_SIZE,
12    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
13};
14use openvm_recursion_circuit_derive::AlignedBorrow;
15use openvm_stark_backend::{
16    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
17};
18use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
19use p3_field::PrimeCharacteristicRing;
20use p3_matrix::Matrix;
21
22use crate::{
23    circuit::deferral::{
24        inner::bus::{InputOrMerkleCommitBus, InputOrMerkleCommitMessage},
25        DEF_AGG_PVS_AIR_ID, DEF_CIRCUIT_PVS_AIR_ID,
26    },
27    utils::digests_to_poseidon2_input,
28};
29
30#[repr(C)]
31#[derive(AlignedBorrow, StructReflection)]
32pub struct InputCommitCols<F> {
33    pub is_valid: F,
34    pub is_first: F,
35
36    pub proof_idx: F,
37    pub row_in_proof_idx: F,
38    pub has_verifier_pvs: F,
39
40    pub air_idx: F,
41    pub cached_idx: F,
42    pub current_commit: [F; DIGEST_SIZE],
43
44    pub res_left: [F; DIGEST_SIZE],
45    pub res_right: [F; DIGEST_SIZE],
46}
47
48#[derive(ColumnsAir)]
49#[columns_via(InputCommitCols<u8>)]
50pub struct InputCommitAir {
51    pub public_values_bus: PublicValuesBus,
52    pub poseidon2_bus: Poseidon2PermuteBus,
53    pub cached_commit_bus: CachedCommitBus,
54    pub input_or_merkle_commit_bus: InputOrMerkleCommitBus,
55}
56
57impl<F> BaseAir<F> for InputCommitAir {
58    fn width(&self) -> usize {
59        InputCommitCols::<u8>::width()
60    }
61}
62impl<F> BaseAirWithPublicValues<F> for InputCommitAir {}
63impl<F> PartitionedBaseAir<F> for InputCommitAir {}
64
65impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB> for InputCommitAir {
66    fn eval(&self, builder: &mut AB) {
67        let main = builder.main();
68        let (local, next) = (
69            main.row_slice(0).expect("window should have two elements"),
70            main.row_slice(1).expect("window should have two elements"),
71        );
72        let local: &InputCommitCols<AB::Var> = (*local).borrow();
73        let next: &InputCommitCols<AB::Var> = (*next).borrow();
74
75        NestedForLoopSubAir::<2> {}.eval(
76            builder,
77            (
78                NestedForLoopIoCols {
79                    is_enabled: local.is_valid,
80                    counter: [local.proof_idx, local.row_in_proof_idx],
81                    is_first: [local.is_first, local.is_valid],
82                }
83                .map_into(),
84                NestedForLoopIoCols {
85                    is_enabled: next.is_valid,
86                    counter: [next.proof_idx, next.row_in_proof_idx],
87                    is_first: [next.is_first, next.is_valid],
88                }
89                .map_into(),
90            ),
91        );
92
93        /*
94         * Constrain that has_verifier_pvs is consistent over all valid rows.
95         */
96        builder.assert_bool(local.has_verifier_pvs);
97        builder
98            .when(local.is_valid * next.is_valid)
99            .assert_eq(local.has_verifier_pvs, next.has_verifier_pvs);
100        builder
101            .when(local.is_valid * local.has_verifier_pvs)
102            .assert_one(local.is_first);
103
104        /*
105         * Read the input commit from public values on the first row for each proof. This is the
106         * input commit for the leaf verifier, and the deferral aggregation Merkle commit at all
107         * internal levels.
108         */
109        let is_leaf = not(local.has_verifier_pvs);
110        let is_internal = local.has_verifier_pvs;
111        let air_idx = is_leaf.clone() * AB::Expr::from_usize(DEF_CIRCUIT_PVS_AIR_ID)
112            + is_internal * AB::Expr::from_usize(DEF_AGG_PVS_AIR_ID);
113
114        for (pv_idx, value) in local.current_commit.iter().enumerate() {
115            self.public_values_bus.receive(
116                builder,
117                local.proof_idx,
118                PublicValuesBusMessage {
119                    air_idx: air_idx.clone(),
120                    pv_idx: AB::Expr::from_usize(pv_idx),
121                    value: (*value).into(),
122                },
123                local.is_first,
124            );
125        }
126
127        /*
128         * Receive cached trace commits and fold them into the sponge on all non-first rows.
129         */
130        self.cached_commit_bus.receive(
131            builder,
132            local.proof_idx,
133            CachedCommitBusMessage {
134                air_idx: local.air_idx.into(),
135                cached_idx: local.cached_idx.into(),
136                global_cached_idx: local.row_in_proof_idx - AB::Expr::ONE,
137                cached_commit: local.current_commit.map(Into::into),
138            },
139            is_leaf.clone() * (local.is_valid - local.is_first),
140        );
141
142        /*
143         * Fold the received current_commit into the sponge. On the first row, the capacity
144         * should be all zeros. On subsequent rows, the capacity should match the res_right
145         * from the previous row.
146         */
147        let is_transition =
148            NestedForLoopSubAir::<1>::local_is_transition(next.is_valid, next.is_first);
149        let is_last =
150            NestedForLoopSubAir::<1>::local_is_last(local.is_valid, next.is_valid, next.is_first);
151
152        self.poseidon2_bus.lookup_key(
153            builder,
154            Poseidon2PermuteMessage {
155                input: digests_to_poseidon2_input(
156                    local.current_commit.map(Into::into),
157                    [AB::Expr::ZERO; _],
158                ),
159                output: digests_to_poseidon2_input(local.res_left, local.res_right).map(Into::into),
160            },
161            local.is_first * is_leaf.clone(),
162        );
163
164        self.poseidon2_bus.lookup_key(
165            builder,
166            Poseidon2PermuteMessage {
167                input: digests_to_poseidon2_input(next.current_commit, local.res_right),
168                output: digests_to_poseidon2_input(next.res_left, next.res_right),
169            },
170            is_transition * is_leaf.clone(),
171        );
172
173        /*
174         * Finally, on the last row for this proof we send the input commit.
175         */
176        self.input_or_merkle_commit_bus.send(
177            builder,
178            local.proof_idx,
179            InputOrMerkleCommitMessage {
180                has_verifier_pvs: local.has_verifier_pvs.into(),
181                commit: from_fn(|i| {
182                    is_internal * local.current_commit[i] + is_leaf.clone() * local.res_left[i]
183                }),
184            },
185            is_last,
186        );
187    }
188}
openvm_continuations/circuit/deferral/inner/input/air.rs

openvm_continuations/circuit/deferral/inner/input/
air.rs
