openvm_recursion_circuit/transcript/transcript/
air.rs

1use core::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    utils::{and, assert_array_eq, not, or},
5    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
6};
7use openvm_recursion_circuit_derive::AlignedBorrow;
8use openvm_stark_backend::{
9    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
10};
11use p3_air::{Air, AirBuilder, BaseAir};
12use p3_field::{Field, PrimeCharacteristicRing};
13use p3_matrix::Matrix;
14
15use crate::{
16    bus::{
17        FinalTranscriptStateBus, FinalTranscriptStateMessage, Poseidon2PermuteBus,
18        Poseidon2PermuteMessage, TranscriptBus, TranscriptBusMessage,
19    },
20    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
21    transcript::poseidon2::{CHUNK, POSEIDON2_WIDTH},
22};
23
24#[repr(C)]
25#[derive(AlignedBorrow, Debug, StructReflection)]
26pub struct TranscriptCols<T> {
27    pub proof_idx: T,
28    pub is_proof_start: T,
29
30    pub tidx: T,
31    /// Indicator for sample/observe.
32    pub is_sample: T,
33    /// 0/1 indicators for the positions that we are absorbing/squeezing (i.e., are in
34    /// the transcript). Constrained to be "decreasing". Because transcript_bus messages
35    /// are sent with multiplicity 0 or 1, this also functions as the lookup count.
36    pub mask: [T; CHUNK],
37
38    /// The poseidon2 state.
39    pub prev_state: [T; POSEIDON2_WIDTH],
40    pub post_state: [T; POSEIDON2_WIDTH],
41}
42
43#[derive(ColumnsAir)]
44#[columns_via(TranscriptCols<u8>)]
45pub struct TranscriptAir {
46    pub transcript_bus: TranscriptBus,
47    pub poseidon2_permute_bus: Poseidon2PermuteBus,
48    pub final_state_bus: Option<FinalTranscriptStateBus>,
49}
50
51impl<F: Field> BaseAir<F> for TranscriptAir {
52    fn width(&self) -> usize {
53        TranscriptCols::<F>::width()
54    }
55}
56
57impl<F: Field> BaseAirWithPublicValues<F> for TranscriptAir {}
58impl<F: Field> PartitionedBaseAir<F> for TranscriptAir {}
59
60impl<AB: AirBuilder + InteractionBuilder> Air<AB> for TranscriptAir {
61    fn eval(&self, builder: &mut AB) {
62        let main = builder.main();
63        let (local, next) = (
64            main.row_slice(0).expect("window should have two elements"),
65            main.row_slice(1).expect("window should have two elements"),
66        );
67        let local: &TranscriptCols<AB::Var> = (*local).borrow();
68        let next: &TranscriptCols<AB::Var> = (*next).borrow();
69
70        ///////////////////////////////////////////////////////////////////////
71        // Constraints
72        ///////////////////////////////////////////////////////////////////////
73        let is_valid = local.mask[0];
74        let next_valid = next.mask[0];
75
76        NestedForLoopSubAir::<1> {}.eval(
77            builder,
78            (
79                NestedForLoopIoCols {
80                    is_enabled: is_valid,
81                    counter: [local.proof_idx],
82                    is_first: [local.is_proof_start],
83                }
84                .map_into(),
85                NestedForLoopIoCols {
86                    is_enabled: next_valid,
87                    counter: [next.proof_idx],
88                    is_first: [next.is_proof_start],
89                }
90                .map_into(),
91            ),
92        );
93
94        builder.when(local.is_proof_start).assert_zero(local.tidx);
95        builder.when(local.is_proof_start).assert_one(is_valid);
96        builder.assert_bool(local.is_sample);
97
98        // Initial state constraints
99        for i in 0..CHUNK {
100            builder
101                .when(local.is_proof_start)
102                .assert_eq(local.prev_state[i + CHUNK], AB::Expr::ZERO);
103
104            builder
105                .when(local.is_proof_start * (AB::Expr::ONE - local.mask[i]))
106                .assert_eq(local.prev_state[i], AB::Expr::ZERO);
107        }
108
109        let mut count = AB::Expr::ZERO;
110        let local_next_same_proof = next_valid - next.is_proof_start;
111        for i in 0..CHUNK {
112            builder.assert_bool(local.mask[i]);
113            count += local.mask[i].into();
114
115            let skip = local.mask[i] - AB::Expr::ONE;
116            if i < CHUNK - 1 {
117                // if mask[i] = 0, then mask[i+1] = 0
118                builder.when(skip.clone()).assert_zero(local.mask[i + 1]);
119            }
120
121            // The state after permutation of this round, should check against next round's input
122            // if next.mask[i] = 0 --> i-th not touched --> it should stay the same (if next is
123            // valid)
124            builder
125                .when((AB::Expr::ONE - next.mask[i]) * local_next_same_proof.clone())
126                .assert_eq(local.post_state[i], next.prev_state[i]);
127            // When it's squeeze(sample), the state always remains the same
128            builder
129                .when(next.is_sample * local_next_same_proof.clone())
130                .assert_eq(local.post_state[i], next.prev_state[i]);
131
132            // The capacity part should always be the same
133            builder
134                .when(local_next_same_proof.clone()) // if next is valid
135                .assert_eq(local.post_state[i + CHUNK], next.prev_state[i + CHUNK]);
136        }
137
138        let mut when_same_proof = builder.when(local_next_same_proof.clone());
139        when_same_proof.assert_eq(next.tidx, local.tidx + count.clone());
140
141        // If local.is_sample = next.is_sample, there have to be CHUNK operations
142        when_same_proof
143            .when_ne(local.is_sample, not(next.is_sample))
144            .assert_eq(count, AB::Expr::from_usize(CHUNK));
145
146        ///////////////////////////////////////////////////////////////////////
147        // Interactions
148        ///////////////////////////////////////////////////////////////////////
149        for i in 0..CHUNK {
150            // When absorb, it's normal order (0 -> RATE)
151            let observe_message = TranscriptBusMessage {
152                tidx: local.tidx + AB::Expr::from_usize(i),
153                value: local.prev_state[i].into(),
154                is_sample: AB::Expr::ZERO,
155            };
156            // When squeeze, it's reverse RATE -> 0, so i means RATE - 1 - i
157            let sample_message = TranscriptBusMessage {
158                tidx: local.tidx + AB::Expr::from_usize(i),
159                value: local.prev_state[CHUNK - 1 - i].into(),
160                is_sample: AB::Expr::ONE,
161            };
162            self.transcript_bus.send(
163                builder,
164                local.proof_idx,
165                observe_message,
166                local.mask[i] * (AB::Expr::ONE - local.is_sample),
167            );
168            self.transcript_bus.send(
169                builder,
170                local.proof_idx,
171                sample_message,
172                local.mask[i] * local.is_sample,
173            );
174        }
175
176        // Permute on all non-final rows except when going from sample to observe,
177        // and never on the final row.
178        let permuted =
179            local_next_same_proof * not::<AB::Expr>(and(local.is_sample, not(next.is_sample)));
180        self.poseidon2_permute_bus.lookup_key(
181            builder,
182            Poseidon2PermuteMessage {
183                input: local.prev_state,
184                output: local.post_state,
185            },
186            permuted.clone(),
187        );
188
189        assert_array_eq(
190            &mut builder.when(not::<AB::Expr>(permuted)),
191            local.prev_state,
192            local.post_state,
193        );
194
195        if let Some(final_state_bus) = self.final_state_bus {
196            final_state_bus.send(
197                builder,
198                local.proof_idx,
199                FinalTranscriptStateMessage {
200                    state: local.post_state,
201                },
202                and(is_valid, or(not(next_valid), next.is_proof_start)),
203            );
204        }
205    }
206}
openvm_recursion_circuit/transcript/transcript/air.rs

openvm_recursion_circuit/transcript/transcript/
air.rs
