From 4072edc11111be5d26ca442862d021f5cd53e779 Mon Sep 17 00:00:00 2001
From: Christian Costa <titan.costa@wanadoo.fr>
Date: Mon, 1 Mar 2010 09:39:29 +0100
Subject: [PATCH] d3dx9_36: Find, store and enable retreival of CTAB comment
 data in shader (based on code from Luis Busquets).

---
 dlls/d3dx9_36/shader.c       | 30 ++++++++++++++++++++++----
 dlls/d3dx9_36/tests/shader.c | 42 ++++++++++++++++++++++++++++++++++++
 include/d3dx9shader.h        |  7 ++++++
 3 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/dlls/d3dx9_36/shader.c b/dlls/d3dx9_36/shader.c
index 22c297e86b9..5de3cadb08b 100644
--- a/dlls/d3dx9_36/shader.c
+++ b/dlls/d3dx9_36/shader.c
@@ -277,6 +277,8 @@ static const struct ID3DXConstantTableVtbl ID3DXConstantTable_Vtbl;
 typedef struct ID3DXConstantTableImpl {
     const ID3DXConstantTableVtbl *lpVtbl;
     LONG ref;
+    LPVOID ctab;
+    DWORD size;
 } ID3DXConstantTableImpl;
 
 /*** IUnknown methods ***/
@@ -316,7 +318,10 @@ static ULONG WINAPI ID3DXConstantTableImpl_Release(ID3DXConstantTable* iface)
     TRACE("(%p)->(): Release from %d\n", This, ref + 1);
 
     if (!ref)
+    {
+        HeapFree(GetProcessHeap(), 0, This->ctab);
         HeapFree(GetProcessHeap(), 0, This);
+    }
 
     return ref;
 }
@@ -326,18 +331,18 @@ static LPVOID WINAPI ID3DXConstantTableImpl_GetBufferPointer(ID3DXConstantTable*
 {
     ID3DXConstantTableImpl *This = (ID3DXConstantTableImpl *)iface;
 
-    FIXME("(%p)->(): stub\n", This);
+    TRACE("(%p)->()\n", This);
 
-    return NULL;
+    return This->ctab;
 }
 
 static DWORD WINAPI ID3DXConstantTableImpl_GetBufferSize(ID3DXConstantTable* iface)
 {
     ID3DXConstantTableImpl *This = (ID3DXConstantTableImpl *)iface;
 
-    FIXME("(%p)->(): stub\n", This);
+    TRACE("(%p)->()\n", This);
 
-    return 0;
+    return This->size;
 }
 
 /*** ID3DXConstantTable methods ***/
@@ -583,12 +588,19 @@ HRESULT WINAPI D3DXGetShaderConstantTableEx(CONST DWORD* pFunction,
                                             LPD3DXCONSTANTTABLE* ppConstantTable)
 {
     ID3DXConstantTableImpl* object;
+    HRESULT hr;
+    LPCVOID data;
+    UINT size;
 
     FIXME("(%p, %x, %p): semi-stub\n", pFunction, flags, ppConstantTable);
 
     if (!pFunction || !ppConstantTable)
         return D3DERR_INVALIDCALL;
 
+    hr = D3DXFindShaderComment(pFunction, MAKEFOURCC('C','T','A','B'), &data, &size);
+    if (hr != D3D_OK)
+        return D3DXERR_INVALIDDATA;
+
     object = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(ID3DXConstantTableImpl));
     if (!object)
     {
@@ -599,6 +611,16 @@ HRESULT WINAPI D3DXGetShaderConstantTableEx(CONST DWORD* pFunction,
     object->lpVtbl = &ID3DXConstantTable_Vtbl;
     object->ref = 1;
 
+    object->ctab = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, size);
+    if (!object->ctab)
+    {
+        HeapFree(GetProcessHeap(), 0, object);
+        ERR("Out of memory\n");
+        return E_OUTOFMEMORY;
+    }
+    object->size = size;
+    memcpy(object->ctab, data, object->size);
+
     *ppConstantTable = (LPD3DXCONSTANTTABLE)object;
 
     return D3D_OK;
diff --git a/dlls/d3dx9_36/tests/shader.c b/dlls/d3dx9_36/tests/shader.c
index 03c77c2c6e4..8a2a5b1d809 100644
--- a/dlls/d3dx9_36/tests/shader.c
+++ b/dlls/d3dx9_36/tests/shader.c
@@ -48,6 +48,13 @@ static const DWORD shader_with_ctab[] = {
     0x0004fffe, FCC_TEXT,   0x00000000, 0x00000000, 0x00000000,             /* TEXT comment                 */
     0x0000ffff};                                                            /* END                          */
 
+static const DWORD shader_with_invalid_ctab[] = {
+    0xfffe0300,                                                             /* vs_3_0                       */
+    0x0005fffe, FCC_CTAB,                                                   /* CTAB comment                 */
+                0x0000001c, 0x000000a9, 0xfffe0300,
+                0x00000000, 0x00000000,
+    0x0000ffff};                                                            /* END                          */
+
 static void test_get_shader_size(void)
 {
     UINT shader_size, expected;
@@ -107,9 +114,44 @@ static void test_find_shader_comment(void)
     ok(size == 20, "Got result %d, expected 20\n", size);
 }
 
+static void test_get_shader_constant_table_ex(void)
+{
+    LPD3DXCONSTANTTABLE constant_table = NULL;
+    HRESULT hr;
+    LPVOID data;
+    DWORD size;
+
+    hr = D3DXGetShaderConstantTableEx(NULL, 0, &constant_table);
+    ok(hr == D3DERR_INVALIDCALL, "Got result %x, expected %x (D3DERR_INVALIDCALL)\n", hr, D3DERR_INVALIDCALL);
+
+    /* No CTAB data */
+    hr = D3DXGetShaderConstantTableEx(simple_ps, 0, &constant_table);
+    ok(hr == D3DXERR_INVALIDDATA, "Got result %x, expected %x (D3DERR_INVALIDCALL)\n", hr, D3DXERR_INVALIDDATA);
+
+    /* With invalid CTAB data */
+    hr = D3DXGetShaderConstantTableEx(shader_with_invalid_ctab, 0, &constant_table);
+    todo_wine ok(hr == D3DXERR_INVALIDDATA, "Got result %x, expected %x (D3DXERR_INVALIDDATA)\n", hr, D3DXERR_INVALIDDATA);
+    if (constant_table) ID3DXConstantTable_Release(constant_table);
+
+    hr = D3DXGetShaderConstantTableEx(shader_with_ctab, 0, &constant_table);
+    ok(hr == D3D_OK, "Got result %x, expected 0 (D3D_OK)\n", hr);
+
+    if (constant_table)
+    {
+        size = ID3DXConstantTable_GetBufferSize(constant_table);
+        ok(size == 20, "Got result %x, expected 20\n", size);
+
+        data = ID3DXConstantTable_GetBufferPointer(constant_table);
+        ok(!memcmp(data, shader_with_ctab + 6, size), "Retreived wrong CTAB data\n");
+
+        ID3DXConstantTable_Release(constant_table);
+    }
+}
+
 START_TEST(shader)
 {
     test_get_shader_size();
     test_get_shader_version();
     test_find_shader_comment();
+    test_get_shader_constant_table_ex();
 }
diff --git a/include/d3dx9shader.h b/include/d3dx9shader.h
index 6d11b3d198b..766d4ca1089 100644
--- a/include/d3dx9shader.h
+++ b/include/d3dx9shader.h
@@ -288,6 +288,13 @@ HRESULT WINAPI D3DXAssembleShader(LPCSTR data,
                                   LPD3DXBUFFER* shader,
                                   LPD3DXBUFFER* error_messages);
 
+HRESULT WINAPI D3DXGetShaderConstantTableEx(CONST DWORD* byte_code,
+                                            DWORD flags,
+                                            LPD3DXCONSTANTTABLE* constant_table);
+
+HRESULT WINAPI D3DXGetShaderConstantTable(CONST DWORD* byte_code,
+                                          LPD3DXCONSTANTTABLE* constant_table);
+
 #ifdef __cplusplus
 }
 #endif