diff --git a/src/lib.rs b/src/lib.rs index 1dcb781..bf56388 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,7 +8,7 @@ //! let pool = BytePool::>::new(); //! //! // Allocate a buffer with capacity 1024. -//! let mut buf = pool.alloc(1024); +//! let mut buf = pool.alloc_and_fill(1024); //! //! // write some data into it //! for i in 0..100 { diff --git a/src/pool.rs b/src/pool.rs index 05e413e..49a6002 100644 --- a/src/pool.rs +++ b/src/pool.rs @@ -56,6 +56,14 @@ impl BytePool { /// The returned `Block` contains arbitrary data, and must be zeroed or overwritten, /// in cases this is needed. pub fn alloc(&self, size: usize) -> Block<'_, T> { + self.alloc_internal(size, false) + } + + pub fn alloc_and_fill(&self, size: usize) -> Block<'_, T> { + self.alloc_internal(size, true) + } + + pub fn alloc_internal(&self, size: usize, fill: bool) -> Block<'_, T> { assert!(size > 0, "Can not allocate empty blocks"); // check the last 4 blocks @@ -64,9 +72,12 @@ impl BytePool { } else { &self.list_large }; - if let Some(el) = list.pop() { - if el.capacity() == size { + if let Some(mut el) = list.pop() { + if el.capacity() >= size && el.capacity() < size + 1024 { // found one, reuse it + if fill { + el.resize(size) + } return Block::new(el, self); } else { // put it back @@ -75,7 +86,11 @@ impl BytePool { } // allocate a new block - let data = T::alloc(size); + let data = if fill { + T::alloc_and_fill(size) + } else { + T::alloc(size) + }; Block::new(data, self) } @@ -90,7 +105,8 @@ impl BytePool { impl<'a, T: Poolable> Drop for Block<'a, T> { fn drop(&mut self) { - let data = mem::ManuallyDrop::into_inner(unsafe { ptr::read(&self.data) }); + let mut data = mem::ManuallyDrop::into_inner(unsafe { ptr::read(&self.data) }); + data.reset(); self.pool.push_raw_block(data); } } @@ -138,6 +154,41 @@ unsafe impl<'a, T: StableDeref + Poolable> StableDeref for Block<'a, T> {} #[cfg(test)] mod tests { use super::*; + #[test] + fn append() { + let pool = BytePool::>::new(); + let mut buf = pool.alloc(4); + assert_eq!(0, buf.len()); + assert_eq!(4, buf.capacity()); + buf.push(12u8); + assert_eq!(1, buf.len()); + buf.extend_from_slice("hello".as_bytes()); + assert_eq!(6, buf.len()); + buf.clear(); + assert_eq!(0, buf.len()); + assert!(buf.capacity() > 0); + } + + #[test] + fn len_and_capacity() { + let pool = BytePool::>::new(); + for i in 1..10 { + let buf = pool.alloc_and_fill(i); + assert_eq!(buf.len(), i) + } + for i in 1..10 { + let buf = pool.alloc(i); + assert_eq!(buf.len(), 0) + } + for i in 1..10 { + let buf = pool.alloc_and_fill(i * 10000); + assert_eq!(buf.len(), i * 10000) + } + for i in 1..10 { + let buf = pool.alloc(i * 10000); + assert_eq!(buf.len(), 0) + } + } #[test] fn basics_vec_u8() { @@ -174,6 +225,7 @@ mod tests { let _slice: &[u8] = &buf; assert_eq!(buf.capacity(), 10); + buf.resize(10, 0); for i in 0..10 { buf[i] = 1; } @@ -198,7 +250,7 @@ mod tests { let pool1 = pool.clone(); let h1 = std::thread::spawn(move || { for _ in 0..100 { - let mut buf = pool1.alloc(64); + let mut buf = pool1.alloc_and_fill(64); buf[10] = 10; } }); @@ -206,7 +258,7 @@ mod tests { let pool2 = pool.clone(); let h2 = std::thread::spawn(move || { for _ in 0..100 { - let mut buf = pool2.alloc(64); + let mut buf = pool2.alloc_and_fill(64); buf[10] = 10; } }); diff --git a/src/poolable.rs b/src/poolable.rs index 23fd0cc..b851e97 100644 --- a/src/poolable.rs +++ b/src/poolable.rs @@ -3,16 +3,41 @@ use std::hash::{BuildHasher, Hash}; /// The trait required to be able to use a type in `BytePool`. pub trait Poolable { + fn empty(&self) -> bool; + fn len(&self) -> usize; fn capacity(&self) -> usize; + fn resize(&mut self, count: usize); + fn reset(&mut self); fn alloc(size: usize) -> Self; + fn alloc_and_fill(size: usize) -> Self; } impl Poolable for Vec { - fn capacity(&self) -> usize { + fn empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { self.len() } + fn capacity(&self) -> usize { + self.capacity() + } + + fn resize(&mut self, count: usize) { + self.resize(count, T::default()); + } + + fn reset(&mut self) { + self.clear(); + } + fn alloc(size: usize) -> Self { + Vec::::with_capacity(size) + } + + fn alloc_and_fill(size: usize) -> Self { vec![T::default(); size] } } @@ -22,11 +47,32 @@ where K: Eq + Hash, S: BuildHasher + Default, { - fn capacity(&self) -> usize { + fn empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { self.len() } + fn capacity(&self) -> usize { + self.capacity() + } + + fn resize(&mut self, _count: usize) { + // do thing + } + + fn reset(&mut self) { + self.clear(); + } + fn alloc(size: usize) -> Self { + Self::alloc_and_fill(size) + } + + fn alloc_and_fill(size: usize) -> Self { + // not actually filling the HaspMap though HashMap::with_capacity_and_hasher(size, Default::default()) } } @@ -42,7 +88,7 @@ impl Realloc for Vec { assert!(new_size > 0); match new_size.cmp(&self.capacity()) { - Greater => self.resize(new_size, T::default()), + Greater => self.reserve(new_size - self.capacity()), Less => { self.truncate(new_size); self.shrink_to_fit();