ab_riscv_benchmarks/
host_utils.rs

1extern crate alloc;
2
3use ab_blake3::{CHUNK_LEN, OUT_LEN};
4use ab_contract_file::instruction::{ContractInstruction, ContractRegister};
5use ab_core_primitives::ed25519::{Ed25519PublicKey, Ed25519Signature};
6use ab_io_type::bool::Bool;
7use ab_riscv_interpreter::prelude::*;
8use ab_riscv_primitives::prelude::*;
9use alloc::alloc::{alloc, dealloc, handle_alloc_error};
10use alloc::vec::Vec;
11use core::alloc::Layout;
12use core::fmt;
13use core::hint::cold_path;
14use core::marker::PhantomData;
15use core::mem::offset_of;
16use core::ops::ControlFlow;
17use core::ptr::NonNull;
18
19/// Contract file bytes
20pub const RISCV_CONTRACT_BYTES: &[u8] = cfg_select! {
21    target_env = "abundance" => &[],
22    _ => {
23        include_bytes!(env!("CONTRACT_PATH"))
24    }
25};
26
27// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
28//  `SimpleInternalArgs` data trait for this or something
29/// Helper data structure for [`Benchmarks::blake3_hash_chunk()`] method
30///
31/// [`Benchmarks::blake3_hash_chunk()`]: crate::Benchmarks::blake3_hash_chunk
32#[derive(Debug, Copy, Clone)]
33#[repr(C)]
34pub struct Blake3HashChunkInternalArgs {
35    chunk_ptr: u64,
36    chunk_size: u32,
37    chunk_capacity: u32,
38    result_ptr: u64,
39    chunk: [u8; CHUNK_LEN],
40    result: [u8; OUT_LEN],
41}
42
43const _: () = {
44    assert!(
45        size_of::<Blake3HashChunkInternalArgs>()
46            == offset_of!(Blake3HashChunkInternalArgs, result) + size_of::<[u8; OUT_LEN]>(),
47        "`Blake3HashChunkInternalArgs` must not have implicit padding"
48    );
49};
50
51impl Blake3HashChunkInternalArgs {
52    /// Create a new instance
53    pub fn new(internal_args_addr: u64, chunk: [u8; CHUNK_LEN]) -> Self {
54        Self {
55            chunk_ptr: internal_args_addr + offset_of!(Self, chunk) as u64,
56            chunk_size: CHUNK_LEN as u32,
57            chunk_capacity: CHUNK_LEN as u32,
58            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
59            chunk,
60            result: [0; _],
61        }
62    }
63
64    /// Extract result
65    pub fn result(&self) -> [u8; OUT_LEN] {
66        self.result
67    }
68}
69
70// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
71//  `SimpleInternalArgs` data trait for this or something
72/// Helper data structure for [`Benchmarks::ed25519_verify()`] method
73///
74/// [`Benchmarks::ed25519_verify()`]: crate::Benchmarks::ed25519_verify
75#[derive(Debug, Copy, Clone)]
76#[repr(C)]
77pub struct Ed25519VerifyInternalArgs {
78    pub public_key_ptr: u64,
79    pub public_key_size: u32,
80    pub public_key_capacity: u32,
81    pub signature_ptr: u64,
82    pub signature_size: u32,
83    pub signature_capacity: u32,
84    pub message_ptr: u64,
85    pub message_size: u32,
86    pub message_capacity: u32,
87    pub result_ptr: u64,
88    pub public_key: Ed25519PublicKey,
89    pub signature: Ed25519Signature,
90    pub message: [u8; OUT_LEN],
91    pub result: Bool,
92    /// Explicit trailing padding.
93    ///
94    /// The host copies the byte representation of this data structure into guest memory, which is
95    /// only sound if every byte of it is initialized, hence implicit padding must not exist here.
96    pub padding: [u8; 7],
97}
98
99const _: () = {
100    assert!(
101        size_of::<Ed25519VerifyInternalArgs>()
102            == offset_of!(Ed25519VerifyInternalArgs, padding) + size_of::<[u8; 7]>(),
103        "`Ed25519VerifyInternalArgs` must not have implicit padding"
104    );
105};
106
107impl Ed25519VerifyInternalArgs {
108    /// Create a new instance
109    pub fn new(
110        internal_args_addr: u64,
111        public_key: Ed25519PublicKey,
112        signature: Ed25519Signature,
113        message: [u8; OUT_LEN],
114    ) -> Self {
115        Self {
116            public_key_ptr: internal_args_addr + offset_of!(Self, public_key) as u64,
117            public_key_size: Ed25519PublicKey::SIZE as u32,
118            public_key_capacity: Ed25519PublicKey::SIZE as u32,
119            signature_ptr: internal_args_addr + offset_of!(Self, signature) as u64,
120            signature_size: Ed25519Signature::SIZE as u32,
121            signature_capacity: Ed25519Signature::SIZE as u32,
122            message_ptr: internal_args_addr + offset_of!(Self, message) as u64,
123            message_size: OUT_LEN as u32,
124            message_capacity: OUT_LEN as u32,
125            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
126            public_key,
127            signature,
128            message,
129            result: Bool::new(false),
130            padding: [0; _],
131        }
132    }
133
134    /// Extract result
135    pub fn result(&self) -> Bool {
136        self.result
137    }
138}
139
140/// Simple test memory implementation
141#[derive(Debug, Copy, Clone)]
142#[repr(align(16))]
143pub struct TestMemory<const BASE_ADDR: u64, const SIZE: usize> {
144    data: [u8; SIZE],
145}
146
147impl<const BASE_ADDR: u64, const SIZE: usize> VirtualMemory for TestMemory<BASE_ADDR, SIZE> {
148    #[inline(always)]
149    fn read<T>(&self, address: u64) -> Result<T, VirtualMemoryError>
150    where
151        T: BasicInt,
152    {
153        let offset = address.wrapping_sub(BASE_ADDR);
154
155        if offset.saturating_add(size_of::<T>() as u64) > self.data.len() as u64 {
156            cold_path();
157            return Err(VirtualMemoryError::OutOfBoundsRead { address });
158        }
159
160        // SAFETY: Only reading basic integers from initialized memory
161        unsafe {
162            Ok(self
163                .data
164                .as_ptr()
165                .cast::<T>()
166                .byte_add(offset as usize)
167                .read_unaligned())
168        }
169    }
170
171    #[inline(always)]
172    unsafe fn read_unchecked<T>(&self, address: u64) -> T
173    where
174        T: BasicInt,
175    {
176        // SAFETY: Guaranteed by function contract
177        unsafe {
178            let offset = address.unchecked_sub(BASE_ADDR) as usize;
179            self.data
180                .as_ptr()
181                .cast::<T>()
182                .byte_add(offset)
183                .read_unaligned()
184        }
185    }
186
187    fn read_slice(&self, address: u64, len: u32) -> Result<&[u8], VirtualMemoryError> {
188        let offset = address.wrapping_sub(BASE_ADDR);
189
190        if offset > self.data.len() as u64 {
191            cold_path();
192            return Err(VirtualMemoryError::OutOfBoundsRead { address });
193        }
194
195        self.data
196            .get(offset as usize..)
197            .and_then(|data| data.get(..len as usize))
198            .ok_or(VirtualMemoryError::OutOfBoundsRead { address })
199    }
200
201    fn read_slice_up_to(&self, address: u64, len: u32) -> &[u8] {
202        let offset = address.wrapping_sub(BASE_ADDR);
203
204        if offset > self.data.len() as u64 {
205            cold_path();
206            return &[];
207        }
208
209        let remaining = self.data.get(offset as usize..).unwrap_or_default();
210        remaining.get(..len as usize).unwrap_or(remaining)
211    }
212
213    #[inline(always)]
214    fn write<T>(&mut self, address: u64, value: T) -> Result<(), VirtualMemoryError>
215    where
216        T: BasicInt,
217    {
218        let offset = address.wrapping_sub(BASE_ADDR);
219
220        if offset.saturating_add(size_of::<T>() as u64) > self.data.len() as u64 {
221            cold_path();
222            return Err(VirtualMemoryError::OutOfBoundsWrite { address });
223        }
224
225        // SAFETY: Only writing basic integers to initialized memory
226        unsafe {
227            self.data
228                .as_mut_ptr()
229                .cast::<T>()
230                .byte_add(offset as usize)
231                .write_unaligned(value);
232        }
233
234        Ok(())
235    }
236
237    fn write_slice(&mut self, address: u64, data: &[u8]) -> Result<(), VirtualMemoryError> {
238        let offset = address.wrapping_sub(BASE_ADDR);
239
240        if offset > self.data.len() as u64 {
241            cold_path();
242            return Err(VirtualMemoryError::OutOfBoundsWrite { address });
243        }
244
245        let len = data.len();
246        let Some(target_data) = self
247            .data
248            .get_mut(offset as usize..)
249            .and_then(|data| data.get_mut(..len))
250        else {
251            cold_path();
252            return Err(VirtualMemoryError::OutOfBoundsWrite { address });
253        };
254
255        target_data.copy_from_slice(data);
256
257        Ok(())
258    }
259}
260
261impl<const BASE_ADDR: u64, const SIZE: usize> Default for TestMemory<BASE_ADDR, SIZE> {
262    fn default() -> Self {
263        Self { data: [0; SIZE] }
264    }
265}
266
267impl<const BASE_ADDR: u64, const SIZE: usize> TestMemory<BASE_ADDR, SIZE> {
268    /// Get a mutable slice of memory
269    pub fn get_mut_bytes(
270        &mut self,
271        address: u64,
272        size: usize,
273    ) -> Result<&mut [u8], VirtualMemoryError> {
274        let Some(offset) = address.checked_sub(BASE_ADDR) else {
275            cold_path();
276            return Err(VirtualMemoryError::OutOfBoundsRead { address });
277        };
278        let offset = offset as usize;
279
280        let Some(slice) = self
281            .data
282            .get_mut(offset..)
283            .and_then(|data| data.get_mut(..size))
284        else {
285            cold_path();
286            return Err(VirtualMemoryError::OutOfBoundsRead { address });
287        };
288
289        Ok(slice)
290    }
291}
292
293/// Lazy instruction fetcher implementation
294#[derive(Debug, Copy, Clone)]
295pub struct LazyInstructionFetcher {
296    return_trap_address: u64,
297    pc: u64,
298}
299
300impl<Memory> ProgramCounter<u64, Memory> for LazyInstructionFetcher
301where
302    Memory: VirtualMemory,
303{
304    #[inline(always)]
305    fn get_pc(&self) -> u64 {
306        self.pc
307    }
308
309    #[inline(always)]
310    unsafe fn try_set_pc_relative(&mut self, instruction_size: u8, offset: i32) -> bool {
311        let old_pc = <Self as ProgramCounter<_, Memory>>::old_pc(self, instruction_size);
312        let pc = old_pc.wrapping_add_signed(i64::from(offset));
313        // Stored either way: on the way out it is what `failed_branch()` reports on, and until then
314        // nothing else is allowed to look at it
315        self.pc = pc;
316
317        pc != self.return_trap_address
318            && pc.is_multiple_of(u64::from(
319                ContractInstruction::<ContractRegister>::alignment(),
320            ))
321    }
322
323    #[cold]
324    #[inline(never)]
325    unsafe fn failed_branch(
326        &mut self,
327        memory: &Memory,
328    ) -> Result<ControlFlow<()>, ExecutionError<u64>> {
329        // The program counter holds the refused target, and `set_pc()` is what says what is wrong
330        // with it
331        self.set_pc(memory, self.pc)
332    }
333
334    #[inline]
335    fn set_pc(&mut self, memory: &Memory, pc: u64) -> Result<ControlFlow<()>, ExecutionError<u64>> {
336        if pc == self.return_trap_address {
337            cold_path();
338            return Ok(ControlFlow::Break(()));
339        }
340
341        if !pc.is_multiple_of(u64::from(
342            ContractInstruction::<ContractRegister>::alignment(),
343        )) {
344            cold_path();
345            return Err(ExecutionError::UnalignedInstruction {
346                address: PackedAddress::new(pc),
347            });
348        }
349
350        // Note: This will not allow reading a 16-bit instruction at the very end of memory range,
351        // but that is going to be the case here anyway since code is followed by read-write memory
352        // anyway
353        if let Err(error) = memory.read::<u32>(pc) {
354            cold_path();
355            return Err(error.into());
356        }
357
358        self.pc = pc;
359
360        Ok(ControlFlow::Continue(()))
361    }
362}
363
364impl<Memory> InstructionFetcher<ContractInstruction, Memory> for LazyInstructionFetcher
365where
366    Memory: VirtualMemory,
367{
368    #[inline]
369    fn peek_instruction(&mut self, memory: &Memory) -> FetchInstructionResult<ContractInstruction> {
370        // SAFETY: Constructor guarantees that the last instruction is a jump, which means going
371        // through `Self::set_pc()` method does the necessary bounds check, so the program counter
372        // always sits on an instruction.
373        let instruction = unsafe { memory.read_unchecked(self.pc) };
374        // SAFETY: All instructions are valid, according to the constructor contract
375        let instruction =
376            unsafe { ContractInstruction::try_decode(instruction).unwrap_unchecked() };
377
378        FetchInstructionResult::Instruction(instruction)
379    }
380
381    #[inline]
382    unsafe fn advance(&mut self, instruction_size: u8) {
383        self.pc = self.pc.wrapping_add(u64::from(instruction_size));
384    }
385
386    #[inline]
387    fn fetch_instruction(
388        &mut self,
389        memory: &Memory,
390    ) -> FetchInstructionResult<ContractInstruction> {
391        let result =
392            InstructionFetcher::<ContractInstruction, Memory>::peek_instruction(self, memory);
393
394        if let FetchInstructionResult::Instruction(instruction) = result {
395            // SAFETY: The instruction was just peeked successfully, and this is the only place that
396            // moves past it
397            unsafe {
398                InstructionFetcher::<ContractInstruction, Memory>::advance(
399                    self,
400                    instruction.size(),
401                );
402            }
403        }
404
405        result
406    }
407}
408
409impl LazyInstructionFetcher {
410    /// Create a new instance.
411    ///
412    /// `return_trap_address` is the address at which the interpreter will stop execution
413    /// (gracefully).
414    ///
415    /// # Safety
416    /// The program counter must be valid and aligned, the instructions processed must be valid and
417    /// end with a jump instruction.
418    #[inline(always)]
419    pub unsafe fn new(return_trap_address: u64, pc: u64) -> Self {
420        Self {
421            return_trap_address,
422            pc,
423        }
424    }
425}
426
427/// Everything [`EagerTestInstructionFetcher`] needs besides its position within the decoded
428/// instruction stream.
429///
430/// This lives in a single heap allocation whose tail holds the decoded instructions themselves,
431/// starting [`EagerTestInstructions::INSTRUCTIONS_OFFSET`] bytes from the beginning of it.
432/// That is what keeps the fetcher itself down to two pointers, so it fits into two argument
433/// registers when threaded through tail-called instruction handlers by value.
434#[derive(Debug)]
435#[repr(C)]
436struct EagerTestInstructionFetcherState {
437    /// Number of decoded instructions stored right after this header
438    instructions_len: usize,
439    /// Guest address that corresponds to the first decoded instruction
440    base_addr: u64,
441    /// Guest address at which execution stops gracefully
442    return_trap_address: u64,
443}
444
445/// Instructions decoded upfront, which [`EagerTestInstructionFetcher`] walks.
446///
447/// Ownership of the allocation lives here rather than in the fetcher because the fetcher is moved
448/// through tail-called instruction handlers by value. A destructor on it would make every handler
449/// that can fail (every load, store, branch and jump) responsible for dropping it on the way out,
450/// which costs a stack frame, callee-saved register spills and a reload in the hot path of each of
451/// them, even though the failing path is never taken.
452pub struct EagerTestInstructions {
453    /// State header, together with the decoded instructions themselves, in a single heap
454    /// allocation.
455    ///
456    /// This is a raw pointer rather than a `Box` on purpose: fetchers point into the same
457    /// allocation, and going through a `Box` would assert unique access to that allocation on
458    /// every use, invalidating pointers that must survive across all of them.
459    state: NonNull<EagerTestInstructionFetcherState>,
460}
461
462impl fmt::Debug for EagerTestInstructions {
463    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
464        f.debug_struct("EagerTestInstructions")
465            .field("instructions_len", &self.instructions_len())
466            .field("base_addr", &self.base_addr())
467            .field("return_trap_address", &self.return_trap_address())
468            .finish_non_exhaustive()
469    }
470}
471
472impl Drop for EagerTestInstructions {
473    fn drop(&mut self) {
474        let layout = Self::allocation_layout(self.instructions_len());
475
476        // SAFETY: Allocated with the global allocator using exactly this layout, and this is the
477        // only owner of the allocation
478        unsafe {
479            dealloc(self.state.as_ptr().cast::<u8>(), layout);
480        }
481    }
482}
483
484impl EagerTestInstructions {
485    /// Byte offset of the decoded instructions from the start of the allocation that
486    /// [`Self::state`] points at
487    const INSTRUCTIONS_OFFSET: usize = size_of::<EagerTestInstructionFetcherState>()
488        .next_multiple_of(align_of::<ContractInstruction>());
489
490    /// Layout of the allocation holding [`EagerTestInstructionFetcherState`] followed by
491    /// `instructions_len` decoded instructions
492    fn allocation_layout(instructions_len: usize) -> Layout {
493        let (layout, instructions_offset) = Layout::new::<EagerTestInstructionFetcherState>()
494            .extend(
495                Layout::array::<ContractInstruction>(instructions_len)
496                    .expect("Decoded instructions fit into memory, they were just allocated; qed"),
497            )
498            .expect("Decoded instructions fit into memory, they were just allocated; qed");
499
500        debug_assert_eq!(instructions_offset, Self::INSTRUCTIONS_OFFSET);
501
502        layout.pad_to_align()
503    }
504
505    fn instantiate(
506        instructions: &[ContractInstruction],
507        base_addr: u64,
508        return_trap_address: u64,
509    ) -> Self {
510        let instructions_len = instructions.len();
511        let layout = Self::allocation_layout(instructions_len);
512        #[expect(
513            clippy::cast_ptr_alignment,
514            reason = "Layout is configured to produce correctly aligned memory"
515        )]
516        // SAFETY: The state itself is always there, so the layout has non-zero size
517        let state = unsafe { alloc(layout) }.cast::<EagerTestInstructionFetcherState>();
518        let Some(state) = NonNull::new(state) else {
519            handle_alloc_error(layout);
520        };
521
522        // SAFETY: Freshly allocated for exactly this type, correctly aligned
523        unsafe {
524            state.write(EagerTestInstructionFetcherState {
525                instructions_len,
526                base_addr,
527                return_trap_address,
528            });
529        }
530
531        let instance = Self { state };
532
533        // SAFETY: The allocation was made for exactly this many instructions and is distinct from
534        // the ones being copied in
535        unsafe {
536            instance.instructions().copy_from_nonoverlapping(
537                NonNull::from(instructions).cast::<ContractInstruction>(),
538                instructions_len,
539            );
540        }
541
542        instance
543    }
544
545    /// Pointer to the first decoded instruction
546    #[inline(always)]
547    fn instructions(&self) -> NonNull<ContractInstruction> {
548        // SAFETY: Decoded instructions are stored at this offset of the same allocation as the
549        // state
550        unsafe { self.state.byte_add(Self::INSTRUCTIONS_OFFSET) }.cast::<ContractInstruction>()
551    }
552
553    /// Number of decoded instructions
554    #[inline(always)]
555    fn instructions_len(&self) -> usize {
556        // SAFETY: State is initialized in the constructor and valid for as long as `self` is
557        unsafe { (*self.state.as_ptr()).instructions_len }
558    }
559
560    /// Guest address that corresponds to the first decoded instruction
561    #[inline(always)]
562    fn base_addr(&self) -> u64 {
563        // SAFETY: State is initialized in the constructor and valid for as long as `self` is
564        unsafe { (*self.state.as_ptr()).base_addr }
565    }
566
567    /// Guest address at which execution stops gracefully
568    #[inline(always)]
569    fn return_trap_address(&self) -> u64 {
570        // SAFETY: State is initialized in the constructor and valid for as long as `self` is
571        unsafe { (*self.state.as_ptr()).return_trap_address }
572    }
573
574    /// Create a fetcher positioned at the instruction that guest address `pc` corresponds to
575    ///
576    /// # Safety
577    /// `pc` must be a valid and aligned guest address of one of the decoded instructions.
578    #[inline(always)]
579    pub unsafe fn fetcher(&self, pc: u64) -> EagerTestInstructionFetcher<'_> {
580        let instruction_offset = (pc - self.base_addr()) as usize / size_of::<u16>();
581
582        EagerTestInstructionFetcher {
583            // SAFETY: Guaranteed by function contract, meaning `instruction_offset` is within
584            // bounds of the decoded stream
585            next_instruction: unsafe { self.instructions().add(instruction_offset) },
586            state: self.state,
587            instructions: PhantomData,
588        }
589    }
590
591    /// Decode `instructions` and create a new instance holding the result.
592    ///
593    /// `base_addr` is the guest address of the first instruction, and `return_trap_address` is the
594    /// address at which the interpreter will stop execution (gracefully).
595    ///
596    /// # Safety
597    /// The instructions processed must be valid and end with a jump instruction, and
598    /// `return_trap_address` must not fall inside them. Instruction fetching does not compare
599    /// against the return trap, so an address inside the instructions would stop execution when
600    /// jumped to but not when reached by falling through.
601    #[inline(always)]
602    pub unsafe fn decode(instructions: &[u8], return_trap_address: u64, base_addr: u64) -> Self {
603        let mut decoded_instructions: Vec<ContractInstruction> =
604            Vec::with_capacity(instructions.len() / size_of::<u16>());
605
606        let mut offset = 0;
607        while let Some(instruction_bytes) = instructions.get(offset..offset + size_of::<u32>()) {
608            let decoded_instruction = u32::from_le_bytes([
609                instruction_bytes[0],
610                instruction_bytes[1],
611                instruction_bytes[2],
612                instruction_bytes[3],
613            ]);
614            // Use `Unimp` as a fallback, though contract is expected to only contain legal
615            // instructions
616            let decoded_instruction = Instruction::try_decode(decoded_instruction).unwrap_or(
617                ContractInstruction::Unimp {
618                    rs1: Register::ZERO,
619                    rs2: Register::ZERO,
620                },
621            );
622            decoded_instructions.push(decoded_instruction);
623            match decoded_instruction.size() {
624                2 => {
625                    offset += 2;
626                }
627                4 => {
628                    // The second half of a 32-bit instruction is a valid offset and may or may not
629                    // decode to a valid instruction on its own. Try to decode it but ignore
630                    // decoding failures.
631
632                    offset += 2;
633
634                    // Could be both 16-bit and 32-bit instruction, need to handle end of the
635                    // instruction stream
636                    let instruction_word = if let Some(instruction_bytes) =
637                        instructions.get(offset..offset + size_of::<u32>())
638                    {
639                        u32::from_le_bytes([
640                            instruction_bytes[0],
641                            instruction_bytes[1],
642                            instruction_bytes[2],
643                            instruction_bytes[3],
644                        ])
645                    } else {
646                        u32::from_le_bytes([instruction_bytes[2], instruction_bytes[3], 0, 0])
647                    };
648
649                    decoded_instructions.push(Instruction::try_decode(instruction_word).unwrap_or(
650                        ContractInstruction::Unimp {
651                            rs1: Register::ZERO,
652                            rs2: Register::ZERO,
653                        },
654                    ));
655                    offset += 2;
656                }
657                instruction_size => {
658                    unreachable!("Invalid instruction size {instruction_size}, expected 2 or 4");
659                }
660            }
661        }
662
663        let remainder_bytes = instructions.get(offset..).unwrap_or(&[]);
664
665        if remainder_bytes.len() == size_of::<u16>() {
666            let instruction_word =
667                u32::from_le_bytes([remainder_bytes[0], remainder_bytes[1], 0, 0]);
668            decoded_instructions.push(Instruction::try_decode(instruction_word).unwrap_or(
669                ContractInstruction::Unimp {
670                    rs1: Register::ZERO,
671                    rs2: Register::ZERO,
672                },
673            ));
674        }
675
676        Self::instantiate(&decoded_instructions, base_addr, return_trap_address)
677    }
678}
679
680/// Eager instruction fetcher walks instructions that [`EagerTestInstructions`] decoded upfront.
681///
682/// This is a plain `Copy` cursor without a destructor, see [`EagerTestInstructions`] for why.
683#[derive(Copy, Clone)]
684#[repr(C)]
685pub struct EagerTestInstructionFetcher<'a> {
686    /// The instruction to be returned by the next [`InstructionFetcher::fetch_instruction()`]
687    /// call.
688    ///
689    /// A pointer rather than an offset helps LLVM with SROA and aliasing analysis, so it can
690    /// retain this in a native register instead of recomputing it from an offset on every
691    /// fetch.
692    next_instruction: NonNull<ContractInstruction>,
693    /// Everything else the fetcher needs, borrowed from [`EagerTestInstructions`]
694    state: NonNull<EagerTestInstructionFetcherState>,
695    /// Fetcher borrows the decoded instructions it walks
696    instructions: PhantomData<&'a EagerTestInstructions>,
697}
698
699const {
700    // When fetcher is used with threaded dispatch, it must fit into two argument registers to be
701    // passed used registers through tail calls
702    assert!(size_of::<EagerTestInstructionFetcher<'_>>() == 16);
703    // Drop glue on the fetcher would force a stack frame into every fallible handler, see
704    // `EagerTestInstructions` for details
705    assert!(!core::mem::needs_drop::<EagerTestInstructionFetcher<'_>>());
706}
707
708impl fmt::Debug for EagerTestInstructionFetcher<'_> {
709    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
710        f.debug_struct("EagerTestInstructionFetcher")
711            .field("next_instruction", &self.next_instruction)
712            .field("instructions_len", &self.instructions_len())
713            .field("base_addr", &self.base_addr())
714            .field("return_trap_address", &self.return_trap_address())
715            .finish_non_exhaustive()
716    }
717}
718
719impl<Memory> ProgramCounter<u64, Memory> for EagerTestInstructionFetcher<'_>
720where
721    Memory: VirtualMemory,
722{
723    #[inline(always)]
724    fn get_pc(&self) -> u64 {
725        let decoded_instruction_byte_offset = self
726            .next_instruction
727            .as_ptr()
728            .addr()
729            .wrapping_sub(self.instructions().as_ptr().addr());
730
731        self.base_addr()
732            + decoded_instruction_byte_offset as u64 * size_of::<u16>() as u64
733                / size_of::<ContractInstruction>() as u64
734    }
735
736    /// Moves within the decoded stream instead of resolving an address and converting it back,
737    /// which is what going through [`Self::set_pc()`] would do.
738    ///
739    /// One comparison and one test are all this needs to recognize every target it cannot resolve:
740    /// one past the end of the decoded stream, a backwards branch that ran off its start, and an
741    /// unaligned one. The return trap sits outside the decoded stream, so a branch to it fails the
742    /// bounds check here too and is answered by [`Self::failed_branch()`] like any other target
743    /// this refuses.
744    #[inline(always)]
745    unsafe fn try_set_pc_relative(&mut self, instruction_size: u8, offset: i32) -> bool {
746        // Byte offset from the instruction being executed to the branch target. The program counter
747        // is advanced during instruction fetching, so that instruction starts `instruction_size`
748        // bytes back.
749        let offset = (offset as isize).wrapping_sub(isize::from(instruction_size));
750        // Every `size_of::<u16>()` of guest code owns one decoded instruction, so the target is
751        // reached by moving within the decoded stream
752        let byte_delta =
753            offset * (size_of::<ContractInstruction>() / size_of::<u16>()).cast_signed();
754        // This may land outside the decoded stream (including before its start), which is fine:
755        // `wrapping_byte_offset()` only computes an address, it never dereferences the pointer, and
756        // the bounds check below rejects such a target before it is ever used
757        let new_next_instruction = self
758            .next_instruction
759            .as_ptr()
760            .wrapping_byte_offset(byte_delta);
761        // Stored either way: on the way out it is the target `failed_branch()` reports on, and
762        // until then nothing else is allowed to look at it
763        // SAFETY: A wrapped pointer is never null, and nothing here dereferences it
764        self.next_instruction = unsafe { NonNull::new_unchecked(new_next_instruction) };
765
766        let decoded_instruction_byte_offset = new_next_instruction
767            .addr()
768            .wrapping_sub(self.instructions().as_ptr().addr());
769
770        // A target that does not land on a decoded instruction sits between two guest
771        // instructions, which makes it an unaligned instruction rather than something to round to
772        // the start of one. That rule lives in `set_pc()`, so rather than restating it here, where
773        // it could drift, such a target simply fails to qualify, as does one past the end of the
774        // decoded stream, which a backwards branch that ran off its start wraps around into.
775        decoded_instruction_byte_offset < self.instructions_len() * size_of::<ContractInstruction>()
776            && decoded_instruction_byte_offset.is_multiple_of(size_of::<ContractInstruction>())
777    }
778
779    /// Turns the refused target back into an address and hands it to [`Self::set_pc()`], which is
780    /// where the rules about what is and is not an instruction address live.
781    #[cold]
782    #[inline(never)]
783    unsafe fn failed_branch(
784        &mut self,
785        memory: &Memory,
786    ) -> Result<ControlFlow<()>, ExecutionError<u64>> {
787        // Signed, because a backwards branch that ran off the start of the decoded stream is
788        // exactly one of the targets that gets here
789        let decoded_instruction_byte_offset = self
790            .next_instruction
791            .as_ptr()
792            .addr()
793            .wrapping_sub(self.instructions().as_ptr().addr())
794            .cast_signed();
795        // Every `size_of::<u16>()` of guest code owns one decoded instruction, and the position is
796        // always that many bytes from the start of the stream, so this is exact
797        let address = self.base_addr().wrapping_add_signed(
798            (decoded_instruction_byte_offset
799                / (size_of::<ContractInstruction>() / size_of::<u16>()).cast_signed())
800                as i64,
801        );
802
803        self.set_pc(memory, address)
804    }
805
806    #[inline]
807    fn set_pc(
808        &mut self,
809        _memory: &Memory,
810        pc: u64,
811    ) -> Result<ControlFlow<()>, ExecutionError<u64>> {
812        let address = pc;
813
814        if address == self.return_trap_address() {
815            cold_path();
816            return Ok(ControlFlow::Break(()));
817        }
818
819        if !address.is_multiple_of(size_of::<u16>() as u64) {
820            cold_path();
821            return Err(ExecutionError::UnalignedInstruction {
822                address: PackedAddress::new(address),
823            });
824        }
825
826        let Some(offset) = address.checked_sub(self.base_addr()) else {
827            cold_path();
828            return Err(ExecutionError::OutOfBoundsRead {
829                address: PackedAddress::new(address),
830            });
831        };
832        let offset = offset as usize;
833        let instruction_offset = offset / size_of::<u16>();
834
835        if instruction_offset >= self.instructions_len() {
836            cold_path();
837            return Err(VirtualMemoryError::OutOfBoundsRead { address }.into());
838        }
839
840        // SAFETY: `instruction_offset` was just checked to be within bounds of the decoded stream
841        self.next_instruction = unsafe { self.instructions().add(instruction_offset) };
842
843        Ok(ControlFlow::Continue(()))
844    }
845}
846
847impl<Memory> InstructionFetcher<ContractInstruction, Memory> for EagerTestInstructionFetcher<'_>
848where
849    Memory: VirtualMemory,
850{
851    #[inline(always)]
852    fn peek_instruction(
853        &mut self,
854        _memory: &Memory,
855    ) -> FetchInstructionResult<ContractInstruction> {
856        // SAFETY: Constructor guarantees that the last instruction is a jump, which means going
857        // through `Self::set_pc()` method does the necessary bounds check, so the position always
858        // points at a decoded instruction.
859        let instruction = unsafe { self.next_instruction.read() };
860
861        FetchInstructionResult::Instruction(instruction)
862    }
863
864    #[inline(always)]
865    unsafe fn advance(&mut self, instruction_size: u8) {
866        let byte_advance =
867            usize::from(instruction_size) / size_of::<u16>() * size_of::<ContractInstruction>();
868        // Wrapping because nothing here dereferences the pointer: the contract of this method is
869        // what makes the resulting position a decoded instruction, and the bounds check that
870        // matters lives in `set_pc()`
871        // SAFETY: A wrapped pointer is never null, and nothing here dereferences it
872        self.next_instruction = unsafe {
873            NonNull::new_unchecked(
874                self.next_instruction
875                    .as_ptr()
876                    .wrapping_byte_add(byte_advance),
877            )
878        };
879    }
880
881    #[inline(always)]
882    fn fetch_instruction(
883        &mut self,
884        memory: &Memory,
885    ) -> FetchInstructionResult<ContractInstruction> {
886        let result =
887            InstructionFetcher::<ContractInstruction, Memory>::peek_instruction(self, memory);
888
889        if let FetchInstructionResult::Instruction(instruction) = result {
890            // SAFETY: The instruction was just peeked successfully, and this is the only place that
891            // moves past it
892            unsafe {
893                InstructionFetcher::<ContractInstruction, Memory>::advance(
894                    self,
895                    instruction.size(),
896                );
897            }
898        }
899
900        result
901    }
902}
903
904impl EagerTestInstructionFetcher<'_> {
905    /// Pointer to the first decoded instruction
906    #[inline(always)]
907    fn instructions(&self) -> NonNull<ContractInstruction> {
908        // SAFETY: Decoded instructions are stored at this offset of the same allocation as the
909        // state
910        unsafe {
911            self.state
912                .byte_add(EagerTestInstructions::INSTRUCTIONS_OFFSET)
913        }
914        .cast::<ContractInstruction>()
915    }
916
917    /// Number of decoded instructions
918    #[inline(always)]
919    fn instructions_len(&self) -> usize {
920        // SAFETY: State is initialized by `EagerTestInstructions` and borrowed for as long as
921        // `self` is alive
922        unsafe { (*self.state.as_ptr()).instructions_len }
923    }
924
925    /// Guest address that corresponds to the first decoded instruction
926    #[inline(always)]
927    fn base_addr(&self) -> u64 {
928        // SAFETY: State is initialized by `EagerTestInstructions` and borrowed for as long as
929        // `self` is alive
930        unsafe { (*self.state.as_ptr()).base_addr }
931    }
932
933    /// Guest address at which execution stops gracefully
934    #[inline(always)]
935    fn return_trap_address(&self) -> u64 {
936        // SAFETY: State is initialized by `EagerTestInstructions` and borrowed for as long as
937        // `self` is alive
938        unsafe { (*self.state.as_ptr()).return_trap_address }
939    }
940}
ab_riscv_benchmarks/host_utils.rs

ab_riscv_benchmarks/
host_utils.rs
