openvm_circuit/system/public_values/
execution.rs

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