diff --git a/src/lib.rs b/src/lib.rs index 4ffc16b..7252d0e 100755 --- a/src/lib.rs +++ b/src/lib.rs @@ -465,7 +465,19 @@ impl Arena { #[inline(never)] fn insert_slow_path(&mut self, value: T) -> Index { - let len = self.items.len(); + let len = if self.capacity() == 0 { + // `drain()` sets the capacity to 0 and if the capacity is 0, the + // next `try_insert() `will refer to an out-of-range index because + // the next `reserve()` does not add element, resulting in a panic. + // So ensure that `self` have at least 1 capacity here. + // + // Ideally, this problem should be handled within `drain()`,but + // this problem cannot be handled within `drain()` because `drain()` + // returns an iterator that borrows `self` mutably. + 1 + } else { + self.items.len() + }; self.reserve(len); self.try_insert(value) .map_err(|_| ()) @@ -899,8 +911,16 @@ impl Arena { /// assert!(arena.get(idx_2).is_none()); /// ``` pub fn drain(&mut self) -> Drain { + let old_len = self.len; + if !self.is_empty() { + // Increment generation, but if there are no elements, do nothing to + // avoid unnecessary incrementing generation. + self.generation += 1; + } + self.free_list_head = None; + self.len = 0; Drain { - len: self.len, + len: old_len, inner: self.items.drain(..).enumerate(), } } diff --git a/tests/tests.rs b/tests/tests.rs index b296bdf..c9e0cd7 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -238,6 +238,23 @@ fn drain() { } assert!(arena.get(idx_1).is_none()); assert!(arena.get(idx_2).is_none()); + + assert_eq!(arena.capacity(), 0); + assert_eq!(arena.len(), 0); + + let idx_3 = arena.insert("a"); + assert_ne!(idx_1, idx_3); + assert_eq!(arena.capacity(), 1); + assert_eq!(arena.len(), 1); + + // If there are no elements, do not increment generation. + let mut arena_2 = Arena::with_capacity(1); + arena_2.drain(); + arena_2.drain(); + arena_2.drain(); + let idx_1 = arena_2.insert(1); + let gen = idx_1.into_raw_parts().1; + assert_eq!(gen, 0); } #[test]