[PATCH] mfreadwrite: Simplify iteration through streams.

[Giovanni Mascellani]

This also fixes a bug happening when the reader has only one stream,
in which case the iteration body was not even executed once.

Signed-off-by: Giovanni Mascellani <gmascellani at codeweavers.com>
---
 dlls/mfreadwrite/reader.c       | 14 ++++-----
 dlls/mfreadwrite/tests/mfplat.c | 50 ++++++++++++++++++++++++++++-----
 2 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/dlls/mfreadwrite/reader.c b/dlls/mfreadwrite/reader.c
index 2819b800eb9..ab9f6c83674 100644
--- a/dlls/mfreadwrite/reader.c
+++ b/dlls/mfreadwrite/reader.c
@@ -1087,13 +1087,10 @@ static BOOL source_reader_get_read_result(struct source_reader *reader, struct m
 
 static HRESULT source_reader_get_next_selected_stream(struct source_reader *reader, unsigned int *stream_index)
 {
-    unsigned int i, start_idx, stop_idx, first_selected = ~0u, requests = ~0u;
+    unsigned int i, first_selected = ~0u, requests = ~0u;
     BOOL selected, stream_drained;
 
-    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))
+    for (i = (reader->last_read_index + 1) % reader->stream_count; ; i = (i + 1) % reader->stream_count)
     {
         stream_drained = reader->streams[i].state == STREAM_STATE_EOS && !reader->streams[i].responses;
         selected = SUCCEEDED(source_reader_get_stream_selection(reader, i, &selected)) && selected;
@@ -1110,6 +1107,9 @@ static HRESULT source_reader_get_next_selected_stream(struct source_reader *read
                 *stream_index = i;
             }
         }
+
+        if (i == reader->last_read_index)
+            break;
     }
 
     /* If all selected streams reached EOS, use first selected. */
@@ -1477,7 +1477,7 @@ static HRESULT WINAPI src_reader_SetStreamSelection(IMFSourceReader *iface, DWOR
     }
 
     if (selection_changed)
-        reader->last_read_index = ~0u;
+        reader->last_read_index = reader->stream_count - 1;
 
     LeaveCriticalSection(&reader->cs);
 
@@ -2360,7 +2360,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;
+    object->last_read_index = object->stream_count - 1;
 
     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 25aa29ac272..c27806f699e 100644
--- a/dlls/mfreadwrite/tests/mfplat.c
+++ b/dlls/mfreadwrite/tests/mfplat.c
@@ -501,7 +501,7 @@ static struct test_media_stream *create_test_stream(DWORD stream_index, IMFMedia
     return stream;
 }
 
-static IMFMediaSource *create_test_source(void)
+static IMFMediaSource *create_test_source(int stream_num)
 {
     struct test_source *source;
     int i;
@@ -511,7 +511,7 @@ static IMFMediaSource *create_test_source(void)
     source->refcount = 1;
     MFCreateEventQueue(&source->event_queue);
     InitializeCriticalSection(&source->cs);
-    for (i = 0; i < ARRAY_SIZE(source->streams); ++i)
+    for (i = 0; i < stream_num; ++i)
         source->streams[i] = create_test_stream(i, &source->IMFMediaSource_iface);
 
     return &source->IMFMediaSource_iface;
@@ -904,7 +904,7 @@ static void test_source_reader_from_media_source(void)
     int i;
     PROPVARIANT pos;
 
-    source = create_test_source();
+    source = create_test_source(3);
     ok(!!source, "Failed to create test source.\n");
 
     hr = MFCreateSourceReaderFromMediaSource(source, NULL, &reader);
@@ -970,8 +970,44 @@ static void test_source_reader_from_media_source(void)
     IMFSourceReader_Release(reader);
     IMFMediaSource_Release(source);
 
+    source = create_test_source(1);
+    ok(!!source, "Failed to create test source.\n");
+
+    hr = MFCreateSourceReaderFromMediaSource(source, NULL, &reader);
+    ok(hr == S_OK, "Failed to create source reader, hr %#x.\n", hr);
+
+    /* MF_SOURCE_READER_ANY_STREAM with a single stream */
+    hr = IMFSourceReader_SetStreamSelection(reader, 0, TRUE);
+    ok(hr == S_OK, "Failed to select a stream, hr %#x.\n", hr);
+
+    pos.vt = VT_I8;
+    pos.hVal.QuadPart = 0;
+    hr = IMFSourceReader_SetCurrentPosition(reader, &GUID_NULL, &pos);
+    ok(hr == S_OK, "Failed to seek to beginning of stream, hr %#x.\n", hr);
+
+    hr = IMFSourceReader_ReadSample(reader, MF_SOURCE_READER_ANY_STREAM, 0, &actual_index, &stream_flags,
+            &timestamp, &sample);
+    ok(hr == S_OK, "Failed to get a sample, hr %#x.\n", hr);
+    ok(actual_index == 0, "Unexpected stream index %u\n", actual_index);
+    ok(!stream_flags, "Unexpected stream flags %#x.\n", stream_flags);
+    ok(timestamp == 123, "Unexpected timestamp.\n");
+    ok(!!sample, "Expected sample object.\n");
+    IMFSample_Release(sample);
+
+    hr = IMFSourceReader_ReadSample(reader, MF_SOURCE_READER_ANY_STREAM, 0, &actual_index, &stream_flags,
+            &timestamp, &sample);
+    ok(hr == S_OK, "Failed to get a sample, hr %#x.\n", hr);
+    ok(actual_index == 0, "Unexpected stream index %u\n", actual_index);
+    ok(!stream_flags, "Unexpected stream flags %#x.\n", stream_flags);
+    ok(timestamp == 123, "Unexpected timestamp.\n");
+    ok(!!sample, "Expected sample object.\n");
+    IMFSample_Release(sample);
+
+    IMFSourceReader_Release(reader);
+    IMFMediaSource_Release(source);
+
     /* Request from stream 0. */
-    source = create_test_source();
+    source = create_test_source(3);
     ok(!!source, "Failed to create test source.\n");
 
     hr = MFCreateSourceReaderFromMediaSource(source, NULL, &reader);
@@ -1005,7 +1041,7 @@ static void test_source_reader_from_media_source(void)
     IMFMediaSource_Release(source);
 
     /* Async mode. */
-    source = create_test_source();
+    source = create_test_source(3);
     ok(!!source, "Failed to create test source.\n");
 
     callback = create_async_callback();
@@ -1050,7 +1086,7 @@ static void test_source_reader_from_media_source(void)
     IMFMediaSource_Release(source);
 
     /* RequestSample failure. */
-    source = create_test_source();
+    source = create_test_source(3);
     ok(!!source, "Failed to create test source.\n");
 
     fail_request_sample = TRUE;
@@ -1120,7 +1156,7 @@ static void test_reader_d3d9(void)
     hr = IDirect3DDeviceManager9_ResetDevice(d3d9_manager, d3d9_device, token);
     ok(hr == S_OK, "Unexpected hr %#x.\n", hr);
 
-    source = create_test_source();
+    source = create_test_source(3);
     ok(!!source, "Failed to create test source.\n");
 
     hr = MFCreateAttributes(&attributes, 1);
-- 
2.31.0