From a11c3419ce9b3bf06989b74957366d5e977133ea Mon Sep 17 00:00:00 2001
From: Zack Gomez <zack.gomez@gmail.com>
Date: Thu, 10 Oct 2024 16:49:40 -0400
Subject: [PATCH] flesh out timestamp query support

---
 ffi/wgpu.h  |  7 +++++--
 src/conv.rs | 14 ++++++++------
 src/lib.rs  | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/ffi/wgpu.h b/ffi/wgpu.h
index 15197a4d..12c1a865 100644
--- a/ffi/wgpu.h
+++ b/ffi/wgpu.h
@@ -31,8 +31,6 @@ typedef enum WGPUNativeFeature {
     WGPUNativeFeature_PartiallyBoundBindingArray = 0x0003000A,
     WGPUNativeFeature_TextureFormat16bitNorm = 0x0003000B,
     WGPUNativeFeature_TextureCompressionAstcHdr = 0x0003000C,
-    // TODO: requires wgpu.h api change
-    // WGPUNativeFeature_TimestampQueryInsidePasses = 0x0003000D,
     WGPUNativeFeature_MappablePrimaryBuffers = 0x0003000E,
     WGPUNativeFeature_BufferBindingArray = 0x0003000F,
     WGPUNativeFeature_UniformBufferAndStorageTextureArrayNonUniformIndexing = 0x00030010,
@@ -56,6 +54,8 @@ typedef enum WGPUNativeFeature {
     WGPUNativeFeature_Subgroup = 0x00030021,
     WGPUNativeFeature_SubgroupVertex = 0x00030022,
     WGPUNativeFeature_SubgroupBarrier = 0x00030023,
+    WGPUNativeFeature_TimestampQueryInsideEncoders = 0x00030024,
+    WGPUNativeFeature_TimestampQueryInsidePasses = 0x00030025,
     WGPUNativeFeature_Force32 = 0x7FFFFFFF
 } WGPUNativeFeature;
 
@@ -297,6 +297,9 @@ void wgpuComputePassEncoderEndPipelineStatisticsQuery(WGPUComputePassEncoder com
 void wgpuRenderPassEncoderBeginPipelineStatisticsQuery(WGPURenderPassEncoder renderPassEncoder, WGPUQuerySet querySet, uint32_t queryIndex);
 void wgpuRenderPassEncoderEndPipelineStatisticsQuery(WGPURenderPassEncoder renderPassEncoder);
 
+void wgpuComputePassEncoderWriteTimestamp(WGPUComputePassEncoder computePassEncoder, WGPUQuerySet querySet, uint32_t queryIndex);
+void wgpuRenderPassEncoderWriteTimestamp(WGPURenderPassEncoder renderPassEncoder, WGPUQuerySet querySet, uint32_t queryIndex);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif
diff --git a/src/conv.rs b/src/conv.rs
index 53a5b8b6..6ca97022 100644
--- a/src/conv.rs
+++ b/src/conv.rs
@@ -1153,10 +1153,12 @@ pub fn features_to_native(features: wgt::Features) -> Vec<native::WGPUFeatureNam
     if features.contains(wgt::Features::TEXTURE_COMPRESSION_ASTC_HDR) {
         temp.push(native::WGPUNativeFeature_TextureCompressionAstcHdr);
     }
-    // TODO: requires wgpu.h api change
-    // if features.contains(wgt::Features::TIMESTAMP_QUERY_INSIDE_PASSES) {
-    //     temp.push(native::WGPUNativeFeature_TimestampQueryInsidePasses);
-    // }
+    if features.contains(wgt::Features::TIMESTAMP_QUERY_INSIDE_PASSES) {
+        temp.push(native::WGPUNativeFeature_TimestampQueryInsidePasses);
+    }
+    if features.contains(wgt::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS) {
+        temp.push(native::WGPUNativeFeature_TimestampQueryInsideEncoders);
+    }
     if features.contains(wgt::Features::MAPPABLE_PRIMARY_BUFFERS) {
         temp.push(native::WGPUNativeFeature_MappablePrimaryBuffers);
     }
@@ -1261,8 +1263,8 @@ pub fn map_feature(feature: native::WGPUFeatureName) -> Option<wgt::Features> {
         native::WGPUNativeFeature_PartiallyBoundBindingArray => Some(Features::PARTIALLY_BOUND_BINDING_ARRAY),
         native::WGPUNativeFeature_TextureFormat16bitNorm => Some(Features::TEXTURE_FORMAT_16BIT_NORM),
         native::WGPUNativeFeature_TextureCompressionAstcHdr => Some(Features::TEXTURE_COMPRESSION_ASTC_HDR),
-        // TODO: requires wgpu.h api change
-        // native::WGPUNativeFeature_TimestampQueryInsidePasses => Some(Features::TIMESTAMP_QUERY_INSIDE_PASSES),
+        native::WGPUNativeFeature_TimestampQueryInsidePasses => Some(Features::TIMESTAMP_QUERY_INSIDE_PASSES),
+        native::WGPUNativeFeature_TimestampQueryInsideEncoders => Some(Features::TIMESTAMP_QUERY_INSIDE_ENCODERS),
         native::WGPUNativeFeature_MappablePrimaryBuffers => Some(Features::MAPPABLE_PRIMARY_BUFFERS),
         native::WGPUNativeFeature_BufferBindingArray => Some(Features::BUFFER_BINDING_ARRAY),
         native::WGPUNativeFeature_UniformBufferAndStorageTextureArrayNonUniformIndexing => Some(Features::UNIFORM_BUFFER_AND_STORAGE_TEXTURE_ARRAY_NON_UNIFORM_INDEXING),
diff --git a/src/lib.rs b/src/lib.rs
index 7aae51a2..7d22084f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4442,3 +4442,45 @@ pub unsafe extern "C" fn wgpuRenderPassEncoderEndPipelineStatisticsQuery(
         ),
     }
 }
+
+#[no_mangle]
+pub unsafe extern "C" fn wgpuComputePassEncoderWriteTimestamp(
+    pass: native::WGPUComputePassEncoder,
+    query_set: native::WGPUQuerySet,
+    query_index: u32,
+) {
+    let pass = pass.as_ref().expect("invalid compute pass");
+    let query_set_id = query_set.as_ref().expect("invalid query set").id;
+    let encoder = pass.encoder.as_mut().unwrap();
+
+    match encoder.write_timestamp(&pass.context, query_set_id, query_index) {
+        Ok(()) => (),
+        Err(cause) => handle_error(
+            &pass.error_sink,
+            cause,
+            None,
+            "wgpuComputePassEncoderWriteTimestamp",
+        ),
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn wgpuRenderPassEncoderWriteTimestamp(
+    pass: native::WGPURenderPassEncoder,
+    query_set: native::WGPUQuerySet,
+    query_index: u32,
+) {
+    let pass = pass.as_ref().expect("invalid render pass");
+    let query_set_id = query_set.as_ref().expect("invalid query set").id;
+    let encoder = pass.encoder.as_mut().unwrap();
+
+    match encoder.write_timestamp(&pass.context, query_set_id, query_index) {
+        Ok(()) => (),
+        Err(cause) => handle_error(
+            &pass.error_sink,
+            cause,
+            None,
+            "wgpuRenderPassEncoderWriteTimestamp",
+        ),
+    }
+}