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 #[cfg(not(feature = "tco"))]
122 #[inline(always)]
123 fn pre_compute<Ctx: ExecutionCtxTrait>(
124 &self,
125 pc: u32,
126 inst: &Instruction<F>,
127 data: &mut [u8],
128 ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError> {
129 let pre_compute: &mut LoadStorePreCompute = data.borrow_mut();
130 let (local_opcode, enabled, is_native_store) =
131 self.pre_compute_impl(pc, inst, pre_compute)?;
132 dispatch!(execute_e1_handler, local_opcode, enabled, is_native_store)
133 }
134
135 #[cfg(feature = "tco")]
136 fn handler<Ctx>(
137 &self,
138 pc: u32,
139 inst: &Instruction<F>,
140 data: &mut [u8],
141 ) -> Result<Handler<F, Ctx>, StaticProgramError>
142 where
143 Ctx: ExecutionCtxTrait,
144 {
145 let pre_compute: &mut LoadStorePreCompute = data.borrow_mut();
146 let (local_opcode, enabled, is_native_store) =
147 self.pre_compute_impl(pc, inst, pre_compute)?;
148 dispatch!(execute_e1_handler, local_opcode, enabled, is_native_store)
149 }
150}
151
152impl<F, A, const NUM_CELLS: usize> MeteredExecutor<F> for LoadStoreExecutor<A, NUM_CELLS>
153where
154 F: PrimeField32,
155{
156 fn metered_pre_compute_size(&self) -> usize {
157 size_of::<E2PreCompute<LoadStorePreCompute>>()
158 }
159
160 #[cfg(not(feature = "tco"))]
161 fn metered_pre_compute<Ctx>(
162 &self,
163 chip_idx: usize,
164 pc: u32,
165 inst: &Instruction<F>,
166 data: &mut [u8],
167 ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
168 where
169 Ctx: MeteredExecutionCtxTrait,
170 {
171 let pre_compute: &mut E2PreCompute<LoadStorePreCompute> = data.borrow_mut();
172 pre_compute.chip_idx = chip_idx as u32;
173 let (local_opcode, enabled, is_native_store) =
174 self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
175 dispatch!(execute_e2_handler, local_opcode, enabled, is_native_store)
176 }
177
178 #[cfg(feature = "tco")]
179 fn metered_handler<Ctx>(
180 &self,
181 chip_idx: usize,
182 pc: u32,
183 inst: &Instruction<F>,
184 data: &mut [u8],
185 ) -> Result<Handler<F, Ctx>, StaticProgramError>
186 where
187 Ctx: MeteredExecutionCtxTrait,
188 {
189 let pre_compute: &mut E2PreCompute<LoadStorePreCompute> = data.borrow_mut();
190 pre_compute.chip_idx = chip_idx as u32;
191 let (local_opcode, enabled, is_native_store) =
192 self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
193 dispatch!(execute_e2_handler, local_opcode, enabled, is_native_store)
194 }
195}
196
197#[inline(always)]
198unsafe fn execute_e12_impl<
199 F: PrimeField32,
200 CTX: ExecutionCtxTrait,
201 T: Copy + Debug + Default,
202 OP: LoadStoreOp<T>,
203 const ENABLED: bool,
204>(
205 pre_compute: &LoadStorePreCompute,
206 instret: &mut u64,
207 pc: &mut u32,
208 exec_state: &mut VmExecState<F, GuestMemory, CTX>,
209) -> Result<(), ExecutionError> {
210 let rs1_bytes: [u8; RV32_REGISTER_NUM_LIMBS] =
211 exec_state.vm_read(RV32_REGISTER_AS, pre_compute.b as u32);
212 let rs1_val = u32::from_le_bytes(rs1_bytes);
213 let ptr_val = rs1_val.wrapping_add(pre_compute.imm_extended);
214 debug_assert!(ptr_val < (1 << POINTER_MAX_BITS));
216 let shift_amount = ptr_val % 4;
217 let ptr_val = ptr_val - shift_amount; let read_data: [u8; RV32_REGISTER_NUM_LIMBS] = if OP::IS_LOAD {
220 exec_state.vm_read(pre_compute.e as u32, ptr_val)
221 } else {
222 exec_state.vm_read(RV32_REGISTER_AS, pre_compute.a as u32)
223 };
224
225 let mut write_data: [T; RV32_REGISTER_NUM_LIMBS] = if OP::HOST_READ {
227 exec_state.host_read(pre_compute.e as u32, ptr_val)
228 } else {
229 [T::default(); RV32_REGISTER_NUM_LIMBS]
230 };
231
232 if !OP::compute_write_data(&mut write_data, read_data, shift_amount as usize) {
233 let err = ExecutionError::Fail {
234 pc: *pc,
235 msg: "Invalid LoadStoreOp",
236 };
237 return Err(err);
238 }
239
240 if ENABLED {
241 if OP::IS_LOAD {
242 exec_state.vm_write(RV32_REGISTER_AS, pre_compute.a as u32, &write_data);
243 } else {
244 exec_state.vm_write(pre_compute.e as u32, ptr_val, &write_data);
245 }
246 }
247
248 *pc += DEFAULT_PC_STEP;
249 *instret += 1;
250
251 Ok(())
252}
253
254#[create_handler]
255#[inline(always)]
256unsafe fn execute_e1_impl<
257 F: PrimeField32,
258 CTX: ExecutionCtxTrait,
259 T: Copy + Debug + Default,
260 OP: LoadStoreOp<T>,
261 const ENABLED: bool,
262>(
263 pre_compute: &[u8],
264 instret: &mut u64,
265 pc: &mut u32,
266 _instret_end: u64,
267 exec_state: &mut VmExecState<F, GuestMemory, CTX>,
268) -> Result<(), ExecutionError> {
269 let pre_compute: &LoadStorePreCompute = pre_compute.borrow();
270 execute_e12_impl::<F, CTX, T, OP, ENABLED>(pre_compute, instret, pc, exec_state)
271}
272
273#[create_handler]
274#[inline(always)]
275unsafe fn execute_e2_impl<
276 F: PrimeField32,
277 CTX: MeteredExecutionCtxTrait,
278 T: Copy + Debug + Default,
279 OP: LoadStoreOp<T>,
280 const ENABLED: bool,
281>(
282 pre_compute: &[u8],
283 instret: &mut u64,
284 pc: &mut u32,
285 _arg: u64,
286 exec_state: &mut VmExecState<F, GuestMemory, CTX>,
287) -> Result<(), ExecutionError> {
288 let pre_compute: &E2PreCompute<LoadStorePreCompute> = pre_compute.borrow();
289 exec_state
290 .ctx
291 .on_height_change(pre_compute.chip_idx as usize, 1);
292 execute_e12_impl::<F, CTX, T, OP, ENABLED>(&pre_compute.data, instret, pc, exec_state)
293}
294
295trait LoadStoreOp<T> {
296 const IS_LOAD: bool;
297 const HOST_READ: bool;
298
299 fn compute_write_data(
301 write_data: &mut [T; RV32_REGISTER_NUM_LIMBS],
302 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
303 shift_amount: usize,
304 ) -> bool;
305}
306#[allow(dead_code)]
309#[derive(Copy, Clone, Debug, Default)]
310struct U8(u8);
311struct LoadWOp;
312struct LoadHUOp;
313struct LoadBUOp;
314struct StoreWOp;
315struct StoreHOp;
316struct StoreBOp;
317impl LoadStoreOp<U8> for LoadWOp {
318 const IS_LOAD: bool = true;
319 const HOST_READ: bool = false;
320
321 #[inline(always)]
322 fn compute_write_data(
323 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
324 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
325 _shift_amount: usize,
326 ) -> bool {
327 *write_data = read_data.map(U8);
328 true
329 }
330}
331
332impl LoadStoreOp<U8> for LoadHUOp {
333 const IS_LOAD: bool = true;
334 const HOST_READ: bool = false;
335 #[inline(always)]
336 fn compute_write_data(
337 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
338 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
339 shift_amount: usize,
340 ) -> bool {
341 if shift_amount != 0 && shift_amount != 2 {
342 return false;
343 }
344 write_data[0] = U8(read_data[shift_amount]);
345 write_data[1] = U8(read_data[shift_amount + 1]);
346 true
347 }
348}
349impl LoadStoreOp<U8> for LoadBUOp {
350 const IS_LOAD: bool = true;
351 const HOST_READ: bool = false;
352 #[inline(always)]
353 fn compute_write_data(
354 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
355 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
356 shift_amount: usize,
357 ) -> bool {
358 write_data[0] = U8(read_data[shift_amount]);
359 true
360 }
361}
362
363impl LoadStoreOp<U8> for StoreWOp {
364 const IS_LOAD: bool = false;
365 const HOST_READ: bool = false;
366 #[inline(always)]
367 fn compute_write_data(
368 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
369 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
370 _shift_amount: usize,
371 ) -> bool {
372 *write_data = read_data.map(U8);
373 true
374 }
375}
376impl LoadStoreOp<U8> for StoreHOp {
377 const IS_LOAD: bool = false;
378 const HOST_READ: bool = true;
379
380 #[inline(always)]
381 fn compute_write_data(
382 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
383 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
384 shift_amount: usize,
385 ) -> bool {
386 if shift_amount != 0 && shift_amount != 2 {
387 return false;
388 }
389 write_data[shift_amount] = U8(read_data[0]);
390 write_data[shift_amount + 1] = U8(read_data[1]);
391 true
392 }
393}
394impl LoadStoreOp<U8> for StoreBOp {
395 const IS_LOAD: bool = false;
396 const HOST_READ: bool = true;
397 #[inline(always)]
398 fn compute_write_data(
399 write_data: &mut [U8; RV32_REGISTER_NUM_LIMBS],
400 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
401 shift_amount: usize,
402 ) -> bool {
403 write_data[shift_amount] = U8(read_data[0]);
404 true
405 }
406}
407
408impl<F: PrimeField32> LoadStoreOp<F> for StoreWOp {
409 const IS_LOAD: bool = false;
410 const HOST_READ: bool = false;
411 #[inline(always)]
412 fn compute_write_data(
413 write_data: &mut [F; RV32_REGISTER_NUM_LIMBS],
414 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
415 _shift_amount: usize,
416 ) -> bool {
417 *write_data = read_data.map(F::from_canonical_u8);
418 true
419 }
420}
421impl<F: PrimeField32> LoadStoreOp<F> for StoreHOp {
422 const IS_LOAD: bool = false;
423 const HOST_READ: bool = true;
424
425 #[inline(always)]
426 fn compute_write_data(
427 write_data: &mut [F; RV32_REGISTER_NUM_LIMBS],
428 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
429 shift_amount: usize,
430 ) -> bool {
431 if shift_amount != 0 && shift_amount != 2 {
432 return false;
433 }
434 write_data[shift_amount] = F::from_canonical_u8(read_data[0]);
435 write_data[shift_amount + 1] = F::from_canonical_u8(read_data[1]);
436 true
437 }
438}
439impl<F: PrimeField32> LoadStoreOp<F> for StoreBOp {
440 const IS_LOAD: bool = false;
441 const HOST_READ: bool = true;
442 #[inline(always)]
443 fn compute_write_data(
444 write_data: &mut [F; RV32_REGISTER_NUM_LIMBS],
445 read_data: [u8; RV32_REGISTER_NUM_LIMBS],
446 shift_amount: usize,
447 ) -> bool {
448 write_data[shift_amount] = F::from_canonical_u8(read_data[0]);
449 true
450 }
451}