diff --git a/src/lib.rs b/src/lib.rs index 681d1de..42320b4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -411,27 +411,20 @@ impl Bump { // Takes `&mut self` so `self` must be unique and there can't be any // borrows active that would get invalidated by resetting. unsafe { - let mut footer = Some(self.all_chunk_footers.get()); - - // Reset the pointer in each of our chunks. - while let Some(f) = footer { - footer = f.as_ref().next.get(); + let mut cur_chunk = self.all_chunk_footers.get(); - if f == self.current_chunk_footer.get() { - // If this is the current chunk, then reset the bump finger - // to the start of the chunk. - f.as_ref() - .ptr - .set(NonNull::new_unchecked(f.as_ref().data.as_ptr() as *mut u8)); - f.as_ref().next.set(None); - self.all_chunk_footers.set(f); - } else { - // If this is not the current chunk, return it to the global - // allocator. - dealloc(f.as_ref().data.as_ptr(), f.as_ref().layout); - } + // Free all chunks except the last one + while let Some(next_chunk) = cur_chunk.as_ref().next.get() { + dealloc(cur_chunk.as_ref().data.as_ptr(), cur_chunk.as_ref().layout); + cur_chunk = next_chunk; } + // Reset the bump finger to the start of the chunk. + cur_chunk.as_ref().ptr.set(cur_chunk.as_ref().data); + + self.all_chunk_footers.set(cur_chunk); + self.current_chunk_footer.set(cur_chunk); + debug_assert_eq!( self.all_chunk_footers.get(), self.current_chunk_footer.get(), diff --git a/tests/tests.rs b/tests/tests.rs index f9f4343..316d300 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -158,3 +158,19 @@ fn with_capacity_test() { with_capacity_helper(0u64..10000); with_capacity_helper(0u128..10000); } + +#[test] +fn test_reset() { + let mut b = Bump::new(); + + for i in 0u64..10_000 { + b.alloc(i); + } + + assert!(b.iter_allocated_chunks().count() > 1); + + let ptr = b.iter_allocated_chunks().last().unwrap().as_ptr(); + b.reset(); + assert_eq!(ptr as usize, b.alloc(0u64) as *const u64 as usize); + assert_eq!(b.iter_allocated_chunks().count(), 1); +}