From 687c1c73581284abd1fe13e2c1f5577a839090b5 Mon Sep 17 00:00:00 2001
From: Henri Verbeet <hverbeet@codeweavers.com>
Date: Thu, 13 Feb 2014 10:14:25 +0100
Subject: [PATCH] d3d10core: Implement d3d10_vertex_shader_GetDevice().

---
 dlls/d3d10core/d3d10core_private.h |  1 +
 dlls/d3d10core/shader.c            | 20 +++++++++++++++++++-
 dlls/d3d10core/tests/device.c      | 17 +++++++++++++----
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/dlls/d3d10core/d3d10core_private.h b/dlls/d3d10core/d3d10core_private.h
index 4439254f4bf..381ddf853d5 100644
--- a/dlls/d3d10core/d3d10core_private.h
+++ b/dlls/d3d10core/d3d10core_private.h
@@ -181,6 +181,7 @@ struct d3d10_vertex_shader
 
     struct wined3d_shader *wined3d_shader;
     struct wined3d_shader_signature output_signature;
+    ID3D10Device1 *device;
 };
 
 HRESULT d3d10_vertex_shader_init(struct d3d10_vertex_shader *shader, struct d3d10_device *device,
diff --git a/dlls/d3d10core/shader.c b/dlls/d3d10core/shader.c
index 4f189e9e09c..15e884c12f0 100644
--- a/dlls/d3d10core/shader.c
+++ b/dlls/d3d10core/shader.c
@@ -167,7 +167,10 @@ static ULONG STDMETHODCALLTYPE d3d10_vertex_shader_AddRef(ID3D10VertexShader *if
     TRACE("%p increasing refcount to %u\n", This, refcount);
 
     if (refcount == 1)
+    {
+        ID3D10Device1_AddRef(This->device);
         wined3d_shader_incref(This->wined3d_shader);
+    }
 
     return refcount;
 }
@@ -180,7 +183,14 @@ static ULONG STDMETHODCALLTYPE d3d10_vertex_shader_Release(ID3D10VertexShader *i
     TRACE("%p decreasing refcount to %u\n", This, refcount);
 
     if (!refcount)
+    {
+        ID3D10Device1 *device = This->device;
+
         wined3d_shader_decref(This->wined3d_shader);
+        /* Release the device last, it may cause the wined3d device to be
+         * destroyed. */
+        ID3D10Device1_Release(device);
+    }
 
     return refcount;
 }
@@ -189,7 +199,12 @@ static ULONG STDMETHODCALLTYPE d3d10_vertex_shader_Release(ID3D10VertexShader *i
 
 static void STDMETHODCALLTYPE d3d10_vertex_shader_GetDevice(ID3D10VertexShader *iface, ID3D10Device **device)
 {
-    FIXME("iface %p, device %p stub!\n", iface, device);
+    struct d3d10_vertex_shader *shader = impl_from_ID3D10VertexShader(iface);
+
+    TRACE("iface %p, device %p.\n", iface, device);
+
+    *device = (ID3D10Device *)shader->device;
+    ID3D10Device_AddRef(*device);
 }
 
 static HRESULT STDMETHODCALLTYPE d3d10_vertex_shader_GetPrivateData(ID3D10VertexShader *iface,
@@ -270,6 +285,9 @@ HRESULT d3d10_vertex_shader_init(struct d3d10_vertex_shader *shader, struct d3d1
         return hr;
     }
 
+    shader->device = &device->ID3D10Device1_iface;
+    ID3D10Device1_AddRef(shader->device);
+
     return S_OK;
 }
 
diff --git a/dlls/d3d10core/tests/device.c b/dlls/d3d10core/tests/device.c
index b2841807246..458e3f520a1 100644
--- a/dlls/d3d10core/tests/device.c
+++ b/dlls/d3d10core/tests/device.c
@@ -549,10 +549,10 @@ float4 main(const float4 color : COLOR) : SV_TARGET
         0x00000000, 0x00000000,
     };
 
+    ULONG refcount, expected_refcount;
     ID3D10VertexShader *vs = NULL;
     ID3D10PixelShader *ps = NULL;
-    ID3D10Device *device;
-    ULONG refcount;
+    ID3D10Device *device, *tmp;
     HRESULT hr;
 
     if (!(device = create_device()))
@@ -561,10 +561,19 @@ float4 main(const float4 color : COLOR) : SV_TARGET
         return;
     }
 
+    expected_refcount = get_refcount((IUnknown *)device) + 1;
     hr = ID3D10Device_CreateVertexShader(device, vs_4_0, sizeof(vs_4_0), &vs);
     ok(SUCCEEDED(hr), "Failed to create SM4 vertex shader, hr %#x\n", hr);
-    if (vs)
-        ID3D10VertexShader_Release(vs);
+    refcount = get_refcount((IUnknown *)device);
+    ok(refcount >= expected_refcount, "Got unexpected refcount %u, expected >= %u.\n", refcount, expected_refcount);
+    tmp = NULL;
+    expected_refcount = refcount + 1;
+    ID3D10VertexShader_GetDevice(vs, &tmp);
+    ok(tmp == device, "Got unexpected device %p, expected %p.\n", tmp, device);
+    refcount = get_refcount((IUnknown *)device);
+    ok(refcount == expected_refcount, "Got unexpected refcount %u, expected %u.\n", refcount, expected_refcount);
+    ID3D10Device_Release(tmp);
+    ID3D10VertexShader_Release(vs);
 
     hr = ID3D10Device_CreateVertexShader(device, vs_2_0, sizeof(vs_2_0), &vs);
     ok(hr == E_INVALIDARG, "Created a SM2 vertex shader, hr %#x\n", hr);