From 988e1a96de51d2c56c71ccec3c8accec66012d9d Mon Sep 17 00:00:00 2001
From: Sven Hesse <shesse@codeweavers.com>
Date: Fri, 25 Aug 2017 15:09:31 +0200
Subject: [PATCH] libs/vkd3d: Implement EndQuery() for QUERY_TYPE_TIMESTAMP.

---
 libs/vkd3d/command.c       | 42 ++++++++++++++++++++++++++++++++++++++
 libs/vkd3d/resource.c      |  8 ++++++++
 libs/vkd3d/vkd3d_private.h |  1 +
 3 files changed, 51 insertions(+)

diff --git a/libs/vkd3d/command.c b/libs/vkd3d/command.c
index b9552915..bd3125c5 100644
--- a/libs/vkd3d/command.c
+++ b/libs/vkd3d/command.c
@@ -3075,6 +3075,31 @@ static void STDMETHODCALLTYPE d3d12_command_list_DiscardResource(ID3D12GraphicsC
     FIXME("iface %p, resource %p, region %p stub!\n", iface, resource, region);
 }
 
+static D3D12_QUERY_HEAP_TYPE vkd3d_query_heap_type_from_query_type(D3D12_QUERY_TYPE type)
+{
+    switch (type)
+    {
+        case D3D12_QUERY_TYPE_OCCLUSION:
+        case D3D12_QUERY_TYPE_BINARY_OCCLUSION:
+            return D3D12_QUERY_HEAP_TYPE_OCCLUSION;
+
+        case D3D12_QUERY_TYPE_TIMESTAMP:
+            return D3D12_QUERY_HEAP_TYPE_TIMESTAMP;
+
+        case D3D12_QUERY_TYPE_PIPELINE_STATISTICS:
+            return D3D12_QUERY_HEAP_TYPE_PIPELINE_STATISTICS;
+
+        case D3D12_QUERY_TYPE_SO_STATISTICS_STREAM0:
+        case D3D12_QUERY_TYPE_SO_STATISTICS_STREAM1:
+        case D3D12_QUERY_TYPE_SO_STATISTICS_STREAM2:
+        case D3D12_QUERY_TYPE_SO_STATISTICS_STREAM3:
+            return D3D12_QUERY_HEAP_TYPE_SO_STATISTICS;
+
+        default:
+            return (D3D12_QUERY_HEAP_TYPE)-1;
+    }
+}
+
 static void STDMETHODCALLTYPE d3d12_command_list_BeginQuery(ID3D12GraphicsCommandList *iface,
         ID3D12QueryHeap *heap, D3D12_QUERY_TYPE type, UINT index)
 {
@@ -3084,6 +3109,23 @@ static void STDMETHODCALLTYPE d3d12_command_list_BeginQuery(ID3D12GraphicsComman
 static void STDMETHODCALLTYPE d3d12_command_list_EndQuery(ID3D12GraphicsCommandList *iface,
         ID3D12QueryHeap *heap, D3D12_QUERY_TYPE type, UINT index)
 {
+    struct d3d12_command_list *list = impl_from_ID3D12GraphicsCommandList(iface);
+    struct d3d12_query_heap *query_heap = unsafe_impl_from_ID3D12QueryHeap(heap);
+    const struct vkd3d_vk_device_procs *vk_procs = &list->device->vk_procs;
+
+    if (query_heap->desc.Type != vkd3d_query_heap_type_from_query_type(type))
+    {
+        WARN("Query type %u not supported by query heap.\n", type);
+        return;
+    }
+
+    if (type == D3D12_QUERY_TYPE_TIMESTAMP)
+    {
+        VK_CALL(vkCmdWriteTimestamp(list->vk_command_buffer, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
+                query_heap->vk_query_pool, index));
+        return;
+    }
+
     FIXME("iface %p, heap %p, type %#x, index %u stub!\n", iface, heap, type, index);
 }
 
diff --git a/libs/vkd3d/resource.c b/libs/vkd3d/resource.c
index f5cc99c9..1dba4ee0 100644
--- a/libs/vkd3d/resource.c
+++ b/libs/vkd3d/resource.c
@@ -1704,6 +1704,14 @@ static const struct ID3D12QueryHeapVtbl d3d12_query_heap_vtbl =
     d3d12_query_heap_GetDevice,
 };
 
+struct d3d12_query_heap *unsafe_impl_from_ID3D12QueryHeap(ID3D12QueryHeap *iface)
+{
+    if (!iface)
+        return NULL;
+    assert(iface->lpVtbl == &d3d12_query_heap_vtbl);
+    return impl_from_ID3D12QueryHeap(iface);
+}
+
 HRESULT d3d12_query_heap_create(struct d3d12_device *device, struct d3d12_query_heap **heap,
         const D3D12_QUERY_HEAP_DESC *desc)
 {
diff --git a/libs/vkd3d/vkd3d_private.h b/libs/vkd3d/vkd3d_private.h
index 1c1ea3f5..2d95debb 100644
--- a/libs/vkd3d/vkd3d_private.h
+++ b/libs/vkd3d/vkd3d_private.h
@@ -267,6 +267,7 @@ struct d3d12_query_heap
 
 HRESULT d3d12_query_heap_create(struct d3d12_device *device,
         struct d3d12_query_heap **heap, const D3D12_QUERY_HEAP_DESC *desc) DECLSPEC_HIDDEN;
+struct d3d12_query_heap *unsafe_impl_from_ID3D12QueryHeap(ID3D12QueryHeap *iface) DECLSPEC_HIDDEN;
 
 struct d3d12_root_descriptor_table_range
 {