ab_riscv_benchmarks/
host_utils.rs

1extern crate alloc;
2
3use ab_blake3::{CHUNK_LEN, OUT_LEN};
4use ab_contract_file::ContractInstruction;
5use ab_core_primitives::ed25519::{Ed25519PublicKey, Ed25519Signature};
6use ab_io_type::IoType;
7use ab_io_type::bool::Bool;
8use ab_riscv_interpreter::rv64::{Rv64InterpreterState, Rv64SystemInstructionHandler};
9use ab_riscv_interpreter::{
10    BasicInt, ExecutableInstruction, ExecutionError, FetchInstructionResult, InstructionFetcher,
11    ProgramCounter, ProgramCounterError, VirtualMemory, VirtualMemoryError,
12};
13use ab_riscv_primitives::instructions::Instruction;
14use ab_riscv_primitives::instructions::rv64::Rv64Instruction;
15use ab_riscv_primitives::registers::general_purpose::{Register, Registers};
16use alloc::vec::Vec;
17use core::marker::PhantomData;
18use core::mem::offset_of;
19use core::ops::ControlFlow;
20
21/// Contract file bytes
22pub const RISCV_CONTRACT_BYTES: &[u8] = cfg_select! {
23    target_env = "abundance" => {
24        &[]
25    }
26    _ => {
27        include_bytes!(env!("CONTRACT_PATH"))
28    }
29};
30
31// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
32//  `SimpleInternalArgs` data trait for this or something
33/// Helper data structure for [`Benchmarks::blake3_hash_chunk()`] method
34///
35/// [`Benchmarks::blake3_hash_chunk()`]: crate::Benchmarks::blake3_hash_chunk
36#[derive(Debug, Copy, Clone)]
37#[repr(C)]
38pub struct Blake3HashChunkInternalArgs {
39    chunk_ptr: u64,
40    chunk_size: u32,
41    chunk_capacity: u32,
42    result_ptr: u64,
43    chunk: [u8; CHUNK_LEN],
44    result: [u8; OUT_LEN],
45}
46
47impl Blake3HashChunkInternalArgs {
48    /// Create a new instance
49    pub fn new(internal_args_addr: u64, chunk: [u8; CHUNK_LEN]) -> Self {
50        Self {
51            chunk_ptr: internal_args_addr + offset_of!(Self, chunk) as u64,
52            chunk_size: CHUNK_LEN as u32,
53            chunk_capacity: CHUNK_LEN as u32,
54            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
55            chunk,
56            result: [0; _],
57        }
58    }
59
60    /// Extract result
61    pub fn result(&self) -> [u8; OUT_LEN] {
62        self.result
63    }
64}
65
66// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
67//  `SimpleInternalArgs` data trait for this or something
68/// Helper data structure for [`Benchmarks::ed25519_verify()`] method
69///
70/// [`Benchmarks::ed25519_verify()`]: crate::Benchmarks::ed25519_verify
71#[derive(Debug, Copy, Clone)]
72#[repr(C)]
73pub struct Ed25519VerifyInternalArgs {
74    pub public_key_ptr: u64,
75    pub public_key_size: u32,
76    pub public_key_capacity: u32,
77    pub signature_ptr: u64,
78    pub signature_size: u32,
79    pub signature_capacity: u32,
80    pub message_ptr: u64,
81    pub message_size: u32,
82    pub message_capacity: u32,
83    pub result_ptr: u64,
84    pub public_key: Ed25519PublicKey,
85    pub signature: Ed25519Signature,
86    pub message: [u8; OUT_LEN],
87    pub result: Bool,
88}
89
90impl Ed25519VerifyInternalArgs {
91    /// Create a new instance
92    pub fn new(
93        internal_args_addr: u64,
94        public_key: Ed25519PublicKey,
95        signature: Ed25519Signature,
96        message: [u8; OUT_LEN],
97    ) -> Self {
98        Self {
99            public_key_ptr: internal_args_addr + offset_of!(Self, public_key) as u64,
100            public_key_size: Ed25519PublicKey::SIZE as u32,
101            public_key_capacity: Ed25519PublicKey::SIZE as u32,
102            signature_ptr: internal_args_addr + offset_of!(Self, signature) as u64,
103            signature_size: Ed25519Signature::SIZE as u32,
104            signature_capacity: Ed25519Signature::SIZE as u32,
105            message_ptr: internal_args_addr + offset_of!(Self, message) as u64,
106            message_size: OUT_LEN as u32,
107            message_capacity: OUT_LEN as u32,
108            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
109            public_key,
110            signature,
111            message,
112            result: Bool::new(false),
113        }
114    }
115
116    /// Extract result
117    pub fn result(&self) -> Bool {
118        self.result
119    }
120}
121
122// Simple test memory implementation
123#[derive(Debug, Copy, Clone)]
124pub struct TestMemory<const MEMORY_SIZE: usize> {
125    data: [u8; MEMORY_SIZE],
126    base_addr: u64,
127}
128
129impl<const MEMORY_SIZE: usize> VirtualMemory for TestMemory<MEMORY_SIZE> {
130    fn read<T>(&self, address: u64) -> Result<T, VirtualMemoryError>
131    where
132        T: BasicInt,
133    {
134        let offset = address
135            .checked_sub(self.base_addr)
136            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })? as usize;
137
138        if offset + size_of::<T>() > self.data.len() {
139            return Err(VirtualMemoryError::OutOfBoundsRead { address });
140        }
141
142        // SAFETY: Only reading basic integers from initialized memory
143        unsafe {
144            Ok(self
145                .data
146                .as_ptr()
147                .cast::<T>()
148                .byte_add(offset)
149                .read_unaligned())
150        }
151    }
152
153    unsafe fn read_unchecked<T>(&self, address: u64) -> T
154    where
155        T: BasicInt,
156    {
157        // SAFETY: Guaranteed by function contract
158        unsafe {
159            let offset = address.unchecked_sub(self.base_addr) as usize;
160            self.data
161                .as_ptr()
162                .cast::<T>()
163                .byte_add(offset)
164                .read_unaligned()
165        }
166    }
167
168    fn write<T>(&mut self, address: u64, value: T) -> Result<(), VirtualMemoryError>
169    where
170        T: BasicInt,
171    {
172        let offset = address
173            .checked_sub(self.base_addr)
174            .ok_or(VirtualMemoryError::OutOfBoundsWrite { address })? as usize;
175
176        if offset + size_of::<T>() > self.data.len() {
177            return Err(VirtualMemoryError::OutOfBoundsWrite { address });
178        }
179
180        // SAFETY: Only writing basic integers to initialized memory
181        unsafe {
182            self.data
183                .as_mut_ptr()
184                .cast::<T>()
185                .byte_add(offset)
186                .write_unaligned(value);
187        }
188
189        Ok(())
190    }
191}
192
193impl<const MEMORY_SIZE: usize> TestMemory<MEMORY_SIZE> {
194    /// Create a new test memory instance with the specified base address
195    pub fn new(base_addr: u64) -> Self {
196        Self {
197            data: [0; _],
198            base_addr,
199        }
200    }
201
202    /// Get a slice of memory
203    pub fn get_bytes(&self, address: u64, size: usize) -> Result<&[u8], VirtualMemoryError> {
204        let offset = address
205            .checked_sub(self.base_addr)
206            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })? as usize;
207
208        if offset + size > self.data.len() {
209            return Err(VirtualMemoryError::OutOfBoundsRead { address });
210        }
211
212        Ok(&self.data[offset..][..size])
213    }
214
215    /// Get a mutable slice of memory
216    pub fn get_mut_bytes(
217        &mut self,
218        address: u64,
219        size: usize,
220    ) -> Result<&mut [u8], VirtualMemoryError> {
221        let offset = address
222            .checked_sub(self.base_addr)
223            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })? as usize;
224
225        if offset + size > self.data.len() {
226            return Err(VirtualMemoryError::OutOfBoundsRead { address });
227        }
228
229        Ok(&mut self.data[offset..][..size])
230    }
231}
232
233/// Eager instruction handler eagerly decodes all instructions upfront
234#[derive(Debug, Default, Clone)]
235pub struct EagerTestInstructionFetcher {
236    instructions: Vec<ContractInstruction>,
237    return_trap_address: u64,
238    base_addr: u64,
239    instruction_offset: usize,
240}
241
242impl<Memory> ProgramCounter<u64, Memory, &'static str> for EagerTestInstructionFetcher
243where
244    Memory: VirtualMemory,
245{
246    #[inline(always)]
247    fn get_pc(&self) -> u64 {
248        self.base_addr + self.instruction_offset as u64 * size_of::<u32>() as u64
249    }
250
251    #[inline]
252    fn set_pc(
253        &mut self,
254        _memory: &mut Memory,
255        pc: u64,
256    ) -> Result<ControlFlow<()>, ProgramCounterError<u64, &'static str>> {
257        let address = pc;
258
259        if address == self.return_trap_address {
260            return Ok(ControlFlow::Break(()));
261        }
262
263        if !address.is_multiple_of(size_of::<u32>() as u64) {
264            return Err(ProgramCounterError::UnalignedInstruction { address });
265        }
266
267        let offset = address
268            .checked_sub(self.base_addr)
269            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })? as usize;
270        let instruction_offset = offset / size_of::<u32>();
271
272        if instruction_offset >= self.instructions.len() {
273            return Err(VirtualMemoryError::OutOfBoundsRead { address }.into());
274        }
275
276        self.instruction_offset = instruction_offset;
277
278        Ok(ControlFlow::Continue(()))
279    }
280}
281
282impl<Memory> InstructionFetcher<ContractInstruction, Memory, &'static str>
283    for EagerTestInstructionFetcher
284where
285    Memory: VirtualMemory,
286{
287    #[inline(always)]
288    fn fetch_instruction(
289        &mut self,
290        _memory: &mut Memory,
291    ) -> Result<
292        FetchInstructionResult<ContractInstruction>,
293        ExecutionError<u64, ContractInstruction, &'static str>,
294    > {
295        // SAFETY: Constructor guarantees that the last instruction is a jump, which means going
296        // through `Self::set_pc()` method that does bound check. Otherwise, advancing forward by
297        // one instruction can't result in out-of-bounds access.
298        let instruction = *unsafe { self.instructions.get_unchecked(self.instruction_offset) };
299        self.instruction_offset += 1;
300
301        Ok(FetchInstructionResult::Instruction(instruction))
302    }
303}
304
305impl EagerTestInstructionFetcher {
306    /// Create a new instance with the specified instructions and base address.
307    ///
308    /// Instructions are in the same order as they appear in the binary, and the base address
309    /// corresponds to the first instruction.
310    ///
311    /// `return_trap_address` is the address at which the interpreter will stop execution
312    /// (gracefully).
313    ///
314    /// # Safety
315    /// The program counter must be valid and aligned, the instructions processed must end with a
316    /// jump instruction.
317    #[inline(always)]
318    pub unsafe fn new(
319        instructions: Vec<ContractInstruction>,
320        return_trap_address: u64,
321        base_addr: u64,
322        pc: u64,
323    ) -> Self {
324        Self {
325            instructions,
326            return_trap_address,
327            base_addr,
328            instruction_offset: (pc - base_addr) as usize / size_of::<u32>(),
329        }
330    }
331}
332
333/// System instruction handler that does nothing
334#[derive(Debug, Clone, Copy)]
335pub struct NoopRv64SystemInstructionHandler<Instruction> {
336    _phantom: PhantomData<Instruction>,
337}
338
339impl<Reg> Default for NoopRv64SystemInstructionHandler<Reg> {
340    #[inline(always)]
341    fn default() -> Self {
342        Self {
343            _phantom: PhantomData,
344        }
345    }
346}
347
348impl<Reg, Memory, PC, CustomError> Rv64SystemInstructionHandler<Reg, Memory, PC, CustomError>
349    for NoopRv64SystemInstructionHandler<Rv64Instruction<Reg>>
350where
351    Reg: Register<Type = u64>,
352    [(); Reg::N]:,
353{
354    #[inline(always)]
355    fn handle_ecall(
356        &mut self,
357        _regs: &mut Registers<Reg>,
358        _memory: &mut Memory,
359        _program_counter: &mut PC,
360    ) -> Result<ControlFlow<()>, ExecutionError<u64, Rv64Instruction<Reg>, CustomError>> {
361        // SAFETY: Contracts are statically known to not contain `ecall` instructions
362        // unsafe { unreachable_unchecked() }
363        // For some known reason this is faster than `unreachable_unchecked()`
364        Ok(ControlFlow::Continue(()))
365    }
366}
367
368/// Execute [`ContractInstruction`]s
369#[expect(clippy::type_complexity)]
370pub fn execute<Memory, IF>(
371    state: &mut Rv64InterpreterState<
372        <ContractInstruction as Instruction>::Reg,
373        Memory,
374        IF,
375        NoopRv64SystemInstructionHandler<
376            Rv64Instruction<<ContractInstruction as Instruction>::Reg>,
377        >,
378        &'static str,
379    >,
380) -> Result<(), ExecutionError<u64, ContractInstruction, &'static str>>
381where
382    Memory: VirtualMemory,
383    IF: InstructionFetcher<ContractInstruction, Memory, &'static str>,
384{
385    loop {
386        let instruction = match state
387            .instruction_fetcher
388            .fetch_instruction(&mut state.memory)?
389        {
390            FetchInstructionResult::Instruction(instruction) => instruction,
391            FetchInstructionResult::ControlFlow(ControlFlow::Continue(())) => {
392                continue;
393            }
394            FetchInstructionResult::ControlFlow(ControlFlow::Break(())) => {
395                break;
396            }
397        };
398
399        match instruction.execute(state)? {
400            ControlFlow::Continue(()) => {
401                continue;
402            }
403            ControlFlow::Break(()) => {
404                break;
405            }
406        }
407    }
408
409    Ok(())
410}
ab_riscv_benchmarks/host_utils.rs

ab_riscv_benchmarks/
host_utils.rs
