diff --git a/wgpu-core/src/weak_vec.rs b/wgpu-core/src/weak_vec.rs index f6034c6e90..e00cc9e599 100644 --- a/wgpu-core/src/weak_vec.rs +++ b/wgpu-core/src/weak_vec.rs @@ -6,28 +6,20 @@ use std::sync::Weak; /// dropping older elements that no longer have strong references to them. #[derive(Debug)] pub(crate) struct WeakVec { - inner: Vec>>, - empty_slots: Vec, - scan_slots_on_next_push: bool, + inner: Vec>, } impl Default for WeakVec { fn default() -> Self { Self { inner: Default::default(), - empty_slots: Default::default(), - scan_slots_on_next_push: false, } } } impl WeakVec { pub(crate) fn new() -> Self { - Self { - inner: Vec::new(), - empty_slots: Vec::default(), - scan_slots_on_next_push: false, - } + Self { inner: Vec::new() } } /// Pushes a new element to this collection. @@ -35,28 +27,25 @@ impl WeakVec { /// If the inner Vec needs to be reallocated, we will first drop older elements that /// no longer have strong references to them. pub(crate) fn push(&mut self, value: Weak) { - if self.scan_slots_on_next_push { - for (i, value) in self.inner.iter_mut().enumerate() { - if let Some(w) = value { - if w.strong_count() == 0 { - *value = None; - self.empty_slots.push(i); - } + if self.inner.len() == self.inner.capacity() { + // Iterating backwards has the advantage that we don't do more work than we have to. + for i in (0..self.inner.len()).rev() { + if self.inner[i].strong_count() == 0 { + self.inner.swap_remove(i); } } + + // Make sure our capacity is twice the number of live elements. + // Leaving some spare capacity ensures that we won't re-scan immediately. + self.inner.reserve_exact(self.inner.len()); } - if let Some(i) = self.empty_slots.pop() { - self.inner[i] = Some(value); - self.scan_slots_on_next_push = self.empty_slots.is_empty(); - } else { - self.inner.push(Some(value)); - self.scan_slots_on_next_push = self.inner.len() == self.inner.capacity(); - } + + self.inner.push(value); } } pub(crate) struct WeakVecIter { - inner: std::iter::Flatten>>>, + inner: std::vec::IntoIter>, } impl Iterator for WeakVecIter { @@ -71,7 +60,7 @@ impl IntoIterator for WeakVec { type IntoIter = WeakVecIter; fn into_iter(self) -> Self::IntoIter { WeakVecIter { - inner: self.inner.into_iter().flatten(), + inner: self.inner.into_iter(), } } }