1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
use parking_lot::Mutex;
use wgt::Backend;

use crate::{
    id::{self},
    Epoch, FastHashMap, Index,
};
use std::{fmt::Debug, marker::PhantomData, sync::Arc};

/// A simple structure to allocate [`Id`] identifiers.
///
/// Calling [`alloc`] returns a fresh, never-before-seen id. Calling [`free`]
/// marks an id as dead; it will never be returned again by `alloc`.
///
/// Use `IdentityManager::default` to construct new instances.
///
/// `IdentityManager` returns `Id`s whose index values are suitable for use as
/// indices into a `Storage<T>` that holds those ids' referents:
///
/// - Every live id has a distinct index value. Each live id's index selects a
///   distinct element in the vector.
///
/// - `IdentityManager` prefers low index numbers. If you size your vector to
///   accommodate the indices produced here, the vector's length will reflect
///   the highwater mark of actual occupancy.
///
/// - `IdentityManager` reuses the index values of freed ids before returning
///   ids with new index values. Freed vector entries get reused.
///
/// See the module-level documentation for an overview of how this
/// fits together.
///
/// [`Id`]: crate::id::Id
/// [`Backend`]: wgt::Backend;
/// [`alloc`]: IdentityManager::alloc
/// [`free`]: IdentityManager::free
#[derive(Debug, Default)]
pub(super) struct IdentityValues {
    free: Vec<(Index, Epoch)>,
    //sorted by Index
    used: FastHashMap<Epoch, Vec<Index>>,
    count: usize,
}

impl IdentityValues {
    /// Allocate a fresh, never-before-seen id with the given `backend`.
    ///
    /// The backend is incorporated into the id, so that ids allocated with
    /// different `backend` values are always distinct.
    pub fn alloc<I: id::TypedId>(&mut self, backend: Backend) -> I {
        self.count += 1;
        match self.free.pop() {
            Some((index, epoch)) => I::zip(index, epoch + 1, backend),
            None => {
                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);
                I::zip(index, epoch, backend)
            }
        }
    }

    pub fn mark_as_used<I: id::TypedId>(&mut self, id: I) -> I {
        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<I: id::TypedId>(&mut self, id: I) {
        let (index, epoch, _backend) = id.unzip();
        self.free.push((index, epoch));
        self.count -= 1;
    }

    pub fn count(&self) -> usize {
        self.count
    }
}

#[derive(Debug)]
pub struct IdentityManager<I: id::TypedId> {
    pub(super) values: Mutex<IdentityValues>,
    _phantom: PhantomData<I>,
}

impl<I: id::TypedId> IdentityManager<I> {
    pub fn process(&self, backend: Backend) -> I {
        self.values.lock().alloc(backend)
    }
    pub fn mark_as_used(&self, id: I) -> I {
        self.values.lock().mark_as_used(id)
    }
    pub fn free(&self, id: I) {
        self.values.lock().release(id)
    }
}

impl<I: id::TypedId> IdentityManager<I> {
    pub fn new() -> Self {
        Self {
            values: Mutex::new(IdentityValues::default()),
            _phantom: PhantomData,
        }
    }
}

/// A type that can produce [`IdentityManager`] filters for ids of type `I`.
///
/// See the module-level documentation for details.
pub trait IdentityHandlerFactory<I: id::TypedId> {
    type Input: Copy;
    /// Create an [`IdentityManager<I>`] implementation that can
    /// transform proto-ids into ids of type `I`.
    /// It can return None if ids are passed from outside
    /// and are not generated by wgpu
    ///
    /// [`IdentityManager<I>`]: IdentityManager
    fn spawn(&self) -> Arc<IdentityManager<I>> {
        Arc::new(IdentityManager::new())
    }
    fn autogenerate_ids() -> bool;
    fn input_to_id(id_in: Self::Input) -> I;
}

/// A global identity handler factory based on [`IdentityManager`].
///
/// Each of this type's `IdentityHandlerFactory<I>::spawn` methods
/// returns a `Mutex<IdentityManager<I>>`, which allocates fresh `I`
/// ids itself, and takes `()` as its proto-id type.
#[derive(Debug)]
pub struct IdentityManagerFactory;

impl<I: id::TypedId> IdentityHandlerFactory<I> for IdentityManagerFactory {
    type Input = ();
    fn autogenerate_ids() -> bool {
        true
    }

    fn input_to_id(_id_in: Self::Input) -> I {
        unreachable!("It should not be called")
    }
}

/// A factory that can build [`IdentityManager`]s for all resource
/// types.
pub trait GlobalIdentityHandlerFactory:
    IdentityHandlerFactory<id::AdapterId>
    + IdentityHandlerFactory<id::DeviceId>
    + IdentityHandlerFactory<id::PipelineLayoutId>
    + IdentityHandlerFactory<id::ShaderModuleId>
    + IdentityHandlerFactory<id::BindGroupLayoutId>
    + IdentityHandlerFactory<id::BindGroupId>
    + IdentityHandlerFactory<id::CommandBufferId>
    + IdentityHandlerFactory<id::RenderBundleId>
    + IdentityHandlerFactory<id::RenderPipelineId>
    + IdentityHandlerFactory<id::ComputePipelineId>
    + IdentityHandlerFactory<id::QuerySetId>
    + IdentityHandlerFactory<id::BufferId>
    + IdentityHandlerFactory<id::StagingBufferId>
    + IdentityHandlerFactory<id::TextureId>
    + IdentityHandlerFactory<id::TextureViewId>
    + IdentityHandlerFactory<id::SamplerId>
    + IdentityHandlerFactory<id::SurfaceId>
{
}

impl GlobalIdentityHandlerFactory for IdentityManagerFactory {}

pub type Input<G, I> = <G as IdentityHandlerFactory<I>>::Input;

#[test]
fn test_epoch_end_of_life() {
    use id::TypedId as _;
    let man = IdentityManager::<id::BufferId>::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);
    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);
}