openvm_mod_circuit_builder/
core_chip.rs

1use std::{
2    marker::PhantomData,
3    mem::{align_of, size_of},
4    sync::Arc,
5};
6
7use itertools::Itertools;
8use num_bigint::BigUint;
9use num_traits::Zero;
10use openvm_circuit::{
11    arch::*,
12    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
13};
14use openvm_circuit_primitives::{
15    var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerChip},
16    ColumnsAir, SubAir, TraceSubRowGenerator,
17};
18use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP};
19use openvm_stark_backend::{
20    interaction::InteractionBuilder,
21    p3_air::BaseAir,
22    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
23    BaseAirWithPublicValues,
24};
25use openvm_stark_sdk::p3_baby_bear::BabyBear;
26
27use crate::builder::{FieldExpr, FieldExprCols};
28
29#[derive(Clone)]
30pub struct FieldExpressionCoreAir {
31    pub expr: FieldExpr,
32
33    /// The global opcode offset.
34    pub offset: usize,
35
36    /// All the opcode indices (including setup) supported by this Air.
37    /// The last one must be the setup opcode if it's a chip needs setup.
38    pub local_opcode_idx: Vec<usize>,
39    /// Opcode flag idx (indices from builder.new_flag()) for all except setup opcode. Empty if
40    /// single op chip.
41    pub opcode_flag_idx: Vec<usize>,
42    // Example 1: 1-op chip EcAdd that needs setup
43    //   local_opcode_idx = [0, 2], where 0 is EcAdd, 2 is setup
44    //   opcode_flag_idx = [], not needed for single op chip.
45    // Example 2: 1-op chip EvaluateLine that doesn't need setup
46    //   local_opcode_idx = [2], the id within PairingOpcodeEnum
47    //   opcode_flag_idx = [], not needed
48    // Example 3: 2-op chip MulDiv that needs setup
49    //   local_opcode_idx = [2, 3, 4], where 2 is Mul, 3 is Div, 4 is setup
50    //   opcode_flag_idx = [0, 1], where 0 is mul_flag, 1 is div_flag, in the builder
51    // We don't support 2-op chip that doesn't need setup right now.
52}
53
54// No columns provided: wraps `FieldExpr`, whose column layout is built dynamically.
55impl ColumnsAir for FieldExpressionCoreAir {}
56
57impl FieldExpressionCoreAir {
58    pub fn new(
59        expr: FieldExpr,
60        offset: usize,
61        local_opcode_idx: Vec<usize>,
62        opcode_flag_idx: Vec<usize>,
63    ) -> Self {
64        let opcode_flag_idx = if opcode_flag_idx.is_empty() && expr.needs_setup() {
65            // single op chip that needs setup, so there is only one default flag, must be 0.
66            vec![0]
67        } else {
68            // multi ops chip or no-setup chip, use as is.
69            opcode_flag_idx
70        };
71        assert_eq!(opcode_flag_idx.len(), local_opcode_idx.len() - 1);
72        Self {
73            expr,
74            offset,
75            local_opcode_idx,
76            opcode_flag_idx,
77        }
78    }
79
80    pub fn num_inputs(&self) -> usize {
81        self.expr.builder.num_input
82    }
83
84    pub fn num_vars(&self) -> usize {
85        self.expr.builder.num_variables
86    }
87
88    pub fn num_flags(&self) -> usize {
89        self.expr.builder.num_flags
90    }
91
92    pub fn output_indices(&self) -> &[usize] {
93        &self.expr.builder.output_indices
94    }
95}
96
97impl<F: Field> BaseAir<F> for FieldExpressionCoreAir {
98    fn width(&self) -> usize {
99        BaseAir::<F>::width(&self.expr)
100    }
101}
102
103impl<F: Field> BaseAirWithPublicValues<F> for FieldExpressionCoreAir {}
104
105impl<AB: InteractionBuilder, I> VmCoreAir<AB, I> for FieldExpressionCoreAir
106where
107    I: VmAdapterInterface<AB::Expr>,
108    AdapterAirContext<AB::Expr, I>:
109        From<AdapterAirContext<AB::Expr, DynAdapterInterface<AB::Expr>>>,
110{
111    fn eval(
112        &self,
113        builder: &mut AB,
114        local: &[AB::Var],
115        _from_pc: AB::Var,
116    ) -> AdapterAirContext<AB::Expr, I> {
117        assert_eq!(local.len(), BaseAir::<AB::F>::width(&self.expr));
118        self.expr.eval(builder, local);
119        let FieldExprCols {
120            is_valid,
121            inputs,
122            vars,
123            flags,
124            ..
125        } = self.expr.load_vars(local);
126        assert_eq!(inputs.len(), self.num_inputs());
127        assert_eq!(vars.len(), self.num_vars());
128        assert_eq!(flags.len(), self.num_flags());
129        let reads: Vec<AB::Expr> = inputs.concat().iter().map(|x| (*x).into()).collect();
130        let writes: Vec<AB::Expr> = self
131            .output_indices()
132            .iter()
133            .flat_map(|&i| vars[i].clone())
134            .map(Into::into)
135            .collect();
136
137        let opcode_flags_except_last = self.opcode_flag_idx.iter().map(|&i| flags[i]).collect_vec();
138        let last_opcode_flag = is_valid
139            - opcode_flags_except_last
140                .iter()
141                .map(|&v| v.into())
142                .sum::<AB::Expr>();
143        builder.assert_bool(last_opcode_flag.clone());
144        let opcode_flags = opcode_flags_except_last
145            .into_iter()
146            .map(Into::into)
147            .chain(Some(last_opcode_flag));
148        let expected_opcode = opcode_flags
149            .zip(self.local_opcode_idx.iter().map(|&i| i + self.offset))
150            .map(|(flag, global_idx)| flag * AB::Expr::from_usize(global_idx))
151            .sum();
152
153        let instruction = MinimalInstruction {
154            is_valid: is_valid.into(),
155            opcode: expected_opcode,
156        };
157
158        let ctx: AdapterAirContext<_, DynAdapterInterface<_>> = AdapterAirContext {
159            to_pc: None,
160            reads: reads.into(),
161            writes: writes.into(),
162            instruction: instruction.into(),
163        };
164        ctx.into()
165    }
166
167    fn start_offset(&self) -> usize {
168        self.offset
169    }
170}
171
172pub struct FieldExpressionMetadata<F, A> {
173    pub total_input_limbs: usize, // num_inputs * limbs_per_input
174    _phantom: PhantomData<(F, A)>,
175}
176
177impl<F, A> Clone for FieldExpressionMetadata<F, A> {
178    fn clone(&self) -> Self {
179        Self {
180            total_input_limbs: self.total_input_limbs,
181            _phantom: PhantomData,
182        }
183    }
184}
185
186impl<F, A> Default for FieldExpressionMetadata<F, A> {
187    fn default() -> Self {
188        Self {
189            total_input_limbs: 0,
190            _phantom: PhantomData,
191        }
192    }
193}
194
195impl<F, A> FieldExpressionMetadata<F, A> {
196    pub fn new(total_input_limbs: usize) -> Self {
197        Self {
198            total_input_limbs,
199            _phantom: PhantomData,
200        }
201    }
202}
203
204impl<F, A> AdapterCoreMetadata for FieldExpressionMetadata<F, A>
205where
206    A: AdapterTraceExecutor<F>,
207{
208    #[inline(always)]
209    fn get_adapter_width() -> usize {
210        A::WIDTH * size_of::<F>()
211    }
212}
213
214pub type FieldExpressionRecordLayout<F, A> = AdapterCoreLayout<FieldExpressionMetadata<F, A>>;
215
216pub struct FieldExpressionCoreRecordMut<'a> {
217    pub opcode: &'a mut u8,
218    pub input_limbs: &'a mut [u8],
219}
220
221impl<'a, F, A> CustomBorrow<'a, FieldExpressionCoreRecordMut<'a>, FieldExpressionRecordLayout<F, A>>
222    for [u8]
223{
224    fn custom_borrow(
225        &'a mut self,
226        layout: FieldExpressionRecordLayout<F, A>,
227    ) -> FieldExpressionCoreRecordMut<'a> {
228        // SAFETY: The buffer length is the width of the trace which should be at least 1
229        let (opcode_buf, input_limbs_buff) = unsafe { self.split_at_mut_unchecked(1) };
230
231        // SAFETY: opcode_buf has exactly 1 element from split_at_mut_unchecked(1)
232        let opcode_buf = unsafe { opcode_buf.get_unchecked_mut(0) };
233
234        FieldExpressionCoreRecordMut {
235            opcode: opcode_buf,
236            input_limbs: &mut input_limbs_buff[..layout.metadata.total_input_limbs],
237        }
238    }
239
240    unsafe fn extract_layout(&self) -> FieldExpressionRecordLayout<F, A> {
241        panic!("Should get the Layout information from FieldExpressionExecutor");
242    }
243}
244
245impl<F, A> SizedRecord<FieldExpressionRecordLayout<F, A>> for FieldExpressionCoreRecordMut<'_> {
246    fn size(layout: &FieldExpressionRecordLayout<F, A>) -> usize {
247        layout.metadata.total_input_limbs + 1
248    }
249
250    fn alignment(_layout: &FieldExpressionRecordLayout<F, A>) -> usize {
251        align_of::<u8>()
252    }
253}
254
255impl<'a> FieldExpressionCoreRecordMut<'a> {
256    // This method is only used in testing
257    pub fn new_from_execution_data(
258        buffer: &'a mut [u8],
259        inputs: &[BigUint],
260        limbs_per_input: usize,
261    ) -> Self {
262        let record_info = FieldExpressionMetadata::<(), ()>::new(inputs.len() * limbs_per_input);
263
264        let record: Self = buffer.custom_borrow(FieldExpressionRecordLayout {
265            metadata: record_info,
266        });
267        record
268    }
269
270    #[inline(always)]
271    pub fn fill_from_execution_data(&mut self, opcode: u8, data: &[u8]) {
272        // Rust will assert that length of `data` and `self.input_limbs` are the same
273        // That is `data.len() == num_inputs * limbs_per_input`
274        *self.opcode = opcode;
275        self.input_limbs.copy_from_slice(data);
276    }
277}
278
279#[derive(Clone)]
280pub struct FieldExpressionExecutor<A> {
281    adapter: A,
282    pub expr: FieldExpr,
283    pub offset: usize,
284    pub local_opcode_idx: Vec<usize>,
285    pub opcode_flag_idx: Vec<usize>,
286    pub name: String,
287}
288
289impl<A> FieldExpressionExecutor<A> {
290    #[allow(clippy::too_many_arguments)]
291    pub fn new(
292        adapter: A,
293        expr: FieldExpr,
294        offset: usize,
295        local_opcode_idx: Vec<usize>,
296        opcode_flag_idx: Vec<usize>,
297        name: &str,
298    ) -> Self {
299        let opcode_flag_idx = if opcode_flag_idx.is_empty() && expr.needs_setup() {
300            // single op chip that needs setup, so there is only one default flag, must be 0.
301            vec![0]
302        } else {
303            // multi ops chip or no-setup chip, use as is.
304            opcode_flag_idx
305        };
306        assert_eq!(opcode_flag_idx.len(), local_opcode_idx.len() - 1);
307        tracing::debug!(
308            "FieldExpressionCoreExecutor: opcode={name}, main_width={}",
309            BaseAir::<BabyBear>::width(&expr)
310        );
311        Self {
312            adapter,
313            expr,
314            offset,
315            local_opcode_idx,
316            opcode_flag_idx,
317            name: name.to_string(),
318        }
319    }
320
321    pub fn get_record_layout<F>(&self) -> FieldExpressionRecordLayout<F, A> {
322        FieldExpressionRecordLayout {
323            metadata: FieldExpressionMetadata::new(
324                self.expr.builder.num_input * self.expr.canonical_num_limbs(),
325            ),
326        }
327    }
328
329    /// Returns a reference to the adapter for use in custom PreflightExecutor implementations.
330    #[inline]
331    pub fn adapter(&self) -> &A {
332        &self.adapter
333    }
334}
335
336pub struct FieldExpressionFiller<A> {
337    adapter: A,
338    pub expr: FieldExpr,
339    pub local_opcode_idx: Vec<usize>,
340    pub opcode_flag_idx: Vec<usize>,
341    pub range_checker: SharedVariableRangeCheckerChip,
342    pub should_finalize: bool,
343}
344
345impl<A> FieldExpressionFiller<A> {
346    #[allow(clippy::too_many_arguments)]
347    pub fn new(
348        adapter: A,
349        expr: FieldExpr,
350        local_opcode_idx: Vec<usize>,
351        opcode_flag_idx: Vec<usize>,
352        range_checker: SharedVariableRangeCheckerChip,
353        should_finalize: bool,
354    ) -> Self {
355        let opcode_flag_idx = if opcode_flag_idx.is_empty() && expr.needs_setup() {
356            // single op chip that needs setup, so there is only one default flag, must be 0.
357            vec![0]
358        } else {
359            // multi ops chip or no-setup chip, use as is.
360            opcode_flag_idx
361        };
362        assert_eq!(opcode_flag_idx.len(), local_opcode_idx.len() - 1);
363        Self {
364            adapter,
365            expr,
366            local_opcode_idx,
367            opcode_flag_idx,
368            range_checker,
369            should_finalize,
370        }
371    }
372    pub fn num_inputs(&self) -> usize {
373        self.expr.builder.num_input
374    }
375
376    pub fn num_flags(&self) -> usize {
377        self.expr.builder.num_flags
378    }
379
380    pub fn get_record_layout<F>(&self) -> FieldExpressionRecordLayout<F, A> {
381        FieldExpressionRecordLayout {
382            metadata: FieldExpressionMetadata::new(
383                self.num_inputs() * self.expr.canonical_num_limbs(),
384            ),
385        }
386    }
387}
388
389impl<F, A, RA> PreflightExecutor<F, RA> for FieldExpressionExecutor<A>
390where
391    F: PrimeField32,
392    A: 'static
393        + AdapterTraceExecutor<F, ReadData: Into<DynArray<u8>>, WriteData: From<DynArray<u8>>>,
394    for<'buf> RA: RecordArena<
395        'buf,
396        FieldExpressionRecordLayout<F, A>,
397        (A::RecordMut<'buf>, FieldExpressionCoreRecordMut<'buf>),
398    >,
399{
400    fn execute(
401        &self,
402        state: VmStateMut<F, TracingMemory, RA>,
403        instruction: &Instruction<F>,
404    ) -> Result<(), ExecutionError> {
405        let (mut adapter_record, mut core_record) = state.ctx.alloc(self.get_record_layout());
406
407        A::start(*state.pc, state.memory, &mut adapter_record);
408
409        let data: DynArray<_> = self
410            .adapter
411            .read(state.memory, instruction, &mut adapter_record)
412            .into();
413
414        core_record.fill_from_execution_data(
415            instruction.opcode.local_opcode_idx(self.offset) as u8,
416            &data.0,
417        );
418
419        let (writes, _, _) = run_field_expression(
420            &self.expr,
421            &self.local_opcode_idx,
422            &self.opcode_flag_idx,
423            core_record.input_limbs,
424            *core_record.opcode as usize,
425        );
426
427        self.adapter.write(
428            state.memory,
429            instruction,
430            writes.into(),
431            &mut adapter_record,
432        );
433
434        *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
435        Ok(())
436    }
437
438    fn get_opcode_name(&self, _opcode: usize) -> String {
439        self.name.clone()
440    }
441}
442
443#[cfg(feature = "aot")]
444impl<F: PrimeField32, A> AotExecutor<F> for FieldExpressionExecutor<A> {}
445
446impl<F, A> TraceFiller<F> for FieldExpressionFiller<A>
447where
448    F: PrimeField32 + Send + Sync + Clone,
449    A: 'static + AdapterTraceFiller<F>,
450{
451    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
452        // Get the core record from the row slice
453        // SAFETY: Caller guarantees that row_slice has width A::WIDTH + core width
454        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
455
456        self.adapter.fill_trace_row(mem_helper, adapter_row);
457
458        // SAFETY:
459        // - caller ensures `core_row` contains a valid record representation that was previously
460        //   written by the executor
461        // - core_row slice is transmuted to FieldExpressionCoreRecordMut using the specified
462        //   layout, which satisfies CustomBorrow requirements for safe access.
463        let record: FieldExpressionCoreRecordMut =
464            unsafe { get_record_from_slice(&mut core_row, self.get_record_layout::<F>()) };
465
466        let (_, inputs, flags) = run_field_expression(
467            &self.expr,
468            &self.local_opcode_idx,
469            &self.opcode_flag_idx,
470            record.input_limbs,
471            *record.opcode as usize,
472        );
473
474        let range_checker = self.range_checker.as_ref();
475        self.expr
476            .generate_subrow((range_checker, inputs, flags), core_row);
477    }
478
479    fn fill_dummy_trace_row(&self, row_slice: &mut [F]) {
480        if !self.should_finalize {
481            return;
482        }
483
484        let inputs: Vec<BigUint> = vec![BigUint::zero(); self.num_inputs()];
485        let flags: Vec<bool> = vec![false; self.num_flags()];
486        let core_row = &mut row_slice[A::WIDTH..];
487        // We **do not** want this trace row to update the range checker
488        // so we must create a temporary range checker
489        let tmp_range_checker = Arc::new(VariableRangeCheckerChip::new(self.range_checker.bus()));
490        self.expr
491            .generate_subrow((&tmp_range_checker, inputs, flags), core_row);
492        core_row[0] = F::ZERO; // is_valid = 0
493    }
494}
495
496fn run_field_expression(
497    expr: &FieldExpr,
498    local_opcode_flags: &[usize],
499    opcode_flag_idx: &[usize],
500    data: &[u8],
501    local_opcode_idx: usize,
502) -> (DynArray<u8>, Vec<BigUint>, Vec<bool>) {
503    let field_element_limbs = expr.canonical_num_limbs();
504    assert_eq!(data.len(), expr.builder.num_input * field_element_limbs);
505
506    let mut inputs = Vec::with_capacity(expr.builder.num_input);
507    for i in 0..expr.builder.num_input {
508        let start = i * field_element_limbs;
509        let end = start + field_element_limbs;
510        let limb_slice = &data[start..end];
511        let input = BigUint::from_bytes_le(limb_slice);
512        inputs.push(input);
513    }
514
515    let mut flags = vec![];
516    if expr.needs_setup() {
517        flags = vec![false; expr.builder.num_flags];
518
519        // Find which opcode this is in our local_opcode_idx list
520        if let Some(opcode_position) = local_opcode_flags
521            .iter()
522            .position(|&idx| idx == local_opcode_idx)
523        {
524            // If this is NOT the last opcode (setup), set the corresponding flag
525            if opcode_position < opcode_flag_idx.len() {
526                let flag_idx = opcode_flag_idx[opcode_position];
527                flags[flag_idx] = true;
528            }
529            // If opcode_position == step.opcode_flag_idx.len(), it's the setup operation
530            // and all flags should remain false (which they already are)
531        }
532    }
533
534    let vars = expr.execute(&inputs, &flags);
535    assert_eq!(vars.len(), expr.builder.num_variables);
536
537    // Write outputs directly to a pre-allocated buffer to avoid intermediate Vecs
538    let num_outputs = expr.builder.output_indices.len();
539    let total_output_bytes = num_outputs * field_element_limbs;
540    let mut write_buffer = vec![0u8; total_output_bytes];
541    for (i, &var_idx) in expr.builder.output_indices.iter().enumerate() {
542        let start = i * field_element_limbs;
543        let bytes = vars[var_idx].to_bytes_le();
544        let copy_len = bytes.len().min(field_element_limbs);
545        write_buffer[start..start + copy_len].copy_from_slice(&bytes[..copy_len]);
546        // Remaining bytes are already zero from vec![0u8; ...]
547    }
548    let writes: DynArray<_> = write_buffer.into();
549
550    (writes, inputs, flags)
551}
552
553#[inline(always)]
554pub fn run_field_expression_precomputed<const NEEDS_SETUP: bool>(
555    expr: &FieldExpr,
556    flag_idx: usize,
557    data: &[u8],
558) -> DynArray<u8> {
559    let field_element_limbs = expr.canonical_num_limbs();
560    assert_eq!(data.len(), expr.num_inputs() * field_element_limbs);
561
562    let mut inputs = Vec::with_capacity(expr.num_inputs());
563    for i in 0..expr.num_inputs() {
564        let start = i * expr.canonical_num_limbs();
565        let end = start + expr.canonical_num_limbs();
566        let limb_slice = &data[start..end];
567        let input = BigUint::from_bytes_le(limb_slice);
568        inputs.push(input);
569    }
570
571    let flags = if NEEDS_SETUP {
572        let mut flags = vec![false; expr.num_flags()];
573        if flag_idx < expr.num_flags() {
574            flags[flag_idx] = true;
575        }
576        flags
577    } else {
578        vec![]
579    };
580
581    let vars = expr.execute(&inputs, &flags);
582    assert_eq!(vars.len(), expr.num_vars());
583
584    // Write outputs directly to a pre-allocated buffer to avoid intermediate Vecs
585    let num_outputs = expr.output_indices().len();
586    let total_output_bytes = num_outputs * field_element_limbs;
587    let mut write_buffer = vec![0u8; total_output_bytes];
588    for (i, &var_idx) in expr.output_indices().iter().enumerate() {
589        let start = i * field_element_limbs;
590        let bytes = vars[var_idx].to_bytes_le();
591        let copy_len = bytes.len().min(field_element_limbs);
592        write_buffer[start..start + copy_len].copy_from_slice(&bytes[..copy_len]);
593    }
594    write_buffer.into()
595}
openvm_mod_circuit_builder/core_chip.rs

openvm_mod_circuit_builder/
core_chip.rs
