openvm_rv32im_circuit/branch_lt/
core.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{
4    arch::*,
5    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
6};
7use openvm_circuit_primitives::{
8    bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
9    utils::not,
10    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
11};
12use openvm_circuit_primitives_derive::AlignedBorrow;
13use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP, LocalOpcode};
14use openvm_rv32im_transpiler::BranchLessThanOpcode;
15use openvm_stark_backend::{
16    interaction::InteractionBuilder,
17    p3_air::{AirBuilder, BaseAir},
18    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
19    BaseAirWithPublicValues,
20};
21use strum::IntoEnumIterator;
22
23#[repr(C)]
24#[derive(AlignedBorrow, StructReflection)]
25pub struct BranchLessThanCoreCols<T, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
26    pub a: [T; NUM_LIMBS],
27    pub b: [T; NUM_LIMBS],
28
29    // Boolean result of a op b. Should branch if and only if cmp_result = 1.
30    pub cmp_result: T,
31    pub imm: T,
32
33    pub opcode_blt_flag: T,
34    pub opcode_bltu_flag: T,
35    pub opcode_bge_flag: T,
36    pub opcode_bgeu_flag: T,
37
38    // Most significant limb of a and b respectively as a field element, will be range
39    // checked to be within [-128, 127) if signed and [0, 256) if unsigned.
40    pub a_msb_f: T,
41    pub b_msb_f: T,
42
43    // 1 if a < b, 0 otherwise.
44    pub cmp_lt: T,
45
46    // 1 at the most significant index i such that a[i] != b[i], otherwise 0. If such
47    // an i exists, diff_val = b[i] - a[i].
48    pub diff_marker: [T; NUM_LIMBS],
49    pub diff_val: T,
50}
51
52#[derive(Copy, Clone, Debug, derive_new::new, ColumnsAir)]
53#[columns_via(BranchLessThanCoreCols<u8, NUM_LIMBS, LIMB_BITS>)]
54pub struct BranchLessThanCoreAir<const NUM_LIMBS: usize, const LIMB_BITS: usize> {
55    pub bus: BitwiseOperationLookupBus,
56    offset: usize,
57}
58
59impl<F: Field, const NUM_LIMBS: usize, const LIMB_BITS: usize> BaseAir<F>
60    for BranchLessThanCoreAir<NUM_LIMBS, LIMB_BITS>
61{
62    fn width(&self) -> usize {
63        BranchLessThanCoreCols::<F, NUM_LIMBS, LIMB_BITS>::width()
64    }
65}
66impl<F: Field, const NUM_LIMBS: usize, const LIMB_BITS: usize> BaseAirWithPublicValues<F>
67    for BranchLessThanCoreAir<NUM_LIMBS, LIMB_BITS>
68{
69}
70
71impl<AB, I, const NUM_LIMBS: usize, const LIMB_BITS: usize> VmCoreAir<AB, I>
72    for BranchLessThanCoreAir<NUM_LIMBS, LIMB_BITS>
73where
74    AB: InteractionBuilder,
75    I: VmAdapterInterface<AB::Expr>,
76    I::Reads: From<[[AB::Expr; NUM_LIMBS]; 2]>,
77    I::Writes: Default,
78    I::ProcessedInstruction: From<ImmInstruction<AB::Expr>>,
79{
80    fn eval(
81        &self,
82        builder: &mut AB,
83        local_core: &[AB::Var],
84        from_pc: AB::Var,
85    ) -> AdapterAirContext<AB::Expr, I> {
86        let cols: &BranchLessThanCoreCols<_, NUM_LIMBS, LIMB_BITS> = local_core.borrow();
87        let flags = [
88            cols.opcode_blt_flag,
89            cols.opcode_bltu_flag,
90            cols.opcode_bge_flag,
91            cols.opcode_bgeu_flag,
92        ];
93
94        let is_valid = flags.iter().fold(AB::Expr::ZERO, |acc, &flag| {
95            builder.assert_bool(flag);
96            acc + flag.into()
97        });
98        builder.assert_bool(is_valid.clone());
99        builder.assert_bool(cols.cmp_result);
100
101        let lt = cols.opcode_blt_flag + cols.opcode_bltu_flag;
102        let ge = cols.opcode_bge_flag + cols.opcode_bgeu_flag;
103        let signed = cols.opcode_blt_flag + cols.opcode_bge_flag;
104        builder.assert_eq(
105            cols.cmp_lt,
106            cols.cmp_result * lt.clone() + not(cols.cmp_result) * ge.clone(),
107        );
108
109        let a = &cols.a;
110        let b = &cols.b;
111        let marker = &cols.diff_marker;
112        let mut prefix_sum = AB::Expr::ZERO;
113
114        // Check if a_msb_f and b_msb_f are signed values of a[NUM_LIMBS - 1] and b[NUM_LIMBS - 1]
115        // in prime field F.
116        let a_diff = a[NUM_LIMBS - 1] - cols.a_msb_f;
117        let b_diff = b[NUM_LIMBS - 1] - cols.b_msb_f;
118        builder.assert_zero(a_diff.clone() * (AB::Expr::from_u32(1 << LIMB_BITS) - a_diff));
119        builder.assert_zero(b_diff.clone() * (AB::Expr::from_u32(1 << LIMB_BITS) - b_diff));
120
121        for i in (0..NUM_LIMBS).rev() {
122            let diff = (if i == NUM_LIMBS - 1 {
123                cols.b_msb_f - cols.a_msb_f
124            } else {
125                b[i] - a[i]
126            }) * (AB::Expr::from_u8(2) * cols.cmp_lt - AB::Expr::ONE);
127            prefix_sum += marker[i].into();
128            builder.assert_bool(marker[i]);
129            builder.assert_zero(not::<AB::Expr>(prefix_sum.clone()) * diff.clone());
130            builder.when(marker[i]).assert_eq(cols.diff_val, diff);
131        }
132        // - If x != y, then prefix_sum = 1 so marker[i] must be 1 iff i is the first index where
133        //   diff != 0. Constrains that diff == diff_val where diff_val is non-zero.
134        // - If x == y, then prefix_sum = 0 and cmp_lt = 0. Here, prefix_sum cannot be 1 because all
135        //   diff are zero, making diff == diff_val fails.
136
137        builder.assert_bool(prefix_sum.clone());
138        builder
139            .when(not::<AB::Expr>(prefix_sum.clone()))
140            .assert_zero(cols.cmp_lt);
141
142        // Check if a_msb_f and b_msb_f are in [-128, 127) if signed, [0, 256) if unsigned.
143        self.bus
144            .send_range(
145                cols.a_msb_f + AB::Expr::from_u32(1 << (LIMB_BITS - 1)) * signed.clone(),
146                cols.b_msb_f + AB::Expr::from_u32(1 << (LIMB_BITS - 1)) * signed.clone(),
147            )
148            .eval(builder, is_valid.clone());
149
150        // Range check to ensure diff_val is non-zero.
151        self.bus
152            .send_range(cols.diff_val - AB::Expr::ONE, AB::F::ZERO)
153            .eval(builder, prefix_sum);
154
155        let expected_opcode = flags
156            .iter()
157            .zip(BranchLessThanOpcode::iter())
158            .fold(AB::Expr::ZERO, |acc, (flag, opcode)| {
159                acc + (*flag).into() * AB::Expr::from_u8(opcode as u8)
160            })
161            + AB::Expr::from_usize(self.offset);
162
163        let to_pc = from_pc
164            + cols.cmp_result * cols.imm
165            + not(cols.cmp_result) * AB::Expr::from_u32(DEFAULT_PC_STEP);
166
167        AdapterAirContext {
168            to_pc: Some(to_pc),
169            reads: [cols.a.map(Into::into), cols.b.map(Into::into)].into(),
170            writes: Default::default(),
171            instruction: ImmInstruction {
172                is_valid,
173                opcode: expected_opcode,
174                immediate: cols.imm.into(),
175            }
176            .into(),
177        }
178    }
179
180    fn start_offset(&self) -> usize {
181        self.offset
182    }
183}
184
185#[repr(C)]
186#[derive(AlignedBytesBorrow, Debug)]
187pub struct BranchLessThanCoreRecord<const NUM_LIMBS: usize, const LIMB_BITS: usize> {
188    pub a: [u8; NUM_LIMBS],
189    pub b: [u8; NUM_LIMBS],
190    pub imm: u32,
191    pub local_opcode: u8,
192}
193
194#[derive(Clone, Copy, derive_new::new)]
195pub struct BranchLessThanExecutor<A, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
196    adapter: A,
197    pub offset: usize,
198}
199
200#[derive(Clone, derive_new::new)]
201pub struct BranchLessThanFiller<A, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
202    adapter: A,
203    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<LIMB_BITS>,
204    pub offset: usize,
205}
206
207impl<F, A, RA, const NUM_LIMBS: usize, const LIMB_BITS: usize> PreflightExecutor<F, RA>
208    for BranchLessThanExecutor<A, NUM_LIMBS, LIMB_BITS>
209where
210    F: PrimeField32,
211    A: 'static + AdapterTraceExecutor<F, ReadData: Into<[[u8; NUM_LIMBS]; 2]>, WriteData = ()>,
212    for<'buf> RA: RecordArena<
213        'buf,
214        EmptyAdapterCoreLayout<F, A>,
215        (
216            A::RecordMut<'buf>,
217            &'buf mut BranchLessThanCoreRecord<NUM_LIMBS, LIMB_BITS>,
218        ),
219    >,
220{
221    fn get_opcode_name(&self, opcode: usize) -> String {
222        format!(
223            "{:?}",
224            BranchLessThanOpcode::from_usize(opcode - self.offset)
225        )
226    }
227
228    fn execute(
229        &self,
230        state: VmStateMut<F, TracingMemory, RA>,
231        instruction: &Instruction<F>,
232    ) -> Result<(), ExecutionError> {
233        let &Instruction { opcode, c: imm, .. } = instruction;
234
235        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
236
237        A::start(*state.pc, state.memory, &mut adapter_record);
238
239        let [rs1, rs2] = self
240            .adapter
241            .read(state.memory, instruction, &mut adapter_record)
242            .into();
243
244        core_record.a = rs1;
245        core_record.b = rs2;
246        core_record.imm = imm.as_canonical_u32();
247        core_record.local_opcode = opcode.local_opcode_idx(self.offset) as u8;
248
249        if run_cmp::<NUM_LIMBS, LIMB_BITS>(core_record.local_opcode, &rs1, &rs2).0 {
250            *state.pc = (F::from_u32(*state.pc) + imm).as_canonical_u32();
251        } else {
252            *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
253        }
254
255        Ok(())
256    }
257}
258
259impl<F, A, const NUM_LIMBS: usize, const LIMB_BITS: usize> TraceFiller<F>
260    for BranchLessThanFiller<A, NUM_LIMBS, LIMB_BITS>
261where
262    F: PrimeField32,
263    A: 'static + AdapterTraceFiller<F>,
264{
265    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
266        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
267        // BranchLessThanCoreCols::width() elements
268        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
269
270        // SAFETY: core_row contains a valid BranchLessThanCoreRecord written by the executor
271        // during trace generation
272        let record: &BranchLessThanCoreRecord<NUM_LIMBS, LIMB_BITS> =
273            unsafe { get_record_from_slice(&mut core_row, ()) };
274
275        self.adapter.fill_trace_row(mem_helper, adapter_row);
276        let core_row: &mut BranchLessThanCoreCols<F, NUM_LIMBS, LIMB_BITS> = core_row.borrow_mut();
277
278        let signed = record.local_opcode == BranchLessThanOpcode::BLT as u8
279            || record.local_opcode == BranchLessThanOpcode::BGE as u8;
280        let ge_op = record.local_opcode == BranchLessThanOpcode::BGE as u8
281            || record.local_opcode == BranchLessThanOpcode::BGEU as u8;
282
283        let (cmp_result, diff_idx, a_sign, b_sign) =
284            run_cmp::<NUM_LIMBS, LIMB_BITS>(record.local_opcode, &record.a, &record.b);
285
286        let cmp_lt = cmp_result ^ ge_op;
287
288        // We range check (a_msb_f + 128) and (b_msb_f + 128) if signed,
289        // a_msb_f and b_msb_f if not
290        let (a_msb_f, a_msb_range) = if a_sign {
291            (
292                -F::from_u32((1 << LIMB_BITS) - record.a[NUM_LIMBS - 1] as u32),
293                record.a[NUM_LIMBS - 1] as u32 - (1 << (LIMB_BITS - 1)),
294            )
295        } else {
296            (
297                F::from_u32(record.a[NUM_LIMBS - 1] as u32),
298                record.a[NUM_LIMBS - 1] as u32 + ((signed as u32) << (LIMB_BITS - 1)),
299            )
300        };
301        let (b_msb_f, b_msb_range) = if b_sign {
302            (
303                -F::from_u32((1 << LIMB_BITS) - record.b[NUM_LIMBS - 1] as u32),
304                record.b[NUM_LIMBS - 1] as u32 - (1 << (LIMB_BITS - 1)),
305            )
306        } else {
307            (
308                F::from_u32(record.b[NUM_LIMBS - 1] as u32),
309                record.b[NUM_LIMBS - 1] as u32 + ((signed as u32) << (LIMB_BITS - 1)),
310            )
311        };
312
313        core_row.diff_val = if diff_idx == NUM_LIMBS {
314            F::ZERO
315        } else if diff_idx == (NUM_LIMBS - 1) {
316            if cmp_lt {
317                b_msb_f - a_msb_f
318            } else {
319                a_msb_f - b_msb_f
320            }
321        } else if cmp_lt {
322            F::from_u8(record.b[diff_idx] - record.a[diff_idx])
323        } else {
324            F::from_u8(record.a[diff_idx] - record.b[diff_idx])
325        };
326
327        self.bitwise_lookup_chip
328            .request_range(a_msb_range, b_msb_range);
329
330        core_row.diff_marker = [F::ZERO; NUM_LIMBS];
331
332        if diff_idx != NUM_LIMBS {
333            self.bitwise_lookup_chip
334                .request_range(core_row.diff_val.as_canonical_u32() - 1, 0);
335            core_row.diff_marker[diff_idx] = F::ONE;
336        }
337
338        core_row.cmp_lt = F::from_bool(cmp_lt);
339        core_row.b_msb_f = b_msb_f;
340        core_row.a_msb_f = a_msb_f;
341        core_row.opcode_bgeu_flag =
342            F::from_bool(record.local_opcode == BranchLessThanOpcode::BGEU as u8);
343        core_row.opcode_bge_flag =
344            F::from_bool(record.local_opcode == BranchLessThanOpcode::BGE as u8);
345        core_row.opcode_bltu_flag =
346            F::from_bool(record.local_opcode == BranchLessThanOpcode::BLTU as u8);
347        core_row.opcode_blt_flag =
348            F::from_bool(record.local_opcode == BranchLessThanOpcode::BLT as u8);
349
350        core_row.imm = F::from_u32(record.imm);
351        core_row.cmp_result = F::from_bool(cmp_result);
352        core_row.b = record.b.map(F::from_u8);
353        core_row.a = record.a.map(F::from_u8);
354    }
355}
356
357// Returns (cmp_result, diff_idx, x_sign, y_sign)
358#[inline(always)]
359pub(super) fn run_cmp<const NUM_LIMBS: usize, const LIMB_BITS: usize>(
360    local_opcode: u8,
361    x: &[u8; NUM_LIMBS],
362    y: &[u8; NUM_LIMBS],
363) -> (bool, usize, bool, bool) {
364    let signed = local_opcode == BranchLessThanOpcode::BLT as u8
365        || local_opcode == BranchLessThanOpcode::BGE as u8;
366    let ge_op = local_opcode == BranchLessThanOpcode::BGE as u8
367        || local_opcode == BranchLessThanOpcode::BGEU as u8;
368    let x_sign = (x[NUM_LIMBS - 1] >> (LIMB_BITS - 1) == 1) && signed;
369    let y_sign = (y[NUM_LIMBS - 1] >> (LIMB_BITS - 1) == 1) && signed;
370    for i in (0..NUM_LIMBS).rev() {
371        if x[i] != y[i] {
372            return ((x[i] < y[i]) ^ x_sign ^ y_sign ^ ge_op, i, x_sign, y_sign);
373        }
374    }
375    (ge_op, NUM_LIMBS, x_sign, y_sign)
376}
openvm_rv32im_circuit/branch_lt/core.rs

openvm_rv32im_circuit/branch_lt/
core.rs
