openvm_circuit/system/public_values/
core.rs

1use std::marker::PhantomData;
2
3use getset::Setters;
4use openvm_circuit_primitives::{encoder::Encoder, AlignedBytesBorrow, SubAir};
5use openvm_instructions::{
6    instruction::Instruction,
7    program::DEFAULT_PC_STEP,
8    LocalOpcode,
9    PublishOpcode::{self, PUBLISH},
10};
11use openvm_stark_backend::{
12    interaction::InteractionBuilder,
13    p3_air::{AirBuilder, AirBuilderWithPublicValues, BaseAir},
14    p3_field::{Field, FieldAlgebra, PrimeField32},
15    rap::BaseAirWithPublicValues,
16};
17
18use crate::{
19    arch::{
20        get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller,
21        BasicAdapterInterface, EmptyAdapterCoreLayout, ExecutionError, MinimalInstruction,
22        PreflightExecutor, RecordArena, TraceFiller, VmCoreAir, VmStateMut,
23    },
24    system::{
25        memory::{online::TracingMemory, MemoryAuxColsFactory},
26        native_adapter::NativeAdapterExecutor,
27        public_values::columns::PublicValuesCoreColsView,
28    },
29};
30
31pub(crate) type AdapterInterface<F> = BasicAdapterInterface<F, MinimalInstruction<F>, 2, 0, 1, 1>;
32
33#[derive(Clone, Debug)]
34pub struct PublicValuesCoreAir {
35    /// Number of custom public values to publish.
36    pub num_custom_pvs: usize,
37    encoder: Encoder,
38}
39
40impl PublicValuesCoreAir {
41    pub fn new(num_custom_pvs: usize, max_degree: u32) -> Self {
42        Self {
43            num_custom_pvs,
44            encoder: Encoder::new(num_custom_pvs, max_degree, true),
45        }
46    }
47}
48
49impl<F: Field> BaseAir<F> for PublicValuesCoreAir {
50    fn width(&self) -> usize {
51        3 + self.encoder.width()
52    }
53}
54
55impl<F: Field> BaseAirWithPublicValues<F> for PublicValuesCoreAir {
56    fn num_public_values(&self) -> usize {
57        self.num_custom_pvs
58    }
59}
60
61impl<AB: InteractionBuilder + AirBuilderWithPublicValues> VmCoreAir<AB, AdapterInterface<AB::Expr>>
62    for PublicValuesCoreAir
63{
64    fn eval(
65        &self,
66        builder: &mut AB,
67        local_core: &[AB::Var],
68        _from_pc: AB::Var,
69    ) -> AdapterAirContext<AB::Expr, AdapterInterface<AB::Expr>> {
70        let cols = PublicValuesCoreColsView::<_, &AB::Var>::borrow(local_core);
71        debug_assert_eq!(cols.width(), BaseAir::<AB::F>::width(self));
72        let is_valid = *cols.is_valid;
73        let value = *cols.value;
74        let index = *cols.index;
75
76        let vars = cols.custom_pv_vars.iter().map(|&&x| x).collect::<Vec<_>>();
77        self.encoder.eval(builder, &vars);
78
79        let flags = self.encoder.flags::<AB>(&vars);
80
81        let mut match_public_value_index = AB::Expr::ZERO;
82        let mut match_public_value = AB::Expr::ZERO;
83        for (i, flag) in flags.iter().enumerate() {
84            match_public_value_index += flag.clone() * AB::F::from_canonical_usize(i);
85            match_public_value += flag.clone() * builder.public_values()[i].into();
86        }
87        builder.assert_eq(is_valid, self.encoder.is_valid::<AB>(&vars));
88
89        let mut when_publish = builder.when(is_valid);
90        when_publish.assert_eq(index, match_public_value_index);
91        when_publish.assert_eq(value, match_public_value);
92
93        AdapterAirContext {
94            to_pc: None,
95            reads: [[value.into()], [index.into()]],
96            writes: [],
97            instruction: MinimalInstruction {
98                is_valid: is_valid.into(),
99                opcode: AB::Expr::from_canonical_usize(PUBLISH.global_opcode().as_usize()),
100            },
101        }
102    }
103
104    fn start_offset(&self) -> usize {
105        PublishOpcode::CLASS_OFFSET
106    }
107}
108
109#[repr(C)]
110#[derive(AlignedBytesBorrow, Debug)]
111pub struct PublicValuesRecord<F> {
112    pub value: F,
113    pub index: F,
114}
115
116/// ATTENTION: If a specific public value is not provided, a default 0 will be used when generating
117/// the proof but in the perspective of constraints, it could be any value.
118#[derive(Clone)]
119pub struct PublicValuesExecutor<F, A = NativeAdapterExecutor<F, 2, 0>> {
120    adapter: A,
121    phantom: PhantomData<F>,
122}
123
124#[derive(Clone, Setters)]
125pub struct PublicValuesFiller<F, A = NativeAdapterExecutor<F, 2, 0>> {
126    adapter: A,
127    encoder: Encoder,
128    num_custom_pvs: usize,
129    public_values: Vec<F>,
130}
131
132impl<F: Clone, A> PublicValuesExecutor<F, A> {
133    pub fn new(adapter: A) -> Self {
134        Self {
135            adapter,
136            phantom: PhantomData,
137        }
138    }
139}
140
141impl<F: Clone, A> PublicValuesFiller<F, A> {
142    /// **Note:** `max_degree` is the maximum degree of the constraint polynomials to represent the
143    /// flags. If you want the overall AIR's constraint degree to be `<= max_constraint_degree`,
144    /// then typically you should set `max_degree` to `max_constraint_degree - 1`.
145    pub fn new(adapter: A, num_custom_pvs: usize, max_degree: u32) -> Self {
146        Self {
147            adapter,
148            encoder: Encoder::new(num_custom_pvs, max_degree, true),
149            num_custom_pvs,
150            public_values: Vec::new(),
151        }
152    }
153
154    pub fn set_public_values(&mut self, public_values: Vec<F>)
155    where
156        F: Field,
157    {
158        assert_eq!(public_values.len(), self.num_custom_pvs);
159        self.public_values = public_values;
160    }
161}
162
163impl<F, A, RA> PreflightExecutor<F, RA> for PublicValuesExecutor<F, A>
164where
165    F: PrimeField32,
166    A: 'static + Clone + AdapterTraceExecutor<F, ReadData = [[F; 1]; 2], WriteData = [[F; 1]; 0]>,
167    for<'buf> RA: RecordArena<
168        'buf,
169        EmptyAdapterCoreLayout<F, A>,
170        (A::RecordMut<'buf>, &'buf mut PublicValuesRecord<F>),
171    >,
172{
173    fn get_opcode_name(&self, opcode: usize) -> String {
174        format!(
175            "{:?}",
176            PublishOpcode::from_usize(opcode - PublishOpcode::CLASS_OFFSET)
177        )
178    }
179
180    fn execute(
181        &self,
182        state: VmStateMut<F, TracingMemory, RA>,
183        instruction: &Instruction<F>,
184    ) -> Result<(), ExecutionError> {
185        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
186
187        A::start(*state.pc, state.memory, &mut adapter_record);
188
189        [[core_record.value], [core_record.index]] =
190            self.adapter
191                .read(state.memory, instruction, &mut adapter_record);
192        {
193            let idx: usize = core_record.index.as_canonical_u32() as usize;
194            let custom_pvs = state.custom_pvs;
195
196            if custom_pvs[idx].is_none() {
197                custom_pvs[idx] = Some(core_record.value);
198            } else {
199                // Not a hard constraint violation when publishing the same value twice but the
200                // program should avoid that.
201                panic!("Custom public value {} already set", idx);
202            }
203        }
204
205        *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
206
207        Ok(())
208    }
209}
210
211impl<F, A> TraceFiller<F> for PublicValuesFiller<F, A>
212where
213    F: PrimeField32,
214    A: 'static + AdapterTraceFiller<F>,
215{
216    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
217        // SAFETY:
218        // - row_slice is guaranteed by the caller to have at least A::WIDTH +
219        //   PublicValuesCoreColsView::width() elements
220        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
221        self.adapter.fill_trace_row(mem_helper, adapter_row);
222        // SAFETY:
223        // - caller ensures core_row contains a valid record written by the executor during trace
224        //   generation
225        let record: &PublicValuesRecord<F> = unsafe { get_record_from_slice(&mut core_row, ()) };
226        let cols = PublicValuesCoreColsView::<_, &mut F>::borrow_mut(core_row);
227
228        let idx: usize = record.index.as_canonical_u32() as usize;
229        let pt = self.encoder.get_flag_pt(idx);
230
231        cols.custom_pv_vars
232            .into_iter()
233            .zip(pt.iter())
234            .for_each(|(var, &val)| {
235                *var = F::from_canonical_u32(val);
236            });
237
238        *cols.index = record.index;
239        *cols.value = record.value;
240        *cols.is_valid = F::ONE;
241    }
242
243    fn generate_public_values(&self) -> Vec<F> {
244        assert_eq!(
245            self.public_values.len(),
246            self.num_custom_pvs,
247            "Did not set public values"
248        );
249        self.public_values.clone()
250    }
251}