openvm_rv32im_circuit/less_than/
core.rs

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