openvm_rv32im_circuit/branch_eq/
core.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{
4    arch::*,
5    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
6};
7use openvm_circuit_primitives::{utils::not, ColumnsAir, StructReflection, StructReflectionHelper};
8use openvm_circuit_primitives_derive::{AlignedBorrow, AlignedBytesBorrow};
9use openvm_instructions::{instruction::Instruction, LocalOpcode};
10use openvm_rv32im_transpiler::BranchEqualOpcode;
11use openvm_stark_backend::{
12    interaction::InteractionBuilder,
13    p3_air::{AirBuilder, BaseAir},
14    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
15    BaseAirWithPublicValues,
16};
17use strum::IntoEnumIterator;
18
19#[repr(C)]
20#[derive(AlignedBorrow, StructReflection)]
21pub struct BranchEqualCoreCols<T, const NUM_LIMBS: usize> {
22    pub a: [T; NUM_LIMBS],
23    pub b: [T; NUM_LIMBS],
24
25    // Boolean result of a op b. Should branch if and only if cmp_result = 1.
26    pub cmp_result: T,
27    pub imm: T,
28
29    pub opcode_beq_flag: T,
30    pub opcode_bne_flag: T,
31
32    pub diff_inv_marker: [T; NUM_LIMBS],
33}
34
35#[derive(Copy, Clone, Debug, derive_new::new, ColumnsAir)]
36#[columns_via(BranchEqualCoreCols<u8, NUM_LIMBS>)]
37pub struct BranchEqualCoreAir<const NUM_LIMBS: usize> {
38    offset: usize,
39    pc_step: u32,
40}
41
42impl<F: Field, const NUM_LIMBS: usize> BaseAir<F> for BranchEqualCoreAir<NUM_LIMBS> {
43    fn width(&self) -> usize {
44        BranchEqualCoreCols::<F, NUM_LIMBS>::width()
45    }
46}
47impl<F: Field, const NUM_LIMBS: usize> BaseAirWithPublicValues<F>
48    for BranchEqualCoreAir<NUM_LIMBS>
49{
50}
51
52impl<AB, I, const NUM_LIMBS: usize> VmCoreAir<AB, I> for BranchEqualCoreAir<NUM_LIMBS>
53where
54    AB: InteractionBuilder,
55    I: VmAdapterInterface<AB::Expr>,
56    I::Reads: From<[[AB::Expr; NUM_LIMBS]; 2]>,
57    I::Writes: Default,
58    I::ProcessedInstruction: From<ImmInstruction<AB::Expr>>,
59{
60    fn eval(
61        &self,
62        builder: &mut AB,
63        local: &[AB::Var],
64        from_pc: AB::Var,
65    ) -> AdapterAirContext<AB::Expr, I> {
66        let cols: &BranchEqualCoreCols<_, NUM_LIMBS> = local.borrow();
67        let flags = [cols.opcode_beq_flag, cols.opcode_bne_flag];
68
69        let is_valid = flags.iter().fold(AB::Expr::ZERO, |acc, &flag| {
70            builder.assert_bool(flag);
71            acc + flag.into()
72        });
73        builder.assert_bool(is_valid.clone());
74        builder.assert_bool(cols.cmp_result);
75
76        let a = &cols.a;
77        let b = &cols.b;
78        let inv_marker = &cols.diff_inv_marker;
79
80        // 1 if cmp_result indicates a and b are equal, 0 otherwise
81        let cmp_eq =
82            cols.cmp_result * cols.opcode_beq_flag + not(cols.cmp_result) * cols.opcode_bne_flag;
83        let mut sum = cmp_eq.clone();
84
85        // For BEQ, inv_marker is used to check equality of a and b:
86        // - If a == b, all inv_marker values must be 0 (sum = 0)
87        // - If a != b, inv_marker contains 0s for all positions except ONE position i where a[i] !=
88        //   b[i]
89        // - At this position, inv_marker[i] contains the multiplicative inverse of (a[i] - b[i])
90        // - This ensures inv_marker[i] * (a[i] - b[i]) = 1, making the sum = 1
91        // Note: There might be multiple valid inv_marker if a != b.
92        // But as long as the trace can provide at least one, that’s sufficient to prove a != b.
93        //
94        // Note:
95        // - If cmp_eq == 0, then it is impossible to have sum != 0 if a == b.
96        // - If cmp_eq == 1, then it is impossible for a[i] - b[i] == 0 to pass for all i if a != b.
97        for i in 0..NUM_LIMBS {
98            sum += (a[i] - b[i]) * inv_marker[i];
99            builder.assert_zero(cmp_eq.clone() * (a[i] - b[i]));
100        }
101        builder.when(is_valid.clone()).assert_one(sum);
102
103        let expected_opcode = flags
104            .iter()
105            .zip(BranchEqualOpcode::iter())
106            .fold(AB::Expr::ZERO, |acc, (flag, opcode)| {
107                acc + (*flag).into() * AB::Expr::from_u8(opcode as u8)
108            })
109            + AB::Expr::from_usize(self.offset);
110
111        let to_pc = from_pc
112            + cols.cmp_result * cols.imm
113            + not(cols.cmp_result) * AB::Expr::from_u32(self.pc_step);
114
115        AdapterAirContext {
116            to_pc: Some(to_pc),
117            reads: [cols.a.map(Into::into), cols.b.map(Into::into)].into(),
118            writes: Default::default(),
119            instruction: ImmInstruction {
120                is_valid,
121                opcode: expected_opcode,
122                immediate: cols.imm.into(),
123            }
124            .into(),
125        }
126    }
127
128    fn start_offset(&self) -> usize {
129        self.offset
130    }
131}
132
133#[repr(C)]
134#[derive(AlignedBytesBorrow, Debug)]
135pub struct BranchEqualCoreRecord<const NUM_LIMBS: usize> {
136    pub a: [u8; NUM_LIMBS],
137    pub b: [u8; NUM_LIMBS],
138    pub imm: u32,
139    pub local_opcode: u8,
140}
141
142#[derive(Clone, Copy, derive_new::new)]
143pub struct BranchEqualExecutor<A, const NUM_LIMBS: usize> {
144    adapter: A,
145    pub offset: usize,
146    pub pc_step: u32,
147}
148
149#[derive(Clone, Copy, derive_new::new)]
150pub struct BranchEqualFiller<A, const NUM_LIMBS: usize> {
151    adapter: A,
152    pub offset: usize,
153    pub pc_step: u32,
154}
155
156impl<F, A, RA, const NUM_LIMBS: usize> PreflightExecutor<F, RA>
157    for BranchEqualExecutor<A, NUM_LIMBS>
158where
159    F: PrimeField32,
160    A: 'static + AdapterTraceExecutor<F, ReadData: Into<[[u8; NUM_LIMBS]; 2]>, WriteData = ()>,
161    for<'buf> RA: RecordArena<
162        'buf,
163        EmptyAdapterCoreLayout<F, A>,
164        (
165            A::RecordMut<'buf>,
166            &'buf mut BranchEqualCoreRecord<NUM_LIMBS>,
167        ),
168    >,
169{
170    fn get_opcode_name(&self, opcode: usize) -> String {
171        format!("{:?}", BranchEqualOpcode::from_usize(opcode - self.offset))
172    }
173
174    fn execute(
175        &self,
176        state: VmStateMut<F, TracingMemory, RA>,
177        instruction: &Instruction<F>,
178    ) -> Result<(), ExecutionError> {
179        let &Instruction { opcode, c: imm, .. } = instruction;
180
181        let branch_eq_opcode = BranchEqualOpcode::from_usize(opcode.local_opcode_idx(self.offset));
182
183        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
184
185        A::start(*state.pc, state.memory, &mut adapter_record);
186
187        let [rs1, rs2] = self
188            .adapter
189            .read(state.memory, instruction, &mut adapter_record)
190            .into();
191
192        core_record.a = rs1;
193        core_record.b = rs2;
194        core_record.imm = imm.as_canonical_u32();
195        core_record.local_opcode = branch_eq_opcode as u8;
196
197        if fast_run_eq(branch_eq_opcode, &rs1, &rs2) {
198            *state.pc = (F::from_u32(*state.pc) + imm).as_canonical_u32();
199        } else {
200            *state.pc = state.pc.wrapping_add(self.pc_step);
201        }
202
203        Ok(())
204    }
205}
206
207impl<F, A, const NUM_LIMBS: usize> TraceFiller<F> for BranchEqualFiller<A, NUM_LIMBS>
208where
209    F: PrimeField32,
210    A: 'static + AdapterTraceFiller<F>,
211{
212    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
213        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
214        // BranchEqualCoreCols::width() elements
215        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
216        self.adapter.fill_trace_row(mem_helper, adapter_row);
217        // SAFETY: core_row contains a valid BranchEqualCoreRecord written by the executor
218        // during trace generation
219        let record: &BranchEqualCoreRecord<NUM_LIMBS> =
220            unsafe { get_record_from_slice(&mut core_row, ()) };
221        let core_row: &mut BranchEqualCoreCols<F, NUM_LIMBS> = core_row.borrow_mut();
222
223        let (cmp_result, diff_idx, diff_inv_val) = run_eq::<F, NUM_LIMBS>(
224            record.local_opcode == BranchEqualOpcode::BEQ as u8,
225            &record.a,
226            &record.b,
227        );
228        core_row.diff_inv_marker = [F::ZERO; NUM_LIMBS];
229        core_row.diff_inv_marker[diff_idx] = diff_inv_val;
230
231        core_row.opcode_bne_flag =
232            F::from_bool(record.local_opcode == BranchEqualOpcode::BNE as u8);
233        core_row.opcode_beq_flag =
234            F::from_bool(record.local_opcode == BranchEqualOpcode::BEQ as u8);
235
236        core_row.imm = F::from_u32(record.imm);
237        core_row.cmp_result = F::from_bool(cmp_result);
238
239        core_row.b = record.b.map(F::from_u8);
240        core_row.a = record.a.map(F::from_u8);
241    }
242}
243
244// Returns (cmp_result, diff_idx, x[diff_idx] - y[diff_idx])
245#[inline(always)]
246pub(super) fn fast_run_eq<const NUM_LIMBS: usize>(
247    local_opcode: BranchEqualOpcode,
248    x: &[u8; NUM_LIMBS],
249    y: &[u8; NUM_LIMBS],
250) -> bool {
251    match local_opcode {
252        BranchEqualOpcode::BEQ => x == y,
253        BranchEqualOpcode::BNE => x != y,
254    }
255}
256
257// Returns (cmp_result, diff_idx, x[diff_idx] - y[diff_idx])
258#[inline(always)]
259pub(super) fn run_eq<F, const NUM_LIMBS: usize>(
260    is_beq: bool,
261    x: &[u8; NUM_LIMBS],
262    y: &[u8; NUM_LIMBS],
263) -> (bool, usize, F)
264where
265    F: PrimeField32,
266{
267    for i in 0..NUM_LIMBS {
268        if x[i] != y[i] {
269            return (!is_beq, i, (F::from_u8(x[i]) - F::from_u8(y[i])).inverse());
270        }
271    }
272    (is_beq, 0, F::ZERO)
273}
openvm_rv32im_circuit/branch_eq/core.rs

openvm_rv32im_circuit/branch_eq/
core.rs
