From 39568b59352627db2c0ec1235c71cfcba05141a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Michael=20M=C3=BCller?= <michael@fds-team.de>
Date: Fri, 9 Jan 2015 07:42:32 +0100
Subject: nvcuda: Properly wrap undocumented 'ContextStorage' interface and add
 tests.

---
 dlls/nvcuda/internal.c     | 103 ++++++++++++++-----
 dlls/nvcuda/tests/nvcuda.c | 247 +++++++++++++++++++++++++++++++++++++++++++--
 include/cuda.h             |   2 +
 3 files changed, 319 insertions(+), 33 deletions(-)

diff --git a/dlls/nvcuda/internal.c b/dlls/nvcuda/internal.c
index 611e310..8a2dffd 100644
--- a/dlls/nvcuda/internal.c
+++ b/dlls/nvcuda/internal.c
@@ -88,7 +88,7 @@ static const CUuuid UUID_Unknown2                   = {{0xA0, 0x94, 0x79, 0x8C,
                                                         0x93, 0xF2, 0x08, 0x00, 0x20, 0x0C, 0x0A, 0x66}};
 static const CUuuid UUID_Unknown3                   = {{0x42, 0xD8, 0x5A, 0x81, 0x23, 0xF6, 0xCB, 0x47,
                                                         0x82, 0x98, 0xF6, 0xE7, 0x8A, 0x3A, 0xEC, 0xDC}};
-static const CUuuid UUID_Unknown4                   = {{0xC6, 0x93, 0x33, 0x6E, 0x11, 0x21, 0xDF, 0x11,
+static const CUuuid UUID_ContextStorage             = {{0xC6, 0x93, 0x33, 0x6E, 0x11, 0x21, 0xDF, 0x11,
                                                         0xA8, 0xC3, 0x68, 0xF3, 0x55, 0xD8, 0x95, 0x93}};
 static const CUuuid UUID_Unknown5                   = {{0x0C, 0xA5, 0x0B, 0x8C, 0x10, 0x04, 0x92, 0x9A,
                                                         0x89, 0xA7, 0xD0, 0xDF, 0x10, 0xE7, 0x72, 0x86}};
@@ -164,20 +164,20 @@ static const struct
 } *Unknown3_orig = NULL;
 
 /*
- * Unknown4
+ * ContextStorage
  */
-struct Unknown4_table
+struct ContextStorage_table
 {
-    void* (WINAPI *func0)(void *param0, void *param1, void *param2, void *param3);
-    void* (WINAPI *func1)(void *param0, void *param1);
-    void* (WINAPI *func2)(void *param0, void *param1, void *param2);
+    CUresult (WINAPI *Set)(CUcontext ctx, void *key, void *value, void *callback);
+    CUresult (WINAPI *Remove)(CUcontext ctx, void *key);
+    CUresult (WINAPI *Get)(void **value, CUcontext ctx, void *key);
 };
 static const struct
 {
-    void* (*func0)(void *param0, void *param1, void *param2, void *param3);
-    void* (*func1)(void *param0, void *param1);
-    void* (*func2)(void *param0, void *param1, void *param2);
-} *Unknown4_orig = NULL;
+    CUresult (*Set)(CUcontext ctx, void *key, void *value, void *callback);
+    CUresult (*Remove)(CUcontext ctx, void *key);
+    CUresult (*Get)(void **value, CUcontext ctx, void *key);
+} *ContextStorage_orig = NULL;
 
 /*
  * TlsNotifyInterface
@@ -310,29 +310,80 @@ static struct Unknown3_table Unknown3_Impl =
     Unknown3_func1_relay,
 };
 
-static void* WINAPI Unknown4_func0_relay(void *param0, void *param1, void *param2, void *param3)
+struct context_storage
 {
-    TRACE("(%p, %p, %p, %p)\n", param0, param1, param2, param3);
-    return Unknown4_orig->func0(param0, param1, param2, param3);
+    void *value;
+    void (WINAPI *callback)(CUcontext ctx, void *key, void *value);
+};
+
+static void storage_destructor_callback(CUcontext ctx, void *key, void *value)
+{
+    struct context_storage *storage = value;
+
+    TRACE("(%p, %p, %p)\n", ctx, key, value);
+
+    if (storage->callback)
+    {
+        TRACE("calling destructor callback %p(%p, %p, %p)\n",
+              storage->callback, ctx, key, storage->value);
+        storage->callback(ctx, key, storage->value);
+        TRACE("destructor callback %p returned\n", storage->callback);
+    }
+
+    HeapFree( GetProcessHeap(), 0, storage );
 }
 
-static void* WINAPI Unknown4_func1_relay(void *param0, void *param1)
+static CUresult WINAPI ContextStorage_Set(CUcontext ctx, void *key, void *value, void *callback)
 {
-    TRACE("(%p, %p)\n", param0, param1);
-    return Unknown4_orig->func1(param0, param1);
+    struct context_storage *storage;
+    CUresult ret;
+
+    TRACE("(%p, %p, %p, %p)\n", ctx, key, value, callback);
+
+    storage = HeapAlloc( GetProcessHeap(), 0, sizeof(*storage) );
+    if (!storage)
+        return CUDA_ERROR_OUT_OF_MEMORY;
+
+    storage->callback = callback;
+    storage->value = value;
+
+    ret = ContextStorage_orig->Set(ctx, key, storage, storage_destructor_callback);
+    if (ret) HeapFree( GetProcessHeap(), 0, storage );
+    return ret;
 }
 
-static void* WINAPI Unknown4_func2_relay(void *param0, void *param1, void *param2)
+static CUresult WINAPI ContextStorage_Remove(CUcontext ctx, void *key)
 {
-    TRACE("(%p, %p, %p)\n", param0, param1, param2);
-    return Unknown4_orig->func2(param0, param1, param2);
+    struct context_storage *storage;
+
+    TRACE("(%p, %p)\n", ctx, key);
+
+    /* FIXME: This is not completely race-condition save, but using a mutex
+     * could have a relatively big overhead. Can still be added later when it
+     * turns out to be necessary. */
+    if (!ContextStorage_orig->Get((void **)&storage, ctx, key))
+        HeapFree( GetProcessHeap(), 0, storage );
+
+    return ContextStorage_orig->Remove(ctx, key);
+}
+
+static CUresult WINAPI ContextStorage_Get(void **value, CUcontext ctx, void *key)
+{
+    struct context_storage *storage;
+    CUresult ret;
+
+    TRACE("(%p, %p, %p)\n", value, ctx, key);
+
+    ret = ContextStorage_orig->Get((void **)&storage, ctx, key);
+    if (!ret) *value = storage->value;
+    return ret;
 }
 
-struct Unknown4_table Unknown4_Impl =
+struct ContextStorage_table ContextStorage_Impl =
 {
-    Unknown4_func0_relay,
-    Unknown4_func1_relay,
-    Unknown4_func2_relay,
+    ContextStorage_Set,
+    ContextStorage_Remove,
+    ContextStorage_Get,
 };
 
 static void* WINAPI Unknown5_func0_relay(void *param0, void *param1, void *param2)
@@ -480,15 +531,15 @@ CUresult cuda_get_table(const void **table, const CUuuid *uuid, const void *orig
         *table = (void *)&Unknown3_Impl;
         return CUDA_SUCCESS;
     }
-    else if (cuda_equal_uuid(uuid, &UUID_Unknown4))
+    else if (cuda_equal_uuid(uuid, &UUID_ContextStorage))
     {
         if (orig_result)
             return orig_result;
         if (!orig_table)
             return CUDA_ERROR_UNKNOWN;
 
-        Unknown4_orig = orig_table;
-        *table = (void *)&Unknown4_Impl;
+        ContextStorage_orig = orig_table;
+        *table = (void *)&ContextStorage_Impl;
         return CUDA_SUCCESS;
     }
     else if (cuda_equal_uuid(uuid, &UUID_Unknown5))
diff --git a/dlls/nvcuda/tests/nvcuda.c b/dlls/nvcuda/tests/nvcuda.c
index fc8f300..9c30c37 100644
--- a/dlls/nvcuda/tests/nvcuda.c
+++ b/dlls/nvcuda/tests/nvcuda.c
@@ -27,10 +27,24 @@
 
 static CUresult (WINAPI *pcuInit)(unsigned int);
 static CUresult (WINAPI *pcuGetExportTable)(const void**, const CUuuid*);
+static CUresult (WINAPI *pcuCtxCreate)(CUcontext *pctx, unsigned int flags, CUdevice dev);
+static CUresult (WINAPI *pcuCtxAttach)(CUcontext *pctx, unsigned int flags);
+static CUresult (WINAPI *pcuCtxDetach)(CUcontext ctx);
 
+static const CUuuid UUID_ContextStorage     = {{0xC6, 0x93, 0x33, 0x6E, 0x11, 0x21, 0xDF, 0x11,
+                                                0xA8, 0xC3, 0x68, 0xF3, 0x55, 0xD8, 0x95, 0x93}};
 static const CUuuid UUID_TlsNotifyInterface = {{0x19, 0x5B, 0xCB, 0xF4, 0xD6, 0x7D, 0x02, 0x4A,
                                                 0xAC, 0xC5, 0x1D, 0x29, 0xCE, 0xA6, 0x31, 0xAE}};
 
+#define CHECK_FUNCPTR(f) if(!p##f){win_skip("Failed to get entry point for %s\n", #f); return;}
+
+struct ContextStorage_table
+{
+    CUresult (WINAPI *Set)(CUcontext ctx, void *key, void *value, void *callback);
+    CUresult (WINAPI *Remove)(CUcontext ctx, void *key);
+    CUresult (WINAPI *Get)(void **value, CUcontext ctx, void *key);
+};
+
 struct TlsNotifyInterface_table
 {
     int size;
@@ -47,8 +61,13 @@ static BOOL init(void)
         return FALSE;
     }
 
-    pcuInit           = (void *)GetProcAddress(nvcuda, "cuInit");
-    pcuGetExportTable = (void *)GetProcAddress(nvcuda, "cuGetExportTable");
+    #define LOAD_FUNCPTR(f) p##f = (void*)GetProcAddress(nvcuda, #f)
+    LOAD_FUNCPTR(cuInit);
+    LOAD_FUNCPTR(cuGetExportTable);
+    LOAD_FUNCPTR(cuCtxCreate);
+    LOAD_FUNCPTR(cuCtxAttach);
+    LOAD_FUNCPTR(cuCtxDetach);
+    #undef LOAD_FUNCPTR
 
     if (!pcuInit)
     {
@@ -100,11 +119,7 @@ static void test_TlsNotifyInterface(void)
     HANDLE thread;
     CUresult res;
 
-    if (!pcuGetExportTable)
-    {
-        win_skip("cuGetExportTable export not found.\n");
-        return;
-    }
+    CHECK_FUNCPTR(cuGetExportTable);
 
     if (pcuGetExportTable((const void **)&iface, &UUID_TlsNotifyInterface))
     {
@@ -159,10 +174,228 @@ static void test_TlsNotifyInterface(void)
     }
 }
 
+struct storage_test_data
+{
+    int count;
+    CUcontext ctx;
+    void *key;
+};
+
+static void WINAPI storage_destructor_callback(CUcontext ctx, void *key, void *value)
+{
+    struct storage_test_data *test_data = value;
+    trace("(%p, %p, %p)\n", ctx, key, value);
+
+    test_data->count++;
+    test_data->ctx = ctx;
+    test_data->key = key;
+}
+
+#define STORAGE_KEY_1 (void *)0xdeadbeef
+#define STORAGE_KEY_2 (void *)0xcafebabe
+
+static void test_ContextStorage(void)
+{
+    const struct ContextStorage_table *iface;
+    struct storage_test_data test_data;
+    CUcontext ctx, ctx2;
+    CUresult res;
+    void *value;
+
+    CHECK_FUNCPTR(cuGetExportTable);
+    CHECK_FUNCPTR(cuCtxCreate);
+    CHECK_FUNCPTR(cuCtxAttach);
+    CHECK_FUNCPTR(cuCtxDetach);
+
+    if (pcuGetExportTable((const void **)&iface, &UUID_ContextStorage))
+    {
+        win_skip("Unknown4 interface not supported.\n");
+        return;
+    }
+
+    /* Call without current context */
+    res = iface->Set(NULL, STORAGE_KEY_1, (void *)42, NULL);
+    ok(res == CUDA_ERROR_INVALID_CONTEXT, "Expected CUDA_ERROR_INVALID_CONTEXT, got %d\n", res);
+
+    /* Create a context, then test Set/Get/Remove */
+    res = pcuCtxCreate(&ctx, 0, 0);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_2);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+    ok(value == (void *)0x55555555, "Value was modified\n");
+
+    res = iface->Set(NULL, STORAGE_KEY_1, (void *)42, NULL);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_2);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+    ok(value == (void *)0x55555555, "Value was modified\n");
+
+    res = iface->Set(NULL, STORAGE_KEY_2, (void *)43, NULL);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_2);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)43, "Unexpected value %p\n", value);
+
+    res = iface->Set(NULL, STORAGE_KEY_1, (void *)42, NULL);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+
+    res = iface->Set(NULL, STORAGE_KEY_1, (void *)44, NULL);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    res = iface->Remove(NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_1);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+    ok(value == (void *)0x55555555, "Value was modified\n");
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_2);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)43, "Unexpected value %p\n", value);
+
+    res = iface->Set(NULL, STORAGE_KEY_1, (void *)44, NULL);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)44, "Unexpected value %p\n", value);
+
+    res = pcuCtxDetach(ctx);
+    ok(!res, "Expected error code 0, got %d\n", res);
+
+    /* Now test with multiple contexts */
+    res = pcuCtxCreate(&ctx, 0, 0);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    res = pcuCtxCreate(&ctx2, 0, 0);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    res = iface->Set(NULL, STORAGE_KEY_1, (void *)42, NULL);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, ctx, STORAGE_KEY_1);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+    ok(value == (void *)0x55555555, "Value was modified\n");
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, ctx2, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    res = iface->Set(ctx, STORAGE_KEY_1, (void *)43, NULL);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, ctx, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)43, "Value was modified\n");
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, ctx2, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    res = iface->Remove(ctx, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, ctx, STORAGE_KEY_1);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+    ok(value == (void *)0x55555555, "Value was modified\n");
+
+    value = (void *)0x55555555;
+    res = iface->Get(&value, ctx2, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(value == (void *)42, "Unexpected value %p\n", value);
+
+    res = pcuCtxDetach(ctx2);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    res = pcuCtxDetach(ctx);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    /* Test with destructor callback */
+    res = pcuCtxCreate(&ctx, 0, 0);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    memset(&test_data, 0, sizeof(test_data));
+    res = iface->Set(NULL, STORAGE_KEY_1, &test_data, &storage_destructor_callback);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 0, "Expected count to be 0, got %d\n", test_data.count);
+
+    res = iface->Set(NULL, STORAGE_KEY_1, &test_data, NULL);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+
+    res = iface->Set(NULL, STORAGE_KEY_1, NULL, &storage_destructor_callback);
+    ok(res == CUDA_ERROR_INVALID_HANDLE, "Expected CUDA_ERROR_INVALID_HANDLE, got %d\n", res);
+
+    ctx2 = NULL;
+    res = pcuCtxAttach(&ctx2, 0);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 0, "Expected count to be 0, got %d\n", test_data.count);
+    ok(ctx2 == ctx, "Expected ctx2 and ctx to be identical\n");
+
+    res = pcuCtxDetach(ctx2);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 0, "Expected count to be 0, got %d\n", test_data.count);
+
+    res = pcuCtxDetach(ctx);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 1, "Expected count to be 1, got %d\n", test_data.count);
+    ok(test_data.ctx == ctx, "Expected ctx to be %p, got %p\n", ctx, test_data.ctx);
+    ok(test_data.key == STORAGE_KEY_1, "Expected key to be %p, got %p\n", STORAGE_KEY_1, test_data.key);
+
+    /* Test if destructor callback is called when removing an element */
+    res = pcuCtxCreate(&ctx, 0, 0);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+
+    memset(&test_data, 0, sizeof(test_data));
+    res = iface->Set(NULL, STORAGE_KEY_1, &test_data, &storage_destructor_callback);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 0, "Expected count to be 0, got %d\n", test_data.count);
+
+    res = iface->Remove(NULL, STORAGE_KEY_1);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 0, "Expected count to be 0, got %d\n", test_data.count);
+
+    res = pcuCtxDetach(ctx);
+    ok(!res, "Expected CUDA_SUCCESS, got %d\n", res);
+    ok(test_data.count == 0, "Expected count to be 0, got %d\n", test_data.count);
+}
+
 START_TEST( nvcuda )
 {
     if (!init())
         return;
 
     test_TlsNotifyInterface();
+    test_ContextStorage();
 }
diff --git a/include/cuda.h b/include/cuda.h
index 0f7b11d..327fe4d 100644
--- a/include/cuda.h
+++ b/include/cuda.h
@@ -22,6 +22,8 @@
 #define CUDA_SUCCESS                0
 #define CUDA_ERROR_INVALID_VALUE    1
 #define CUDA_ERROR_OUT_OF_MEMORY    2
+#define CUDA_ERROR_INVALID_CONTEXT  201
+#define CUDA_ERROR_INVALID_HANDLE   400
 #define CUDA_ERROR_NOT_SUPPORTED    801
 #define CUDA_ERROR_UNKNOWN          999
 
-- 
2.2.1

