[PATCH] mfreadwrite/reader: Alternate between selected streams for MF_SOURCE_READER_ANY_STREAM requests.

[Nikolay Sivov]

Signed-off-by: Nikolay Sivov <nsivov(a)codeweavers.com>
---
 dlls/mfreadwrite/reader.c       | 68 ++++++++++++++-------------------
 dlls/mfreadwrite/tests/mfplat.c | 10 ++---
 2 files changed, 32 insertions(+), 46 deletions(-)

diff --git a/dlls/mfreadwrite/reader.c b/dlls/mfreadwrite/reader.c
index cc1b29a2b6a..7aa8f65df7a 100644
--- a/dlls/mfreadwrite/reader.c
+++ b/dlls/mfreadwrite/reader.c
@@ -101,7 +101,6 @@ enum media_stream_flags
     STREAM_FLAG_SAMPLE_REQUESTED = 0x1, /* Protects from making multiple sample requests. */
     STREAM_FLAG_SELECTED = 0x2,         /* Mirrors descriptor, used to simplify tests when starting the source. */
     STREAM_FLAG_PRESENTED = 0x4,        /* Set if stream was selected last time Start() was called. */
-    STREAM_FLAG_REQUESTED_ONCE = 0x8,   /* Used for MF_SOURCE_READER_ANY_STREAM in synchronous mode. */
 };
 
 struct media_stream
@@ -109,7 +108,7 @@ struct media_stream
     IMFMediaStream *stream;
     IMFMediaType *current;
     IMFTransform *decoder;
-    DWORD id;
+    unsigned int id;
     unsigned int index;
     enum media_stream_state state;
     unsigned int flags;
@@ -169,13 +168,14 @@ struct source_reader
     LONG refcount;
     IMFMediaSource *source;
     IMFPresentationDescriptor *descriptor;
-    DWORD first_audio_stream_index;
-    DWORD first_video_stream_index;
     IMFSourceReaderCallback *async_callback;
+    unsigned int first_audio_stream_index;
+    unsigned int first_video_stream_index;
+    unsigned int last_read_index;
+    unsigned int stream_count;
     unsigned int flags;
     enum media_source_state source_state;
     struct media_stream *streams;
-    DWORD stream_count;
     struct list responses;
     CRITICAL_SECTION cs;
     CONDITION_VARIABLE sample_event;
@@ -393,7 +393,7 @@ static HRESULT source_reader_request_sample(struct source_reader *reader, struct
             WARN("Sample request failed, hr %#x.\n", hr);
         else
         {
-            stream->flags |= (STREAM_FLAG_SAMPLE_REQUESTED | STREAM_FLAG_REQUESTED_ONCE);
+            stream->flags |= STREAM_FLAG_SAMPLE_REQUESTED;
         }
     }
 
@@ -988,39 +988,42 @@ static BOOL source_reader_get_read_result(struct source_reader *reader, struct m
     return !request_sample;
 }
 
-static HRESULT source_reader_get_first_selected_stream(struct source_reader *reader, unsigned int flags,
-        unsigned int *stream_index)
+static HRESULT source_reader_get_next_selected_stream(struct source_reader *reader, unsigned int *stream_index)
 {
-    unsigned int i, first_selected = ~0u;
+    unsigned int i, start_idx, stop_idx, first_selected = ~0u, requests = ~0u;
     BOOL selected, stream_drained;
 
-    for (i = 0; i < reader->stream_count; ++i)
+    start_idx = (reader->last_read_index + 1) % reader->stream_count;
+    stop_idx = reader->last_read_index == ~0u ? reader->stream_count : reader->last_read_index;
+
+    for (i = start_idx; i < reader->stream_count && i != stop_idx; i = (i + 1) % (reader->stream_count + 1))
     {
         stream_drained = reader->streams[i].state == STREAM_STATE_EOS && !reader->streams[i].responses;
         selected = SUCCEEDED(source_reader_get_stream_selection(reader, i, &selected)) && selected;
 
-        if (selected && !(reader->streams[i].flags & flags))
+        if (selected)
         {
             if (first_selected == ~0u)
                 first_selected = i;
 
-            if (!stream_drained)
+            /* Try to balance pending reads. */
+            if (!stream_drained && reader->streams[i].requests < requests)
             {
+                requests = reader->streams[i].requests;
                 *stream_index = i;
-                break;
             }
         }
     }
 
-    /* If all selected streams reached EOS, use first selected. This fallback only applies after reader went through all
-       selected streams once. */
-    if (i == reader->stream_count && first_selected != ~0u && !flags)
+    /* If all selected streams reached EOS, use first selected. */
+    if (first_selected != ~0u)
     {
-        *stream_index = first_selected;
-        i = first_selected;
+        if (requests == ~0u)
+            *stream_index = first_selected;
+        reader->last_read_index = *stream_index;
     }
 
-    return i == reader->stream_count ? MF_E_MEDIA_SOURCE_NO_STREAMS_SELECTED : S_OK;
+    return first_selected == ~0u ? MF_E_MEDIA_SOURCE_NO_STREAMS_SELECTED : S_OK;
 }
 
 static HRESULT source_reader_get_stream_read_index(struct source_reader *reader, unsigned int index, unsigned int *stream_index)
@@ -1037,18 +1040,7 @@ static HRESULT source_reader_get_stream_read_index(struct source_reader *reader,
             *stream_index = reader->first_audio_stream_index;
             break;
         case MF_SOURCE_READER_ANY_STREAM:
-            if (reader->async_callback)
-            {
-                /* Pick first selected stream. */
-                hr = source_reader_get_first_selected_stream(reader, 0, stream_index);
-            }
-            else
-            {
-                /* Cycle through all selected streams once, next pick first selected. */
-                if (FAILED(hr = source_reader_get_first_selected_stream(reader, STREAM_FLAG_REQUESTED_ONCE, stream_index)))
-                    hr = source_reader_get_first_selected_stream(reader, 0, stream_index);
-            }
-            return hr;
+            return source_reader_get_next_selected_stream(reader, stream_index);
         default:
             *stream_index = index;
     }
@@ -1339,8 +1331,7 @@ static HRESULT WINAPI src_reader_SetStreamSelection(IMFSourceReader *iface, DWOR
             if (!selection_changed)
             {
                 source_reader_get_stream_selection(reader, i, &selected);
-                if (selected ^ selection)
-                    selection_changed = TRUE;
+                selection_changed = !!(selected ^ selection);
             }
 
             if (selection)
@@ -1364,8 +1355,7 @@ static HRESULT WINAPI src_reader_SetStreamSelection(IMFSourceReader *iface, DWOR
         }
 
         source_reader_get_stream_selection(reader, index, &selected);
-        if (selected ^ selection)
-            selection_changed = TRUE;
+        selection_changed = !!(selected ^ selection);
 
         if (selection)
             hr = IMFPresentationDescriptor_SelectStream(reader->descriptor, index);
@@ -1373,11 +1363,8 @@ static HRESULT WINAPI src_reader_SetStreamSelection(IMFSourceReader *iface, DWOR
             hr = IMFPresentationDescriptor_DeselectStream(reader->descriptor, index);
     }
 
-    if (SUCCEEDED(hr) && selection_changed)
-    {
-        for (i = 0; i < reader->stream_count; ++i)
-            reader->streams[i].flags &= ~STREAM_FLAG_REQUESTED_ONCE;
-    }
+    if (selection_changed)
+        reader->last_read_index = ~0u;
 
     LeaveCriticalSection(&reader->cs);
 
@@ -2135,6 +2122,7 @@ static HRESULT create_source_reader_from_source(IMFMediaSource *source, IMFAttri
     /* At least one major type has to be set. */
     object->first_audio_stream_index = reader_get_first_stream_index(object->descriptor, &MFMediaType_Audio);
     object->first_video_stream_index = reader_get_first_stream_index(object->descriptor, &MFMediaType_Video);
+    object->last_read_index = ~0u;
 
     if (object->first_audio_stream_index == MF_SOURCE_READER_INVALID_STREAM_INDEX &&
             object->first_video_stream_index == MF_SOURCE_READER_INVALID_STREAM_INDEX)
diff --git a/dlls/mfreadwrite/tests/mfplat.c b/dlls/mfreadwrite/tests/mfplat.c
index 2a6e3a1cc36..8455ac64688 100644
--- a/dlls/mfreadwrite/tests/mfplat.c
+++ b/dlls/mfreadwrite/tests/mfplat.c
@@ -886,13 +886,12 @@ static void test_source_reader_from_media_source(void)
     hr = IMFSourceReader_SetStreamSelection(reader, 2, TRUE);
     ok(hr == S_OK, "Failed to select a stream, hr %#x.\n", hr);
 
-    for (i = 0; i < 2 * TEST_SOURCE_NUM_STREAMS; ++i)
+    for (i = 0; i < TEST_SOURCE_NUM_STREAMS + 1; ++i)
     {
         hr = IMFSourceReader_ReadSample(reader, MF_SOURCE_READER_ANY_STREAM, 0, &actual_index, &stream_flags,
                 ×tamp, &sample);
         ok(hr == S_OK, "Failed to get a sample, hr %#x.\n", hr);
-        ok(actual_index == (i < TEST_SOURCE_NUM_STREAMS ? i : 0), "%d: Unexpected stream index %u\n",
-                i, actual_index);
+        ok(actual_index == i % TEST_SOURCE_NUM_STREAMS, "%d: Unexpected stream index %u\n", i, actual_index);
         ok(!stream_flags, "Unexpected stream flags %#x.\n", stream_flags);
         ok(timestamp == 123, "Unexpected timestamp.\n");
         ok(!!sample, "Expected sample object.\n");
@@ -905,13 +904,12 @@ static void test_source_reader_from_media_source(void)
     hr = IMFSourceReader_SetStreamSelection(reader, 0, TRUE);
     ok(hr == S_OK, "Failed to select a stream, hr %#x.\n", hr);
 
-    for (i = 0; i < 2 * TEST_SOURCE_NUM_STREAMS; ++i)
+    for (i = 0; i < TEST_SOURCE_NUM_STREAMS + 1; ++i)
     {
         hr = IMFSourceReader_ReadSample(reader, MF_SOURCE_READER_ANY_STREAM, 0, &actual_index, &stream_flags,
                 ×tamp, &sample);
         ok(hr == S_OK, "Failed to get a sample, hr %#x.\n", hr);
-        ok(actual_index == (i < TEST_SOURCE_NUM_STREAMS ? i : 0), "%d: Unexpected stream index %u\n",
-                i, actual_index);
+        ok(actual_index == i % TEST_SOURCE_NUM_STREAMS, "%d: Unexpected stream index %u\n", i, actual_index);
         ok(!stream_flags, "Unexpected stream flags %#x.\n", stream_flags);
         ok(timestamp == 123, "Unexpected timestamp.\n");
         ok(!!sample, "Expected sample object.\n");
-- 
2.29.2