1use std::{
2 borrow::{Borrow, BorrowMut},
3 fmt::Debug,
4 mem::size_of,
5};
6
7use openvm_circuit::{
8 arch::*,
9 system::memory::{online::GuestMemory, POINTER_MAX_BITS},
10};
11use openvm_circuit_primitives::AlignedBytesBorrow;
12use openvm_instructions::{
13 instruction::Instruction,
14 program::DEFAULT_PC_STEP,
15 riscv::{RV32_IMM_AS, RV32_REGISTER_AS, RV32_REGISTER_NUM_LIMBS},
16 LocalOpcode, NATIVE_AS,
17};
18use openvm_rv32im_transpiler::Rv32LoadStoreOpcode::{self, *};
19use openvm_stark_backend::p3_field::PrimeField32;
20
21use super::core::LoadStoreExecutor;
22
23#[derive(AlignedBytesBorrow, Clone)]
24#[repr(C)]
25struct LoadStorePreCompute {
26 imm_extended: u32,
27 a: u8,
28 b: u8,
29 e: u8,
30}
31
32impl<A, const NUM_CELLS: usize> LoadStoreExecutor<A, NUM_CELLS> {
33 fn pre_compute_impl<F: PrimeField32>(
35 &self,
36 pc: u32,
37 inst: &Instruction<F>,
38 data: &mut LoadStorePreCompute,
39 ) -> Result<(Rv32LoadStoreOpcode, bool, bool), StaticProgramError> {
40 let Instruction {
41 opcode,
42 a,
43 b,
44 c,
45 d,
46 e,
47 f,
48 g,
49 ..
50 } = inst;
51 let enabled = !f.is_zero();
52
53 let e_u32 = e.as_canonical_u32();
54 if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 == RV32_IMM_AS {
55 return Err(StaticProgramError::InvalidInstruction(pc));
56 }
57
58 let local_opcode = Rv32LoadStoreOpcode::from_usize(
59 opcode.local_opcode_idx(Rv32LoadStoreOpcode::CLASS_OFFSET),
60 );
61 match local_opcode {
62 LOADW | LOADBU | LOADHU => {}
63 STOREW | STOREH | STOREB => {
64 if !enabled {
65 return Err(StaticProgramError::InvalidInstruction(pc));
66 }
67 }
68 _ => unreachable!("LoadStoreExecutor should not handle LOADB/LOADH opcodes"),
69 }
70
71 let imm = c.as_canonical_u32();
72 let imm_sign = g.as_canonical_u32();
73 let imm_extended = imm + imm_sign * 0xffff0000;
74 let is_native_store = e_u32 == NATIVE_AS;
75
76 *data = LoadStorePreCompute {
77 imm_extended,
78 a: a.as_canonical_u32() as u8,
79 b: b.as_canonical_u32() as u8,
80 e: e_u32 as u8,
81 };
82 Ok((local_opcode, enabled, is_native_store))
83 }
84}
85
86macro_rules! dispatch {
87 ($execute_impl:ident, $local_opcode:ident, $enabled:ident, $is_native_store:ident) => {
88 match ($local_opcode, $enabled, $is_native_store) {
89 (LOADW, true, _) => Ok($execute_impl::<_, _, U8, LoadWOp, true>),
90 (LOADW, false, _) => Ok($execute_impl::<_, _, U8, LoadWOp, false>),
91 (LOADHU, true, _) => Ok($execute_impl::<_, _, U8, LoadHUOp, true>),
92 (LOADHU, false, _) => Ok($execute_impl::<_, _, U8, LoadHUOp, false>),
93 (LOADBU, true, _) => Ok($execute_impl::<_, _, U8, LoadBUOp, true>),
94 (LOADBU, false, _) => Ok($execute_impl::<_, _, U8, LoadBUOp, false>),
95 (STOREW, true, false) => Ok($execute_impl::<_, _, U8, StoreWOp, true>),
96 (STOREW, false, false) => Ok($execute_impl::<_, _, U8, StoreWOp, false>),
97 (STOREW, true, true) => Ok($execute_impl::<_, _, F, StoreWOp, true>),
98 (STOREW, false, true) => Ok($execute_impl::<_, _, F, StoreWOp, false>),
99 (STOREH, true, false) => Ok($execute_impl::<_, _, U8, StoreHOp, true>),
100 (STOREH, false, false) => Ok($execute_impl::<_, _, U8, StoreHOp, false>),
101 (STOREH, true, true) => Ok($execute_impl::<_, _, F, StoreHOp, true>),
102 (STOREH, false, true) => Ok($execute_impl::<_, _, F, StoreHOp, false>),
103 (STOREB, true, false) => Ok($execute_impl::<_, _, U8, StoreBOp, true>),
104 (STOREB, false, false) => Ok($execute_impl::<_, _, U8, StoreBOp, false>),
105 (STOREB, true, true) => Ok($execute_impl::<_, _, F, StoreBOp, true>),
106 (STOREB, false, true) => Ok($execute_impl::<_, _, F, StoreBOp, false>),
107 (_, _, _) => unreachable!(),
108 }
109 };
110}
111
112impl<F, A, const NUM_CELLS: usize> Executor<F> for LoadStoreExecutor<A, NUM_CELLS>
113where
114 F: PrimeField32,
115{
116 #[inline(always)]
117 fn pre_compute_size(&self) -> usize {
118 size_of::<LoadStorePreCompute>()
119 }
120
121 #[inline(always)]
122 fn pre_compute<Ctx: ExecutionCtxTrait>(
123 &self,
124 pc: u32,
125 inst: &Instruction<F>,
126 data: &mut [u8],
127 ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError> {
128 let pre_compute: &mut LoadStorePreCompute = data.borrow_mut();
129 let (local_opcode, enabled, is_native_store) =
130 self.pre_compute_impl(pc, inst, pre_compute)?;
131 dispatch!(execute_e1_impl, local_opcode, enabled, is_native_store)
132 }
133
134 #[cfg(feature = "tco")]
135 fn handler<Ctx>(
136 &self,
137 pc: u32,
138 inst: &Instruction<F>,
139 data: &mut [u8],
140 ) -> Result<Handler<F, Ctx>, StaticProgramError>
141 where
142 Ctx: ExecutionCtxTrait,
143 {
144 let pre_compute: &mut LoadStorePreCompute = data.borrow_mut();
145 let (local_opcode, enabled, is_native_store) =
146 self.pre_compute_impl(pc, inst, pre_compute)?;
147 dispatch!(
148 execute_e1_tco_handler,
149 local_opcode,
150 enabled,
151 is_native_store
152 )
153 }
154}
155
156impl<F, A, const NUM_CELLS: usize> MeteredExecutor<F> for LoadStoreExecutor<A, NUM_CELLS>
157where
158 F: PrimeField32,
159{
160 fn metered_pre_compute_size(&self) -> usize {
161 size_of::<E2PreCompute<LoadStorePreCompute>>()
162 }
163
164 fn metered_pre_compute<Ctx>(
165 &self,
166 chip_idx: usize,
167 pc: u32,
168 inst: &Instruction<F>,
169 data: &mut [u8],
170 ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
171 where
172 Ctx: MeteredExecutionCtxTrait,
173 {
174 let pre_compute: &mut E2PreCompute<LoadStorePreCompute> = data.borrow_mut();
175 pre_compute.chip_idx = chip_idx as u32;
176 let (local_opcode, enabled, is_native_store) =
177 self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
178 dispatch!(execute_e2_impl, local_opcode, enabled, is_native_store)
179 }
180
181 #[cfg(feature = "tco")]
182 fn metered_handler<Ctx>(
183 &self,
184 chip_idx: usize,
185 pc: u32,
186 inst: &Instruction<F>,
187 data: &mut [u8],
188 ) -> Result<Handler<F, Ctx>, StaticProgramError>
189 where
190 Ctx: MeteredExecutionCtxTrait,
191 {
192 let pre_compute: &mut E2PreCompute<LoadStorePreCompute> = data.borrow_mut();
193 pre_compute.chip_idx = chip_idx as u32;
194 let (local_opcode, enabled, is_native_store) =
195 self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
196 dispatch!(
197 execute_e2_tco_handler,
198 local_opcode,
199 enabled,
200 is_native_store
201 )
202 }
203}
204
205#[inline(always)]
206unsafe fn execute_e12_impl<
207 F: PrimeField32,
208 CTX: ExecutionCtxTrait,
209 T: Copy + Debug + Default,
210 OP: LoadStoreOp<T>,
211 const ENABLED: bool,
212>(
213 pre_compute: &LoadStorePreCompute,
214 vm_state: &mut VmExecState<F, GuestMemory, CTX>,
215) {
216 let rs1_bytes: [u8; RV32_REGISTER_NUM_LIMBS] =
217 vm_state.vm_read(RV32_REGISTER_AS, pre_compute.b as u32);
218 let rs1_val = u32::from_le_bytes(rs1_bytes);
219 let ptr_val = rs1_val.wrapping_add(pre_compute.imm_extended);
220 debug_assert!(ptr_val < (1 << POINTER_MAX_BITS));
222 let shift_amount = ptr_val % 4;
223 let ptr_val = ptr_val - shift_amount; let read_data: [u8; RV32_REGISTER_NUM_LIMBS] = if OP::IS_LOAD {
226 vm_state.vm_read(pre_compute.e as u32, ptr_val)
227 } else {
228 vm_state.vm_read(RV32_REGISTER_AS, pre_compute.a as u32)
229 };
230
231 let mut write_data: [T; RV32_REGISTER_NUM_LIMBS] = if OP::HOST_READ {
233 vm_state.host_read(pre_compute.e as u32, ptr_val)
234 } else {
235 [T::default(); RV32_REGISTER_NUM_LIMBS]
236 };
237
238 if !OP::compute_write_data(&mut write_data, read_data, shift_amount as usize) {
239 vm_state.exit_code = Err(ExecutionError::Fail {
240 pc: vm_state.pc,
241 msg: "Invalid LoadStoreOp",
242 });
243 return;
244 }
245
246 if ENABLED {
247 if OP::IS_LOAD {
248 vm_state.vm_write(RV32_REGISTER_AS, pre_compute.a as u32, &write_data);
249 } else {
250 vm_state.vm_write(pre_compute.e as u32, ptr_val, &write_data);
251 }
252 }
253
254 vm_state.pc += DEFAULT_PC_STEP;
255 vm_state.instret += 1;
256}
257
258#[create_tco_handler]
259unsafe fn execute_e1_impl<
260 F: PrimeField32,
261 CTX: ExecutionCtxTrait,
262 T: Copy + Debug + Default,
263 OP: LoadStoreOp<T>,
264 const ENABLED: bool,
265>(
266 pre_compute: &[u8],
267 vm_state: &mut VmExecState<F, GuestMemory, CTX>,
268) {
269 let pre_compute: &LoadStorePreCompute = pre_compute.borrow();
270 execute_e12_impl::<F, CTX, T, OP, ENABLED>(pre_compute, vm_state);
271}
272
273#[create_tco_handler]
274unsafe fn execute_e2_impl<
275 F: PrimeField32,
276 CTX: MeteredExecutionCtxTrait,
277 T: Copy + Debug + Default,
278 OP: LoadStoreOp<T>,
279 const ENABLED: bool,
280>(
281 pre_compute: &[u8],
282 vm_state: &mut VmExecState<F, GuestMemory, CTX>,
283) {
284 let pre_compute: &E2PreCompute<LoadStorePreCompute> = pre_compute.borrow();
285 vm_state
286 .ctx
287 .on_height_change(pre_compute.chip_idx as usize, 1);
288 execute_e12_impl::<F, CTX, T, OP, ENABLED>(&pre_compute.data, vm_state);
289}
290
291trait LoadStoreOp<T> {
292 const IS_LOAD: bool;
293 const HOST_READ: bool;
294
295 fn compute_write_data(
297 write_data: &mut [T; RV32_REGISTER_NUM_LIMBS],
298 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
299 shift_amount: usize,
300 ) -> bool;
301}
302#[allow(dead_code)]
305#[derive(Copy, Clone, Debug, Default)]
306struct U8(u8);
307struct LoadWOp;
308struct LoadHUOp;
309struct LoadBUOp;
310struct StoreWOp;
311struct StoreHOp;
312struct StoreBOp;
313impl LoadStoreOp<U8> for LoadWOp {
314 const IS_LOAD: bool = true;
315 const HOST_READ: bool = false;
316
317 #[inline(always)]
318 fn compute_write_data(
319 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
320 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
321 _shift_amount: usize,
322 ) -> bool {
323 *write_data = read_data.map(U8);
324 true
325 }
326}
327
328impl LoadStoreOp<U8> for LoadHUOp {
329 const IS_LOAD: bool = true;
330 const HOST_READ: bool = false;
331 #[inline(always)]
332 fn compute_write_data(
333 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
334 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
335 shift_amount: usize,
336 ) -> bool {
337 if shift_amount != 0 && shift_amount != 2 {
338 return false;
339 }
340 write_data[0] = U8(read_data[shift_amount]);
341 write_data[1] = U8(read_data[shift_amount + 1]);
342 true
343 }
344}
345impl LoadStoreOp<U8> for LoadBUOp {
346 const IS_LOAD: bool = true;
347 const HOST_READ: bool = false;
348 #[inline(always)]
349 fn compute_write_data(
350 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
351 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
352 shift_amount: usize,
353 ) -> bool {
354 write_data[0] = U8(read_data[shift_amount]);
355 true
356 }
357}
358
359impl LoadStoreOp<U8> for StoreWOp {
360 const IS_LOAD: bool = false;
361 const HOST_READ: bool = false;
362 #[inline(always)]
363 fn compute_write_data(
364 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
365 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
366 _shift_amount: usize,
367 ) -> bool {
368 *write_data = read_data.map(U8);
369 true
370 }
371}
372impl LoadStoreOp<U8> for StoreHOp {
373 const IS_LOAD: bool = false;
374 const HOST_READ: bool = true;
375
376 #[inline(always)]
377 fn compute_write_data(
378 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
379 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
380 shift_amount: usize,
381 ) -> bool {
382 if shift_amount != 0 && shift_amount != 2 {
383 return false;
384 }
385 write_data[shift_amount] = U8(read_data[0]);
386 write_data[shift_amount + 1] = U8(read_data[1]);
387 true
388 }
389}
390impl LoadStoreOp<U8> for StoreBOp {
391 const IS_LOAD: bool = false;
392 const HOST_READ: bool = true;
393 #[inline(always)]
394 fn compute_write_data(
395 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
396 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
397 shift_amount: usize,
398 ) -> bool {
399 write_data[shift_amount] = U8(read_data[0]);
400 true
401 }
402}
403
404impl<F: PrimeField32> LoadStoreOp<F> for StoreWOp {
405 const IS_LOAD: bool = false;
406 const HOST_READ: bool = false;
407 #[inline(always)]
408 fn compute_write_data(
409 write_data: &mut [F; RV32_REGISTER_NUM_LIMBS],
410 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
411 _shift_amount: usize,
412 ) -> bool {
413 *write_data = read_data.map(F::from_canonical_u8);
414 true
415 }
416}
417impl<F: PrimeField32> LoadStoreOp<F> for StoreHOp {
418 const IS_LOAD: bool = false;
419 const HOST_READ: bool = true;
420
421 #[inline(always)]
422 fn compute_write_data(
423 write_data: &mut [F; RV32_REGISTER_NUM_LIMBS],
424 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
425 shift_amount: usize,
426 ) -> bool {
427 if shift_amount != 0 && shift_amount != 2 {
428 return false;
429 }
430 write_data[shift_amount] = F::from_canonical_u8(read_data[0]);
431 write_data[shift_amount + 1] = F::from_canonical_u8(read_data[1]);
432 true
433 }
434}
435impl<F: PrimeField32> LoadStoreOp<F> for StoreBOp {
436 const IS_LOAD: bool = false;
437 const HOST_READ: bool = true;
438 #[inline(always)]
439 fn compute_write_data(
440 write_data: &mut [F; RV32_REGISTER_NUM_LIMBS],
441 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
442 shift_amount: usize,
443 ) -> bool {
444 write_data[shift_amount] = F::from_canonical_u8(read_data[0]);
445 true
446 }
447}