Skip to content

Commit

Permalink
malloc toggle for nockstack
Browse files Browse the repository at this point in the history
Null checking, documentation
  • Loading branch information
bitemyapp committed Jan 11, 2025
1 parent 6e1b1c1 commit 54e4a3e
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 5 deletions.
3 changes: 3 additions & 0 deletions rust/sword/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,9 @@ cc = "1.0"

# run with e.g. 'cargo build --features check_forwarding,check_acyclic'
[features]
default = ["malloc"]
malloc = []
mmap = []
# FOR DEBUGGING MEMORY ISSUES ONLY
check_all = ["check_acyclic", "check_forwarding", "check_junior"]
check_acyclic = []
Expand Down
4 changes: 2 additions & 2 deletions rust/sword/benches/cue_pill.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ fn main() -> io::Result<()> {
let in_len = f.metadata()?.len();
let mut stack = NockStack::new(1 << 10 << 10 << 10, 0);
let jammed_input = unsafe {
let in_map = memmap::Mmap::map(&f)?;
let in_map = memmap2::Mmap::map(&f)?;
let word_len = (in_len + 7) >> 3;
let (mut atom, dest) = IndirectAtom::new_raw_mut(&mut stack, word_len as usize);
write_bytes(dest.add(word_len as usize - 1), 0, 8);
Expand Down Expand Up @@ -61,7 +61,7 @@ fn main() -> io::Result<()> {
.open(output_filename)?;
f_out.set_len((jammed_output.size() << 3) as u64)?;
unsafe {
let mut out_map = memmap::MmapMut::map_mut(&f_out)?;
let mut out_map = memmap2::MmapMut::map_mut(&f_out)?;
copy_nonoverlapping(
jammed_output.data_pointer() as *mut u8,
out_map.as_mut_ptr(),
Expand Down
85 changes: 82 additions & 3 deletions rust/sword/src/mem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use either::Either::{self, Left, Right};
use ibig::Stack;
use memmap2::MmapMut;
use std::alloc::Layout;
use std::ops::{Deref, DerefMut};
use std::panic::panic_any;
use std::ptr::copy_nonoverlapping;
use std::{mem, ptr};
Expand Down Expand Up @@ -152,6 +153,80 @@ pub enum Direction {
IncreasingDeref,
}

pub enum AllocType {
Mmap,
Malloc,
}

pub(crate) enum Memory {
Mmap(MmapMut),
Malloc(*mut u8, usize),
}

impl Deref for Memory {
type Target = [u8];

#[inline]
fn deref(&self) -> &[u8] {
match self {
Memory::Mmap(mmap) => mmap.deref(),
Memory::Malloc(ptr, size) => {
unsafe { core::slice::from_raw_parts(*ptr, *size) }
},
}
}
}

impl DerefMut for Memory {
#[inline]
fn deref_mut(&mut self) -> &mut [u8] {
match self {
Memory::Mmap(mmap) => mmap.deref_mut(),
Memory::Malloc(ptr, size) => unsafe {
core::slice::from_raw_parts_mut(*ptr, *size)
},
}
}
}

impl AsRef<[u8]> for Memory {
#[inline]
fn as_ref(&self) -> &[u8] {
self.deref()
}
}

impl AsMut<[u8]> for Memory {
#[inline]
fn as_mut(&mut self) -> &mut [u8] {
self.deref_mut()
}
}

impl Memory {
/// Layout and MmapMut::map_anon take their sizes/lengths in bytes but we speak in terms
/// of machine words which are u64 for our purposes so we're 8x'ing them with a cutesy shift.
pub fn allocate(alloc_type: AllocType, size: usize) -> Result<Self, NewStackError> {
let memory = match alloc_type {
AllocType::Mmap => {
let mmap_mut = MmapMut::map_anon(size << 3)?;
Self::Mmap(mmap_mut)
},
AllocType::Malloc => {
// Align is in terms of bytes so I'm aligning it to 64-bits / 8 bytes, word size.
let layout = Layout::from_size_align(size << 3, std::mem::size_of::<u64>()).expect("Invalid layout");
let alloc = unsafe { std::alloc::alloc(layout) };
if alloc.is_null() {
// std promises that std::alloc::handle_alloc_error will diverge
std::alloc::handle_alloc_error(layout);
}
Self::Malloc(alloc, size)
},
};
Ok(memory)
}
}

/// A stack for Nock computation, which supports stack allocation and delimited copying collection
/// for returned nouns
#[allow(dead_code)] // We need the memory field to keep our memory from being unmapped
Expand All @@ -166,8 +241,9 @@ pub struct NockStack {
stack_pointer: *mut u64,
/// Alloc pointer for the current stack frame.
alloc_pointer: *mut u64,
/// MMap which must be kept alive as long as this [NockStack] is
memory: MmapMut,
// /// MMap which must be kept alive as long as this [NockStack] is
// memory: MmapMut,
memory: Memory,
/// PMA from which we will copy into the [NockStack]
/// Whether or not [`Self::pre_copy()`] has been called on the current stack frame.
pc: bool,
Expand Down Expand Up @@ -196,7 +272,10 @@ impl NockStack {
return Err(NewStackError::StackTooSmall);
}
let free = size - (top_slots + RESERVED);
let mut memory = MmapMut::map_anon(size << 3)?;
#[cfg(feature = "mmap")]
let mut memory = Memory::allocate(AllocType::Mmap, size)?;
#[cfg(feature = "malloc")]
let mut memory = Memory::allocate(AllocType::Malloc, size)?;
let start = memory.as_mut_ptr() as *mut u64;
// Here, frame_pointer < alloc_pointer, so the initial frame is West
let frame_pointer = unsafe { start.add(RESERVED + top_slots) } as *mut u64;
Expand Down

0 comments on commit 54e4a3e

Please sign in to comment.