[PATCH] mfreadwrite/reader: Add support for reading to SA-allocated samples.

[Nikolay Sivov]

Signed-off-by: Nikolay Sivov <nsivov(a)codeweavers.com>
---
 dlls/mfreadwrite/reader.c | 141 +++++++++++++++++++++++++++++---------
 1 file changed, 110 insertions(+), 31 deletions(-)

diff --git a/dlls/mfreadwrite/reader.c b/dlls/mfreadwrite/reader.c
index ec7d3fd62e6..9903898a155 100644
--- a/dlls/mfreadwrite/reader.c
+++ b/dlls/mfreadwrite/reader.c
@@ -85,6 +85,7 @@ struct stream_response
     DWORD stream_flags;
     LONGLONG timestamp;
     IMFSample *sample;
+    unsigned int sa_pending : 1;
 };
 
 enum media_stream_state
@@ -365,12 +366,83 @@ static HRESULT WINAPI source_reader_callback_GetParameters(IMFAsyncCallback *ifa
     return E_NOTIMPL;
 }
 
+static void source_reader_response_ready(struct source_reader *reader, struct stream_response *response)
+{
+    struct source_reader_async_command *command;
+    struct media_stream *stream = &reader->streams[response->stream_index];
+    HRESULT hr;
+
+    if (!stream->requests || response->sa_pending)
+        return;
+
+    if (reader->async_callback)
+    {
+        if (SUCCEEDED(source_reader_create_async_op(SOURCE_READER_ASYNC_SAMPLE_READY, &command)))
+        {
+            command->u.sample.stream_index = stream->index;
+            if (FAILED(hr = MFPutWorkItem(MFASYNC_CALLBACK_QUEUE_STANDARD, &reader->async_commands_callback,
+                    &command->IUnknown_iface)))
+                WARN("Failed to submit async result, hr %#x.\n", hr);
+            IUnknown_Release(&command->IUnknown_iface);
+        }
+    }
+    else
+        WakeAllConditionVariable(&reader->sample_event);
+
+    stream->requests--;
+}
+
+static void source_reader_copy_sample_buffer(IMFSample *src, IMFSample *dst)
+{
+    IMFMediaBuffer *buffer;
+    unsigned int flags;
+    LONGLONG time;
+    HRESULT hr;
+
+    IMFSample_CopyAllItems(src, (IMFAttributes *)dst);
+
+    IMFSample_SetSampleDuration(dst, 0);
+    IMFSample_SetSampleTime(dst, 0);
+    IMFSample_SetSampleFlags(dst, 0);
+
+    if (SUCCEEDED(IMFSample_GetSampleDuration(src, &time)))
+        IMFSample_SetSampleDuration(dst, time);
+
+    if (SUCCEEDED(IMFSample_GetSampleTime(src, &time)))
+        IMFSample_SetSampleTime(dst, time);
+
+    if (SUCCEEDED(IMFSample_GetSampleFlags(src, &flags)))
+        IMFSample_SetSampleFlags(dst, flags);
+
+    if (SUCCEEDED(IMFSample_ConvertToContiguousBuffer(src, NULL)))
+    {
+        if (SUCCEEDED(IMFSample_GetBufferByIndex(dst, 0, &buffer)))
+        {
+            if (FAILED(hr = IMFSample_CopyToBuffer(src, buffer)))
+                WARN("Failed to copy a buffer, hr %#x.\n", hr);
+            IMFMediaBuffer_Release(buffer);
+        }
+    }
+}
+
+static void source_reader_set_sa_response(struct source_reader *reader, struct stream_response *response)
+{
+    struct media_stream *stream = &reader->streams[response->stream_index];
+    IMFSample *sample;
+
+    if (SUCCEEDED(IMFVideoSampleAllocatorEx_AllocateSample(stream->allocator, &sample)))
+    {
+        source_reader_copy_sample_buffer(response->sample, sample);
+        response->sa_pending = 0;
+        IMFSample_Release(response->sample);
+        response->sample = sample;
+    }
+}
+
 static void source_reader_queue_response(struct source_reader *reader, struct media_stream *stream, HRESULT status,
         DWORD stream_flags, LONGLONG timestamp, IMFSample *sample)
 {
-    struct source_reader_async_command *command;
     struct stream_response *response;
-    HRESULT hr;
 
     response = heap_alloc_zero(sizeof(*response));
     response->status = status;
@@ -381,27 +453,16 @@ static void source_reader_queue_response(struct source_reader *reader, struct me
     if (response->sample)
         IMFSample_AddRef(response->sample);
 
+    if (response->sample && stream->allocator)
+    {
+        response->sa_pending = 1;
+        source_reader_set_sa_response(reader, response);
+    }
+
     list_add_tail(&reader->responses, &response->entry);
     stream->responses++;
 
-    if (stream->requests)
-    {
-        if (reader->async_callback)
-        {
-            if (SUCCEEDED(source_reader_create_async_op(SOURCE_READER_ASYNC_SAMPLE_READY, &command)))
-            {
-                command->u.sample.stream_index = stream->index;
-                if (FAILED(hr = MFPutWorkItem(MFASYNC_CALLBACK_QUEUE_STANDARD, &reader->async_commands_callback,
-                        &command->IUnknown_iface)))
-                    WARN("Failed to submit async result, hr %#x.\n", hr);
-                IUnknown_Release(&command->IUnknown_iface);
-            }
-        }
-        else
-            WakeAllConditionVariable(&reader->sample_event);
-
-        stream->requests--;
-    }
+    source_reader_response_ready(reader, response);
 }
 
 static HRESULT source_reader_request_sample(struct source_reader *reader, struct media_stream *stream)
@@ -874,20 +935,26 @@ static struct stream_response * media_stream_detach_response(struct source_reade
 static struct stream_response *media_stream_pop_response(struct source_reader *reader, struct media_stream *stream)
 {
     struct stream_response *response;
-    struct list *head;
 
-    if (stream)
+    LIST_FOR_EACH_ENTRY(response, &reader->responses, struct stream_response, entry)
     {
-        LIST_FOR_EACH_ENTRY(response, &reader->responses, struct stream_response, entry)
-        {
-            if (response->stream_index == stream->index)
-                return media_stream_detach_response(reader, response);
-        }
+        if ((stream && response->stream_index != stream->index) || response->sa_pending)
+            continue;
+
+        return media_stream_detach_response(reader, response);
     }
-    else
+
+    return NULL;
+}
+
+static struct stream_response *media_stream_pick_pending_response(struct source_reader *reader, unsigned int stream)
+{
+    struct stream_response *response;
+
+    LIST_FOR_EACH_ENTRY(response, &reader->responses, struct stream_response, entry)
     {
-        if ((head = list_head(&reader->responses)))
-            return media_stream_detach_response(reader, LIST_ENTRY(head, struct stream_response, entry));
+        if (response->stream_index == stream && response->sa_pending)
+            return response;
     }
 
     return NULL;
@@ -963,7 +1030,7 @@ static BOOL source_reader_got_response_for_stream(struct source_reader *reader,
 
     LIST_FOR_EACH_ENTRY(response, &reader->responses, struct stream_response, entry)
     {
-        if (response->stream_index == stream->index)
+        if (response->stream_index == stream->index && !response->sa_pending)
             return TRUE;
     }
 
@@ -1201,6 +1268,18 @@ static HRESULT WINAPI source_reader_async_commands_callback_Invoke(IMFAsyncCallb
 
             break;
 
+        case SOURCE_READER_ASYNC_SA_READY:
+
+            EnterCriticalSection(&reader->cs);
+            if ((response = media_stream_pick_pending_response(reader, command->u.sa.stream_index)))
+            {
+                source_reader_set_sa_response(reader, response);
+                source_reader_response_ready(reader, response);
+            }
+            LeaveCriticalSection(&reader->cs);
+
+            break;
+
         case SOURCE_READER_ASYNC_SAMPLE_READY:
 
             EnterCriticalSection(&reader->cs);
-- 
2.30.0