polkavm_common/
utils.rs

1use core::mem::MaybeUninit;
2
3#[cfg(feature = "alloc")]
4use alloc::{borrow::Cow, vec::Vec};
5
6use crate::program::Reg;
7
8/// A replacement for `alloc::borrow::Cow<[u8]>` which also works in pure no_std.
9#[derive(Clone, PartialEq, Eq, Debug, Default)]
10#[repr(transparent)]
11pub struct CowBytes<'a>(CowBytesImpl<'a>);
12
13#[cfg(feature = "alloc")]
14type CowBytesImpl<'a> = Cow<'a, [u8]>;
15
16#[cfg(not(feature = "alloc"))]
17type CowBytesImpl<'a> = &'a [u8];
18
19impl<'a> CowBytes<'a> {
20    #[cfg(feature = "alloc")]
21    pub fn into_owned(self) -> CowBytes<'static> {
22        match self.0 {
23            Cow::Borrowed(data) => CowBytes(Cow::Owned(data.into())),
24            Cow::Owned(data) => CowBytes(Cow::Owned(data)),
25        }
26    }
27}
28
29impl<'a> core::ops::Deref for CowBytes<'a> {
30    type Target = [u8];
31
32    fn deref(&self) -> &Self::Target {
33        &self.0
34    }
35}
36
37impl<'a> From<&'a [u8]> for CowBytes<'a> {
38    fn from(slice: &'a [u8]) -> Self {
39        CowBytes(slice.into())
40    }
41}
42
43impl<'a> From<&'a str> for CowBytes<'a> {
44    fn from(slice: &'a str) -> Self {
45        CowBytes(slice.as_bytes().into())
46    }
47}
48
49#[cfg(feature = "alloc")]
50impl<'a> From<Vec<u8>> for CowBytes<'a> {
51    fn from(vec: Vec<u8>) -> Self {
52        CowBytes(vec.into())
53    }
54}
55
56#[cfg(feature = "alloc")]
57impl<'a> From<Cow<'a, [u8]>> for CowBytes<'a> {
58    fn from(cow: Cow<'a, [u8]>) -> Self {
59        CowBytes(cow)
60    }
61}
62
63macro_rules! define_align_to_next_page {
64    ($name:ident, $type:ty) => {
65        /// Aligns the `value` to the next `page_size`, or returns the `value` as-is if it's already aligned.
66        #[inline]
67        pub const fn $name(page_size: $type, value: $type) -> Option<$type> {
68            assert!(
69                page_size != 0 && (page_size & (page_size - 1)) == 0,
70                "page size is not a power of two"
71            );
72            if value & page_size - 1 == 0 {
73                Some(value)
74            } else {
75                if value <= <$type>::MAX - page_size {
76                    Some((value + page_size) & !(page_size - 1))
77                } else {
78                    None
79                }
80            }
81        }
82    };
83}
84
85define_align_to_next_page!(align_to_next_page_u32, u32);
86define_align_to_next_page!(align_to_next_page_u64, u64);
87define_align_to_next_page!(align_to_next_page_usize, usize);
88
89#[test]
90fn test_align_to_next_page() {
91    assert_eq!(align_to_next_page_u64(4096, 0), Some(0));
92    assert_eq!(align_to_next_page_u64(4096, 1), Some(4096));
93    assert_eq!(align_to_next_page_u64(4096, 4095), Some(4096));
94    assert_eq!(align_to_next_page_u64(4096, 4096), Some(4096));
95    assert_eq!(align_to_next_page_u64(4096, 4097), Some(8192));
96    let max = (0x10000000000000000_u128 - 4096) as u64;
97    assert_eq!(align_to_next_page_u64(4096, max), Some(max));
98    assert_eq!(align_to_next_page_u64(4096, max + 1), None);
99}
100
101pub trait AsUninitSliceMut {
102    fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>];
103}
104
105impl AsUninitSliceMut for [MaybeUninit<u8>] {
106    fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>] {
107        self
108    }
109}
110
111impl AsUninitSliceMut for [u8] {
112    fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>] {
113        #[allow(unsafe_code)]
114        // SAFETY: `MaybeUnunit<T>` is guaranteed to have the same representation as `T`,
115        //         so casting `[T]` into `[MaybeUninit<T>]` is safe.
116        unsafe {
117            core::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), self.len())
118        }
119    }
120}
121
122impl<const N: usize> AsUninitSliceMut for MaybeUninit<[u8; N]> {
123    fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>] {
124        #[allow(unsafe_code)]
125        // SAFETY: `MaybeUnunit<T>` is guaranteed to have the same representation as `T`,
126        //         so casting `[T; N]` into `[MaybeUninit<T>]` is safe.
127        unsafe {
128            core::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), N)
129        }
130    }
131}
132
133impl<const N: usize> AsUninitSliceMut for [u8; N] {
134    fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>] {
135        let slice: &mut [u8] = &mut self[..];
136        slice.as_uninit_slice_mut()
137    }
138}
139
140/// A gas value used for gas metering.
141#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Default)]
142#[repr(transparent)]
143pub struct Gas(u64);
144
145impl core::fmt::Display for Gas {
146    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
147        self.0.fmt(fmt)
148    }
149}
150
151impl Gas {
152    /// The maximum possible available gas.
153    pub const MAX: Self = Gas(i64::MAX as u64);
154
155    /// The minimum possible available gas.
156    pub const MIN: Self = Gas(0);
157
158    /// Constructs a new gas value. Alias for [`Gas::from_u64`].
159    pub const fn new(gas: u64) -> Option<Self> {
160        Self::from_u64(gas)
161    }
162
163    /// Constructs a new gas value from an `u64`, checking whether it's in range between [`Gas::MIN`] and [`Gas::MAX`].
164    pub const fn from_u64(gas: u64) -> Option<Self> {
165        let gas = Self(gas);
166        if gas.0 > Self::MAX.0 {
167            None
168        } else {
169            Some(gas)
170        }
171    }
172
173    /// Constructs a new gas value from an `i64`, checking whether it's in range between [`Gas::MIN`] and [`Gas::MAX`].
174    pub const fn from_i64(gas: i64) -> Option<Self> {
175        Self::from_u64(gas as u64)
176    }
177
178    /// Gets the raw gas value.
179    pub const fn get(self) -> u64 {
180        self.0
181    }
182
183    /// Checks whether there is no gas remaining.
184    pub const fn is_empty(self) -> bool {
185        self.0 == 0
186    }
187}
188
189impl From<u32> for Gas {
190    fn from(gas: u32) -> Self {
191        Gas(u64::from(gas))
192    }
193}
194
195impl TryFrom<u64> for Gas {
196    type Error = &'static str;
197    fn try_from(gas: u64) -> Result<Self, Self::Error> {
198        Self::from_u64(gas).ok_or("out of range gas")
199    }
200}
201
202impl TryFrom<i64> for Gas {
203    type Error = &'static str;
204    fn try_from(gas: i64) -> Result<Self, Self::Error> {
205        Self::from_i64(gas).ok_or("out of range gas")
206    }
207}
208
209pub trait Access<'a> {
210    type Error: core::fmt::Display;
211
212    fn get_reg(&self, reg: Reg) -> u32;
213    fn set_reg(&mut self, reg: Reg, value: u32);
214    fn read_memory_into_slice<'slice, T>(&self, address: u32, buffer: &'slice mut T) -> Result<&'slice mut [u8], Self::Error>
215    where
216        T: ?Sized + AsUninitSliceMut;
217    fn write_memory(&mut self, address: u32, data: &[u8]) -> Result<(), Self::Error>;
218
219    fn sbrk(&mut self, size: u32) -> Option<u32>;
220
221    /// Returns the current size of the program's heap.
222    fn heap_size(&self) -> u32;
223
224    fn program_counter(&self) -> Option<u32>;
225    fn native_program_counter(&self) -> Option<u64>;
226
227    /// Gets the amount of gas remaining, or `None` if gas metering is not enabled for this instance.
228    ///
229    /// Note that this being zero doesn't necessarily mean that the execution ran out of gas,
230    /// if the program ended up consuming *exactly* the amount of gas that it was provided with!
231    fn gas_remaining(&self) -> Option<Gas>;
232
233    fn consume_gas(&mut self, gas: u64);
234
235    #[cfg(feature = "alloc")]
236    fn read_memory_into_vec(&self, address: u32, length: u32) -> Result<Vec<u8>, Self::Error> {
237        let mut buffer = Vec::new();
238        buffer.reserve_exact(length as usize);
239
240        let pointer = buffer.as_ptr();
241        let slice = self.read_memory_into_slice(address, buffer.spare_capacity_mut())?;
242
243        // Since `read_memory_into_slice` returns a `&mut [u8]` we can be sure it initialized the buffer
244        // we've passed to it, as long as it's actually the same buffer we gave it.
245        assert_eq!(slice.as_ptr(), pointer);
246        assert_eq!(slice.len(), length as usize);
247
248        #[allow(unsafe_code)]
249        // SAFETY: `read_memory_into_slice` initialized this buffer, and we've verified this with `assert`s.
250        unsafe {
251            buffer.set_len(length as usize);
252        }
253
254        Ok(buffer)
255    }
256}
257
258// Copied from `MaybeUninit::slice_assume_init_mut`.
259// TODO: Remove this once this API is stabilized.
260#[allow(clippy::missing_safety_doc)]
261#[allow(unsafe_code)]
262pub unsafe fn slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
263    // SAFETY: The caller is responsible for making sure the `slice` was properly initialized.
264    unsafe { &mut *(slice as *mut [MaybeUninit<T>] as *mut [T]) }
265}
266
267#[allow(unsafe_code)]
268pub fn byte_slice_init<'dst>(dst: &'dst mut [MaybeUninit<u8>], src: &[u8]) -> &'dst mut [u8] {
269    assert_eq!(dst.len(), src.len());
270
271    let length = dst.len();
272    let src_ptr: *const u8 = src.as_ptr();
273    let dst_ptr: *mut u8 = dst.as_mut_ptr().cast::<u8>();
274
275    // SAFETY: Both pointers are valid and are guaranteed to point to a region of memory
276    // at least `length` bytes big.
277    unsafe {
278        core::ptr::copy_nonoverlapping(src_ptr, dst_ptr, length);
279    }
280
281    // SAFETY: We've just initialized this slice.
282    unsafe { slice_assume_init_mut(dst) }
283}
polkavm_common/utils.rs

polkavm_common/
utils.rs
