solana_sbpf/
aligned_memory.rs

1//! Aligned memory
2
3use std::{
4    alloc::{alloc, alloc_zeroed, dealloc, handle_alloc_error, Layout},
5    mem,
6    ptr::NonNull,
7};
8
9use crate::memory_region::HostMemoryObject;
10
11/// Scalar types, aka "plain old data"
12pub trait Pod: Copy {}
13
14impl Pod for u8 {}
15impl Pod for u16 {}
16impl Pod for u32 {}
17impl Pod for u64 {}
18impl Pod for i8 {}
19impl Pod for i16 {}
20impl Pod for i32 {}
21impl Pod for i64 {}
22
23/// Provides u8 slices at a specified alignment
24#[derive(Debug, PartialEq, Eq)]
25pub struct AlignedMemory<const ALIGN: usize> {
26    mem: AlignedVec<ALIGN>,
27    zero_up_to_max_len: bool,
28}
29
30impl<const ALIGN: usize> AlignedMemory<ALIGN> {
31    /// Returns a filled AlignedMemory by copying the given slice
32    pub fn from_slice(data: &[u8]) -> Self {
33        let max_len = data.len();
34        let mut mem = AlignedVec::new(max_len, false);
35        unsafe {
36            // SAFETY: `mem` was allocated with `max_len` bytes
37            core::ptr::copy_nonoverlapping(data.as_ptr(), mem.as_mut_ptr(), max_len);
38            mem.set_len(max_len);
39        }
40        Self {
41            mem,
42            zero_up_to_max_len: false,
43        }
44    }
45
46    /// Returns a new empty AlignedMemory with uninitialized preallocated memory
47    pub fn with_capacity(max_len: usize) -> Self {
48        let mem = AlignedVec::new(max_len, false);
49        Self {
50            mem,
51            zero_up_to_max_len: false,
52        }
53    }
54
55    /// Returns a new empty AlignedMemory with zero initialized preallocated memory
56    pub fn with_capacity_zeroed(max_len: usize) -> Self {
57        let mem = AlignedVec::new(max_len, true);
58        Self {
59            mem,
60            zero_up_to_max_len: true,
61        }
62    }
63
64    /// Returns a new filled AlignedMemory with zero initialized preallocated memory
65    pub fn zero_filled(max_len: usize) -> Self {
66        let mut mem = AlignedVec::new(max_len, true);
67        // SAFETY: Bytes were zeroed
68        unsafe {
69            mem.set_len(max_len);
70        }
71        Self {
72            mem,
73            zero_up_to_max_len: true,
74        }
75    }
76
77    /// Calculate memory size (allocated memory block and the size of [`AlignedMemory`] itself).
78    pub fn mem_size(&self) -> usize {
79        self.mem.capacity().saturating_add(mem::size_of::<Self>())
80    }
81
82    /// Get the length of the data
83    pub fn len(&self) -> usize {
84        self.mem.len()
85    }
86
87    /// Is the memory empty
88    pub fn is_empty(&self) -> bool {
89        self.mem.is_empty()
90    }
91
92    /// Get the current write index
93    pub fn write_index(&self) -> usize {
94        self.mem.len()
95    }
96
97    /// Get an aligned slice
98    pub fn as_slice(&self) -> &[u8] {
99        self.mem.as_slice()
100    }
101
102    /// Get an aligned mutable slice
103    pub fn as_slice_mut(&mut self) -> &mut [u8] {
104        self.mem.as_slice_mut()
105    }
106
107    /// Grows memory with `value` repeated `num` times starting at the `write_index`
108    pub fn fill_write(&mut self, num: usize, value: u8) -> std::io::Result<()> {
109        let (ptr, new_len) = self.mem.write_ptr_for(num).ok_or_else(|| {
110            std::io::Error::new(
111                std::io::ErrorKind::InvalidInput,
112                "aligned memory fill_write failed",
113            )
114        })?;
115
116        if self.zero_up_to_max_len && value == 0 {
117            // No action needed because up to `max_len` is zeroed and no shrinking is allowed
118        } else {
119            unsafe {
120                core::ptr::write_bytes(ptr, value, num);
121            }
122        }
123        unsafe {
124            self.mem.set_len(new_len);
125        }
126        Ok(())
127    }
128
129    /// Write a generic type T into the memory.
130    ///
131    /// # Safety
132    ///
133    /// Unsafe since it assumes that there is enough capacity.
134    pub unsafe fn write_unchecked<T: Pod>(&mut self, value: T) {
135        let pos = self.mem.len();
136        let new_len = pos.saturating_add(mem::size_of::<T>());
137        debug_assert!(new_len <= self.mem.capacity());
138        unsafe {
139            self.mem.write_ptr().cast::<T>().write_unaligned(value);
140            self.mem.set_len(new_len);
141        }
142    }
143
144    /// Write a slice of bytes into the memory.
145    ///
146    /// # Safety
147    ///
148    /// Unsafe since it assumes that there is enough capacity.
149    pub unsafe fn write_all_unchecked(&mut self, value: &[u8]) {
150        let pos = self.mem.len();
151        let new_len = pos.saturating_add(value.len());
152        debug_assert!(new_len <= self.mem.capacity());
153        core::ptr::copy_nonoverlapping(value.as_ptr(), self.mem.write_ptr(), value.len());
154        self.mem.set_len(new_len);
155    }
156}
157
158// Custom Clone impl is needed to ensure alignment. Derived clone would just
159// clone self.mem and there would be no guarantee that the clone allocation is
160// aligned.
161impl<const ALIGN: usize> Clone for AlignedMemory<ALIGN> {
162    fn clone(&self) -> Self {
163        AlignedMemory::from_slice(self.as_slice())
164    }
165}
166
167impl<const ALIGN: usize> std::io::Write for AlignedMemory<ALIGN> {
168    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
169        let (ptr, new_len) = self.mem.write_ptr_for(buf.len()).ok_or_else(|| {
170            std::io::Error::new(
171                std::io::ErrorKind::InvalidInput,
172                "aligned memory fill_write failed",
173            )
174        })?;
175        unsafe {
176            core::ptr::copy_nonoverlapping(buf.as_ptr(), ptr, buf.len());
177            self.mem.set_len(new_len);
178        }
179        Ok(buf.len())
180    }
181    fn flush(&mut self) -> std::io::Result<()> {
182        Ok(())
183    }
184}
185
186impl<const ALIGN: usize, T: AsRef<[u8]>> From<T> for AlignedMemory<ALIGN> {
187    fn from(bytes: T) -> Self {
188        AlignedMemory::from_slice(bytes.as_ref())
189    }
190}
191
192unsafe impl<const A: usize> HostMemoryObject for &AlignedMemory<A> {
193    const WRITABLE: bool = false;
194    fn host_address(self) -> usize {
195        self.mem.ptr.as_ptr().expose_provenance()
196    }
197    fn byte_length(&self) -> usize {
198        self.len()
199    }
200}
201
202unsafe impl<const A: usize> HostMemoryObject for &mut AlignedMemory<A> {
203    const WRITABLE: bool = true;
204    fn host_address(self) -> usize {
205        self.mem.ptr.as_ptr().expose_provenance()
206    }
207    fn byte_length(&self) -> usize {
208        self.len()
209    }
210}
211
212/// Returns true if `ptr` is aligned to `align`.
213pub fn is_memory_aligned(ptr: usize, align: usize) -> bool {
214    ptr.checked_rem(align)
215        .map(|remainder| remainder == 0)
216        .unwrap_or(false)
217}
218
219/// Provides backing storage for [`AlignedMemory`]. Allocates a block of bytes with the
220/// requested alignment, and can be increased in length up to the requested capacity.
221struct AlignedVec<const ALIGN: usize> {
222    ptr: NonNull<u8>,
223    length: usize,
224    capacity: usize,
225}
226
227impl<const ALIGN: usize> Drop for AlignedVec<ALIGN> {
228    fn drop(&mut self) {
229        if self.capacity == 0 {
230            return;
231        }
232        let ptr = self.ptr.as_ptr();
233        unsafe {
234            // SAFETY: Layout is checked on construction
235            let layout = Layout::from_size_align_unchecked(self.capacity, ALIGN);
236            dealloc(ptr, layout);
237        }
238    }
239}
240
241impl<const A: usize> std::fmt::Debug for AlignedVec<A> {
242    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
243        f.debug_list().entries(self.as_slice()).finish()
244    }
245}
246
247impl<const A: usize> PartialEq for AlignedVec<A> {
248    fn eq(&self, other: &Self) -> bool {
249        self.as_slice() == other.as_slice()
250    }
251}
252
253impl<const A: usize> Eq for AlignedVec<A> {}
254
255impl<const ALIGN: usize> AlignedVec<ALIGN> {
256    /// Allocates a [`Vec<u8>`] with the requested alignment.
257    /// Ensure that the Vec is only dropped with the correct layout
258    ///
259    /// # Panics
260    /// Panics if the requested size is incompatible with the requested alignment or if allocation fails.
261    fn new(max_len: usize, zeroed: bool) -> Self {
262        assert!(ALIGN != 0, "Alignment must not be zero");
263        if max_len == 0 {
264            return Self::empty();
265        }
266        unsafe {
267            let layout = Layout::from_size_align(max_len, ALIGN).expect("invalid layout");
268            // SAFETY: Layout is non-zero, and allocation errors are handled
269            let ptr = if zeroed {
270                alloc_zeroed(layout)
271            } else {
272                alloc(layout)
273            };
274            if ptr.is_null() {
275                handle_alloc_error(layout);
276            }
277            Self {
278                ptr: NonNull::new(ptr).unwrap_or_else(|| handle_alloc_error(layout)),
279                length: 0,
280                capacity: max_len,
281            }
282        }
283    }
284
285    fn as_slice(&self) -> &[u8] {
286        unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.length) }
287    }
288
289    fn as_slice_mut(&mut self) -> &mut [u8] {
290        unsafe { core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.length) }
291    }
292
293    fn empty() -> Self {
294        Self {
295            // Create a dangling pointer
296            // FIXME: Use `Layout::dangling_ptr` once Rust 1.95.0 is released
297            ptr: NonNull::new(ALIGN as *mut u8).expect("alignment may not be zero"),
298            length: 0,
299            capacity: 0,
300        }
301    }
302
303    fn as_mut_ptr(&mut self) -> *mut u8 {
304        self.ptr.as_ptr()
305    }
306
307    /// Returns a pointer to the end of the current initialized length, i.e.
308    /// `mem.as_mut_ptr().mem(self.len())`.
309    /// Users must ensure that any writes to this pointer are in bounds of `capacity`
310    fn write_ptr(&mut self) -> *mut u8 {
311        unsafe { self.as_mut_ptr().add(self.len()) }
312    }
313
314    /// Similar to [`write_ptr`], but checks that there is room for the write.
315    /// Returns (pointer, new_length)
316    fn write_ptr_for(&mut self, bytes: usize) -> Option<(*mut u8, usize)> {
317        let ptr = self.write_ptr();
318        let new_len = self
319            .len()
320            .checked_add(bytes)
321            .filter(|l| *l <= self.capacity())?;
322        Some((ptr, new_len))
323    }
324
325    fn len(&self) -> usize {
326        self.length
327    }
328
329    fn capacity(&self) -> usize {
330        self.capacity
331    }
332
333    fn is_empty(&self) -> bool {
334        self.len() == 0
335    }
336
337    /// Set the length of the `AlignedVec`. The new length must be less than or equal to
338    /// the capacity, and the memory must be initialized up to that length.
339    /// The new length must not be less than the previous length.
340    unsafe fn set_len(&mut self, new_len: usize) {
341        debug_assert!(
342            new_len <= self.capacity,
343            "attempted to grow AlignedVec beyond capacity"
344        );
345        debug_assert!(new_len >= self.length, "attempted to shrink AlignedVec");
346        self.length = new_len;
347    }
348}
349
350/// `AlignedVec` is [`Send`] as `u8` is `Send` and the data behind the pointer is uniquely owned.
351unsafe impl<const N: usize> Send for AlignedVec<N> {}
352
353/// `AlignedVec` is [`Sync`] as `u8` is `Send` and the data behind the pointer is uniquely owned.
354unsafe impl<const N: usize> Sync for AlignedVec<N> {}
355
356#[allow(clippy::arithmetic_side_effects)]
357#[cfg(test)]
358mod tests {
359    use {super::*, std::io::Write};
360
361    fn do_test<const ALIGN: usize>() {
362        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity(10);
363        let ptr = aligned_memory.mem.as_mut_ptr();
364        assert_eq!(
365            ptr.addr() & (ALIGN - 1),
366            0,
367            "memory is not correctly aligned"
368        );
369
370        assert_eq!(aligned_memory.write(&[42u8; 1]).unwrap(), 1);
371        assert_eq!(aligned_memory.write(&[42u8; 9]).unwrap(), 9);
372        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
373        assert_eq!(aligned_memory.write(&[42u8; 0]).unwrap(), 0);
374        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
375        aligned_memory.write(&[42u8; 1]).unwrap_err();
376        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
377        aligned_memory.as_slice_mut().copy_from_slice(&[84u8; 10]);
378        assert_eq!(aligned_memory.as_slice(), &[84u8; 10]);
379
380        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(10);
381        aligned_memory.fill_write(5, 0).unwrap();
382        aligned_memory.fill_write(2, 1).unwrap();
383        assert_eq!(aligned_memory.write(&[2u8; 3]).unwrap(), 3);
384        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
385        aligned_memory.fill_write(1, 3).unwrap_err();
386        aligned_memory.write(&[4u8; 1]).unwrap_err();
387        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
388
389        let aligned_memory = AlignedMemory::<ALIGN>::zero_filled(10);
390        assert_eq!(aligned_memory.len(), 10);
391        assert_eq!(aligned_memory.as_slice(), &[0u8; 10]);
392
393        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(15);
394        unsafe {
395            aligned_memory.write_unchecked::<u8>(42);
396            assert_eq!(aligned_memory.len(), 1);
397            aligned_memory.write_unchecked::<u64>(0xCAFEBADDDEADCAFE);
398            assert_eq!(aligned_memory.len(), 9);
399            aligned_memory.fill_write(3, 0).unwrap();
400            aligned_memory.write_all_unchecked(b"foo");
401            assert_eq!(aligned_memory.len(), 15);
402        }
403        let mem = aligned_memory.as_slice();
404        assert_eq!(mem[0], 42);
405        assert_eq!(
406            unsafe {
407                core::ptr::read_unaligned::<u64>(mem[1..1 + mem::size_of::<u64>()].as_ptr().cast())
408            },
409            0xCAFEBADDDEADCAFE
410        );
411        assert_eq!(&mem[1 + mem::size_of::<u64>()..][..3], &[0, 0, 0]);
412        assert_eq!(&mem[1 + mem::size_of::<u64>() + 3..], b"foo");
413    }
414
415    #[test]
416    fn test_aligned_memory() {
417        do_test::<1>();
418        do_test::<16>();
419        do_test::<32768>();
420    }
421
422    #[cfg(debug_assertions)]
423    #[test]
424    #[should_panic(expected = "<= self.mem.capacity()")]
425    fn test_write_unchecked_debug_assert() {
426        let mut aligned_memory = AlignedMemory::<8>::with_capacity(15);
427        unsafe {
428            aligned_memory.write_unchecked::<u64>(42);
429            aligned_memory.write_unchecked::<u64>(24);
430        }
431    }
432
433    #[cfg(debug_assertions)]
434    #[test]
435    #[should_panic(expected = "<= self.mem.capacity()")]
436    fn test_write_all_unchecked_debug_assert() {
437        let mut aligned_memory = AlignedMemory::<8>::with_capacity(5);
438        unsafe {
439            aligned_memory.write_all_unchecked(b"foo");
440            aligned_memory.write_all_unchecked(b"bar");
441        }
442    }
443
444    const fn assert_send<T: Send>() {}
445    const fn assert_sync<T: Sync>() {}
446    const fn assert_unpin<T: Unpin>() {}
447    const _: () = assert_send::<AlignedMemory<8>>();
448    const _: () = assert_sync::<AlignedMemory<8>>();
449    const _: () = assert_unpin::<AlignedMemory<8>>();
450}
solana_sbpf/aligned_memory.rs

solana_sbpf/
aligned_memory.rs
