From 4fa62d40dcde49557b306c731fb534673ec1036f Mon Sep 17 00:00:00 2001
From: Nikolay Sivov <nsivov@codeweavers.com>
Date: Thu, 8 Oct 2020 14:44:17 +0300
Subject: [PATCH] evr/allocator: Implement SetCallback().

Signed-off-by: Nikolay Sivov <nsivov@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>
---
 dlls/evr/sample.c    | 21 +++++++++++++++++--
 dlls/evr/tests/evr.c | 48 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/dlls/evr/sample.c b/dlls/evr/sample.c
index f371d920762..31e3ff43334 100644
--- a/dlls/evr/sample.c
+++ b/dlls/evr/sample.c
@@ -30,6 +30,9 @@ struct sample_allocator
     IMFVideoSampleAllocator IMFVideoSampleAllocator_iface;
     IMFVideoSampleAllocatorCallback IMFVideoSampleAllocatorCallback_iface;
     LONG refcount;
+
+    IMFVideoSampleAllocatorNotify *callback;
+    CRITICAL_SECTION cs;
 };
 
 static struct sample_allocator *impl_from_IMFVideoSampleAllocator(IMFVideoSampleAllocator *iface)
@@ -87,6 +90,9 @@ static ULONG WINAPI sample_allocator_Release(IMFVideoSampleAllocator *iface)
 
     if (!refcount)
     {
+        if (allocator->callback)
+            IMFVideoSampleAllocatorNotify_Release(allocator->callback);
+        DeleteCriticalSection(&allocator->cs);
         heap_free(allocator);
     }
 
@@ -156,9 +162,19 @@ static ULONG WINAPI sample_allocator_callback_Release(IMFVideoSampleAllocatorCal
 static HRESULT WINAPI sample_allocator_callback_SetCallback(IMFVideoSampleAllocatorCallback *iface,
         IMFVideoSampleAllocatorNotify *callback)
 {
-    FIXME("%p, %p.\n", iface, callback);
+    struct sample_allocator *allocator = impl_from_IMFVideoSampleAllocatorCallback(iface);
 
-    return E_NOTIMPL;
+    TRACE("%p, %p.\n", iface, callback);
+
+    EnterCriticalSection(&allocator->cs);
+    if (allocator->callback)
+        IMFVideoSampleAllocatorNotify_Release(allocator->callback);
+    allocator->callback = callback;
+    if (allocator->callback)
+        IMFVideoSampleAllocatorNotify_AddRef(allocator->callback);
+    LeaveCriticalSection(&allocator->cs);
+
+    return S_OK;
 }
 
 static HRESULT WINAPI sample_allocator_callback_GetFreeSampleCount(IMFVideoSampleAllocatorCallback *iface,
@@ -191,6 +207,7 @@ HRESULT WINAPI MFCreateVideoSampleAllocator(REFIID riid, void **obj)
     object->IMFVideoSampleAllocator_iface.lpVtbl = &sample_allocator_vtbl;
     object->IMFVideoSampleAllocatorCallback_iface.lpVtbl = &sample_allocator_callback_vtbl;
     object->refcount = 1;
+    InitializeCriticalSection(&object->cs);
 
     hr = IMFVideoSampleAllocator_QueryInterface(&object->IMFVideoSampleAllocator_iface, riid, obj);
     IMFVideoSampleAllocator_Release(&object->IMFVideoSampleAllocator_iface);
diff --git a/dlls/evr/tests/evr.c b/dlls/evr/tests/evr.c
index 3ae36d9b5df..3951b813a02 100644
--- a/dlls/evr/tests/evr.c
+++ b/dlls/evr/tests/evr.c
@@ -1147,8 +1147,47 @@ static void test_MFCreateVideoMixerAndPresenter(void)
     IUnknown_Release(presenter);
 }
 
+static HRESULT WINAPI test_notify_callback_QueryInterface(IMFVideoSampleAllocatorNotify *iface,
+        REFIID riid, void **obj)
+{
+    if (IsEqualIID(riid, &IID_IMFVideoSampleAllocatorNotify) ||
+            IsEqualIID(riid, &IID_IUnknown))
+    {
+        *obj = iface;
+        IMFVideoSampleAllocatorNotify_AddRef(iface);
+        return S_OK;
+    }
+
+    *obj = NULL;
+    return E_NOINTERFACE;
+}
+
+static ULONG WINAPI test_notify_callback_AddRef(IMFVideoSampleAllocatorNotify *iface)
+{
+    return 2;
+}
+
+static ULONG WINAPI test_notify_callback_Release(IMFVideoSampleAllocatorNotify *iface)
+{
+    return 1;
+}
+
+static HRESULT WINAPI test_notify_callback_NotifyRelease(IMFVideoSampleAllocatorNotify *iface)
+{
+    return E_NOTIMPL;
+}
+
+static const IMFVideoSampleAllocatorNotifyVtbl test_notify_callback_vtbl =
+{
+    test_notify_callback_QueryInterface,
+    test_notify_callback_AddRef,
+    test_notify_callback_Release,
+    test_notify_callback_NotifyRelease,
+};
+
 static void test_MFCreateVideoSampleAllocator(void)
 {
+    IMFVideoSampleAllocatorNotify test_notify = { &test_notify_callback_vtbl };
     IMFVideoSampleAllocatorCallback *allocator_cb;
     IMFVideoSampleAllocator *allocator;
     IMFVideoMediaType *video_type;
@@ -1171,6 +1210,15 @@ static void test_MFCreateVideoSampleAllocator(void)
     hr = IMFVideoSampleAllocator_QueryInterface(allocator, &IID_IMFVideoSampleAllocatorCallback, (void **)&allocator_cb);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
 
+    hr = IMFVideoSampleAllocatorCallback_SetCallback(allocator_cb, NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = IMFVideoSampleAllocatorCallback_SetCallback(allocator_cb, &test_notify);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = IMFVideoSampleAllocatorCallback_SetCallback(allocator_cb, NULL);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
     count = 10;
     hr = IMFVideoSampleAllocatorCallback_GetFreeSampleCount(allocator_cb, &count);
 todo_wine {