LibMedia: Implement buffered time range scanning for Matroska

Libraries/LibMedia/Containers/Matroska/MatroskaDemuxer.cpp

@@ -25,8 +25,10 @@ DecoderErrorOr<NonnullRefPtr<MatroskaDemuxer>> MatroskaDemuxer::from_stream(Nonn
 
 MatroskaDemuxer::MatroskaDemuxer(NonnullRefPtr<MediaStream> const& stream, Reader&& reader)
     : m_stream(stream)
+    , m_buffered_scan_cursor(stream->create_cursor())
     , m_reader(move(reader))
 {
+    m_buffered_scan_cursor->set_is_blocking(false);
 }
 
 MatroskaDemuxer::~MatroskaDemuxer() = default;
@@ -177,12 +179,10 @@ DecoderErrorOr<AK::Duration> MatroskaDemuxer::total_duration()
 
 TimeRanges MatroskaDemuxer::buffered_time_ranges() const
 {
-    // FIXME: Scan the stream for buffered ranges.
-    TimeRanges ranges;
-    auto duration = m_reader.duration();
-    if (duration.has_value())
-        ranges.add_range(AK::Duration::zero(), duration.value());
-    return ranges;
+    auto byte_ranges = m_stream->available_byte_ranges();
+    if (byte_ranges.is_empty())
+        return {};
+    return m_reader.buffered_time_ranges(m_buffered_scan_cursor, byte_ranges);
 }
 
 DecoderErrorOr<AK::Duration> MatroskaDemuxer::duration_of_track(Track const&)

Libraries/LibMedia/Containers/Matroska/MatroskaDemuxer.h

@@ -63,6 +63,7 @@ private:
     TrackStatus& get_track_status(Track const&);
 
     NonnullRefPtr<MediaStream> m_stream;
+    NonnullRefPtr<MediaStreamCursor> m_buffered_scan_cursor;
     Reader m_reader;
 
     mutable Sync::Mutex m_track_statuses_mutex;

Libraries/LibMedia/Containers/Matroska/Reader.cpp

@@ -866,7 +866,33 @@ DecoderErrorOr<Block> Reader::parse_block_group(Streamer& streamer, AK::Duration
 
 DecoderErrorOr<SampleIterator> Reader::create_sample_iterator(NonnullRefPtr<MediaStreamCursor> const& cursor, Optional<u64> track_number) const
 {
-    dbgln_if(MATROSKA_DEBUG, "Creating sample iterator starting at {} relative to segment at {}", m_first_cluster_position, m_segment_contents_position);
+    return create_sample_iterator_at_byte_position(cursor, 0, track_number);
+}
+
+DecoderErrorOr<SampleIterator> Reader::create_sample_iterator_at_byte_position(NonnullRefPtr<MediaStreamCursor> const& cursor, size_t position, Optional<u64> track_number) const
+{
+    Optional<size_t> cluster_position;
+
+    if (m_first_cluster_position >= position) {
+        cluster_position = m_first_cluster_position;
+    } else {
+        for (auto const& [track_number_with_cues, cue_points] : m_cues) {
+            if (track_number.has_value() && track_number != track_number_with_cues)
+                continue;
+            for (auto const& cue_point : cue_points) {
+                auto cue_cluster_position = m_segment_contents_position + cue_point.position.cluster_position();
+                if (cue_cluster_position < position)
+                    continue;
+                if (!cluster_position.has_value() || cue_cluster_position < cluster_position.value())
+                    cluster_position = cue_cluster_position;
+            }
+        }
+    }
+
+    if (!cluster_position.has_value())
+        return DecoderError::format(DecoderErrorCategory::EndOfStream, "Could not find a Cluster element after {}", position);
+
+    dbgln_if(MATROSKA_DEBUG, "Creating sample iterator starting at {} relative to segment at {}", cluster_position, m_segment_contents_position);
     TrackBlockContexts track_contexts;
     if (track_number.has_value()) {
         auto track = TRY(track_for_track_number(track_number.value()));
@@ -875,7 +901,7 @@ DecoderErrorOr<SampleIterator> Reader::create_sample_iterator(NonnullRefPtr<Medi
         for (auto const& [number, track_entry] : m_tracks)
             track_contexts.set(number, TrackBlockContext::from_track_entry(*track_entry));
     }
-    return SampleIterator(cursor, track_number, move(track_contexts), m_segment_information.timestamp_scale(), m_segment_contents_position, m_first_cluster_position);
+    return SampleIterator(cursor, track_number, move(track_contexts), m_segment_information.timestamp_scale(), m_segment_contents_position, cluster_position.value());
 }
 
 static DecoderErrorOr<CueTrackPosition> parse_cue_track_position(Streamer& streamer)
@@ -1121,4 +1147,131 @@ Optional<Vector<TrackCuePoint> const&> Reader::cue_points_for_track(u64 track_nu
     return m_cues.get(track_number);
 }
 
+TimeRanges Reader::buffered_time_ranges(NonnullRefPtr<MediaStreamCursor> const& cursor, Vector<MediaStream::ByteRange> const& byte_ranges) const
+{
+    auto create_iterator = [&](size_t position) -> Optional<SampleIterator> {
+        auto iterator = create_sample_iterator_at_byte_position(cursor, position);
+        if (iterator.is_error())
+            return {};
+        return iterator.release_value();
+    };
+
+    size_t cached_range_index = 0;
+    size_t byte_range_index = 0;
+    while (byte_range_index < byte_ranges.size()) {
+        auto cached_range = m_buffered_ranges.get(cached_range_index);
+        auto const& byte_range = byte_ranges[byte_range_index];
+        VERIFY(byte_range.start < byte_range.end);
+
+        auto previous_byte_range = byte_ranges.get(byte_range_index - 1);
+        if (previous_byte_range.has_value())
+            VERIFY(previous_byte_range->end < byte_range.start);
+
+        // If the current byte range precedes the current cached range, insert a new one for that byte range.
+        // Restart the loop with the same cached range.
+        if (!cached_range.has_value() || byte_range.start < cached_range->start) {
+            auto new_cached_range = BufferedRange {
+                .start = byte_range.start,
+                .end = byte_range.end,
+                .iterator = create_iterator(byte_range.start),
+            };
+            m_buffered_ranges.insert(cached_range_index, move(new_cached_range));
+            cached_range_index++;
+            byte_range_index++;
+            continue;
+        }
+
+        VERIFY(cached_range.has_value());
+
+        // If the current range is an exact match to the byte range, we can just update the end byte and advance.
+        if (byte_range.start == cached_range->start) {
+            cached_range->end = byte_range.end;
+            cached_range_index++;
+            byte_range_index++;
+            continue;
+        }
+
+        auto previous_cached_range = m_buffered_ranges.get(cached_range_index - 1);
+        if (previous_cached_range.has_value()) {
+            VERIFY(previous_cached_range->start < cached_range->start);
+
+            // If the current cached range is entirely encompassed by the previous one, then remove that range. We'll
+            // need to rescan its contents.
+            if (previous_cached_range->end >= cached_range->end) {
+                m_buffered_ranges.remove(cached_range_index);
+                continue;
+            }
+        }
+
+        // The range has shifted forward. We'll need to re-read from the new start position.
+        auto new_iterator = create_iterator(byte_range.start);
+
+        // If the cached range's last read is still contained in the new byte range, we can keep using its end time.
+        // Just grab the first frame at the new byte range's start and update the cached range's start from it.
+        auto& cached_iterator = cached_range->iterator;
+        if (cached_iterator.has_value() && new_iterator.has_value()) {
+            auto last_cached_position = cached_iterator->position();
+            if (byte_range.start <= last_cached_position && last_cached_position <= byte_range.end) {
+                auto first_block = new_iterator->next_block();
+
+                if (!first_block.is_error() && first_block.value().timestamp().has_value()) {
+                    cached_range->start = byte_range.start;
+                    cached_range->end = byte_range.end;
+                    cached_range->time_start = first_block.value().timestamp().value();
+                    cached_range_index++;
+                    byte_range_index++;
+                    continue;
+                }
+            }
+        }
+
+        // Otherwise, we have to reset everything for this range.
+        *cached_range = {
+            .start = byte_range.start,
+            .end = byte_range.end,
+            .iterator = move(new_iterator),
+        };
+        cached_range_index++;
+        byte_range_index++;
+    }
+
+    // Remove any leftover ranges. We should be left with only the exact ranges provided to us.
+    m_buffered_ranges.remove(cached_range_index, m_buffered_ranges.size() - cached_range_index);
+
+    // All previously known buffered ranges are now matched up or discarded. Iterate the blocks to update the ranges'
+    // end times and append the ranges.
+    VERIFY(m_buffered_ranges.size() == byte_ranges.size());
+    TimeRanges result;
+
+    for (size_t i = 0; i < byte_ranges.size(); i++) {
+        auto& cached_range = m_buffered_ranges[i];
+        auto const& byte_range = byte_ranges[i];
+        VERIFY(cached_range.start == byte_range.start);
+        VERIFY(cached_range.end == byte_range.end);
+
+        if (cached_range.iterator.has_value() && cached_range.iterator->position() < cached_range.end) {
+            auto& iterator = cached_range.iterator;
+
+            while (iterator->position() < byte_range.end) {
+                auto block_or_error = iterator->next_block();
+                if (block_or_error.is_error())
+                    break;
+                auto block = block_or_error.release_value();
+                if (block.timestamp().has_value() && block.duration().has_value()) {
+                    if (!cached_range.time_start.has_value())
+                        cached_range.time_start = block.timestamp().value();
+
+                    auto block_end = block.timestamp().value() + block.duration().value();
+                    cached_range.time_end = block_end;
+                }
+            }
+        }
+
+        if (cached_range.time_start.has_value())
+            result.add_range(max(AK::Duration::zero(), cached_range.time_start.value()), cached_range.time_end);
+    }
+
+    return result;
+}
+
 }

Libraries/LibMedia/Containers/Matroska/Reader.h

@@ -13,6 +13,8 @@
 #include <LibMedia/DecoderError.h>
 #include <LibMedia/Export.h>
 #include <LibMedia/Forward.h>
+#include <LibMedia/MediaStream.h>
+#include <LibMedia/TimeRanges.h>
 
 #include "Document.h"
 #include "SampleIterator.h"
@@ -51,12 +53,15 @@ public:
     DecoderErrorOr<size_t> track_count() const;
 
     DecoderErrorOr<SampleIterator> create_sample_iterator(NonnullRefPtr<MediaStreamCursor> const& cursor, Optional<u64> track_number = {}) const;
+    DecoderErrorOr<SampleIterator> create_sample_iterator_at_byte_position(NonnullRefPtr<MediaStreamCursor> const& cursor, size_t position, Optional<u64> track_number = {}) const;
     DecoderErrorOr<SampleIterator> seek_to_random_access_point(SampleIterator, AK::Duration) const;
 
     Optional<Vector<TrackCuePoint> const&> cue_points_for_track(u64 track_number) const;
 
     static size_t find_cue_point_index_at_or_before(Vector<TrackCuePoint> const&, Optional<AK::Duration> total_duration, AK::Duration target);
 
+    TimeRanges buffered_time_ranges(NonnullRefPtr<MediaStreamCursor> const&, Vector<MediaStream::ByteRange> const& byte_ranges) const;
+
 private:
     Reader() = default;
 
@@ -90,6 +95,15 @@ private:
 
     // The vectors must be sorted by timestamp at all times.
     HashMap<u64, Vector<TrackCuePoint>> m_cues;
+
+    struct BufferedRange {
+        size_t start { 0 };
+        size_t end { 0 };
+        Optional<SampleIterator> iterator;
+        Optional<AK::Duration> time_start { OptionalNone() };
+        AK::Duration time_end { AK::Duration::zero() };
+    };
+    mutable Vector<BufferedRange> m_buffered_ranges;
 };
 
 }

Tests/LibMedia/FILE_LICENSES.md

@@ -1,4 +1,4 @@
-## vp9_clamp_reference_mvs.webm & master_elements_containing_crc32.mkv
+## vp9_clamp_reference_mvs.webm, master_elements_containing_crc32.mkv & big_buck_bunny_5s.webm
 
 Licensed under [Creative Commons Attribution 3.0](http://creativecommons.org/licenses/by/3.0/)\
 (c) copyright 2008, Blender Foundation / [www.bigbuckbunny.org](https://www.bigbuckbunny.org)

Tests/LibMedia/TestParseMatroska.cpp

@@ -483,3 +483,204 @@ TEST_CASE(opus_frame_duration)
         EXPECT_EQ(block.duration().value(), expected_duration);
     }
 }
+
+static ByteBuffer load_test_file_data(StringView path)
+{
+    auto file = MUST(Core::File::open(path, Core::File::OpenMode::Read));
+    return MUST(file->read_until_eof());
+}
+
+static constexpr size_t CUES_START = 298382;
+
+static auto create_incremental_demuxer(ByteBuffer const& file_data, NonnullRefPtr<Media::IncrementallyPopulatedStream>& stream, size_t initial_end)
+{
+    stream = Media::IncrementallyPopulatedStream::create_empty();
+    stream->add_chunk_at(0, file_data.bytes().slice(0, initial_end));
+    stream->add_chunk_at(CUES_START, file_data.bytes().slice(CUES_START));
+    return MUST(Media::Matroska::MatroskaDemuxer::from_stream(stream));
+}
+
+TEST_CASE(buffered_time_ranges_full_file)
+{
+    auto file_data = load_test_file_data("./vp9_in_webm.webm"sv);
+    auto stream = Media::IncrementallyPopulatedStream::create_from_buffer(file_data);
+    auto demuxer = MUST(Media::Matroska::MatroskaDemuxer::from_stream(stream));
+
+    auto ranges = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges.size(), 1u);
+    EXPECT_EQ(ranges[0].start, AK::Duration::from_microseconds(500));
+    EXPECT_EQ(ranges[0].end, AK::Duration::from_microseconds(1021500));
+}
+
+TEST_CASE(buffered_time_ranges_incremental_thirds)
+{
+    auto file_data = load_test_file_data("./vp9_in_webm.webm"sv);
+    auto start = AK::Duration::from_microseconds(500);
+    size_t one_third = file_data.size() / 3;
+    size_t two_thirds = one_third * 2;
+
+    NonnullRefPtr<Media::IncrementallyPopulatedStream> stream = Media::IncrementallyPopulatedStream::create_empty();
+    auto demuxer = create_incremental_demuxer(file_data, stream, one_third);
+
+    // Stage 1: first third.
+    auto ranges_1 = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_1.size(), 1u);
+    EXPECT_EQ(ranges_1[0].start, start);
+    EXPECT_EQ(ranges_1[0].end, AK::Duration::from_microseconds(91000));
+
+    // Stage 2: extend to two thirds.
+    stream->add_chunk_at(one_third, file_data.bytes().slice(one_third, two_thirds - one_third));
+    auto ranges_2 = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_2.size(), 1u);
+    EXPECT_EQ(ranges_2[0].start, start);
+    EXPECT(ranges_2[0].end > ranges_1[0].end);
+    EXPECT_EQ(ranges_2[0].end, AK::Duration::from_microseconds(571000));
+
+    // Stage 3: complete the file.
+    stream->add_chunk_at(two_thirds, file_data.bytes().slice(two_thirds, CUES_START - two_thirds));
+    stream->close();
+    auto ranges_3 = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_3.size(), 1u);
+    EXPECT_EQ(ranges_3[0].start, start);
+    EXPECT(ranges_3[0].end > ranges_2[0].end);
+    EXPECT_EQ(ranges_3[0].end, AK::Duration::from_microseconds(1021500));
+}
+
+// big_buck_bunny_5s.webm cluster layout:
+//   Cluster 0 (0ms):    byte 482
+//   Cluster 1 (500ms):  byte 6128
+//   Cluster 2 (1000ms): byte 13177
+//   Cluster 3 (1500ms): byte 21628
+//   Cluster 4 (2000ms): byte 31913
+//   Cluster 5 (2500ms): byte 49246
+//   Cluster 6 (3000ms): byte 59860
+//   Cluster 7 (3500ms): byte 71977
+//   Cluster 8 (4000ms): byte 91687
+//   Cluster 9 (4500ms): byte 102297
+//   Cues:               byte 113303
+//   File size:          113491
+
+static constexpr size_t BBB_CUES_START = 113303;
+
+TEST_CASE(buffered_time_ranges_gap_then_fill)
+{
+    auto file_data = load_test_file_data("./big_buck_bunny_5s.webm"sv);
+
+    // Buffer clusters 0-3 (0-2s) and clusters 7-9 (3.5-5s), leaving a gap at clusters 4-6 (2-3.5s).
+    constexpr size_t first_chunk_end = 31913;
+    constexpr size_t second_chunk_start = 71977;
+    auto stream = Media::IncrementallyPopulatedStream::create_empty();
+    stream->add_chunk_at(0, file_data.bytes().slice(0, first_chunk_end));
+    stream->add_chunk_at(BBB_CUES_START, file_data.bytes().slice(BBB_CUES_START));
+    stream->close();
+    stream->add_chunk_at(second_chunk_start, file_data.bytes().slice(second_chunk_start, BBB_CUES_START - second_chunk_start));
+    auto demuxer = MUST(Media::Matroska::MatroskaDemuxer::from_stream(stream));
+
+    auto ranges_gap = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_gap.size(), 2u);
+    EXPECT_EQ(ranges_gap[0].start, AK::Duration::zero());
+    EXPECT_EQ(ranges_gap[1].start, AK::Duration::from_milliseconds(3500));
+
+    // Fill the gap with clusters 4-6.
+    stream->add_chunk_at(first_chunk_end, file_data.bytes().slice(first_chunk_end, second_chunk_start - first_chunk_end));
+    auto ranges_filled = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_filled.size(), 1u);
+    EXPECT_EQ(ranges_filled[0].start, AK::Duration::zero());
+    EXPECT_EQ(ranges_filled[0].end, AK::Duration::from_nanoseconds(4999666666));
+}
+
+TEST_CASE(buffered_time_ranges_reverse_order_chunks)
+{
+    auto file_data = load_test_file_data("./big_buck_bunny_5s.webm"sv);
+
+    // Buffer init + clusters 0-2 (0-1.5s) and clusters 7-9 (3.5-5s), then fill the middle.
+    constexpr size_t first_chunk_end = 21628;
+    constexpr size_t second_chunk_start = 71977;
+    NonnullRefPtr<Media::IncrementallyPopulatedStream> stream = Media::IncrementallyPopulatedStream::create_empty();
+    stream = Media::IncrementallyPopulatedStream::create_empty();
+    stream->add_chunk_at(0, file_data.bytes().slice(0, first_chunk_end));
+    stream->add_chunk_at(second_chunk_start, file_data.bytes().slice(second_chunk_start));
+    stream->close();
+    auto demuxer = MUST(Media::Matroska::MatroskaDemuxer::from_stream(stream));
+
+    auto ranges_1 = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_1.size(), 2u);
+    EXPECT_EQ(ranges_1[0].start, AK::Duration::zero());
+    EXPECT_EQ(ranges_1[1].start, AK::Duration::from_milliseconds(3500));
+
+    // Fill the gap with clusters 3-6.
+    stream->add_chunk_at(first_chunk_end, file_data.bytes().slice(first_chunk_end, second_chunk_start - first_chunk_end));
+    auto ranges_2 = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_2.size(), 1u);
+    EXPECT_EQ(ranges_2[0].start, AK::Duration::zero());
+    EXPECT(ranges_2[0].end > AK::Duration::from_milliseconds(4900));
+}
+
+TEST_CASE(buffered_time_ranges_repeated_query)
+{
+    auto file_data = load_test_file_data("./vp9_in_webm.webm"sv);
+    auto stream = Media::IncrementallyPopulatedStream::create_from_buffer(file_data);
+    auto demuxer = MUST(Media::Matroska::MatroskaDemuxer::from_stream(stream));
+
+    // Query multiple times — results should be identical and stable.
+    auto ranges_1 = demuxer->buffered_time_ranges();
+    auto ranges_2 = demuxer->buffered_time_ranges();
+    auto ranges_3 = demuxer->buffered_time_ranges();
+    EXPECT_EQ(ranges_1, ranges_2);
+    EXPECT_EQ(ranges_2, ranges_3);
+}
+
+TEST_CASE(buffered_time_ranges_evicted_start)
+{
+    // Simulate data eviction by calling buffered_time_ranges with the full file first,
+    // then with a byte range whose start is later (as if the beginning was evicted).
+    auto file_data = load_test_file_data("./big_buck_bunny_5s.webm"sv);
+    auto stream = Media::IncrementallyPopulatedStream::create_from_buffer(file_data);
+    auto reader = MUST(Media::Matroska::Reader::from_stream(stream->create_cursor()));
+    auto cursor = stream->create_cursor();
+
+    // Get buffered ranges for the full file.
+    Vector<Media::MediaStream::ByteRange> byte_ranges;
+    byte_ranges.append({ 0, file_data.size() });
+    auto time_ranges = reader.buffered_time_ranges(cursor, byte_ranges);
+    EXPECT_EQ(time_ranges.size(), 1u);
+    EXPECT_EQ(time_ranges[0].start, AK::Duration::zero());
+    EXPECT_EQ(time_ranges[0].end, AK::Duration::from_nanoseconds(4999666666));
+
+    // Simulate eviction of the first four clusters.
+    byte_ranges[0] = { 31913, file_data.size() };
+    time_ranges = reader.buffered_time_ranges(cursor, byte_ranges);
+    EXPECT_EQ(time_ranges.size(), 1u);
+    EXPECT_EQ(time_ranges[0].start, AK::Duration::from_milliseconds(2000));
+    EXPECT_EQ(time_ranges[0].end, AK::Duration::from_nanoseconds(4999666666));
+}
+
+TEST_CASE(buffered_time_ranges_evicted_start_appended_end)
+{
+    // Simulate data eviction by calling buffered_time_ranges with an early portion of the file,
+    // then again with a new range that does not overlap.
+    auto file_data = load_test_file_data("./big_buck_bunny_5s.webm"sv);
+    auto stream = Media::IncrementallyPopulatedStream::create_from_buffer(file_data);
+    auto reader = MUST(Media::Matroska::Reader::from_stream(stream->create_cursor()));
+    auto cursor = stream->create_cursor();
+
+    // Get buffered ranges with only the first two clusters available.
+    Vector<Media::MediaStream::ByteRange> byte_ranges;
+    byte_ranges.append({ 0, 21628 });
+    auto time_ranges = reader.buffered_time_ranges(cursor, byte_ranges);
+    EXPECT_EQ(time_ranges.size(), 1u);
+    EXPECT_EQ(time_ranges[0].start, AK::Duration::zero());
+    EXPECT_EQ(time_ranges[0].end, AK::Duration::from_nanoseconds(1499666666));
+
+    // Get buffered ranges with only clusters 8 and 9 available.
+    byte_ranges[0] = { 91687, 113303 };
+    time_ranges = reader.buffered_time_ranges(cursor, byte_ranges);
+    EXPECT_EQ(time_ranges.size(), 1u);
+    EXPECT_EQ(time_ranges[0].start, AK::Duration::from_milliseconds(4000));
+    EXPECT_EQ(time_ranges[0].end, AK::Duration::from_nanoseconds(4999666666));
+
+    // Get buffered ranges with a byte range containing no clusters.
+    byte_ranges[0] = { 113303, file_data.size() };
+    time_ranges = reader.buffered_time_ranges(cursor, byte_ranges);
+    EXPECT_EQ(time_ranges.size(), 0u);
+}