openvm_recursion_circuit/proof_shape/pvs/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    utils::not, AlignedBorrow, ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
5};
6use openvm_stark_backend::{
7    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
8};
9use p3_air::{Air, AirBuilder, BaseAir};
10use p3_field::{PrimeCharacteristicRing, PrimeField32};
11use p3_matrix::Matrix;
12
13use crate::{
14    bus::{PublicValuesBus, PublicValuesBusMessage, TranscriptBus, TranscriptBusMessage},
15    proof_shape::bus::{NumPublicValuesBus, NumPublicValuesMessage},
16    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
17};
18
19#[repr(C)]
20#[derive(AlignedBorrow, Debug, StructReflection)]
21pub struct PublicValuesCols<F> {
22    pub is_valid: F,
23
24    pub proof_idx: F,
25    pub air_idx: F,
26    pub pv_idx: F,
27
28    pub is_first_in_proof: F,
29    pub is_first_in_air: F,
30
31    pub tidx: F,
32    pub value: F,
33}
34
35#[derive(ColumnsAir)]
36#[columns_via(PublicValuesCols<u8>)]
37pub struct PublicValuesAir {
38    pub public_values_bus: PublicValuesBus,
39    pub num_pvs_bus: NumPublicValuesBus,
40    pub transcript_bus: TranscriptBus,
41    pub(crate) continuations_enabled: bool,
42}
43
44impl<F> BaseAir<F> for PublicValuesAir {
45    fn width(&self) -> usize {
46        PublicValuesCols::<F>::width()
47    }
48}
49impl<F> BaseAirWithPublicValues<F> for PublicValuesAir {}
50impl<F> PartitionedBaseAir<F> for PublicValuesAir {}
51
52impl<AB: AirBuilder + InteractionBuilder> Air<AB> for PublicValuesAir
53where
54    AB::F: PrimeField32,
55{
56    fn eval(&self, builder: &mut AB) {
57        let main = builder.main();
58
59        let (local, next) = (
60            main.row_slice(0).expect("window should have two elements"),
61            main.row_slice(1).expect("window should have two elements"),
62        );
63        let local: &PublicValuesCols<AB::Var> = (*local).borrow();
64        let next: &PublicValuesCols<AB::Var> = (*next).borrow();
65
66        NestedForLoopSubAir::<1> {}.eval(
67            builder,
68            (
69                NestedForLoopIoCols {
70                    is_enabled: local.is_valid,
71                    counter: [local.proof_idx],
72                    is_first: [local.is_first_in_proof],
73                }
74                .map_into(),
75                NestedForLoopIoCols {
76                    is_enabled: next.is_valid,
77                    counter: [next.proof_idx],
78                    is_first: [next.is_first_in_proof],
79                }
80                .map_into(),
81            ),
82        );
83        // Constrain is_first_for_air, send NumPublicValuesBus message when true
84        builder.assert_bool(local.is_first_in_air);
85        builder
86            .when(local.is_first_in_proof)
87            .assert_one(local.is_first_in_air);
88        builder
89            .when(local.is_first_in_air)
90            .assert_one(local.is_valid);
91        builder
92            .when(next.is_valid * (next.air_idx - local.air_idx))
93            .assert_one(next.is_first_in_air);
94        builder
95            .when(local.is_first_in_air)
96            .assert_zero(local.pv_idx);
97
98        let is_same_air = local.is_valid * next.is_valid * not(next.is_first_in_air);
99        self.num_pvs_bus.receive(
100            builder,
101            local.proof_idx,
102            NumPublicValuesMessage {
103                air_idx: local.air_idx.into(),
104                tidx: local.tidx - local.pv_idx,
105                num_pvs: local.pv_idx + AB::Expr::ONE,
106            },
107            local.is_valid - is_same_air.clone(),
108        );
109
110        let mut when_same_air = builder.when(is_same_air);
111        when_same_air.assert_eq(local.air_idx, next.air_idx);
112        when_same_air.assert_eq(next.pv_idx, local.pv_idx + AB::Expr::ONE);
113        when_same_air.assert_eq(next.tidx, local.tidx + AB::Expr::ONE);
114
115        self.public_values_bus.send(
116            builder,
117            local.proof_idx,
118            PublicValuesBusMessage {
119                air_idx: local.air_idx,
120                pv_idx: local.pv_idx,
121                value: local.value,
122            },
123            local.is_valid,
124        );
125        if self.continuations_enabled {
126            self.public_values_bus.send(
127                builder,
128                local.proof_idx,
129                PublicValuesBusMessage {
130                    air_idx: local.air_idx,
131                    pv_idx: local.pv_idx,
132                    value: local.value,
133                },
134                local.is_valid,
135            );
136        }
137
138        // Receive transcript read of public values
139        self.transcript_bus.receive(
140            builder,
141            local.proof_idx,
142            TranscriptBusMessage {
143                tidx: local.tidx.into(),
144                value: local.value.into(),
145                is_sample: AB::Expr::ZERO,
146            },
147            local.is_valid,
148        );
149    }
150}
openvm_recursion_circuit/proof_shape/pvs/air.rs

openvm_recursion_circuit/proof_shape/pvs/
air.rs
