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::{HostBuffer, 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    fn host(self) -> HostBuffer {
194        HostBuffer::Immutable(std::ptr::slice_from_raw_parts(
195            self.mem.ptr.as_ptr(),
196            self.len(),
197        ))
198    }
199}
200
201unsafe impl<const A: usize> HostMemoryObject for &mut AlignedMemory<A> {
202    fn host(self) -> HostBuffer {
203        HostBuffer::Mutable(std::ptr::slice_from_raw_parts_mut(
204            self.mem.ptr.as_ptr(),
205            self.len(),
206        ))
207    }
208}
209
210/// Returns true if `ptr` is aligned to `align`.
211pub fn is_memory_aligned(ptr: usize, align: usize) -> bool {
212    ptr.checked_rem(align)
213        .map(|remainder| remainder == 0)
214        .unwrap_or(false)
215}
216
217/// Provides backing storage for [`AlignedMemory`]. Allocates a block of bytes with the
218/// requested alignment, and can be increased in length up to the requested capacity.
219struct AlignedVec<const ALIGN: usize> {
220    ptr: NonNull<u8>,
221    length: usize,
222    capacity: usize,
223}
224
225impl<const ALIGN: usize> Drop for AlignedVec<ALIGN> {
226    fn drop(&mut self) {
227        if self.capacity == 0 {
228            return;
229        }
230        let ptr = self.ptr.as_ptr();
231        unsafe {
232            // SAFETY: Layout is checked on construction
233            let layout = Layout::from_size_align_unchecked(self.capacity, ALIGN);
234            dealloc(ptr, layout);
235        }
236    }
237}
238
239impl<const A: usize> std::fmt::Debug for AlignedVec<A> {
240    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
241        f.debug_list().entries(self.as_slice()).finish()
242    }
243}
244
245impl<const A: usize> PartialEq for AlignedVec<A> {
246    fn eq(&self, other: &Self) -> bool {
247        self.as_slice() == other.as_slice()
248    }
249}
250
251impl<const A: usize> Eq for AlignedVec<A> {}
252
253impl<const ALIGN: usize> AlignedVec<ALIGN> {
254    /// Allocates a [`Vec<u8>`] with the requested alignment.
255    /// Ensure that the Vec is only dropped with the correct layout
256    ///
257    /// # Panics
258    /// Panics if the requested size is incompatible with the requested alignment or if allocation fails.
259    fn new(max_len: usize, zeroed: bool) -> Self {
260        assert!(ALIGN != 0, "Alignment must not be zero");
261        if max_len == 0 {
262            return Self::empty();
263        }
264        unsafe {
265            let layout = Layout::from_size_align(max_len, ALIGN).expect("invalid layout");
266            // SAFETY: Layout is non-zero, and allocation errors are handled
267            let ptr = if zeroed {
268                alloc_zeroed(layout)
269            } else {
270                alloc(layout)
271            };
272            if ptr.is_null() {
273                handle_alloc_error(layout);
274            }
275            Self {
276                ptr: NonNull::new(ptr).unwrap_or_else(|| handle_alloc_error(layout)),
277                length: 0,
278                capacity: max_len,
279            }
280        }
281    }
282
283    fn as_slice(&self) -> &[u8] {
284        unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.length) }
285    }
286
287    fn as_slice_mut(&mut self) -> &mut [u8] {
288        unsafe { core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.length) }
289    }
290
291    fn empty() -> Self {
292        let layout = Layout::from_size_align(0, ALIGN).expect("invalid layout");
293        Self {
294            ptr: layout.dangling_ptr(),
295            length: 0,
296            capacity: 0,
297        }
298    }
299
300    fn as_mut_ptr(&mut self) -> *mut u8 {
301        self.ptr.as_ptr()
302    }
303
304    /// Returns a pointer to the end of the current initialized length, i.e.
305    /// `mem.as_mut_ptr().mem(self.len())`.
306    /// Users must ensure that any writes to this pointer are in bounds of `capacity`
307    fn write_ptr(&mut self) -> *mut u8 {
308        unsafe { self.as_mut_ptr().add(self.len()) }
309    }
310
311    /// Similar to [`write_ptr`], but checks that there is room for the write.
312    /// Returns (pointer, new_length)
313    fn write_ptr_for(&mut self, bytes: usize) -> Option<(*mut u8, usize)> {
314        let ptr = self.write_ptr();
315        let new_len = self
316            .len()
317            .checked_add(bytes)
318            .filter(|l| *l <= self.capacity())?;
319        Some((ptr, new_len))
320    }
321
322    fn len(&self) -> usize {
323        self.length
324    }
325
326    fn capacity(&self) -> usize {
327        self.capacity
328    }
329
330    fn is_empty(&self) -> bool {
331        self.len() == 0
332    }
333
334    /// Set the length of the `AlignedVec`. The new length must be less than or equal to
335    /// the capacity, and the memory must be initialized up to that length.
336    /// The new length must not be less than the previous length.
337    unsafe fn set_len(&mut self, new_len: usize) {
338        debug_assert!(
339            new_len <= self.capacity,
340            "attempted to grow AlignedVec beyond capacity"
341        );
342        debug_assert!(new_len >= self.length, "attempted to shrink AlignedVec");
343        self.length = new_len;
344    }
345}
346
347/// `AlignedVec` is [`Send`] as `u8` is `Send` and the data behind the pointer is uniquely owned.
348unsafe impl<const N: usize> Send for AlignedVec<N> {}
349
350/// `AlignedVec` is [`Sync`] as `u8` is `Send` and the data behind the pointer is uniquely owned.
351unsafe impl<const N: usize> Sync for AlignedVec<N> {}
352
353#[allow(clippy::arithmetic_side_effects)]
354#[cfg(test)]
355mod tests {
356    use {super::*, std::io::Write};
357
358    fn do_test<const ALIGN: usize>() {
359        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity(10);
360        let ptr = aligned_memory.mem.as_mut_ptr();
361        assert_eq!(
362            ptr.addr() & (ALIGN - 1),
363            0,
364            "memory is not correctly aligned"
365        );
366
367        assert_eq!(aligned_memory.write(&[42u8; 1]).unwrap(), 1);
368        assert_eq!(aligned_memory.write(&[42u8; 9]).unwrap(), 9);
369        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
370        assert_eq!(aligned_memory.write(&[42u8; 0]).unwrap(), 0);
371        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
372        aligned_memory.write(&[42u8; 1]).unwrap_err();
373        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
374        aligned_memory.as_slice_mut().copy_from_slice(&[84u8; 10]);
375        assert_eq!(aligned_memory.as_slice(), &[84u8; 10]);
376
377        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(10);
378        aligned_memory.fill_write(5, 0).unwrap();
379        aligned_memory.fill_write(2, 1).unwrap();
380        assert_eq!(aligned_memory.write(&[2u8; 3]).unwrap(), 3);
381        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
382        aligned_memory.fill_write(1, 3).unwrap_err();
383        aligned_memory.write(&[4u8; 1]).unwrap_err();
384        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
385
386        let aligned_memory = AlignedMemory::<ALIGN>::zero_filled(10);
387        assert_eq!(aligned_memory.len(), 10);
388        assert_eq!(aligned_memory.as_slice(), &[0u8; 10]);
389
390        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(15);
391        unsafe {
392            aligned_memory.write_unchecked::<u8>(42);
393            assert_eq!(aligned_memory.len(), 1);
394            aligned_memory.write_unchecked::<u64>(0xCAFEBADDDEADCAFE);
395            assert_eq!(aligned_memory.len(), 9);
396            aligned_memory.fill_write(3, 0).unwrap();
397            aligned_memory.write_all_unchecked(b"foo");
398            assert_eq!(aligned_memory.len(), 15);
399        }
400        let mem = aligned_memory.as_slice();
401        assert_eq!(mem[0], 42);
402        assert_eq!(
403            unsafe {
404                core::ptr::read_unaligned::<u64>(mem[1..1 + mem::size_of::<u64>()].as_ptr().cast())
405            },
406            0xCAFEBADDDEADCAFE
407        );
408        assert_eq!(&mem[1 + mem::size_of::<u64>()..][..3], &[0, 0, 0]);
409        assert_eq!(&mem[1 + mem::size_of::<u64>() + 3..], b"foo");
410    }
411
412    #[test]
413    fn test_aligned_memory() {
414        do_test::<1>();
415        do_test::<16>();
416        do_test::<32768>();
417    }
418
419    #[cfg(debug_assertions)]
420    #[test]
421    #[should_panic(expected = "<= self.mem.capacity()")]
422    fn test_write_unchecked_debug_assert() {
423        let mut aligned_memory = AlignedMemory::<8>::with_capacity(15);
424        unsafe {
425            aligned_memory.write_unchecked::<u64>(42);
426            aligned_memory.write_unchecked::<u64>(24);
427        }
428    }
429
430    #[cfg(debug_assertions)]
431    #[test]
432    #[should_panic(expected = "<= self.mem.capacity()")]
433    fn test_write_all_unchecked_debug_assert() {
434        let mut aligned_memory = AlignedMemory::<8>::with_capacity(5);
435        unsafe {
436            aligned_memory.write_all_unchecked(b"foo");
437            aligned_memory.write_all_unchecked(b"bar");
438        }
439    }
440
441    const fn assert_send<T: Send>() {}
442    const fn assert_sync<T: Sync>() {}
443    const fn assert_unpin<T: Unpin>() {}
444    const _: () = assert_send::<AlignedMemory<8>>();
445    const _: () = assert_sync::<AlignedMemory<8>>();
446    const _: () = assert_unpin::<AlignedMemory<8>>();
447}
solana_sbpf/aligned_memory.rs

solana_sbpf/
aligned_memory.rs
