openvm_circuit/system/public_values/
execution.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    mem::size_of,
4};
5
6use openvm_circuit_primitives::AlignedBytesBorrow;
7use openvm_instructions::{
8    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_IMM_AS, NATIVE_AS,
9};
10use openvm_stark_backend::p3_field::PrimeField32;
11
12use super::PublicValuesExecutor;
13#[cfg(not(feature = "tco"))]
14use crate::arch::ExecuteFunc;
15#[cfg(feature = "tco")]
16use crate::arch::Handler;
17use crate::{
18    arch::{
19        create_handler,
20        execution_mode::{ExecutionCtxTrait, MeteredExecutionCtxTrait},
21        E2PreCompute, Executor, MeteredExecutor, StaticProgramError, VmExecState,
22    },
23    system::memory::online::GuestMemory,
24    utils::{transmute_field_to_u32, transmute_u32_to_field},
25};
26
27#[derive(AlignedBytesBorrow)]
28#[repr(C)]
29struct PublicValuesPreCompute {
30    b_or_imm: u32,
31    c_or_imm: u32,
32}
33
34impl<F, A> PublicValuesExecutor<F, A>
35where
36    F: PrimeField32,
37{
38    fn pre_compute_impl(
39        &self,
40        inst: &Instruction<F>,
41        data: &mut PublicValuesPreCompute,
42    ) -> (bool, bool) {
43        let &Instruction { b, c, e, f, .. } = inst;
44
45        let e = e.as_canonical_u32();
46        let f = f.as_canonical_u32();
47
48        let b_is_imm = e == RV32_IMM_AS;
49        let c_is_imm = f == RV32_IMM_AS;
50
51        let b_or_imm = if b_is_imm {
52            transmute_field_to_u32(&b)
53        } else {
54            b.as_canonical_u32()
55        };
56        let c_or_imm = if c_is_imm {
57            transmute_field_to_u32(&c)
58        } else {
59            c.as_canonical_u32()
60        };
61
62        *data = PublicValuesPreCompute { b_or_imm, c_or_imm };
63
64        (b_is_imm, c_is_imm)
65    }
66}
67
68macro_rules! dispatch {
69    ($execute_impl:ident, $b_is_imm:ident, $c_is_imm:ident) => {
70        match ($b_is_imm, $c_is_imm) {
71            (true, true) => Ok($execute_impl::<_, _, true, true>),
72            (true, false) => Ok($execute_impl::<_, _, true, false>),
73            (false, true) => Ok($execute_impl::<_, _, false, true>),
74            (false, false) => Ok($execute_impl::<_, _, false, false>),
75        }
76    };
77}
78
79impl<F, A> Executor<F> for PublicValuesExecutor<F, A>
80where
81    F: PrimeField32,
82{
83    #[inline(always)]
84    fn pre_compute_size(&self) -> usize {
85        size_of::<PublicValuesPreCompute>()
86    }
87
88    #[cfg(not(feature = "tco"))]
89    #[inline(always)]
90    fn pre_compute<Ctx>(
91        &self,
92        _pc: u32,
93        inst: &Instruction<F>,
94        data: &mut [u8],
95    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
96    where
97        Ctx: ExecutionCtxTrait,
98    {
99        let data: &mut PublicValuesPreCompute = data.borrow_mut();
100        let (b_is_imm, c_is_imm) = self.pre_compute_impl(inst, data);
101
102        dispatch!(execute_e1_handler, b_is_imm, c_is_imm)
103    }
104
105    #[cfg(feature = "tco")]
106    fn handler<Ctx>(
107        &self,
108        _pc: u32,
109        inst: &Instruction<F>,
110        data: &mut [u8],
111    ) -> Result<Handler<F, Ctx>, StaticProgramError>
112    where
113        Ctx: ExecutionCtxTrait,
114    {
115        let data: &mut PublicValuesPreCompute = data.borrow_mut();
116        let (b_is_imm, c_is_imm) = self.pre_compute_impl(inst, data);
117
118        dispatch!(execute_e1_handler, b_is_imm, c_is_imm)
119    }
120}
121
122impl<F, A> MeteredExecutor<F> for PublicValuesExecutor<F, A>
123where
124    F: PrimeField32,
125{
126    fn metered_pre_compute_size(&self) -> usize {
127        size_of::<E2PreCompute<PublicValuesPreCompute>>()
128    }
129
130    #[cfg(not(feature = "tco"))]
131    fn metered_pre_compute<Ctx>(
132        &self,
133        chip_idx: usize,
134        _pc: u32,
135        inst: &Instruction<F>,
136        data: &mut [u8],
137    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
138    where
139        Ctx: MeteredExecutionCtxTrait,
140    {
141        let data: &mut E2PreCompute<PublicValuesPreCompute> = data.borrow_mut();
142        data.chip_idx = chip_idx as u32;
143        let (b_is_imm, c_is_imm) = self.pre_compute_impl(inst, &mut data.data);
144
145        dispatch!(execute_e2_handler, b_is_imm, c_is_imm)
146    }
147
148    #[cfg(feature = "tco")]
149    fn metered_handler<Ctx>(
150        &self,
151        chip_idx: usize,
152        _pc: u32,
153        inst: &Instruction<F>,
154        data: &mut [u8],
155    ) -> Result<Handler<F, Ctx>, StaticProgramError>
156    where
157        Ctx: MeteredExecutionCtxTrait,
158    {
159        let data: &mut E2PreCompute<PublicValuesPreCompute> = data.borrow_mut();
160        data.chip_idx = chip_idx as u32;
161        let (b_is_imm, c_is_imm) = self.pre_compute_impl(inst, &mut data.data);
162
163        dispatch!(execute_e2_handler, b_is_imm, c_is_imm)
164    }
165}
166
167#[inline(always)]
168unsafe fn execute_e12_impl<F: PrimeField32, CTX, const B_IS_IMM: bool, const C_IS_IMM: bool>(
169    pre_compute: &PublicValuesPreCompute,
170    instret: &mut u64,
171    pc: &mut u32,
172    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
173) where
174    CTX: ExecutionCtxTrait,
175{
176    let value = if B_IS_IMM {
177        transmute_u32_to_field(&pre_compute.b_or_imm)
178    } else {
179        exec_state.vm_read::<F, 1>(NATIVE_AS, pre_compute.b_or_imm)[0]
180    };
181    let index = if C_IS_IMM {
182        transmute_u32_to_field(&pre_compute.c_or_imm)
183    } else {
184        exec_state.vm_read::<F, 1>(NATIVE_AS, pre_compute.c_or_imm)[0]
185    };
186
187    let idx: usize = index.as_canonical_u32() as usize;
188    {
189        let custom_pvs = &mut exec_state.vm_state.custom_pvs;
190
191        if custom_pvs[idx].is_none() {
192            custom_pvs[idx] = Some(value);
193        } else {
194            // Not a hard constraint violation when publishing the same value twice but the
195            // program should avoid that.
196            panic!("Custom public value {} already set", idx);
197        }
198    }
199    *pc = pc.wrapping_add(DEFAULT_PC_STEP);
200    *instret += 1;
201}
202
203#[create_handler]
204#[inline(always)]
205unsafe fn execute_e1_impl<F: PrimeField32, CTX, const B_IS_IMM: bool, const C_IS_IMM: bool>(
206    pre_compute: &[u8],
207    instret: &mut u64,
208    pc: &mut u32,
209    _instret_end: u64,
210    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
211) where
212    CTX: ExecutionCtxTrait,
213{
214    let pre_compute: &PublicValuesPreCompute = pre_compute.borrow();
215    execute_e12_impl::<_, _, B_IS_IMM, C_IS_IMM>(pre_compute, instret, pc, exec_state);
216}
217
218#[create_handler]
219#[inline(always)]
220unsafe fn execute_e2_impl<F: PrimeField32, CTX, const B_IS_IMM: bool, const C_IS_IMM: bool>(
221    pre_compute: &[u8],
222    instret: &mut u64,
223    pc: &mut u32,
224    _arg: u64,
225    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
226) where
227    CTX: MeteredExecutionCtxTrait,
228{
229    let pre_compute: &E2PreCompute<PublicValuesPreCompute> = pre_compute.borrow();
230    exec_state
231        .ctx
232        .on_height_change(pre_compute.chip_idx as usize, 1);
233    execute_e12_impl::<_, _, B_IS_IMM, C_IS_IMM>(&pre_compute.data, instret, pc, exec_state);
234}
openvm_circuit/system/public_values/execution.rs

openvm_circuit/system/public_values/
execution.rs
