openvm_circuit/system/connector/
mod.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    marker::PhantomData,
4};
5
6use openvm_circuit_primitives::{
7    var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus},
8    ColumnsAir, StructReflection, StructReflectionHelper,
9};
10use openvm_circuit_primitives_derive::AlignedBorrow;
11use openvm_cpu_backend::CpuBackend;
12use openvm_instructions::LocalOpcode;
13use openvm_stark_backend::{
14    interaction::InteractionBuilder,
15    p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder},
16    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
17    p3_matrix::{dense::RowMajorMatrix, Matrix},
18    prover::AirProvingContext,
19    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
20};
21use serde::{Deserialize, Serialize};
22
23use crate::{
24    arch::{instructions::SystemOpcode::TERMINATE, ExecutionBus, ExecutionState},
25    primitives::Chip,
26    system::program::ProgramBus,
27};
28
29#[cfg(test)]
30mod tests;
31
32/// When a program hasn't terminated. There is no constraints on the exit code.
33/// But we will use this value when generating the proof.
34pub const DEFAULT_SUSPEND_EXIT_CODE: u32 = 42;
35
36#[derive(Debug, Clone, Copy, ColumnsAir)]
37#[columns_via(ConnectorCols<u8>)]
38pub struct VmConnectorAir {
39    pub execution_bus: ExecutionBus,
40    pub program_bus: ProgramBus,
41    pub range_bus: VariableRangeCheckerBus,
42    /// The final timestamp will be constrained to be in the range [0, 2^timestamp_max_bits).
43    timestamp_max_bits: usize,
44}
45
46#[derive(Debug, Clone, Copy, AlignedBorrow, StructReflection)]
47#[repr(C)]
48pub struct VmConnectorPvs<F> {
49    /// The initial PC of this segment.
50    pub initial_pc: F,
51    /// The final PC of this segment.
52    pub final_pc: F,
53    /// The exit code of the whole program. 0 means exited normally. This is only meaningful when
54    /// `is_terminate` is 1.
55    pub exit_code: F,
56    /// Whether the whole program is terminated. 0 means not terminated. 1 means terminated.
57    /// Only the last segment of an execution can have `is_terminate` = 1.
58    pub is_terminate: F,
59}
60
61impl<F: PrimeField32> VmConnectorPvs<F> {
62    pub fn is_terminate(&self) -> bool {
63        self.is_terminate == F::from_bool(true)
64    }
65
66    pub fn exit_code(&self) -> Option<u32> {
67        if self.is_terminate() && self.exit_code == F::ZERO {
68            Some(self.exit_code.as_canonical_u32())
69        } else {
70            None
71        }
72    }
73}
74
75impl<F: Field> BaseAirWithPublicValues<F> for VmConnectorAir {
76    fn num_public_values(&self) -> usize {
77        VmConnectorPvs::<F>::width()
78    }
79}
80impl<F: Field> PartitionedBaseAir<F> for VmConnectorAir {}
81impl<F: Field> BaseAir<F> for VmConnectorAir {
82    fn width(&self) -> usize {
83        ConnectorCols::<F>::width()
84    }
85}
86
87impl VmConnectorAir {
88    pub fn new(
89        execution_bus: ExecutionBus,
90        program_bus: ProgramBus,
91        range_bus: VariableRangeCheckerBus,
92        timestamp_max_bits: usize,
93    ) -> Self {
94        assert!(
95            range_bus.range_max_bits * 2 >= timestamp_max_bits,
96            "Range checker not large enough: range_max_bits={}, timestamp_max_bits={}",
97            range_bus.range_max_bits,
98            timestamp_max_bits
99        );
100        Self {
101            execution_bus,
102            program_bus,
103            range_bus,
104            timestamp_max_bits,
105        }
106    }
107
108    /// Returns (low_bits, high_bits) to range check.
109    fn timestamp_limb_bits(&self) -> (usize, usize) {
110        let range_max_bits = self.range_bus.range_max_bits;
111        if self.timestamp_max_bits <= range_max_bits {
112            (self.timestamp_max_bits, 0)
113        } else {
114            (range_max_bits, self.timestamp_max_bits - range_max_bits)
115        }
116    }
117}
118
119#[derive(Debug, Copy, Clone, AlignedBorrow, StructReflection, Serialize, Deserialize)]
120#[repr(C)]
121pub struct ConnectorCols<T> {
122    pub pc: T,
123    pub timestamp: T,
124    pub is_terminate: T,
125    pub exit_code: T,
126    /// Lowest `range_bus.range_max_bits` bits of the timestamp
127    timestamp_low_limb: T,
128    /// Equals 1 if this is the first row of the segment, 0 if this is the second row of the
129    /// segment. Used to enforce that the trace has exactly two rows.
130    is_begin: T,
131}
132
133impl<T: Copy> ConnectorCols<T> {
134    fn map<F>(self, f: impl Fn(T) -> F) -> ConnectorCols<F> {
135        ConnectorCols {
136            pc: f(self.pc),
137            timestamp: f(self.timestamp),
138            is_terminate: f(self.is_terminate),
139            exit_code: f(self.exit_code),
140            timestamp_low_limb: f(self.timestamp_low_limb),
141            is_begin: f(self.is_begin),
142        }
143    }
144
145    fn flatten(&self) -> [T; 6] {
146        [
147            self.pc,
148            self.timestamp,
149            self.is_terminate,
150            self.exit_code,
151            self.timestamp_low_limb,
152            self.is_begin,
153        ]
154    }
155}
156
157impl<AB: InteractionBuilder + PairBuilder + AirBuilderWithPublicValues> Air<AB> for VmConnectorAir {
158    fn eval(&self, builder: &mut AB) {
159        let main = builder.main();
160        let (local, next) = (
161            main.row_slice(0).expect("window should have two elements"),
162            main.row_slice(1).expect("window should have two elements"),
163        );
164
165        let local: &ConnectorCols<AB::Var> = (*local).borrow();
166        let next: &ConnectorCols<AB::Var> = (*next).borrow();
167
168        let &VmConnectorPvs {
169            initial_pc,
170            final_pc,
171            exit_code,
172            is_terminate,
173        } = builder.public_values().borrow();
174
175        builder.when_transition().assert_eq(local.pc, initial_pc);
176        builder.when_transition().assert_eq(next.pc, final_pc);
177        builder
178            .when_transition()
179            .when(next.is_terminate)
180            .assert_eq(next.exit_code, exit_code);
181        builder
182            .when_transition()
183            .assert_eq(next.is_terminate, is_terminate);
184        // Assert is_terminate is boolean on every row to ensure lookup multiplicity is boolean
185        // below
186        builder.assert_bool(local.is_terminate);
187
188        builder.when_transition().assert_one(local.timestamp);
189
190        // We force the first row to have is_begin = 1 and the last row to have is_begin = 0.
191        // Additionally, we enforce that the is_begin column decreases by exactly 1 per row.
192        // The only way to satisfy this is to have exactly two rows: one with is_begin = 1 and
193        // one with is_begin = 0 (assuming max height < field characteristic)
194        builder.when_first_row().assert_one(local.is_begin);
195        builder
196            .when_transition()
197            .assert_eq(next.is_begin + AB::Expr::ONE, local.is_begin);
198        builder.when_last_row().assert_zero(local.is_begin);
199
200        self.execution_bus.execute(
201            builder,
202            local.is_begin, // 1 only if these are [0th, 1st] and not [1st, 0th]
203            ExecutionState::new(next.pc, next.timestamp),
204            ExecutionState::new(local.pc, local.timestamp),
205        );
206        self.program_bus.lookup_instruction(
207            builder,
208            next.pc,
209            AB::Expr::from_usize(TERMINATE.global_opcode().as_usize()),
210            [AB::Expr::ZERO, AB::Expr::ZERO, next.exit_code.into()],
211            local.is_begin * next.is_terminate,
212        );
213
214        // We decompose and range check `local.timestamp` as `timestamp_low_limb,
215        // timestamp_high_limb` where `timestamp = timestamp_low_limb + timestamp_high_limb
216        // * 2^range_max_bits`.
217        let (low_bits, high_bits) = self.timestamp_limb_bits();
218        let high_limb = (local.timestamp - local.timestamp_low_limb)
219            * AB::F::ONE.div_2exp_u64(self.range_bus.range_max_bits as u64);
220        self.range_bus
221            .range_check(local.timestamp_low_limb, low_bits)
222            .eval(builder, AB::Expr::ONE);
223        self.range_bus
224            .range_check(high_limb, high_bits)
225            .eval(builder, AB::Expr::ONE);
226    }
227}
228
229pub struct VmConnectorChip<F> {
230    pub range_checker: SharedVariableRangeCheckerChip,
231    pub boundary_states: [Option<ConnectorCols<u32>>; 2],
232    timestamp_max_bits: usize,
233    _marker: PhantomData<F>,
234}
235
236impl<F> VmConnectorChip<F> {
237    pub fn new(range_checker: SharedVariableRangeCheckerChip, timestamp_max_bits: usize) -> Self {
238        let range_bus = range_checker.bus();
239        assert!(
240            range_bus.range_max_bits * 2 >= timestamp_max_bits,
241            "Range checker not large enough: range_max_bits={}, timestamp_max_bits={}",
242            range_bus.range_max_bits,
243            timestamp_max_bits
244        );
245        Self {
246            range_checker,
247            boundary_states: [None, None],
248            timestamp_max_bits,
249            _marker: PhantomData,
250        }
251    }
252
253    pub fn begin(&mut self, state: ExecutionState<u32>) {
254        self.boundary_states[0] = Some(ConnectorCols {
255            pc: state.pc,
256            timestamp: state.timestamp,
257            is_terminate: 0,
258            exit_code: 0,
259            timestamp_low_limb: 0, // will be computed during tracegen
260            is_begin: 1,
261        });
262    }
263
264    pub fn end(&mut self, state: ExecutionState<u32>, exit_code: Option<u32>) {
265        self.boundary_states[1] = Some(ConnectorCols {
266            pc: state.pc,
267            timestamp: state.timestamp,
268            is_terminate: exit_code.is_some() as u32,
269            exit_code: exit_code.unwrap_or(DEFAULT_SUSPEND_EXIT_CODE),
270            timestamp_low_limb: 0, // will be computed during tracegen
271            is_begin: 0,
272        });
273    }
274
275    fn timestamp_limb_bits(&self) -> (usize, usize) {
276        let range_max_bits = self.range_checker.bus().range_max_bits;
277        if self.timestamp_max_bits <= range_max_bits {
278            (self.timestamp_max_bits, 0)
279        } else {
280            (range_max_bits, self.timestamp_max_bits - range_max_bits)
281        }
282    }
283}
284
285impl<RA, SC> Chip<RA, CpuBackend<SC>> for VmConnectorChip<Val<SC>>
286where
287    SC: StarkProtocolConfig,
288    Val<SC>: PrimeField32,
289{
290    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<CpuBackend<SC>> {
291        let [initial_state, final_state] = self.boundary_states.map(|state| {
292            let mut state = state.unwrap();
293            // Decompose and range check timestamp
294            let range_max_bits = self.range_checker.range_max_bits();
295            let timestamp_low_limb = state.timestamp & ((1u32 << range_max_bits) - 1);
296            state.timestamp_low_limb = timestamp_low_limb;
297            let (low_bits, high_bits) = self.timestamp_limb_bits();
298            self.range_checker.add_count(timestamp_low_limb, low_bits);
299            self.range_checker
300                .add_count(state.timestamp >> range_max_bits, high_bits);
301
302            state.map(Val::<SC>::from_u32)
303        });
304
305        let trace = RowMajorMatrix::new(
306            [initial_state.flatten(), final_state.flatten()].concat(),
307            ConnectorCols::<Val<SC>>::width(),
308        );
309
310        let mut public_values = Val::<SC>::zero_vec(VmConnectorPvs::<Val<SC>>::width());
311        *public_values.as_mut_slice().borrow_mut() = VmConnectorPvs {
312            initial_pc: initial_state.pc,
313            final_pc: final_state.pc,
314            exit_code: final_state.exit_code,
315            is_terminate: final_state.is_terminate,
316        };
317        AirProvingContext::simple(trace, public_values)
318    }
319}
openvm_circuit/system/connector/mod.rs

openvm_circuit/system/connector/
mod.rs
