Trace pipeline creation before it actually happens

wgpu-core/src/device/mod.rs

@@ -4018,40 +4018,32 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let hub = B::hub(self);
         let mut token = Token::root();
 
+        let fid = hub.render_pipelines.prepare(id_in);
         let (device_guard, mut token) = hub.devices.read(&mut token);
         let error = loop {
             let device = match device_guard.get(device_id) {
                 Ok(device) => device,
                 Err(_) => break DeviceError::Invalid.into(),
             };
-            let (pipeline, derived_bind_group_count, layout_id) = match device
+            #[cfg(feature = "trace")]
+            if let Some(ref trace) = device.trace {
+                trace
+                    .lock()
+                    .add(trace::Action::CreateRenderPipeline(fid.id(), desc.clone()));
+            }
+
+            let (pipeline, derived_bind_group_count, _layout_id) = match device
                 .create_render_pipeline(device_id, desc, implicit_pipeline_ids, &hub, &mut token)
             {
                 Ok(pair) => pair,
                 Err(e) => break e,
             };
 
-            let id = hub
+            let id = fid.assign(pipeline, &mut token);
-                .render_pipelines
-                .register_identity(id_in, pipeline, &mut token);
-
-            #[cfg(feature = "trace")]
-            if let Some(ref trace) = device.trace {
-                trace.lock().add(trace::Action::CreateRenderPipeline(
-                    id.0,
-                    pipeline::RenderPipelineDescriptor {
-                        layout: Some(layout_id),
-                        ..desc.clone()
-                    },
-                ));
-            }
-            let _ = layout_id;
             return (id.0, derived_bind_group_count, None);
         };
 
-        let id =
+        let id = fid.assign_error(desc.label.borrow_or_default(), &mut token);
-            hub.render_pipelines
-                .register_error(id_in, desc.label.borrow_or_default(), &mut token);
         (id, 0, Some(error))
     }
 
@@ -4149,40 +4141,32 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let hub = B::hub(self);
         let mut token = Token::root();
 
+        let fid = hub.compute_pipelines.prepare(id_in);
         let (device_guard, mut token) = hub.devices.read(&mut token);
         let error = loop {
             let device = match device_guard.get(device_id) {
                 Ok(device) => device,
                 Err(_) => break DeviceError::Invalid.into(),
             };
-            let (pipeline, derived_bind_group_count, layout_id) = match device
+            #[cfg(feature = "trace")]
+            if let Some(ref trace) = device.trace {
+                trace
+                    .lock()
+                    .add(trace::Action::CreateComputePipeline(fid.id(), desc.clone()));
+            }
+
+            let (pipeline, derived_bind_group_count, _layout_id) = match device
                 .create_compute_pipeline(device_id, desc, implicit_pipeline_ids, &hub, &mut token)
             {
                 Ok(pair) => pair,
                 Err(e) => break e,
             };
 
-            let id = hub
+            let id = fid.assign(pipeline, &mut token);
-                .compute_pipelines
-                .register_identity(id_in, pipeline, &mut token);
-
-            #[cfg(feature = "trace")]
-            if let Some(ref trace) = device.trace {
-                trace.lock().add(trace::Action::CreateComputePipeline(
-                    id.0,
-                    pipeline::ComputePipelineDescriptor {
-                        layout: Some(layout_id),
-                        ..desc.clone()
-                    },
-                ));
-            }
-            let _ = layout_id;
             return (id.0, derived_bind_group_count, None);
         };
 
-        let id =
+        let id = fid.assign_error(desc.label.borrow_or_default(), &mut token);
-            hub.compute_pipelines
-                .register_error(id_in, desc.label.borrow_or_default(), &mut token);
         (id, 0, Some(error))
     }
 

wgpu-core/src/hub.rs

@@ -439,10 +439,37 @@ impl<T: Resource, I: TypedId, F: IdentityHandlerFactory<I>> Registry<T, I, F> {
     }
 }
 
+#[must_use]
+pub(crate) struct FutureId<'a, I: TypedId, T> {
+    id: I,
+    data: &'a RwLock<Storage<T, I>>,
+}
+
+impl<I: TypedId + Copy, T> FutureId<'_, I, T> {
+    pub fn id(&self) -> I {
+        self.id
+    }
+
+    pub fn assign<'a, A: Access<T>>(self, value: T, _: &'a mut Token<A>) -> Valid<I> {
+        self.data.write().insert(self.id, value);
+        Valid(self.id)
+    }
+
+    pub fn assign_error<'a, A: Access<T>>(self, label: &str, _: &'a mut Token<A>) -> I {
+        self.data.write().insert_error(self.id, label);
+        self.id
+    }
+}
+
 impl<T: Resource, I: TypedId + Copy, F: IdentityHandlerFactory<I>> Registry<T, I, F> {
-    pub fn register<A: Access<T>>(&self, id: I, value: T, _token: &mut Token<A>) {
+    pub(crate) fn prepare(
-        debug_assert_eq!(id.unzip().2, self.backend);
+        &self,
-        self.data.write().insert(id, value);
+        id_in: <F::Filter as IdentityHandler<I>>::Input,
+    ) -> FutureId<I, T> {
+        FutureId {
+            id: self.identity.process(id_in, self.backend),
+            data: &self.data,
+        }
     }
 
     pub fn read<'a, A: Access<T>>(
@@ -459,15 +486,14 @@ impl<T: Resource, I: TypedId + Copy, F: IdentityHandlerFactory<I>> Registry<T, I
         (self.data.write(), Token::new())
     }
 
+    //TODO: consider remove this once everything uses `prepare`
     pub(crate) fn register_identity<A: Access<T>>(
         &self,
         id_in: <F::Filter as IdentityHandler<I>>::Input,
         value: T,
         token: &mut Token<A>,
     ) -> Valid<I> {
-        let id = self.identity.process(id_in, self.backend);
+        self.prepare(id_in).assign(value, token)
-        self.register(id, value, token);
-        Valid(id)
     }
 
     pub(crate) fn register_identity_locked(
@@ -481,16 +507,14 @@ impl<T: Resource, I: TypedId + Copy, F: IdentityHandlerFactory<I>> Registry<T, I
         Valid(id)
     }
 
+    //TODO: consider remove this once everything uses `prepare`
     pub fn register_error<A: Access<T>>(
         &self,
         id_in: <F::Filter as IdentityHandler<I>>::Input,
         label: &str,
-        _token: &mut Token<A>,
+        token: &mut Token<A>,
     ) -> I {
-        let id = self.identity.process(id_in, self.backend);
+        self.prepare(id_in).assign_error(label, token)
-        debug_assert_eq!(id.unzip().2, self.backend);
-        self.data.write().insert_error(id, label);
-        id
     }
 
     pub fn unregister_locked(&self, id: I, guard: &mut Storage<T, I>) -> Option<T> {
@@ -513,10 +537,6 @@ impl<T: Resource, I: TypedId + Copy, F: IdentityHandlerFactory<I>> Registry<T, I
         (value, Token::new())
     }
 
-    pub fn process_id(&self, id_in: <F::Filter as IdentityHandler<I>>::Input) -> I {
-        self.identity.process(id_in, self.backend)
-    }
-
     pub fn free_id(&self, id: I) {
         self.identity.free(id)
     }