diff --git a/src/ext.rs b/src/ext.rs index a7777d4..c7b0aba 100644 --- a/src/ext.rs +++ b/src/ext.rs @@ -1,5 +1,6 @@ //! Extensions and macros use std::cell::RefCell; +use std::ptr; #[macro_export] macro_rules! basic_enum { ($(#[$meta:meta])* $vis:vis $name:ident $(; $tcomment:literal)?: $($var:ident $(=> $comment:literal)?),+ $(,)?) => { @@ -138,16 +139,34 @@ where F: FnOnce(&mut [u8]) -> T static BUFFER: RefCell> = RefCell::new(vec![0u8; STACK_SIZE_LIMIT*2]); } BUFFER.with(move |buf| { - let mut buf = buf.borrow_mut(); - if buf.len() < size { - buf.resize(size, 0); + // If the borrow fails then `f` has recursively called into this function, so for that we allocate a new buffer instead of reusing this static one. + if let Ok(mut buf) = buf.try_borrow_mut() { + if buf.len() < size { + buf.resize(size, 0); + } + let res = f(&mut buf[..size]); + bytes::blank(&mut buf[..size]); + res + } else { + f(&mut vec![0u8; size]) } - f(&mut buf[..size]) }) } else { stackalloc::alloca_zeroed(size, f) - // - // TODO: Is this okay to do? I'm not sure it is.. We'll see + + // I don't think this is okay to do. //stackalloc::alloca(size, move |buf| f(unsafe { stackalloc::helpers::slice_assume_init_mut(buf) })) } } + +pub mod bytes +{ + use super::*; + /// `bzero` this slice + pub fn blank(slice: &mut [u8]) + { + unsafe { + ptr::write_bytes(slice.as_mut_ptr(), 0, slice.len()); + } + } +}