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 pub b_msb_f: T,
39 pub c_msb_f: T,
40
41 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 builder.assert_bool(prefix_sum.clone());
118 builder
119 .when(not::<AB::Expr>(prefix_sum.clone()))
120 .assert_zero(cols.cmp_result);
121
122 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 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 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 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 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#[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}