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::basic::BasicInterpreterState;
9use ab_riscv_interpreter::prelude::*;
10use ab_riscv_primitives::prelude::*;
11use alloc::vec::Vec;
12use core::marker::PhantomData;
13use core::mem::offset_of;
14use core::ops::ControlFlow;
15
16/// Contract file bytes
17pub const RISCV_CONTRACT_BYTES: &[u8] = cfg_select! {
18    target_env = "abundance" => {
19        &[]
20    }
21    _ => {
22        include_bytes!(env!("CONTRACT_PATH"))
23    }
24};
25
26// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
27//  `SimpleInternalArgs` data trait for this or something
28/// Helper data structure for [`Benchmarks::blake3_hash_chunk()`] method
29///
30/// [`Benchmarks::blake3_hash_chunk()`]: crate::Benchmarks::blake3_hash_chunk
31#[derive(Debug, Copy, Clone)]
32#[repr(C)]
33pub struct Blake3HashChunkInternalArgs {
34    chunk_ptr: u64,
35    chunk_size: u32,
36    chunk_capacity: u32,
37    result_ptr: u64,
38    chunk: [u8; CHUNK_LEN],
39    result: [u8; OUT_LEN],
40}
41
42impl Blake3HashChunkInternalArgs {
43    /// Create a new instance
44    pub fn new(internal_args_addr: u64, chunk: [u8; CHUNK_LEN]) -> Self {
45        Self {
46            chunk_ptr: internal_args_addr + offset_of!(Self, chunk) as u64,
47            chunk_size: CHUNK_LEN as u32,
48            chunk_capacity: CHUNK_LEN as u32,
49            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
50            chunk,
51            result: [0; _],
52        }
53    }
54
55    /// Extract result
56    pub fn result(&self) -> [u8; OUT_LEN] {
57        self.result
58    }
59}
60
61// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
62//  `SimpleInternalArgs` data trait for this or something
63/// Helper data structure for [`Benchmarks::ed25519_verify()`] method
64///
65/// [`Benchmarks::ed25519_verify()`]: crate::Benchmarks::ed25519_verify
66#[derive(Debug, Copy, Clone)]
67#[repr(C)]
68pub struct Ed25519VerifyInternalArgs {
69    pub public_key_ptr: u64,
70    pub public_key_size: u32,
71    pub public_key_capacity: u32,
72    pub signature_ptr: u64,
73    pub signature_size: u32,
74    pub signature_capacity: u32,
75    pub message_ptr: u64,
76    pub message_size: u32,
77    pub message_capacity: u32,
78    pub result_ptr: u64,
79    pub public_key: Ed25519PublicKey,
80    pub signature: Ed25519Signature,
81    pub message: [u8; OUT_LEN],
82    pub result: Bool,
83}
84
85impl Ed25519VerifyInternalArgs {
86    /// Create a new instance
87    pub fn new(
88        internal_args_addr: u64,
89        public_key: Ed25519PublicKey,
90        signature: Ed25519Signature,
91        message: [u8; OUT_LEN],
92    ) -> Self {
93        Self {
94            public_key_ptr: internal_args_addr + offset_of!(Self, public_key) as u64,
95            public_key_size: Ed25519PublicKey::SIZE as u32,
96            public_key_capacity: Ed25519PublicKey::SIZE as u32,
97            signature_ptr: internal_args_addr + offset_of!(Self, signature) as u64,
98            signature_size: Ed25519Signature::SIZE as u32,
99            signature_capacity: Ed25519Signature::SIZE as u32,
100            message_ptr: internal_args_addr + offset_of!(Self, message) as u64,
101            message_size: OUT_LEN as u32,
102            message_capacity: OUT_LEN as u32,
103            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
104            public_key,
105            signature,
106            message,
107            result: Bool::new(false),
108        }
109    }
110
111    /// Extract result
112    pub fn result(&self) -> Bool {
113        self.result
114    }
115}
116
117/// Simple test memory implementation
118#[derive(Debug, Copy, Clone)]
119#[repr(align(16))]
120pub struct TestMemory<const BASE_ADDR: u64, const SIZE: usize> {
121    data: [u8; SIZE],
122}
123
124impl<const BASE_ADDR: u64, const SIZE: usize> VirtualMemory for TestMemory<BASE_ADDR, SIZE> {
125    fn read<T>(&self, address: u64) -> Result<T, VirtualMemoryError>
126    where
127        T: BasicInt,
128    {
129        let offset = address.wrapping_sub(BASE_ADDR);
130
131        if offset.saturating_add(size_of::<T>() as u64) > self.data.len() as u64 {
132            return Err(VirtualMemoryError::OutOfBoundsRead { address });
133        }
134
135        // SAFETY: Only reading basic integers from initialized memory
136        unsafe {
137            Ok(self
138                .data
139                .as_ptr()
140                .cast::<T>()
141                .byte_add(offset as usize)
142                .read_unaligned())
143        }
144    }
145
146    unsafe fn read_unchecked<T>(&self, address: u64) -> T
147    where
148        T: BasicInt,
149    {
150        // SAFETY: Guaranteed by function contract
151        unsafe {
152            let offset = address.unchecked_sub(BASE_ADDR) as usize;
153            self.data
154                .as_ptr()
155                .cast::<T>()
156                .byte_add(offset)
157                .read_unaligned()
158        }
159    }
160
161    fn read_slice(&self, address: u64, len: u32) -> Result<&[u8], VirtualMemoryError> {
162        let offset = address.wrapping_sub(BASE_ADDR);
163
164        if offset > self.data.len() as u64 {
165            return Err(VirtualMemoryError::OutOfBoundsRead { address });
166        }
167
168        self.data
169            .get(offset as usize..)
170            .and_then(|data| data.get(..len as usize))
171            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })
172    }
173
174    fn read_slice_up_to(&self, address: u64, len: u32) -> &[u8] {
175        let offset = address.wrapping_sub(BASE_ADDR);
176
177        if offset > self.data.len() as u64 {
178            return &[];
179        }
180
181        let remaining = self.data.get(offset as usize..).unwrap_or_default();
182        remaining.get(..len as usize).unwrap_or(remaining)
183    }
184
185    fn write<T>(&mut self, address: u64, value: T) -> Result<(), VirtualMemoryError>
186    where
187        T: BasicInt,
188    {
189        let offset = address.wrapping_sub(BASE_ADDR);
190
191        if offset.saturating_add(size_of::<T>() as u64) > self.data.len() as u64 {
192            return Err(VirtualMemoryError::OutOfBoundsWrite { address });
193        }
194
195        // SAFETY: Only writing basic integers to initialized memory
196        unsafe {
197            self.data
198                .as_mut_ptr()
199                .cast::<T>()
200                .byte_add(offset as usize)
201                .write_unaligned(value);
202        }
203
204        Ok(())
205    }
206
207    fn write_slice(&mut self, address: u64, data: &[u8]) -> Result<(), VirtualMemoryError> {
208        let offset = address.wrapping_sub(BASE_ADDR);
209
210        if offset > self.data.len() as u64 {
211            return Err(VirtualMemoryError::OutOfBoundsWrite { address });
212        }
213
214        let len = data.len();
215        self.data
216            .get_mut(offset as usize..)
217            .and_then(|data| data.get_mut(..len))
218            .ok_or(VirtualMemoryError::OutOfBoundsWrite { address })?
219            .copy_from_slice(data);
220
221        Ok(())
222    }
223}
224
225impl<const BASE_ADDR: u64, const SIZE: usize> Default for TestMemory<BASE_ADDR, SIZE> {
226    fn default() -> Self {
227        Self { data: [0; SIZE] }
228    }
229}
230
231impl<const BASE_ADDR: u64, const SIZE: usize> TestMemory<BASE_ADDR, SIZE> {
232    /// Get a mutable slice of memory
233    pub fn get_mut_bytes(
234        &mut self,
235        address: u64,
236        size: usize,
237    ) -> Result<&mut [u8], VirtualMemoryError> {
238        let offset = address
239            .checked_sub(BASE_ADDR)
240            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })? as usize;
241
242        if offset + size > self.data.len() {
243            return Err(VirtualMemoryError::OutOfBoundsRead { address });
244        }
245
246        Ok(&mut self.data[offset..][..size])
247    }
248}
249
250/// Lazy instruction fetcher implementation
251#[derive(Debug, Copy, Clone)]
252pub struct LazyInstructionFetcher {
253    return_trap_address: u64,
254    pc: u64,
255}
256
257impl<Memory> ProgramCounter<u64, Memory> for LazyInstructionFetcher
258where
259    Memory: VirtualMemory,
260{
261    #[inline(always)]
262    fn get_pc(&self) -> u64 {
263        self.pc
264    }
265
266    #[inline]
267    fn set_pc(
268        &mut self,
269        memory: &Memory,
270        pc: u64,
271    ) -> Result<ControlFlow<()>, ProgramCounterError<u64>> {
272        if pc == self.return_trap_address {
273            return Ok(ControlFlow::Break(()));
274        }
275
276        if !pc.is_multiple_of(u64::from(ContractInstruction::alignment())) {
277            return Err(ProgramCounterError::UnalignedInstruction { address: pc });
278        }
279
280        // Note: This will not allow reading a 16-bit instruction at the very end of memory range,
281        // but that is going to be the case here anyway since code is followed by read-write memory
282        // anyway
283        memory.read::<u32>(pc)?;
284
285        self.pc = pc;
286
287        Ok(ControlFlow::Continue(()))
288    }
289}
290
291impl<Memory> InstructionFetcher<ContractInstruction, Memory> for LazyInstructionFetcher
292where
293    Memory: VirtualMemory,
294{
295    #[inline]
296    fn fetch_instruction(
297        &mut self,
298        memory: &Memory,
299    ) -> Result<FetchInstructionResult<ContractInstruction>, ExecutionError<u64>> {
300        // SAFETY: Constructor guarantees that the last instruction is a jump, which means going
301        // through `Self::set_pc()` method does the necessary bounds check and advancing forward by
302        // one instruction can't result in out-of-bounds access.
303        let instruction = unsafe { memory.read_unchecked(self.pc) };
304        // SAFETY: All instructions are valid, according to the constructor contract
305        let instruction =
306            unsafe { ContractInstruction::try_decode(instruction).unwrap_unchecked() };
307
308        self.pc += u64::from(instruction.size());
309
310        Ok(FetchInstructionResult::Instruction(instruction))
311    }
312}
313
314impl LazyInstructionFetcher {
315    /// Create a new instance.
316    ///
317    /// `return_trap_address` is the address at which the interpreter will stop execution
318    /// (gracefully).
319    ///
320    /// # Safety
321    /// The program counter must be valid and aligned, the instructions processed must be valid and
322    /// end with a jump instruction.
323    #[inline(always)]
324    pub unsafe fn new(return_trap_address: u64, pc: u64) -> Self {
325        Self {
326            return_trap_address,
327            pc,
328        }
329    }
330}
331
332/// Eager instruction handler eagerly decodes all instructions upfront
333#[derive(Debug, Default, Clone)]
334pub struct EagerTestInstructionFetcher {
335    instructions: Vec<ContractInstruction>,
336    return_trap_address: u64,
337    base_addr: u64,
338    instruction_offset: usize,
339}
340
341impl<Memory> ProgramCounter<u64, Memory> for EagerTestInstructionFetcher
342where
343    Memory: VirtualMemory,
344{
345    #[inline(always)]
346    fn get_pc(&self) -> u64 {
347        self.base_addr + self.instruction_offset as u64 * size_of::<u16>() as u64
348    }
349
350    #[inline]
351    fn set_pc(
352        &mut self,
353        _memory: &Memory,
354        pc: u64,
355    ) -> Result<ControlFlow<()>, ProgramCounterError<u64>> {
356        let address = pc;
357
358        if address == self.return_trap_address {
359            return Ok(ControlFlow::Break(()));
360        }
361
362        if !address.is_multiple_of(size_of::<u16>() as u64) {
363            return Err(ProgramCounterError::UnalignedInstruction { address });
364        }
365
366        let offset = address
367            .checked_sub(self.base_addr)
368            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })? as usize;
369        let instruction_offset = offset / size_of::<u16>();
370
371        if instruction_offset >= self.instructions.len() {
372            return Err(VirtualMemoryError::OutOfBoundsRead { address }.into());
373        }
374
375        self.instruction_offset = instruction_offset;
376
377        Ok(ControlFlow::Continue(()))
378    }
379}
380
381impl<Memory> InstructionFetcher<ContractInstruction, Memory> for EagerTestInstructionFetcher
382where
383    Memory: VirtualMemory,
384{
385    #[inline(always)]
386    fn fetch_instruction(
387        &mut self,
388        _memory: &Memory,
389    ) -> Result<FetchInstructionResult<ContractInstruction>, ExecutionError<u64>> {
390        // SAFETY: Constructor guarantees that the last instruction is a jump, which means going
391        // through `Self::set_pc()` method does the necessary bounds check and advancing forward by
392        // one instruction can't result in out-of-bounds access.
393        let instruction = *unsafe { self.instructions.get_unchecked(self.instruction_offset) };
394        self.instruction_offset += usize::from(instruction.size()) / size_of::<u16>();
395
396        Ok(FetchInstructionResult::Instruction(instruction))
397    }
398}
399
400impl EagerTestInstructionFetcher {
401    /// Create a new instance with the specified instructions and base address.
402    ///
403    /// Instructions are decoded during instantiation of the instruction fetcher, and the base
404    /// address corresponds to the first instruction.
405    ///
406    /// `return_trap_address` is the address at which the interpreter will stop execution
407    /// (gracefully).
408    ///
409    /// # Safety
410    /// The program counter must be valid and aligned, the instructions processed must end with a
411    /// jump instruction.
412    #[inline(always)]
413    pub unsafe fn new(
414        instructions: &[u8],
415        return_trap_address: u64,
416        base_addr: u64,
417        pc: u64,
418    ) -> Self {
419        let mut decoded_instructions = Vec::with_capacity(instructions.len() / size_of::<u16>());
420
421        let mut offset = 0;
422        while let Some(instruction_bytes) = instructions.get(offset..offset + size_of::<u32>()) {
423            let decoded_instruction = u32::from_le_bytes([
424                instruction_bytes[0],
425                instruction_bytes[1],
426                instruction_bytes[2],
427                instruction_bytes[3],
428            ]);
429            // Use `Unimp` as a fallback, though contract is expected to only contain legal
430            // instructions
431            let decoded_instruction =
432                Instruction::try_decode(decoded_instruction).unwrap_or(ContractInstruction::Unimp);
433            decoded_instructions.push(decoded_instruction);
434            match decoded_instruction.size() {
435                2 => {
436                    offset += 2;
437                }
438                4 => {
439                    // The second half of a 32-bit instruction is a valid offset and may or may not
440                    // decode to a valid instruction on its own. Try to decode it but ignore
441                    // decoding failures.
442
443                    offset += 2;
444
445                    // Could be both 16-bit and 32-bit instruction, need to handle end of the
446                    // instruction stream
447                    let instruction_word = if let Some(instruction_bytes) =
448                        instructions.get(offset..offset + size_of::<u32>())
449                    {
450                        u32::from_le_bytes([
451                            instruction_bytes[0],
452                            instruction_bytes[1],
453                            instruction_bytes[2],
454                            instruction_bytes[3],
455                        ])
456                    } else {
457                        u32::from_le_bytes([instruction_bytes[2], instruction_bytes[3], 0, 0])
458                    };
459
460                    decoded_instructions.push(
461                        Instruction::try_decode(instruction_word)
462                            .unwrap_or(ContractInstruction::Unimp),
463                    );
464                    offset += 2;
465                }
466                instruction_size => {
467                    unreachable!("Invalid instruction size {instruction_size}, expected 2 or 4");
468                }
469            }
470        }
471
472        let remainder_bytes = instructions.get(offset..).unwrap_or(&[]);
473
474        if remainder_bytes.len() == size_of::<u16>() {
475            let instruction_word =
476                u32::from_le_bytes([remainder_bytes[0], remainder_bytes[1], 0, 0]);
477            decoded_instructions.push(
478                Instruction::try_decode(instruction_word).unwrap_or(ContractInstruction::Unimp),
479            );
480        };
481
482        Self {
483            instructions: decoded_instructions,
484            return_trap_address,
485            base_addr,
486            instruction_offset: (pc - base_addr) as usize / size_of::<u16>(),
487        }
488    }
489}
490
491/// System instruction handler that does nothing
492#[derive(Debug, Clone, Copy)]
493pub struct NoopRv64SystemInstructionHandler<Instruction> {
494    _phantom: PhantomData<Instruction>,
495}
496
497impl<Reg> Default for NoopRv64SystemInstructionHandler<Reg> {
498    #[inline(always)]
499    fn default() -> Self {
500        Self {
501            _phantom: PhantomData,
502        }
503    }
504}
505
506impl<Reg, Regs, Memory, PC, CustomError>
507    SystemInstructionHandler<Reg, Regs, Memory, PC, CustomError>
508    for NoopRv64SystemInstructionHandler<Rv64Instruction<Reg>>
509where
510    Reg: Register<Type = u64>,
511    Regs: RegisterFile<Reg>,
512{
513    #[inline(always)]
514    fn handle_ecall(
515        &mut self,
516        _regs: &mut Regs,
517        _memory: &mut Memory,
518        _program_counter: &mut PC,
519    ) -> Result<ControlFlow<()>, ExecutionError<u64, CustomError>> {
520        // SAFETY: Contracts are statically known to not contain `ecall` instructions
521        // unsafe { unreachable_unchecked() }
522        // For some known reason this is faster than `unreachable_unchecked()`
523        Ok(ControlFlow::Continue(()))
524    }
525}
526
527/// Execute [`ContractInstruction`]s
528pub fn execute<Regs, Memory, IF>(
529    state: &mut BasicInterpreterState<
530        Regs,
531        (),
532        Memory,
533        IF,
534        NoopRv64SystemInstructionHandler<
535            Rv64Instruction<<ContractInstruction as Instruction>::Reg>,
536        >,
537    >,
538) -> Result<(), ExecutionError<u64>>
539where
540    Regs: RegisterFile<<ContractInstruction as Instruction>::Reg>,
541    Memory: VirtualMemory,
542    IF: InstructionFetcher<ContractInstruction, Memory>,
543{
544    loop {
545        let instruction = match state.instruction_fetcher.fetch_instruction(&state.memory)? {
546            FetchInstructionResult::Instruction(instruction) => instruction,
547            FetchInstructionResult::ControlFlow(ControlFlow::Continue(())) => {
548                continue;
549            }
550            FetchInstructionResult::ControlFlow(ControlFlow::Break(())) => {
551                break;
552            }
553        };
554
555        match instruction.execute(
556            &mut state.regs,
557            &mut state.ext_state,
558            &mut state.memory,
559            &mut state.instruction_fetcher,
560            &mut state.system_instruction_handler,
561        )? {
562            ControlFlow::Continue(()) => {
563                continue;
564            }
565            ControlFlow::Break(()) => {
566                break;
567            }
568        }
569    }
570
571    Ok(())
572}
ab_riscv_benchmarks/host_utils.rs

ab_riscv_benchmarks/
host_utils.rs
