solana_sbpf/
aligned_memory.rs

1//! Aligned memory
2
3use std::{mem, ptr};
4
5/// Scalar types, aka "plain old data"
6pub trait Pod {}
7
8impl Pod for bool {}
9impl Pod for u8 {}
10impl Pod for u16 {}
11impl Pod for u32 {}
12impl Pod for u64 {}
13impl Pod for i8 {}
14impl Pod for i16 {}
15impl Pod for i32 {}
16impl Pod for i64 {}
17
18/// Provides u8 slices at a specified alignment
19#[derive(Debug, PartialEq, Eq)]
20pub struct AlignedMemory<const ALIGN: usize> {
21    max_len: usize,
22    align_offset: usize,
23    mem: Vec<u8>,
24    zero_up_to_max_len: bool,
25}
26
27impl<const ALIGN: usize> AlignedMemory<ALIGN> {
28    fn get_mem(max_len: usize) -> (Vec<u8>, usize) {
29        let mut mem: Vec<u8> = Vec::with_capacity(max_len.saturating_add(ALIGN));
30        let align_offset = mem.as_ptr().align_offset(ALIGN);
31        mem.resize(align_offset, 0);
32        (mem, align_offset)
33    }
34    fn get_mem_zeroed(max_len: usize) -> (Vec<u8>, usize) {
35        // use calloc() to get zeroed memory from the OS instead of using
36        // malloc() + memset(), see
37        // https://github.com/rust-lang/rust/issues/54628
38        let mut mem = vec![0; max_len];
39        let align_offset = mem.as_ptr().align_offset(ALIGN);
40        mem.resize(max_len.saturating_add(align_offset), 0);
41        (mem, align_offset)
42    }
43    /// Returns a filled AlignedMemory by copying the given slice
44    pub fn from_slice(data: &[u8]) -> Self {
45        let max_len = data.len();
46        let (mut mem, align_offset) = Self::get_mem(max_len);
47        mem.extend_from_slice(data);
48        Self {
49            max_len,
50            align_offset,
51            mem,
52            zero_up_to_max_len: false,
53        }
54    }
55    /// Returns a new empty AlignedMemory with uninitialized preallocated memory
56    pub fn with_capacity(max_len: usize) -> Self {
57        let (mem, align_offset) = Self::get_mem(max_len);
58        Self {
59            max_len,
60            align_offset,
61            mem,
62            zero_up_to_max_len: false,
63        }
64    }
65    /// Returns a new empty AlignedMemory with zero initialized preallocated memory
66    pub fn with_capacity_zeroed(max_len: usize) -> Self {
67        let (mut mem, align_offset) = Self::get_mem_zeroed(max_len);
68        mem.truncate(align_offset);
69        Self {
70            max_len,
71            align_offset,
72            mem,
73            zero_up_to_max_len: true,
74        }
75    }
76    /// Returns a new filled AlignedMemory with zero initialized preallocated memory
77    pub fn zero_filled(max_len: usize) -> Self {
78        let (mem, align_offset) = Self::get_mem_zeroed(max_len);
79        Self {
80            max_len,
81            align_offset,
82            mem,
83            zero_up_to_max_len: true,
84        }
85    }
86    /// Calculate memory size
87    pub fn mem_size(&self) -> usize {
88        self.mem.capacity().saturating_add(mem::size_of::<Self>())
89    }
90    /// Get the length of the data
91    pub fn len(&self) -> usize {
92        self.mem.len().saturating_sub(self.align_offset)
93    }
94    /// Is the memory empty
95    pub fn is_empty(&self) -> bool {
96        self.mem.len() == self.align_offset
97    }
98    /// Get the current write index
99    pub fn write_index(&self) -> usize {
100        self.mem.len()
101    }
102    /// Get an aligned slice
103    pub fn as_slice(&self) -> &[u8] {
104        let start = self.align_offset;
105        let end = self.mem.len();
106        &self.mem[start..end]
107    }
108    /// Get an aligned mutable slice
109    pub fn as_slice_mut(&mut self) -> &mut [u8] {
110        let start = self.align_offset;
111        let end = self.mem.len();
112        &mut self.mem[start..end]
113    }
114    /// Grows memory with `value` repeated `num` times starting at the `write_index`
115    pub fn fill_write(&mut self, num: usize, value: u8) -> std::io::Result<()> {
116        let new_len = match (
117            self.mem.len().checked_add(num),
118            self.align_offset.checked_add(self.max_len),
119        ) {
120            (Some(new_len), Some(allocation_end)) if new_len <= allocation_end => new_len,
121            _ => {
122                return Err(std::io::Error::new(
123                    std::io::ErrorKind::InvalidInput,
124                    "aligned memory fill_write failed",
125                ))
126            }
127        };
128        if self.zero_up_to_max_len && value == 0 {
129            // Safe because everything up to `max_len` is zeroed and no shrinking is allowed
130            unsafe {
131                self.mem.set_len(new_len);
132            }
133        } else {
134            self.mem.resize(new_len, value);
135        }
136        Ok(())
137    }
138
139    /// Write a generic type T into the memory.
140    ///
141    /// # Safety
142    ///
143    /// Unsafe since it assumes that there is enough capacity.
144    pub unsafe fn write_unchecked<T: Pod>(&mut self, value: T) {
145        let pos = self.mem.len();
146        let new_len = pos.saturating_add(mem::size_of::<T>());
147        debug_assert!(new_len <= self.align_offset.saturating_add(self.max_len));
148        self.mem.set_len(new_len);
149        ptr::write_unaligned(
150            self.mem.get_unchecked_mut(pos..new_len).as_mut_ptr().cast(),
151            value,
152        );
153    }
154
155    /// Write a slice of bytes into the memory.
156    ///
157    /// # Safety
158    ///
159    /// Unsafe since it assumes that there is enough capacity.
160    pub unsafe fn write_all_unchecked(&mut self, value: &[u8]) {
161        let pos = self.mem.len();
162        let new_len = pos.saturating_add(value.len());
163        debug_assert!(new_len <= self.align_offset.saturating_add(self.max_len));
164        self.mem.set_len(new_len);
165        self.mem
166            .get_unchecked_mut(pos..new_len)
167            .copy_from_slice(value);
168    }
169}
170
171// Custom Clone impl is needed to ensure alignment. Derived clone would just
172// clone self.mem and there would be no guarantee that the clone allocation is
173// aligned.
174impl<const ALIGN: usize> Clone for AlignedMemory<ALIGN> {
175    fn clone(&self) -> Self {
176        AlignedMemory::from_slice(self.as_slice())
177    }
178}
179
180impl<const ALIGN: usize> std::io::Write for AlignedMemory<ALIGN> {
181    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
182        match (
183            self.mem.len().checked_add(buf.len()),
184            self.align_offset.checked_add(self.max_len),
185        ) {
186            (Some(new_len), Some(allocation_end)) if new_len <= allocation_end => {}
187            _ => {
188                return Err(std::io::Error::new(
189                    std::io::ErrorKind::InvalidInput,
190                    "aligned memory write failed",
191                ))
192            }
193        }
194        self.mem.extend_from_slice(buf);
195        Ok(buf.len())
196    }
197    fn flush(&mut self) -> std::io::Result<()> {
198        Ok(())
199    }
200}
201
202impl<const ALIGN: usize, T: AsRef<[u8]>> From<T> for AlignedMemory<ALIGN> {
203    fn from(bytes: T) -> Self {
204        AlignedMemory::from_slice(bytes.as_ref())
205    }
206}
207
208/// Returns true if `ptr` is aligned to `align`.
209pub fn is_memory_aligned(ptr: usize, align: usize) -> bool {
210    ptr.checked_rem(align)
211        .map(|remainder| remainder == 0)
212        .unwrap_or(false)
213}
214
215#[allow(clippy::arithmetic_side_effects)]
216#[cfg(test)]
217mod tests {
218    use {super::*, std::io::Write};
219
220    fn do_test<const ALIGN: usize>() {
221        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity(10);
222
223        assert_eq!(aligned_memory.write(&[42u8; 1]).unwrap(), 1);
224        assert_eq!(aligned_memory.write(&[42u8; 9]).unwrap(), 9);
225        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
226        assert_eq!(aligned_memory.write(&[42u8; 0]).unwrap(), 0);
227        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
228        aligned_memory.write(&[42u8; 1]).unwrap_err();
229        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
230        aligned_memory.as_slice_mut().copy_from_slice(&[84u8; 10]);
231        assert_eq!(aligned_memory.as_slice(), &[84u8; 10]);
232
233        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(10);
234        aligned_memory.fill_write(5, 0).unwrap();
235        aligned_memory.fill_write(2, 1).unwrap();
236        assert_eq!(aligned_memory.write(&[2u8; 3]).unwrap(), 3);
237        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
238        aligned_memory.fill_write(1, 3).unwrap_err();
239        aligned_memory.write(&[4u8; 1]).unwrap_err();
240        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
241
242        let aligned_memory = AlignedMemory::<ALIGN>::zero_filled(10);
243        assert_eq!(aligned_memory.len(), 10);
244        assert_eq!(aligned_memory.as_slice(), &[0u8; 10]);
245
246        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(15);
247        unsafe {
248            aligned_memory.write_unchecked::<u8>(42);
249            assert_eq!(aligned_memory.len(), 1);
250            aligned_memory.write_unchecked::<u64>(0xCAFEBADDDEADCAFE);
251            assert_eq!(aligned_memory.len(), 9);
252            aligned_memory.fill_write(3, 0).unwrap();
253            aligned_memory.write_all_unchecked(b"foo");
254            assert_eq!(aligned_memory.len(), 15);
255        }
256        let mem = aligned_memory.as_slice();
257        assert_eq!(mem[0], 42);
258        assert_eq!(
259            unsafe {
260                ptr::read_unaligned::<u64>(mem[1..1 + mem::size_of::<u64>()].as_ptr().cast())
261            },
262            0xCAFEBADDDEADCAFE
263        );
264        assert_eq!(&mem[1 + mem::size_of::<u64>()..][..3], &[0, 0, 0]);
265        assert_eq!(&mem[1 + mem::size_of::<u64>() + 3..], b"foo");
266    }
267
268    #[test]
269    fn test_aligned_memory() {
270        do_test::<1>();
271        do_test::<32768>();
272    }
273
274    #[cfg(debug_assertions)]
275    #[test]
276    #[should_panic(expected = "<= self.align_offset.saturating_add(self.max_len)")]
277    fn test_write_unchecked_debug_assert() {
278        let mut aligned_memory = AlignedMemory::<8>::with_capacity(15);
279        unsafe {
280            aligned_memory.write_unchecked::<u64>(42);
281            aligned_memory.write_unchecked::<u64>(24);
282        }
283    }
284
285    #[cfg(debug_assertions)]
286    #[test]
287    #[should_panic(expected = "<= self.align_offset.saturating_add(self.max_len)")]
288    fn test_write_all_unchecked_debug_assert() {
289        let mut aligned_memory = AlignedMemory::<8>::with_capacity(5);
290        unsafe {
291            aligned_memory.write_all_unchecked(b"foo");
292            aligned_memory.write_all_unchecked(b"bar");
293        }
294    }
295}
solana_sbpf/aligned_memory.rs

solana_sbpf/
aligned_memory.rs
