Skip to main content

ab_riscv_benchmarks/
host_utils.rs

1use ab_blake3::{CHUNK_LEN, OUT_LEN};
2use ab_contract_file::instruction::{ContractInstruction, ContractRegister};
3use ab_core_primitives::ed25519::{Ed25519PublicKey, Ed25519Signature};
4use ab_io_type::bool::Bool;
5use ab_riscv_interpreter::prelude::*;
6use ab_riscv_primitives::prelude::*;
7use core::hint::cold_path;
8use core::mem::offset_of;
9use core::ops::ControlFlow;
10
11/// Contract file bytes
12pub const RISCV_CONTRACT_BYTES: &[u8] = cfg_select! {
13    target_env = "abundance" => &[],
14    _ => {
15        include_bytes!(env!("CONTRACT_PATH"))
16    }
17};
18
19// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
20//  `SimpleInternalArgs` data trait for this or something
21/// Helper data structure for [`Benchmarks::blake3_hash_chunk()`] method
22///
23/// [`Benchmarks::blake3_hash_chunk()`]: crate::Benchmarks::blake3_hash_chunk
24#[derive(Debug, Copy, Clone)]
25#[repr(C)]
26pub struct Blake3HashChunkInternalArgs {
27    chunk_ptr: u64,
28    chunk_size: u32,
29    chunk_capacity: u32,
30    result_ptr: u64,
31    chunk: [u8; CHUNK_LEN],
32    result: [u8; OUT_LEN],
33}
34
35const _: () = {
36    assert!(
37        size_of::<Blake3HashChunkInternalArgs>()
38            == offset_of!(Blake3HashChunkInternalArgs, result) + size_of::<[u8; OUT_LEN]>(),
39        "`Blake3HashChunkInternalArgs` must not have implicit padding"
40    );
41};
42
43impl Blake3HashChunkInternalArgs {
44    /// Create a new instance
45    pub fn new(internal_args_addr: u64, chunk: [u8; CHUNK_LEN]) -> Self {
46        Self {
47            chunk_ptr: internal_args_addr + offset_of!(Self, chunk) as u64,
48            chunk_size: CHUNK_LEN as u32,
49            chunk_capacity: CHUNK_LEN as u32,
50            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
51            chunk,
52            result: [0; _],
53        }
54    }
55
56    /// Extract result
57    pub fn result(&self) -> [u8; OUT_LEN] {
58        self.result
59    }
60}
61
62// TODO: Generate similar helper data structures in the `#[contract]` macro itself, maybe introduce
63//  `SimpleInternalArgs` data trait for this or something
64/// Helper data structure for [`Benchmarks::ed25519_verify()`] method
65///
66/// [`Benchmarks::ed25519_verify()`]: crate::Benchmarks::ed25519_verify
67#[derive(Debug, Copy, Clone)]
68#[repr(C)]
69pub struct Ed25519VerifyInternalArgs {
70    pub public_key_ptr: u64,
71    pub public_key_size: u32,
72    pub public_key_capacity: u32,
73    pub signature_ptr: u64,
74    pub signature_size: u32,
75    pub signature_capacity: u32,
76    pub message_ptr: u64,
77    pub message_size: u32,
78    pub message_capacity: u32,
79    pub result_ptr: u64,
80    pub public_key: Ed25519PublicKey,
81    pub signature: Ed25519Signature,
82    pub message: [u8; OUT_LEN],
83    pub result: Bool,
84    /// Explicit trailing padding.
85    ///
86    /// The host copies the byte representation of this data structure into guest memory, which is
87    /// only sound if every byte of it is initialized, hence implicit padding must not exist here.
88    pub padding: [u8; 7],
89}
90
91const _: () = {
92    assert!(
93        size_of::<Ed25519VerifyInternalArgs>()
94            == offset_of!(Ed25519VerifyInternalArgs, padding) + size_of::<[u8; 7]>(),
95        "`Ed25519VerifyInternalArgs` must not have implicit padding"
96    );
97};
98
99impl Ed25519VerifyInternalArgs {
100    /// Create a new instance
101    pub fn new(
102        internal_args_addr: u64,
103        public_key: Ed25519PublicKey,
104        signature: Ed25519Signature,
105        message: [u8; OUT_LEN],
106    ) -> Self {
107        Self {
108            public_key_ptr: internal_args_addr + offset_of!(Self, public_key) as u64,
109            public_key_size: Ed25519PublicKey::SIZE as u32,
110            public_key_capacity: Ed25519PublicKey::SIZE as u32,
111            signature_ptr: internal_args_addr + offset_of!(Self, signature) as u64,
112            signature_size: Ed25519Signature::SIZE as u32,
113            signature_capacity: Ed25519Signature::SIZE as u32,
114            message_ptr: internal_args_addr + offset_of!(Self, message) as u64,
115            message_size: OUT_LEN as u32,
116            message_capacity: OUT_LEN as u32,
117            result_ptr: internal_args_addr + offset_of!(Self, result) as u64,
118            public_key,
119            signature,
120            message,
121            result: Bool::new(false),
122            padding: [0; _],
123        }
124    }
125
126    /// Extract result
127    pub fn result(&self) -> Bool {
128        self.result
129    }
130}
131
132/// Instruction stored by [`BasicEagerInstructions::decode()`] in slots whose bytes do not decode
133/// into a valid instruction.
134///
135/// Contract code is only expected to contain legal instructions, so this is only reachable by
136/// jumping into the middle of one.
137///
138/// [`BasicEagerInstructions::decode()`]: ab_riscv_interpreter::basic::BasicEagerInstructions::decode
139pub const UNDECODABLE_INSTRUCTION: ContractInstruction = ContractInstruction::Unimp {
140    rs1: ContractRegister::Zero,
141    rs2: ContractRegister::Zero,
142};
143
144/// Lazy instruction fetcher implementation
145#[derive(Debug, Copy, Clone)]
146pub struct LazyInstructionFetcher {
147    return_trap_address: u64,
148    pc: u64,
149}
150
151impl<Memory> ProgramCounter<u64, Memory> for LazyInstructionFetcher
152where
153    Memory: VirtualMemory,
154{
155    #[inline(always)]
156    fn get_pc(&self) -> u64 {
157        self.pc
158    }
159
160    #[inline(always)]
161    unsafe fn try_set_pc_relative(&mut self, instruction_size: u8, offset: i32) -> bool {
162        let old_pc = <Self as ProgramCounter<_, Memory>>::old_pc(self, instruction_size);
163        let pc = old_pc.wrapping_add_signed(i64::from(offset));
164        // Stored either way: on the way out it is what `failed_branch()` reports on, and until then
165        // nothing else is allowed to look at it
166        self.pc = pc;
167
168        pc != self.return_trap_address
169            && pc.is_multiple_of(u64::from(
170                ContractInstruction::<ContractRegister>::ALIGNMENT,
171            ))
172    }
173
174    #[cold]
175    #[inline(never)]
176    unsafe fn failed_branch(
177        &mut self,
178        memory: &Memory,
179    ) -> Result<ControlFlow<()>, ExecutionError<u64>> {
180        // The program counter holds the refused target, and `set_pc()` is what says what is wrong
181        // with it
182        self.set_pc(memory, self.pc)
183    }
184
185    #[inline]
186    fn set_pc(&mut self, memory: &Memory, pc: u64) -> Result<ControlFlow<()>, ExecutionError<u64>> {
187        if pc == self.return_trap_address {
188            cold_path();
189            return Ok(ControlFlow::Break(()));
190        }
191
192        if !pc.is_multiple_of(u64::from(
193            ContractInstruction::<ContractRegister>::ALIGNMENT,
194        )) {
195            cold_path();
196            return Err(ExecutionError::UnalignedInstruction {
197                address: PackedAddress::new(pc),
198            });
199        }
200
201        // Note: This will not allow reading a 16-bit instruction at the very end of memory range,
202        // but that is going to be the case here anyway since code is followed by read-write memory
203        // anyway
204        if let Err(error) = memory.read::<u32>(pc) {
205            cold_path();
206            return Err(error.into());
207        }
208
209        self.pc = pc;
210
211        Ok(ControlFlow::Continue(()))
212    }
213}
214
215impl<Memory> InstructionFetcher<ContractInstruction, Memory> for LazyInstructionFetcher
216where
217    Memory: VirtualMemory,
218{
219    type Peeked = ContractInstruction;
220
221    #[inline(always)]
222    fn peeked_instruction<'a>(
223        &'a self,
224        peeked: &'a ContractInstruction,
225    ) -> &'a ContractInstruction {
226        peeked
227    }
228
229    #[inline]
230    fn peek_instruction(&mut self, memory: &Memory) -> FetchInstructionResult<ContractInstruction> {
231        // SAFETY: Constructor guarantees that the last instruction is a jump, which means going
232        // through `Self::set_pc()` method does the necessary bounds check, so the program counter
233        // always sits on an instruction.
234        let instruction = unsafe { memory.read_unchecked(self.pc) };
235        // SAFETY: All instructions are valid, according to the constructor contract
236        let instruction =
237            unsafe { ContractInstruction::try_decode(instruction).unwrap_unchecked() };
238
239        FetchInstructionResult::Instruction(instruction)
240    }
241
242    #[inline]
243    unsafe fn advance(&mut self, instruction_size: u8) {
244        self.pc = self.pc.wrapping_add(u64::from(instruction_size));
245    }
246
247    #[inline]
248    fn fetch_instruction(
249        &mut self,
250        memory: &Memory,
251    ) -> FetchInstructionResult<ContractInstruction> {
252        let result =
253            InstructionFetcher::<ContractInstruction, Memory>::peek_instruction(self, memory);
254
255        if let FetchInstructionResult::Instruction(instruction) = result {
256            // SAFETY: The instruction was just peeked successfully, and this is the only place that
257            // moves past it
258            unsafe {
259                InstructionFetcher::<ContractInstruction, Memory>::advance(
260                    self,
261                    instruction.size(),
262                );
263            }
264        }
265
266        result
267    }
268}
269
270impl LazyInstructionFetcher {
271    /// Create a new instance.
272    ///
273    /// `return_trap_address` is the address at which the interpreter will stop execution
274    /// (gracefully).
275    ///
276    /// # Safety
277    /// The program counter must be valid and aligned, the instructions processed must be valid and
278    /// end with a jump instruction.
279    #[inline(always)]
280    pub unsafe fn new(return_trap_address: u64, pc: u64) -> Self {
281        Self {
282            return_trap_address,
283            pc,
284        }
285    }
286}