Add ReferenceType trait for command structures

andyleiserson · andyleiserson · commit 27ee35c70a4b · 2025-10-02T17:03:10.000-07:00
Use it to deduplicate some `Foo` / `ArcFoo` / `TraceFoo` definitions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -142,6 +142,7 @@ libloading = "0.8"
 libm = { version = "0.2.6", default-features = false }
 libtest-mimic = "0.8"
 log = "0.4.21"
+macro_rules_attribute = "0.2"
 nanoserde = "0.2"
 nanorand = { version = "0.8", default-features = false, features = ["wyrand"] }
 noise = "0.9"
diff --git a/player/src/bin/play.rs b/player/src/bin/play.rs
@@ -7,7 +7,7 @@ fn main() {
 
     use player::GlobalPlay as _;
     use wgc::device::trace;
-    use wgpu_core::identity::IdentityManager;
+    use wgpu_core::{command::IdReferences, identity::IdentityManager};
 
     use std::{
         fs,
@@ -52,7 +52,7 @@ fn main() {
 
     log::info!("Loading trace '{trace:?}'");
     let file = fs::File::open(trace).unwrap();
-    let mut actions: Vec<trace::Action> = ron::de::from_reader(file).unwrap();
+    let mut actions: Vec<trace::Action<IdReferences>> = ron::de::from_reader(file).unwrap();
     actions.reverse(); // allows us to pop from the top
     log::info!("Found {} actions", actions.len());
 
diff --git a/player/src/lib.rs b/player/src/lib.rs
@@ -6,7 +6,11 @@
 extern crate wgpu_core as wgc;
 extern crate wgpu_types as wgt;
 
-use wgc::{command::Command, device::trace, identity::IdentityManager};
+use wgc::{
+    command::{Command, IdReferences},
+    device::trace,
+    identity::IdentityManager,
+};
 
 use std::{borrow::Cow, fs, path::Path};
 
@@ -21,7 +25,7 @@ pub trait GlobalPlay {
         &self,
         device: wgc::id::DeviceId,
         queue: wgc::id::QueueId,
-        action: trace::Action,
+        action: trace::Action<IdReferences>,
         dir: &Path,
         command_encoder_id_manager: &mut IdentityManager<wgc::id::markers::CommandEncoder>,
         command_buffer_id_manager: &mut IdentityManager<wgc::id::markers::CommandBuffer>,
@@ -143,7 +147,7 @@ impl GlobalPlay for wgc::global::Global {
                             }
                         };
                         wgc::ray_tracing::BlasBuildEntry {
-                            blas_id: x.blas_id,
+                            blas_id: x.blas,
                             geometries,
                         }
                     });
@@ -153,14 +157,14 @@ impl GlobalPlay for wgc::global::Global {
                             instance
                                 .as_ref()
                                 .map(|instance| wgc::ray_tracing::TlasInstance {
-                                    blas_id: instance.blas_id,
+                                    blas_id: instance.blas,
                                     transform: &instance.transform,
                                     custom_data: instance.custom_data,
                                     mask: instance.mask,
                                 })
                         });
                         wgc::ray_tracing::TlasPackage {
-                            tlas_id: x.tlas_id,
+                            tlas_id: x.tlas,
                             instances: Box::new(instances),
                             lowest_unmodified: x.lowest_unmodified,
                         }
@@ -188,7 +192,7 @@ impl GlobalPlay for wgc::global::Global {
         &self,
         device: wgc::id::DeviceId,
         queue: wgc::id::QueueId,
-        action: trace::Action,
+        action: trace::Action<IdReferences>,
         dir: &Path,
         command_encoder_id_manager: &mut IdentityManager<wgc::id::markers::CommandEncoder>,
         command_buffer_id_manager: &mut IdentityManager<wgc::id::markers::CommandBuffer>,
diff --git a/player/tests/player/main.rs b/player/tests/player/main.rs
@@ -20,7 +20,7 @@ use std::{
     path::{Path, PathBuf},
     slice,
 };
-use wgc::identity::IdentityManager;
+use wgc::{command::IdReferences, identity::IdentityManager};
 
 #[derive(serde::Deserialize)]
 struct RawId {
@@ -57,7 +57,7 @@ struct Expectation {
 struct Test<'a> {
     features: wgt::Features,
     expectations: Vec<Expectation>,
-    actions: Vec<wgc::device::trace::Action<'a>>,
+    actions: Vec<wgc::device::trace::Action<'a, IdReferences>>,
 }
 
 fn map_callback(status: Result<(), wgc::resource::BufferAccessError>) {
diff --git a/wgpu-core/Cargo.toml b/wgpu-core/Cargo.toml
@@ -178,6 +178,7 @@ document-features.workspace = true
 hashbrown.workspace = true
 indexmap.workspace = true
 log.workspace = true
+macro_rules_attribute.workspace = true
 once_cell = { workspace = true, features = ["std"] }
 parking_lot.workspace = true
 profiling = { workspace = true, default-features = false }
diff --git a/wgpu-core/src/command/encoder_command.rs b/wgpu-core/src/command/encoder_command.rs
@@ -1,12 +1,77 @@
 use core::convert::Infallible;
 
 use alloc::{string::String, sync::Arc, vec::Vec};
+#[cfg(feature = "serde")]
+use macro_rules_attribute::attribute_alias;
 
 use crate::{
     id,
     resource::{Buffer, QuerySet, Texture},
 };
 
+pub trait ReferenceType {
+    type Buffer: Clone + core::fmt::Debug;
+    type Texture: Clone + core::fmt::Debug;
+    type TextureView: Clone + core::fmt::Debug;
+    type QuerySet: Clone + core::fmt::Debug;
+    type BindGroup: Clone + core::fmt::Debug;
+    type RenderPipeline: Clone + core::fmt::Debug;
+    type RenderBundle: Clone + core::fmt::Debug;
+    type ComputePipeline: Clone + core::fmt::Debug;
+    type Blas: Clone + core::fmt::Debug;
+    type Tlas: Clone + core::fmt::Debug;
+}
+
+#[derive(Clone, Debug)]
+pub struct IdReferences;
+
+#[derive(Clone, Debug)]
+pub struct ArcReferences;
+
+impl ReferenceType for IdReferences {
+    type Buffer = id::BufferId;
+    type Texture = id::TextureId;
+    type TextureView = id::TextureViewId;
+    type QuerySet = id::QuerySetId;
+    type BindGroup = id::BindGroupId;
+    type RenderPipeline = id::RenderPipelineId;
+    type RenderBundle = id::RenderBundleId;
+    type ComputePipeline = id::ComputePipelineId;
+    type Blas = id::BlasId;
+    type Tlas = id::TlasId;
+}
+
+impl ReferenceType for ArcReferences {
+    type Buffer = Arc<Buffer>;
+    type Texture = Arc<Texture>;
+    type TextureView = Arc<crate::resource::TextureView>;
+    type QuerySet = Arc<QuerySet>;
+    type BindGroup = Arc<crate::binding_model::BindGroup>;
+    type RenderPipeline = Arc<crate::pipeline::RenderPipeline>;
+    type RenderBundle = Arc<crate::command::RenderBundle>;
+    type ComputePipeline = Arc<crate::pipeline::ComputePipeline>;
+    type Blas = Arc<crate::resource::Blas>;
+    type Tlas = Arc<crate::resource::Tlas>;
+}
+
+#[cfg(feature = "serde")]
+attribute_alias! {
+    #[apply(serde_object_reference_struct)] =
+    #[derive(serde::Serialize, serde::Deserialize)]
+    #[serde(bound =
+         "R::Buffer: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::Texture: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::TextureView: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::QuerySet: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::BindGroup: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::RenderPipeline: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::RenderBundle: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::ComputePipeline: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::Blas: serde::Serialize + for<'d> serde::Deserialize<'d>,\
+          R::Tlas: serde::Serialize + for<'d> serde::Deserialize<'d>"
+    )];
+}
+
 #[derive(Clone, Debug)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum Command {
diff --git a/wgpu-core/src/command/mod.rs b/wgpu-core/src/command/mod.rs
@@ -34,13 +34,15 @@ use core::mem::{self, ManuallyDrop};
 use core::ops;
 
 pub(crate) use self::clear::clear_texture;
+#[cfg(feature = "serde")]
+pub(crate) use self::encoder_command::serde_object_reference_struct;
 pub use self::{
     bundle::*,
     clear::ClearError,
     compute::*,
     compute_command::{ArcComputeCommand, ComputeCommand},
     draw::*,
-    encoder_command::{ArcCommand, Command},
+    encoder_command::{ArcCommand, ArcReferences, Command, IdReferences, ReferenceType},
     query::*,
     render::*,
     render_command::{ArcRenderCommand, RenderCommand},
diff --git a/wgpu-core/src/command/ray_tracing.rs b/wgpu-core/src/command/ray_tracing.rs
@@ -8,32 +8,29 @@ use core::{
 use wgt::{math::align_to, BufferUsages, BufferUses, Features};
 
 use crate::{
-    command::CommandBufferMutable,
+    command::encoder::EncodingState,
+    ray_tracing::{AsAction, AsBuild, BlasTriangleGeometryInfo, TlasBuild, ValidateAsActionsError},
+    resource::InvalidResourceError,
+    track::Tracker,
+};
+use crate::{command::EncoderStateError, device::resource::CommandIndices};
+use crate::{
+    command::{ArcCommand, ArcReferences, CommandBufferMutable},
     device::queue::TempResource,
     global::Global,
     id::CommandEncoderId,
     init_tracker::MemoryInitKind,
     ray_tracing::{
-        BlasBuildEntry, BlasGeometries, BuildAccelerationStructureError, TlasPackage,
-        TraceBlasBuildEntry, TraceBlasGeometries, TraceBlasTriangleGeometry, TraceTlasInstance,
-        TraceTlasPackage,
+        ArcBlasBuildEntry, ArcBlasGeometries, ArcBlasTriangleGeometry, ArcTlasInstance,
+        ArcTlasPackage, BlasBuildEntry, BlasGeometries, BuildAccelerationStructureError,
+        OwnedBlasBuildEntry, OwnedTlasPackage, TlasPackage, TraceBlasBuildEntry,
+        TraceBlasGeometries, TraceBlasTriangleGeometry, TraceTlasInstance, TraceTlasPackage,
     },
     resource::{Blas, BlasCompactState, Buffer, Labeled, StagingBuffer, Tlas},
     scratch::ScratchBuffer,
     snatch::SnatchGuard,
     track::PendingTransition,
 };
-use crate::{command::EncoderStateError, device::resource::CommandIndices};
-use crate::{
-    command::{encoder::EncodingState, ArcCommand},
-    ray_tracing::{
-        ArcBlasBuildEntry, ArcBlasGeometries, ArcBlasTriangleGeometry, ArcTlasInstance,
-        ArcTlasPackage, AsAction, AsBuild, BlasTriangleGeometryInfo, TlasBuild,
-        ValidateAsActionsError,
-    },
-    resource::InvalidResourceError,
-    track::Tracker,
-};
 use crate::{lock::RwLockWriteGuard, resource::RawResourceAccess};
 
 use crate::id::{BlasId, TlasId};
@@ -146,7 +143,7 @@ impl Global {
                     }
                 };
                 TraceBlasBuildEntry {
-                    blas_id: blas_entry.blas_id,
+                    blas: blas_entry.blas_id,
                     geometries,
                 }
             })
@@ -158,15 +155,15 @@ impl Global {
                     .instances
                     .map(|instance| {
                         instance.map(|instance| TraceTlasInstance {
-                            blas_id: instance.blas_id,
+                            blas: instance.blas_id,
                             transform: *instance.transform,
                             custom_data: instance.custom_data,
                             mask: instance.mask,
                         })
                     })
                     .collect();
                 TraceTlasPackage {
-                    tlas_id: package.tlas_id,
+                    tlas: package.tlas_id,
                     instances,
                     lowest_unmodified: package.lowest_unmodified,
                 }
@@ -214,7 +211,7 @@ impl Global {
                         }
                     };
                     Ok(ArcBlasBuildEntry {
-                        blas: self.resolve_blas_id(blas_entry.blas_id)?,
+                        blas: self.resolve_blas_id(blas_entry.blas)?,
                         geometries,
                     })
                 })
@@ -231,7 +228,7 @@ impl Global {
                                 .as_ref()
                                 .map(|instance| {
                                     Ok(ArcTlasInstance {
-                                        blas: self.resolve_blas_id(instance.blas_id)?,
+                                        blas: self.resolve_blas_id(instance.blas)?,
                                         transform: instance.transform,
                                         custom_data: instance.custom_data,
                                         mask: instance.mask,
@@ -241,7 +238,7 @@ impl Global {
                         })
                         .collect::<Result<_, BuildAccelerationStructureError>>()?;
                     Ok(ArcTlasPackage {
-                        tlas: self.resolve_tlas_id(tlas_package.tlas_id)?,
+                        tlas: self.resolve_tlas_id(tlas_package.tlas)?,
                         instances,
                         lowest_unmodified: tlas_package.lowest_unmodified,
                     })
@@ -255,8 +252,8 @@ impl Global {
 
 pub(crate) fn build_acceleration_structures(
     state: &mut EncodingState,
-    blas: Vec<ArcBlasBuildEntry>,
-    tlas: Vec<ArcTlasPackage>,
+    blas: Vec<OwnedBlasBuildEntry<ArcReferences>>,
+    tlas: Vec<OwnedTlasPackage<ArcReferences>>,
 ) -> Result<(), BuildAccelerationStructureError> {
     state
         .device
@@ -281,7 +278,7 @@ pub(crate) fn build_acceleration_structures(
         &mut scratch_buffer_blas_size,
         &mut blas_storage,
     )?;
-    let mut tlas_lock_store = Vec::<(Option<ArcTlasPackage>, Arc<Tlas>)>::new();
+    let mut tlas_lock_store = Vec::<(Option<OwnedTlasPackage<ArcReferences>>, Arc<Tlas>)>::new();
 
     for package in tlas.into_iter() {
         let tlas = package.tlas.clone();
@@ -614,7 +611,7 @@ impl CommandBufferMutable {
 
 ///iterates over the blas iterator, and it's geometry, pushing the buffers into a storage vector (and also some validation).
 fn iter_blas(
-    blas_iter: impl Iterator<Item = ArcBlasBuildEntry>,
+    blas_iter: impl Iterator<Item = OwnedBlasBuildEntry<ArcReferences>>,
     tracker: &mut Tracker,
     build_command: &mut AsBuild,
     buf_storage: &mut Vec<TriangleBufferStore>,
diff --git a/wgpu-core/src/device/trace.rs b/wgpu-core/src/device/trace.rs
diff --git a/wgpu-core/src/ray_tracing.rs b/wgpu-core/src/ray_tracing.rs
diff --git a/wgpu-types/src/lib.rs b/wgpu-types/src/lib.rs
