From 8ee9df9eb39c3c8114e23cc07218388a1b12558b Mon Sep 17 00:00:00 2001
From: teoxoy <28601907+teoxoy@users.noreply.github.com>
Date: Tue, 25 Jun 2024 14:45:31 +0200
Subject: [PATCH] move same device check and tracker insertion inside
 `validate_and_begin_pipeline_statistics_query`

---
 wgpu-core/src/command/compute.rs |  9 +++------
 wgpu-core/src/command/query.rs   | 10 +++++++++-
 wgpu-core/src/command/render.rs  |  6 +++---
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/wgpu-core/src/command/compute.rs b/wgpu-core/src/command/compute.rs
index d4da468c0..428783105 100644
--- a/wgpu-core/src/command/compute.rs
+++ b/wgpu-core/src/command/compute.rs
@@ -662,14 +662,11 @@ impl Global {
                     query_index,
                 } => {
                     let scope = PassErrorScope::BeginPipelineStatisticsQuery;
-
-                    query_set.same_device_as(cmd_buf).map_pass_err(scope)?;
-
-                    let query_set = state.tracker.query_sets.insert_single(query_set);
-
                     validate_and_begin_pipeline_statistics_query(
-                        query_set.clone(),
+                        query_set,
                         raw,
+                        &mut state.tracker.query_sets,
+                        cmd_buf,
                         query_index,
                         None,
                         &mut state.active_query,
diff --git a/wgpu-core/src/command/query.rs b/wgpu-core/src/command/query.rs
index 16557a4b4..1da5badb7 100644
--- a/wgpu-core/src/command/query.rs
+++ b/wgpu-core/src/command/query.rs
@@ -10,7 +10,7 @@ use crate::{
     id,
     init_tracker::MemoryInitKind,
     resource::{DestroyedResourceError, ParentDevice, QuerySet},
-    track::TrackerIndex,
+    track::{StatelessTracker, TrackerIndex},
     FastHashMap,
 };
 use std::{iter, marker::PhantomData, sync::Arc};
@@ -124,6 +124,8 @@ impl crate::error::PrettyError for QueryError {
 #[derive(Clone, Debug, Error)]
 #[non_exhaustive]
 pub enum QueryUseError {
+    #[error(transparent)]
+    Device(#[from] DeviceError),
     #[error("Query {query_index} is out of bounds for a query set of size {query_set_size}")]
     OutOfBounds {
         query_index: u32,
@@ -269,10 +271,14 @@ pub(super) fn end_occlusion_query<A: HalApi>(
 pub(super) fn validate_and_begin_pipeline_statistics_query<A: HalApi>(
     query_set: Arc<QuerySet<A>>,
     raw_encoder: &mut A::CommandEncoder,
+    tracker: &mut StatelessTracker<QuerySet<A>>,
+    cmd_buf: &CommandBuffer<A>,
     query_index: u32,
     reset_state: Option<&mut QueryResetMap<A>>,
     active_query: &mut Option<(Arc<QuerySet<A>>, u32)>,
 ) -> Result<(), QueryUseError> {
+    query_set.same_device_as(cmd_buf)?;
+
     let needs_reset = reset_state.is_none();
     query_set.validate_query(
         SimplifiedQueryType::PipelineStatistics,
@@ -280,6 +286,8 @@ pub(super) fn validate_and_begin_pipeline_statistics_query<A: HalApi>(
         reset_state,
     )?;
 
+    tracker.add_single(&query_set);
+
     if let Some((_old, old_idx)) = active_query.take() {
         return Err(QueryUseError::AlreadyStarted {
             active_query_index: old_idx,
diff --git a/wgpu-core/src/command/render.rs b/wgpu-core/src/command/render.rs
index e2721e71d..85f99faf6 100644
--- a/wgpu-core/src/command/render.rs
+++ b/wgpu-core/src/command/render.rs
@@ -2324,11 +2324,11 @@ impl Global {
                         );
                         let scope = PassErrorScope::BeginPipelineStatisticsQuery;
 
-                        let query_set = tracker.query_sets.insert_single(query_set);
-
                         validate_and_begin_pipeline_statistics_query(
-                            query_set.clone(),
+                            query_set,
                             raw,
+                            &mut tracker.query_sets,
+                            cmd_buf.as_ref(),
                             query_index,
                             Some(&mut cmd_buf_data.pending_query_resets),
                             &mut active_query,