Simplify the ID allocation in IdentityValues

CHANGELOG.md

@@ -113,6 +113,7 @@ Bottom level categories:
 - Fix panic when creating a surface while no backend is available. By @wumpf [#5166](https://github.com/gfx-rs/wgpu/pull/5166)
 - Correctly compute minimum buffer size for array-typed `storage` and `uniform` vars. By @jimblandy [#5222](https://github.com/gfx-rs/wgpu/pull/5222)
 - Fix timeout when presenting a surface where no work has been done. By @waywardmonkeys in [#5200](https://github.com/gfx-rs/wgpu/pull/5200)
+- Simplify and speed up the allocation of internal IDs. By @nical in [#5229](https://github.com/gfx-rs/wgpu/pull/5229)
 - Fix an issue where command encoders weren't properly freed if an error occurred during command encoding. By @ErichDonGubler in [#5251](https://github.com/gfx-rs/wgpu/pull/5251).
 
 #### WGL

wgpu-core/src/identity.rs

@@ -3,10 +3,17 @@ use wgt::Backend;
 
 use crate::{
  id::{Id, Marker},
- Epoch, FastHashMap, Index,
+ Epoch, Index,
 };
 use std::{fmt::Debug, marker::PhantomData};
 
+#[derive(Copy, Clone, Debug, PartialEq)]
+enum IdSource {
+ External,
+ Allocated,
+ None,
+}
+
 /// A simple structure to allocate [`Id`] identifiers.
 ///
 /// Calling [`alloc`] returns a fresh, never-before-seen id. Calling [`free`]
@@ -34,12 +41,15 @@ use std::{fmt::Debug, marker::PhantomData};
 /// [`Backend`]: wgt::Backend;
 /// [`alloc`]: IdentityManager::alloc
 /// [`free`]: IdentityManager::free
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub(super) struct IdentityValues {
  free: Vec<(Index, Epoch)>,
- //sorted by Index
- used: FastHashMap<Epoch, Vec<Index>>,
+ next_index: Index,
  count: usize,
+ // Sanity check: The allocation logic works under the assumption that we don't
+ // do a mix of allocating ids from here and providing ids manually for the same
+ // storage container.
+ id_source: IdSource,
 }
 
 impl IdentityValues {
@@ -48,35 +58,41 @@ impl IdentityValues {
  /// The backend is incorporated into the id, so that ids allocated with
  /// different `backend` values are always distinct.
  pub fn alloc<T: Marker>(&mut self, backend: Backend) -> Id<T> {
+ assert!(
+ self.id_source != IdSource::External,
+ "Mix of internally allocated and externally provided IDs"
+ );
+ self.id_source = IdSource::Allocated;
+
  self.count += 1;
  match self.free.pop() {
  Some((index, epoch)) => Id::zip(index, epoch + 1, backend),
  None => {
+ let index = self.next_index;
+ self.next_index += 1;
  let epoch = 1;
- let used = self.used.entry(epoch).or_insert_with(Default::default);
- let index = if let Some(i) = used.iter().max_by_key(|v| *v) {
- i + 1
- } else {
- 0
- };
- used.push(index);
  Id::zip(index, epoch, backend)
  }
  }
  }
 
  pub fn mark_as_used<T: Marker>(&mut self, id: Id<T>) -> Id<T> {
+ assert!(
+ self.id_source != IdSource::Allocated,
+ "Mix of internally allocated and externally provided IDs"
+ );
+ self.id_source = IdSource::External;
+
  self.count += 1;
- let (index, epoch, _backend) = id.unzip();
- let used = self.used.entry(epoch).or_insert_with(Default::default);
- used.push(index);
  id
  }
 
  /// Free `id`. It will never be returned from `alloc` again.
  pub fn release<T: Marker>(&mut self, id: Id<T>) {
- let (index, epoch, _backend) = id.unzip();
- self.free.push((index, epoch));
+ if let IdSource::Allocated = self.id_source {
+ let (index, epoch, _backend) = id.unzip();
+ self.free.push((index, epoch));
+ }
  self.count -= 1;
  }
 
@@ -106,7 +122,12 @@ impl<T: Marker> IdentityManager<T> {
 impl<T: Marker> IdentityManager<T> {
  pub fn new() -> Self {
  Self {
- values: Mutex::new(IdentityValues::default()),
+ values: Mutex::new(IdentityValues {
+ free: Vec::new(),
+ next_index: 0,
+ count: 0,
+ id_source: IdSource::None,
+ }),
  _phantom: PhantomData,
  }
  }
@@ -115,15 +136,11 @@ impl<T: Marker> IdentityManager<T> {
 #[test]
 fn test_epoch_end_of_life() {
  use crate::id;
-
  let man = IdentityManager::<id::markers::Buffer>::new();
- let forced_id = man.mark_as_used(id::BufferId::zip(0, 1, Backend::Empty));
- assert_eq!(forced_id.unzip().0, 0);
  let id1 = man.process(Backend::Empty);
- assert_eq!(id1.unzip().0, 1);
+ assert_eq!(id1.unzip(), (0, 1, Backend::Empty));
  man.free(id1);
  let id2 = man.process(Backend::Empty);
  // confirm that the epoch 1 is no longer re-used
- assert_eq!(id2.unzip().0, 1);
- assert_eq!(id2.unzip().1, 2);
+ assert_eq!(id2.unzip(), (0, 2, Backend::Empty));
 }