ab_riscv_benchmarks/
host_utils.rs

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

ab_riscv_benchmarks/
host_utils.rs
