mfreadwrite/reader: Create sample allocator for video streams.

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

dlls/mfreadwrite/reader.c

@@ -33,6 +33,9 @@
 #include "mferror.h"
 #include "mfidl.h"
 #include "mfreadwrite.h"
+#include "d3d9.h"
+#include "initguid.h"
+#include "dxva2api.h"
 
 #include "wine/debug.h"
 #include "wine/heap.h"
@@ -108,6 +111,7 @@ struct media_stream
     IMFMediaStream *stream;
     IMFMediaType *current;
     IMFTransform *decoder;
+    IMFVideoSampleAllocatorEx *allocator;
     unsigned int id;
     unsigned int index;
     enum media_stream_state state;
@@ -157,6 +161,9 @@ enum source_reader_flags
     SOURCE_READER_FLUSHING = 0x1,
     SOURCE_READER_SEEKING = 0x2,
     SOURCE_READER_SHUTDOWN_ON_RELEASE = 0x4,
+    SOURCE_READER_D3D9_DEVICE_MANAGER = 0x8,
+    SOURCE_READER_DXGI_DEVICE_MANAGER = 0x10,
+    SOURCE_READER_HAS_DEVICE_MANAGER = SOURCE_READER_D3D9_DEVICE_MANAGER | SOURCE_READER_DXGI_DEVICE_MANAGER,
 };
 
 struct source_reader
@@ -170,6 +177,7 @@ struct source_reader
     IMFPresentationDescriptor *descriptor;
     IMFSourceReaderCallback *async_callback;
     IMFAttributes *attributes;
+    IUnknown *device_manager;
     unsigned int first_audio_stream_index;
     unsigned int first_video_stream_index;
     unsigned int last_read_index;
@@ -1283,6 +1291,8 @@ static ULONG WINAPI src_reader_Release(IMFSourceReader *iface)
                 IMFMediaType_Release(stream->current);
             if (stream->decoder)
                 IMFTransform_Release(stream->decoder);
+            if (stream->allocator)
+                IMFVideoSampleAllocatorEx_Release(stream->allocator);
         }
         source_reader_release_responses(reader, NULL);
         heap_free(reader->streams);
@@ -1519,6 +1529,48 @@ static HRESULT source_reader_set_compatible_media_type(struct source_reader *rea
     return type_set ? S_OK : S_FALSE;
 }
 
+static HRESULT source_reader_setup_sample_allocator(struct source_reader *reader, unsigned int index,
+        IMFMediaType *media_type)
+{
+    struct media_stream *stream = &reader->streams[index];
+    GUID major = { 0 };
+    HRESULT hr;
+
+    IMFMediaType_GetMajorType(media_type, &major);
+    if (!IsEqualGUID(&major, &MFMediaType_Video))
+        return S_OK;
+
+    if (!(reader->flags & SOURCE_READER_HAS_DEVICE_MANAGER))
+        return S_OK;
+
+    if (reader->flags & SOURCE_READER_DXGI_DEVICE_MANAGER)
+    {
+        FIXME("DXGI device manager is not supported.\n");
+        return S_OK;
+    }
+
+    if (!stream->allocator)
+    {
+        if (FAILED(hr = MFCreateVideoSampleAllocatorEx(&IID_IMFVideoSampleAllocatorEx, (void **)&stream->allocator)))
+        {
+            WARN("Failed to create sample allocator, hr %#x.\n", hr);
+            return hr;
+        }
+    }
+
+    IMFVideoSampleAllocatorEx_UninitializeSampleAllocator(stream->allocator);
+    if (FAILED(hr = IMFVideoSampleAllocatorEx_SetDirectXManager(stream->allocator, reader->device_manager)))
+    {
+        WARN("Failed to set device manager, hr %#x.\n", hr);
+        return hr;
+    }
+
+    if (FAILED(hr = IMFVideoSampleAllocatorEx_InitializeSampleAllocatorEx(stream->allocator, 2, 8, NULL, media_type)))
+        WARN("Failed to initialize sample allocator, hr %#x.\n", hr);
+
+    return hr;
+}
+
 static HRESULT source_reader_configure_decoder(struct source_reader *reader, DWORD index, const CLSID *clsid,
         IMFMediaType *input_type, IMFMediaType *output_type)
 {
@@ -1668,8 +1720,11 @@ static HRESULT WINAPI src_reader_SetCurrentMediaType(IMFSourceReader *iface, DWO
 
     EnterCriticalSection(&reader->cs);
 
-    if ((hr = source_reader_set_compatible_media_type(reader, index, type)) == S_FALSE)
+    hr = source_reader_set_compatible_media_type(reader, index, type);
+    if (hr == S_FALSE)
         hr = source_reader_create_decoder_for_stream(reader, index, type);
+    if (SUCCEEDED(hr))
+        hr = source_reader_setup_sample_allocator(reader, index, type);
 
     LeaveCriticalSection(&reader->cs);
 
@@ -2172,6 +2227,27 @@ static HRESULT create_source_reader_from_source(IMFMediaSource *source, IMFAttri
                 (void **)&object->async_callback);
         if (object->async_callback)
             TRACE("Using async callback %p.\n", object->async_callback);
+
+        IMFAttributes_GetUnknown(attributes, &MF_SOURCE_READER_D3D_MANAGER, &IID_IUnknown, (void **)&object->device_manager);
+        if (object->device_manager)
+        {
+            IUnknown *unk = NULL;
+
+            if (SUCCEEDED(IUnknown_QueryInterface(object->device_manager, &IID_IMFDXGIDeviceManager, (void **)&unk)))
+                object->flags |= SOURCE_READER_DXGI_DEVICE_MANAGER;
+            else if (SUCCEEDED(IUnknown_QueryInterface(object->device_manager, &IID_IDirect3DDeviceManager9, (void **)&unk)))
+                object->flags |= SOURCE_READER_D3D9_DEVICE_MANAGER;
+
+            if (!(object->flags & (SOURCE_READER_HAS_DEVICE_MANAGER)))
+            {
+                WARN("Unknown device manager.\n");
+                IUnknown_Release(object->device_manager);
+                object->device_manager = NULL;
+            }
+
+            if (unk)
+                IUnknown_Release(unk);
+        }
     }
 
     hr = IMFSourceReader_QueryInterface(&object->IMFSourceReader_iface, riid, out);

dlls/mfreadwrite/tests/Makefile.in

@@ -1,5 +1,5 @@
 TESTDLL   = mfreadwrite.dll
-IMPORTS   = ole32 mfplat mfreadwrite mfuuid
+IMPORTS   = ole32 user32 d3d9 dxva2 mfplat mfreadwrite mfuuid
 
 C_SRCS = \
 	mfplat.c

dlls/mfreadwrite/tests/mfplat.c

@@ -37,6 +37,8 @@ DEFINE_GUID(GUID_NULL, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
 #include "mfidl.h"
 #include "mferror.h"
 #include "mfreadwrite.h"
+#include "d3d9.h"
+#include "dxva2api.h"
 
 #include "wine/heap.h"
 #include "wine/test.h"
@@ -48,6 +50,37 @@ static ULONG get_refcount(void *iface)
     return IUnknown_Release(unknown);
 }
 
+static HWND create_window(void)
+{
+    RECT r = {0, 0, 640, 480};
+
+    AdjustWindowRect(&r, WS_OVERLAPPEDWINDOW | WS_VISIBLE, FALSE);
+
+    return CreateWindowA("static", "mfreadwrite_test", WS_OVERLAPPEDWINDOW | WS_VISIBLE,
+            0, 0, r.right - r.left, r.bottom - r.top, NULL, NULL, NULL, NULL);
+}
+
+static IDirect3DDevice9 *create_d3d9_device(IDirect3D9 *d3d9, HWND focus_window)
+{
+    D3DPRESENT_PARAMETERS present_parameters = {0};
+    IDirect3DDevice9 *device = NULL;
+
+    present_parameters.BackBufferWidth = 640;
+    present_parameters.BackBufferHeight = 480;
+    present_parameters.BackBufferFormat = D3DFMT_X8R8G8B8;
+    present_parameters.SwapEffect = D3DSWAPEFFECT_DISCARD;
+    present_parameters.hDeviceWindow = focus_window;
+    present_parameters.Windowed = TRUE;
+    present_parameters.EnableAutoDepthStencil = TRUE;
+    present_parameters.AutoDepthStencilFormat = D3DFMT_D24S8;
+    present_parameters.Flags = D3DPRESENTFLAG_LOCKABLE_BACKBUFFER;
+
+    IDirect3D9_CreateDevice(d3d9, D3DADAPTER_DEFAULT, D3DDEVTYPE_HAL, focus_window,
+            D3DCREATE_HARDWARE_VERTEXPROCESSING, &present_parameters, &device);
+
+    return device;
+}
+
 static HRESULT (WINAPI *pMFCreateMFByteStreamOnStream)(IStream *stream, IMFByteStream **bytestream);
 
 static void init_functions(void)
@@ -1054,6 +1087,57 @@ static void test_source_reader_from_media_source(void)
     fail_request_sample = FALSE;
 }
 
+static void test_reader_d3d9(void)
+{
+    IDirect3DDeviceManager9 *d3d9_manager;
+    IDirect3DDevice9 *d3d9_device;
+    IMFAttributes *attributes;
+    IMFSourceReader *reader;
+    IMFMediaSource *source;
+    IDirect3D9 *d3d9;
+    HWND window;
+    HRESULT hr;
+    UINT token;
+
+    window = create_window();
+    d3d9 = Direct3DCreate9(D3D_SDK_VERSION);
+    ok(!!d3d9, "Failed to create a D3D9 object.\n");
+    if (!(d3d9_device = create_d3d9_device(d3d9, window)))
+    {
+        skip("Failed to create a D3D9 device, skipping tests.\n");
+        goto done;
+    }
+
+    hr = DXVA2CreateDirect3DDeviceManager9(&token, &d3d9_manager);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    hr = IDirect3DDeviceManager9_ResetDevice(d3d9_manager, d3d9_device, token);
+    ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
+
+    source = create_test_source();
+    ok(!!source, "Failed to create test source.\n");
+
+    hr = MFCreateAttributes(&attributes, 1);
+    ok(hr == S_OK, "Failed to create attributes object, hr %#x.\n", hr);
+
+    hr = IMFAttributes_SetUnknown(attributes, &MF_SOURCE_READER_D3D_MANAGER, (IUnknown *)d3d9_manager);
+    ok(hr == S_OK, "Failed to set attribute, hr %#x.\n", hr);
+
+    hr = MFCreateSourceReaderFromMediaSource(source, attributes, &reader);
+    ok(hr == S_OK, "Failed to create source reader, hr %#x.\n", hr);
+
+    IMFAttributes_Release(attributes);
+
+    IMFSourceReader_Release(reader);
+
+    IDirect3DDeviceManager9_Release(d3d9_manager);
+    IDirect3DDevice9_Release(d3d9_device);
+
+done:
+    IDirect3D9_Release(d3d9);
+    DestroyWindow(window);
+}
+
 START_TEST(mfplat)
 {
     HRESULT hr;
@@ -1066,6 +1150,7 @@ START_TEST(mfplat)
     test_factory();
     test_source_reader();
     test_source_reader_from_media_source();
+    test_reader_d3d9();
 
     hr = MFShutdown();
     ok(hr == S_OK, "Failed to shut down, hr %#x.\n", hr);