Skip to content

Commit

Permalink
Backport, changes taken from #28079 of the solana monorepo. (#389)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Lichtso
Lichtso authored Oct 5, 2022
1 parent 34660e9 commit 8ca225a
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 63 deletions.
125 changes: 63 additions & 62 deletions src/aligned_memory.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,20 @@
#![allow(clippy::integer_arithmetic)]
//! Aligned memory
use std::{mem, ptr};

/// Scalar types, aka "plain old data"
pub trait Pod {}

impl Pod for bool {}
impl Pod for u8 {}
impl Pod for u16 {}
impl Pod for u32 {}
impl Pod for u64 {}
impl Pod for i8 {}
impl Pod for i16 {}
impl Pod for i32 {}
impl Pod for i64 {}

/// Provides u8 slices at a specified alignment
#[derive(Debug, PartialEq, Eq)]
pub struct AlignedMemory<const ALIGN: usize> {
Expand All @@ -14,7 +26,7 @@ pub struct AlignedMemory<const ALIGN: usize> {

impl<const ALIGN: usize> AlignedMemory<ALIGN> {
fn get_mem(max_len: usize) -> (Vec<u8>, usize) {
let mut mem: Vec<u8> = Vec::with_capacity(max_len + ALIGN);
let mut mem: Vec<u8> = Vec::with_capacity(max_len.saturating_add(ALIGN));
mem.push(0);
let align_offset = mem.as_ptr().align_offset(ALIGN);
mem.resize(align_offset, 0);
Expand All @@ -26,7 +38,7 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
// https://github.com/rust-lang/rust/issues/54628
let mut mem = vec![0; max_len];
let align_offset = mem.as_ptr().align_offset(ALIGN);
mem.resize(align_offset + max_len, 0);
mem.resize(max_len.saturating_add(align_offset), 0);
(mem, align_offset)
}
/// Returns a filled AlignedMemory by copying the given slice
Expand Down Expand Up @@ -74,15 +86,15 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
}
/// Calculate memory size
pub fn mem_size(&self) -> usize {
mem::size_of::<Self>() + self.mem.capacity()
self.mem.capacity().saturating_add(mem::size_of::<Self>())
}
/// Get the length of the data
pub fn len(&self) -> usize {
self.mem.len() - self.align_offset
self.mem.len().saturating_sub(self.align_offset)
}
/// Is the memory empty
pub fn is_empty(&self) -> bool {
self.mem.len() - self.align_offset == 0
self.mem.len() == self.align_offset
}
/// Get the current write index
pub fn write_index(&self) -> usize {
Expand All @@ -102,44 +114,38 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
}
/// Grows memory with `value` repeated `num` times starting at the `write_index`
pub fn fill_write(&mut self, num: usize, value: u8) -> std::io::Result<()> {
if self.mem.len() + num > self.align_offset + self.max_len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"aligned memory resize failed",
));
}
let new_len = match (
self.mem.len().checked_add(num),
self.align_offset.checked_add(self.max_len),
) {
(Some(new_len), Some(allocation_end)) if new_len <= allocation_end => new_len,
_ => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"aligned memory resize failed",
))
}
};
if self.zero_up_to_max_len && value == 0 {
// Safe because everything up to `max_len` is zeroed and no shrinking is allowed
unsafe {
self.mem.set_len(self.mem.len() + num);
self.mem.set_len(new_len);
}
} else {
self.mem.resize(self.mem.len() + num, value);
self.mem.resize(new_len, value);
}
Ok(())
}

/// Write one byte into the memory.
/// Write a generic type T into the memory.
///
/// # Safety
///
/// Unsafe since it assumes that there is enough capacity.
pub unsafe fn write_u8_unchecked(&mut self, value: u8) {
pub unsafe fn write_unchecked<T: Pod>(&mut self, value: T) {
let pos = self.mem.len();
debug_assert!(pos < self.align_offset + self.max_len);
self.mem.set_len(pos + 1);
*self.mem.get_unchecked_mut(pos) = value;
}

/// Write one u64 into the memory.
///
/// # Safety
///
/// Unsafe since it assumes that there is enough capacity.
pub unsafe fn write_u64_unchecked(&mut self, value: u64) {
let pos = self.mem.len();
let new_len = pos + mem::size_of::<u64>();
debug_assert!(new_len <= self.align_offset + self.max_len);
let new_len = pos.saturating_add(mem::size_of::<T>());
debug_assert!(new_len <= self.align_offset.saturating_add(self.max_len));
self.mem.set_len(new_len);
ptr::write_unaligned(
self.mem.get_unchecked_mut(pos..new_len).as_mut_ptr().cast(),
Expand All @@ -154,10 +160,11 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
/// Unsafe since it assumes that there is enough capacity.
pub unsafe fn write_all_unchecked(&mut self, value: &[u8]) {
let pos = self.mem.len();
debug_assert!(pos + value.len() <= self.align_offset + self.max_len);
self.mem.set_len(pos + value.len());
let new_len = pos.saturating_add(value.len());
debug_assert!(new_len <= self.align_offset.saturating_add(self.max_len));
self.mem.set_len(new_len);
self.mem
.get_unchecked_mut(pos..pos + value.len())
.get_unchecked_mut(pos..new_len)
.copy_from_slice(value);
}
}
Expand All @@ -173,11 +180,17 @@ impl<const ALIGN: usize> Clone for AlignedMemory<ALIGN> {

impl<const ALIGN: usize> std::io::Write for AlignedMemory<ALIGN> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
if self.mem.len() + buf.len() > self.align_offset + self.max_len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"aligned memory write failed",
));
match (
self.mem.len().checked_add(buf.len()),
self.align_offset.checked_add(self.max_len),
) {
(Some(new_len), Some(allocation_end)) if new_len <= allocation_end => {}
_ => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"aligned memory write failed",
))
}
}
self.mem.extend_from_slice(buf);
Ok(buf.len())
Expand All @@ -193,18 +206,17 @@ impl<const ALIGN: usize, T: AsRef<[u8]>> From<T> for AlignedMemory<ALIGN> {
}
}

/// Returns true if `data` is aligned to `align`.
pub fn is_memory_aligned(data: &[u8], align: usize) -> bool {
(data.as_ptr() as usize)
.checked_rem(align)
/// Returns true if `ptr` is aligned to `align`.
pub fn is_memory_aligned(ptr: usize, align: usize) -> bool {
ptr.checked_rem(align)
.map(|remainder| remainder == 0)
.unwrap_or(false)
}

#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#![allow(clippy::integer_arithmetic)]
use {super::*, std::io::Write};

fn do_test<const ALIGN: usize>() {
let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity(10);
Expand Down Expand Up @@ -234,9 +246,9 @@ mod tests {

let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(15);
unsafe {
aligned_memory.write_u8_unchecked(42);
aligned_memory.write_unchecked::<u8>(42);
assert_eq!(aligned_memory.len(), 1);
aligned_memory.write_u64_unchecked(0xCAFEBADDDEADCAFE);
aligned_memory.write_unchecked::<u64>(0xCAFEBADDDEADCAFE);
assert_eq!(aligned_memory.len(), 9);
aligned_memory.fill_write(3, 0).unwrap();
aligned_memory.write_all_unchecked(b"foo");
Expand All @@ -262,29 +274,18 @@ mod tests {

#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "< self.align_offset + self.max_len")]
fn test_write_u8_unchecked_debug_assert() {
let mut aligned_memory = AlignedMemory::<8>::with_capacity(1);
unsafe {
aligned_memory.write_u8_unchecked(42);
aligned_memory.write_u8_unchecked(24);
}
}

#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "<= self.align_offset + self.max_len")]
fn test_write_u64_unchecked_debug_assert() {
#[should_panic(expected = "<= self.align_offset.saturating_add(self.max_len)")]
fn test_write_unchecked_debug_assert() {
let mut aligned_memory = AlignedMemory::<8>::with_capacity(15);
unsafe {
aligned_memory.write_u64_unchecked(42);
aligned_memory.write_u64_unchecked(24);
aligned_memory.write_unchecked::<u64>(42);
aligned_memory.write_unchecked::<u64>(24);
}
}

#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "<= self.align_offset + self.max_len")]
#[should_panic(expected = "<= self.align_offset.saturating_add(self.max_len)")]
fn test_write_all_unchecked_debug_assert() {
let mut aligned_memory = AlignedMemory::<8>::with_capacity(5);
unsafe {
Expand Down
2 changes: 1 addition & 1 deletion src/elf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -414,7 +414,7 @@ impl<I: InstructionMeter> Executable<I> {
// it must be properly aligned. We assume that HOST_ALIGN is a
// multiple of the ELF "natural" alignment. See test_load_unaligned.
let aligned;
let bytes = if is_memory_aligned(bytes, HOST_ALIGN) {
let bytes = if is_memory_aligned(bytes.as_ptr() as usize, HOST_ALIGN) {
bytes
} else {
aligned = AlignedMemory::<{ HOST_ALIGN }>::from_slice(bytes);
Expand Down

0 comments on commit 8ca225a

Please sign in to comment.