diff --git a/Cargo.toml b/Cargo.toml index 4fa2552..8ce1cd6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "stackalloc" -version = "1.0.1" +version = "1.1.0" homepage = "https://git.flanchan.moe/flanchan/stackalloc-rs" repository = "https://github.com/notflan/stackalloc-rs" keywords = ["alloca", "stack", "stack-allocation", "safe"] diff --git a/src/avec.rs b/src/avec.rs index f1a58c3..1ab49b6 100644 --- a/src/avec.rs +++ b/src/avec.rs @@ -68,8 +68,8 @@ impl<'a, T> Drop for AVec<'a, T> /* for x in slice::from_raw_parts_mut(self.inner.stack.buf_ptr, self.fill_ptr()) { - std::ptr::drop_in_place(x.as_mut_ptr()); - }*/ + std::ptr::drop_in_place(x.as_mut_ptr()); + }*/ } } } @@ -90,7 +90,7 @@ impl<'a, T> AVec<'a, T> /// Have the elements been moved to the heap? pub fn is_allocated(&self) -> bool { - self.fill_ptr() > self.stack_sz + self.fill_ptr() >= self.stack_sz } /// Create a new `AVec` with this backing buffer. @@ -109,4 +109,52 @@ impl<'a, T> AVec<'a, T> _stack: PhantomData } } + + fn move_to_heap(&mut self) + { + let buf: Vec = unsafe { + slice::from_raw_parts(self.inner.stack.buf_ptr as *const MaybeUninit, self.fill_ptr()).iter().map(|x| x.as_ptr().read()).collect() + }; + self.inner = Internal { + heap: ManuallyDrop::new(HeapBuffer { + _fill_ptr: self.stack_sz, + buf, + }), + }; + } + + /// Insert an element into this `AVec`. + pub fn push(&mut self, item: T) + { + if self.is_allocated() + { + unsafe { + (*self.inner.heap).buf.push(item) + } + } else { + unsafe { + let ptr = self.inner.stack.fill_ptr; + *self.inner.stack.buf_ptr.add(ptr) = MaybeUninit::new(item); + self.inner.stack.fill_ptr += 1; + + if self.is_allocated() { + // Move all items to heap + self.move_to_heap(); + } + } + } + } + + /// The number of elements in this `AVec`. + pub fn len(&self) -> usize + { + if self.is_allocated() + { + unsafe { + self.inner.heap.buf.len() + } + } else { + self.fill_ptr() + } + } } diff --git a/src/lib.rs b/src/lib.rs index cdca58b..b35218c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -106,7 +106,7 @@ use std::{ ptr, }; -//TODO: pub mod avec; pub use avec::AVec; +pub mod avec; pub use avec::AVec; mod ffi; /// Allocate a runtime length uninitialised byte buffer on the stack, call `callback` with this buffer, and then deallocate the buffer. @@ -228,12 +228,12 @@ use helpers::*; where F: FnOnce(&mut [u8]) -> T { alloca(size, move |buf| { - // SAFETY: We zero-initialise the backing slice - callback(unsafe { - ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len()); // buf.fill(MaybeUninit::zeroed()); - slice_assume_init_mut(buf) - }) + // SAFETY: We zero-initialise the backing slice + callback(unsafe { + ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len()); // buf.fill(MaybeUninit::zeroed()); + slice_assume_init_mut(buf) }) + }) } @@ -247,12 +247,12 @@ where F: FnOnce(&mut [MaybeUninit]) -> U { let size_bytes = (std::mem::size_of::() * size) + std::mem::align_of::(); alloca(size_bytes, move |buf| { - let abuf = align_buffer_to::>(buf.as_mut_ptr() as *mut u8); - debug_assert!(buf.as_ptr_range().contains(&(abuf as *const _ as *const MaybeUninit))); - unsafe { - callback(slice::from_raw_parts_mut(abuf, size)) - } - }) + let abuf = align_buffer_to::>(buf.as_mut_ptr() as *mut u8); + debug_assert!(buf.as_ptr_range().contains(&(abuf as *const _ as *const MaybeUninit))); + unsafe { + callback(slice::from_raw_parts_mut(abuf, size)) + } + }) } /// Allocate a runtime length slice of `T` on the stack, fill it by calling `init_with`, call `callback` with this buffer, and then drop and deallocate the buffer. @@ -262,22 +262,22 @@ where F: FnOnce(&mut [MaybeUninit]) -> U /// See `alloca()`. #[inline] pub fn stackalloc_with(size: usize, mut init_with: I, callback: F) -> U where F: FnOnce(&mut [T]) -> U, -I: FnMut() -> T + I: FnMut() -> T { stackalloc_uninit(size, move |buf| { - buf.fill_with(move || MaybeUninit::new(init_with())); + buf.fill_with(move || MaybeUninit::new(init_with())); + // SAFETY: We have initialised the buffer above + let buf = unsafe { slice_assume_init_mut(buf) }; + let ret = callback(buf); + if mem::needs_drop::() + { // SAFETY: We have initialised the buffer above - let buf = unsafe { slice_assume_init_mut(buf) }; - let ret = callback(buf); - if mem::needs_drop::() - { - // SAFETY: We have initialised the buffer above - unsafe { - ptr::drop_in_place(buf as *mut _); - } + unsafe { + ptr::drop_in_place(buf as *mut _); } - ret - }) + } + ret + }) } /// Allocate a runtime length slice of `T` on the stack, fill it by cloning `init`, call `callback` with this buffer, and then drop and deallocate the buffer. @@ -287,7 +287,7 @@ I: FnMut() -> T /// See `alloca()`. #[inline] pub fn stackalloc(size: usize, init: T, callback: F) -> U where F: FnOnce(&mut [T]) -> U, -T: Clone + T: Clone { stackalloc_with(size, move || init.clone(), callback) } @@ -300,10 +300,84 @@ T: Clone /// See `alloca()`. #[inline] pub fn stackalloc_with_default(size: usize, callback: F) -> U where F: FnOnce(&mut [T]) -> U, -T: Default + T: Default { stackalloc_with(size, T::default, callback) } + +/// Collect an iterator into a stack allocated buffer up to `size` elements, call `callback` with this buffer, and then drop and deallocate the buffer. +/// +/// See `stackalloc()`. +/// +/// # Size +/// We will only take up to `size` elements from the iterator, the rest of the iterator is dropped. +/// If the iterator yield less elements than `size`, then the slice passed to callback will be smaller than `size` and only contain the elements actually yielded. +#[inline] pub fn stackalloc_with_iter(size: usize, iter: I, callback: F) -> U +where F: FnOnce(&mut [T]) -> U, + I: IntoIterator, +{ + stackalloc_uninit(size, move |buf| { + let mut done = 0; + for (d, s) in buf.iter_mut().zip(iter.into_iter()) + { + *d = MaybeUninit::new(s); + done+=1; + } + // SAFETY: We just initialised `done` elements of `buf` above. + let buf = unsafe { + slice_assume_init_mut(&mut buf[..done]) + }; + let ret = callback(buf); + if mem::needs_drop::() + { + // SAFETY: We have initialised the `buf` above + unsafe { + ptr::drop_in_place(buf as *mut _); + } + } + ret + }) +} + +/// Collect an exact size iterator into a stack allocated slice, call `callback` with this buffer, and then drop and deallocate the buffer. +/// +/// See `stackalloc_with_iter()`. +/// +/// # Size +/// If the implementation of `ExactSizeIterator` on `I` is incorrect and reports a longer length than the iterator actually produces, then the slice passed to `callback` is shortened to the number of elements actually produced. +#[inline] pub fn stackalloc_from_iter_exact(iter: I, callback: F) -> U +where F: FnOnce(&mut [T]) -> U, + I: IntoIterator, + I::IntoIter: ExactSizeIterator, +{ + let iter = iter.into_iter(); + stackalloc_with_iter(iter.len(), iter, callback) +} + +/// Collect an iterator into a stack allocated buffer, call `callback` with this buffer, and then drop and deallocate the buffer. +/// +/// # Safety +/// While the slice passed to `callback` is guaranteed to be safe to use, regardless of if the iterator fills (or tries to overfill) it, this function is still marked as `unsafe` because it trusts the iterator `I` reports an accurate length with its `size_hint()`. +/// It is recommended to instead use `stackalloc_with_iter()` specifying a strict upper bound on the buffer's size, or `stackalloc_from_iter_exact()` for `ExactSizeIterator`s, as this function may allocate far more, or far less (even 0) memory needed to hold all the iterator's elements; therefore this function will very easily not work properly and/or cause stack overflow if used carelessly. +/// +/// If the standard library's `std::iter::TrustedLen` trait becomes stablised, this function will be changed to require that as a bound on `I` and this function will no longer be `unsafe`. +/// +/// # Size +/// The size allocated for the buffer will be the upper bound of the iterator's `size_hint()` if one exists. If not, then the size allocated will be the lower bound of `size_hint()`. +/// This can potentially result in only some of the iterator being present in the buffer, or the buffer allocated being much larger than the iterator itself. +/// If this iterator does not have a good `size_hint()` for this purpose, use `stackalloc_with_iter()`, or `stackalloc_from_iter_exact()` if the iterator has an exact size. +#[inline] pub unsafe fn stackalloc_from_iter_trusted(iter: I, callback: F) -> U +where F: FnOnce(&mut [T]) -> U, + I: IntoIterator, +{ + let iter = iter.into_iter(); + stackalloc_with_iter(match iter.size_hint() { + (_, Some(x)) | + (x, _) => x, + }, iter, callback) +} + + #[cfg(test)] mod tests; diff --git a/src/tests.rs b/src/tests.rs index 22982d2..98a11f2 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,5 +1,37 @@ //! Contains tests and benchmarks +#[test] +fn unknown_size_iter() +{ + let iter = 10..; + let result: u64 = iter.clone().take(10).sum(); + + assert_eq!(result, super::stackalloc_with_iter(10, iter, |buf| buf.iter().copied().sum::())); +} + +fn unknown_size_iter_len() { + assert_eq!(super::stackalloc_with_iter(1024, 0..100, |buf| buf.len()), 100); +} + +#[test] +fn exact_size_iter() +{ + let iter = vec![ + 1, + 2, + 13, + 24, + 100, + ]; + let len = iter.len(); + let result: u64 = iter.iter().copied().sum(); + + assert_eq!(super::stackalloc_from_iter_exact(iter, |buf| { + assert_eq!(buf.len(), len); + buf.iter().copied().sum::() + }), result); +} + #[test] #[should_panic] fn unwinding_over_boundary() @@ -149,7 +181,7 @@ mod bench black_box(crate::alloca(SIZE, |b| {black_box(b);})); }) } - + #[bench] fn stackalloc_of_zeroed_bytes_known(b: &mut Bencher) {