mfplat/allocator: Improve initialization handling.

Signed-off-by: Nikolay Sivov <nsivov@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/mfplat/sample.c

@@ -76,6 +76,7 @@ struct sample_allocator
     } frame_desc;
 
     IMFAttributes *attributes;
+    IMFMediaType *media_type;
 
     unsigned int free_sample_count;
     unsigned int cold_sample_count;
@@ -1114,6 +1115,48 @@ static void sample_allocator_release_samples(struct sample_allocator *allocator)
         list_remove(&iter->entry);
         heap_free(iter);
     }
+
+    allocator->free_sample_count = 0;
+    allocator->cold_sample_count = 0;
+}
+
+static void sample_allocator_set_media_type(struct sample_allocator *allocator, IMFMediaType *media_type)
+{
+    UINT64 frame_size;
+    GUID subtype;
+
+    if (!media_type)
+    {
+        if (allocator->media_type)
+            IMFMediaType_Release(allocator->media_type);
+        allocator->media_type = NULL;
+        return;
+    }
+
+    /* Check if type is the same. */
+    IMFMediaType_GetUINT64(media_type, &MF_MT_FRAME_SIZE, &frame_size);
+    IMFMediaType_GetGUID(media_type, &MF_MT_SUBTYPE, &subtype);
+
+    if (frame_size == ((UINT64) allocator->frame_desc.width << 32 | allocator->frame_desc.height) &&
+            subtype.Data1 == allocator->frame_desc.d3d9_format)
+    {
+        return;
+    }
+
+    if (allocator->media_type)
+        IMFMediaType_Release(allocator->media_type);
+    allocator->media_type = media_type;
+    if (allocator->media_type)
+        IMFMediaType_AddRef(allocator->media_type);
+}
+
+static void sample_allocator_set_attributes(struct sample_allocator *allocator, IMFAttributes *attributes)
+{
+    if (allocator->attributes)
+        IMFAttributes_Release(allocator->attributes);
+    allocator->attributes = attributes;
+    if (allocator->attributes)
+        IMFAttributes_AddRef(allocator->attributes);
 }
 
 static ULONG WINAPI sample_allocator_Release(IMFVideoSampleAllocatorEx *iface)
@@ -1133,6 +1176,8 @@ static ULONG WINAPI sample_allocator_Release(IMFVideoSampleAllocatorEx *iface)
             IMFDXGIDeviceManager_Release(allocator->dxgi_device_manager);
         if (allocator->attributes)
             IMFAttributes_Release(allocator->attributes);
+        sample_allocator_set_media_type(allocator, NULL);
+        sample_allocator_set_attributes(allocator, NULL);
         sample_allocator_release_samples(allocator);
         DeleteCriticalSection(&allocator->cs);
         heap_free(allocator);
@@ -1190,7 +1235,9 @@ static HRESULT WINAPI sample_allocator_UninitializeSampleAllocator(IMFVideoSampl
     EnterCriticalSection(&allocator->cs);
 
     sample_allocator_release_samples(allocator);
-    allocator->free_sample_count = 0;
+    sample_allocator_set_media_type(allocator, NULL);
+    sample_allocator_set_attributes(allocator, NULL);
+    memset(&allocator->frame_desc, 0, sizeof(allocator->frame_desc));
 
     LeaveCriticalSection(&allocator->cs);
 
@@ -1348,16 +1395,14 @@ static HRESULT sample_allocator_initialize(struct sample_allocator *allocator, u
     if (sample_count > max_sample_count)
         return E_INVALIDARG;
 
+    sample_allocator_set_media_type(allocator, media_type);
+    sample_allocator_set_attributes(allocator, attributes);
+
     sample_count = max(1, sample_count);
     max_sample_count = max(1, max_sample_count);
 
     if (attributes)
-    {
-        allocator->attributes = attributes;
-        IMFAttributes_AddRef(allocator->attributes);
-
         IMFAttributes_GetUINT32(attributes, &MF_SA_BUFFERS_PER_SAMPLE, &allocator->frame_desc.buffer_count);
-    }
 
     allocator->frame_desc.d3d9_format = subtype.Data1;
     allocator->frame_desc.dxgi_format = MFMapDX9FormatToDXGIFormat(allocator->frame_desc.d3d9_format);

dlls/mfplat/tests/mfplat.c

@@ -6290,27 +6290,27 @@ static void test_dxgi_surface_buffer(void)
 static void test_sample_allocator(void)
 {
     IMFVideoSampleAllocatorNotify test_notify = { &test_notify_callback_vtbl };
+    IMFMediaType *media_type, *video_type, *video_type2;
     IMFVideoSampleAllocatorCallback *allocator_cb;
     IMFVideoSampleAllocatorEx *allocatorex;
-    IMFMediaType *media_type, *video_type;
+    IDirect3DDeviceManager9 *d3d9_manager;
     IMFVideoSampleAllocator *allocator;
+    unsigned int buffer_count, token;
+    IDirect3DDevice9 *d3d9_device;
     IMFDXGIDeviceManager *manager;
     IMFSample *sample, *sample2;
     IMFDXGIBuffer *dxgi_buffer;
     IMFAttributes *attributes;
     D3D11_TEXTURE2D_DESC desc;
-    unsigned int buffer_count, token;
     ID3D11Texture2D *texture;
     IMFMediaBuffer *buffer;
     ID3D11Device *device;
     LONG refcount, count;
+    IDirect3D9 *d3d9;
     IUnknown *unk;
     HRESULT hr;
     BYTE *data;
-    IDirect3D9 *d3d9;
     HWND window;
-    IDirect3DDeviceManager9 *d3d9_manager;
-    IDirect3DDevice9 *d3d9_device;
 
     if (!pMFCreateVideoSampleAllocatorEx)
     {
@@ -6366,6 +6366,7 @@ static void test_sample_allocator(void)
     ok(hr == MF_E_ATTRIBUTENOTFOUND, "Unexpected hr %#x.\n", hr);
 
     video_type = create_video_type(&MFVideoFormat_RGB32);
+    video_type2 = create_video_type(&MFVideoFormat_RGB32);
 
     hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 2, video_type);
     ok(hr == MF_E_ATTRIBUTENOTFOUND, "Unexpected hr %#x.\n", hr);
@@ -6374,11 +6375,42 @@ static void test_sample_allocator(void)
     hr = IMFMediaType_SetUINT64(video_type, &MF_MT_FRAME_SIZE, (UINT64) 320 << 32 | 240);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
 
+    hr = IMFMediaType_SetUINT64(video_type2, &MF_MT_FRAME_SIZE, (UINT64) 320 << 32 | 240);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
     hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 0, video_type);
     ok(hr == E_INVALIDARG, "Unexpected hr %#x.\n", hr);
 
+    EXPECT_REF(video_type, 1);
     hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 1, video_type);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    EXPECT_REF(video_type, 2);
+
+    hr = IMFMediaType_SetUINT64(video_type2, &IID_IUnknown, (UINT64) 320 << 32 | 240);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    /* Setting identical type does not replace it. */
+    hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 1, video_type2);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    EXPECT_REF(video_type, 2);
+    EXPECT_REF(video_type2, 1);
+
+    hr = IMFMediaType_SetUINT64(video_type2, &MF_MT_FRAME_SIZE, (UINT64) 64 << 32 | 64);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 1, video_type2);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    EXPECT_REF(video_type2, 2);
+    EXPECT_REF(video_type, 1);
+
+    /* Modify referenced type. */
+    hr = IMFMediaType_SetUINT64(video_type2, &MF_MT_FRAME_SIZE, (UINT64) 320 << 32 | 64);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 1, video_type);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    EXPECT_REF(video_type, 2);
+    EXPECT_REF(video_type2, 1);
 
     count = 0;
     hr = IMFVideoSampleAllocatorCallback_GetFreeSampleCount(allocator_cb, &count);
@@ -6401,6 +6433,7 @@ static void test_sample_allocator(void)
     hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 2, video_type);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
     ok(refcount == get_refcount(sample), "Unexpected refcount %u.\n", get_refcount(sample));
+    EXPECT_REF(video_type, 2);
 
     hr = IMFVideoSampleAllocatorCallback_GetFreeSampleCount(allocator_cb, &count);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
@@ -6426,6 +6459,18 @@ todo_wine
 
     IMFSample_Release(sample);
 
+    hr = IMFVideoSampleAllocator_UninitializeSampleAllocator(allocator);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+todo_wine
+    EXPECT_REF(video_type, 2);
+
+    hr = IMFVideoSampleAllocatorCallback_GetFreeSampleCount(allocator_cb, &count);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    ok(!count, "Unexpected count %d.\n", count);
+
+    hr = IMFVideoSampleAllocator_AllocateSample(allocator, &sample);
+    ok(hr == MF_E_NOT_INITIALIZED, "Unexpected hr %#x.\n", hr);
+
     IMFVideoSampleAllocatorCallback_Release(allocator_cb);
     IMFVideoSampleAllocator_Release(allocator);
 
@@ -6468,6 +6513,15 @@ todo_wine
     hr = IMFVideoSampleAllocatorEx_AllocateSample(allocatorex, &sample2);
     ok(hr == MF_E_SAMPLEALLOCATOR_EMPTY, "Unexpected hr %#x.\n", hr);
 
+    /* Reinitialize with already allocated samples. */
+    hr = IMFVideoSampleAllocatorEx_InitializeSampleAllocatorEx(allocatorex, 0, 0, NULL, video_type);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    EXPECT_REF(attributes, 1);
+
+    hr = IMFVideoSampleAllocatorEx_AllocateSample(allocatorex, &sample2);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    IMFSample_Release(sample2);
+
     IMFSample_Release(sample);
 
     IMFVideoSampleAllocatorCallback_Release(allocator_cb);
@@ -6490,8 +6544,10 @@ todo_wine
     hr = pMFCreateVideoSampleAllocatorEx(&IID_IMFVideoSampleAllocator, (void **)&allocator);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
 
+    EXPECT_REF(manager, 1);
     hr = IMFVideoSampleAllocator_SetDirectXManager(allocator, (IUnknown *)manager);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+    EXPECT_REF(manager, 2);
 
     hr = IMFVideoSampleAllocator_InitializeSampleAllocator(allocator, 0, video_type);
     ok(hr == E_INVALIDARG, "Unexpected hr %#x.\n", hr);