From 004e3efe84a320d9331371ed31fa50baa2414911 Mon Sep 17 00:00:00 2001
From: Nicolas Silva <nical@fastmail.com>
Date: Thu, 15 Feb 2024 21:55:55 +0100
Subject: [PATCH] Simplify the ID allocation in IdentityValues (#5229)

Co-authored-by: Connor Fitzgerald <connorwadefitzgerald@gmail.com>
---
 CHANGELOG.md              |  1 +
 wgpu-core/src/identity.rs | 63 +++++++++++++++++++++++++--------------
 2 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0717cf26c..05a4611da 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/wgpu-core/src/identity.rs b/wgpu-core/src/identity.rs
index 0e34055c7..d76d29341 100644
--- a/wgpu-core/src/identity.rs
+++ b/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));
 }