From eef8867d9e6b35c9ad8b472f7597c455dd9db455 Mon Sep 17 00:00:00 2001
From: Zaggy1024 <zaggy1024@gmail.com>
Date: Sat, 12 Nov 2022 04:04:13 -0600
Subject: [PATCH] LibVideo: Implement Matroska keyframe search for when there
 are no Cues

This just searches sequentially through each block in a SampleIterator
until it finds a block after the specified seek timestamp. Once it
finds one, it will try to set the input/output iterator to the most
recent keyframe. If the iterator's original position is closer to the
target, however, it leaves it at that original position, allowing
callers to continue decoding from that position until they reach the
target timestamp.
---
 .../Containers/Matroska/MatroskaDemuxer.cpp   | 13 ++--
 .../LibVideo/Containers/Matroska/Reader.cpp   | 61 +++++++++++++++++--
 .../LibVideo/Containers/Matroska/Reader.h     |  2 +
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp b/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp
index 10984bd6de1..8834857a893 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp
@@ -53,13 +53,16 @@ DecoderErrorOr<MatroskaDemuxer::TrackStatus*> MatroskaDemuxer::get_track_status(
     return &m_track_statuses.get(track).release_value();
 }
 
-DecoderErrorOr<void> MatroskaDemuxer::seek_to_most_recent_keyframe(Track track, Time)
+DecoderErrorOr<void> MatroskaDemuxer::seek_to_most_recent_keyframe(Track track, Time timestamp)
 {
     // Removing the track status will cause us to start from the beginning.
-    // FIXME: We just go back to the beginning always, so that the PlaybackManager seeking
-    //        technology can be tested.
-    m_track_statuses.remove(track);
-    return {};
+    if (timestamp.is_zero()) {
+        m_track_statuses.remove(track);
+        return {};
+    }
+
+    auto& track_status = *TRY(get_track_status(track));
+    return m_reader.seek_to_random_access_point(track_status.iterator, timestamp);
 }
 
 DecoderErrorOr<NonnullOwnPtr<Sample>> MatroskaDemuxer::get_next_sample_for_track(Track track)
diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp
index de3c3cc54a9..7b3ee4fb256 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp
@@ -599,11 +599,62 @@ DecoderErrorOr<SampleIterator> Reader::create_sample_iterator(u64 track_number)
     return SampleIterator(this->m_mapped_file, segment_view, track_number, position, TRY(segment_information()).timestamp_scale());
 }
 
+static DecoderErrorOr<bool> find_keyframe_before_timestamp(SampleIterator& iterator, Time const& timestamp)
+{
+#if MATROSKA_DEBUG
+    size_t inter_frames_count;
+#endif
+    Optional<SampleIterator> last_keyframe;
+
+    while (true) {
+        SampleIterator rewind_iterator = iterator;
+        auto block = TRY(iterator.next_block());
+
+        if (block.only_keyframes()) {
+            last_keyframe.emplace(rewind_iterator);
+#if MATROSKA_DEBUG
+            inter_frames_count = 0;
+#endif
+        }
+
+        if (block.timestamp() > timestamp)
+            break;
+
+#if MATROSKA_DEBUG
+        inter_frames_count++;
+#endif
+    }
+
+    if (last_keyframe.has_value()) {
+#if MATROSKA_DEBUG
+        dbgln("Seeked to a keyframe with {} inter frames to skip", inter_frames_count);
+#endif
+        iterator = last_keyframe.release_value();
+        return true;
+    }
+
+    return false;
+}
+
 DecoderErrorOr<void> Reader::seek_to_random_access_point(SampleIterator& iterator, Time timestamp)
 {
-    (void)iterator;
-    (void)timestamp;
-    return DecoderError::not_implemented();
+    // FIXME: Use Cues to look these up if the element is present.
+
+    // FIXME: This could cache the keyframes it finds. Is it worth doing? Probably not, most files will have Cues :^)
+    if (timestamp < iterator.last_timestamp() || iterator.last_timestamp().is_negative()) {
+        // If the timestamp is before the iterator's current position, then we need to start from the beginning of the Segment.
+        iterator = TRY(create_sample_iterator(iterator.m_track_id));
+        if (!TRY(find_keyframe_before_timestamp(iterator, timestamp)))
+            return DecoderError::corrupted("No random access points found"sv);
+
+        return {};
+    }
+
+    auto seeked_iterator = iterator;
+    if (TRY(find_keyframe_before_timestamp(seeked_iterator, timestamp)))
+        iterator = seeked_iterator;
+    VERIFY(iterator.last_timestamp() <= timestamp);
+    return {};
 }
 
 DecoderErrorOr<Block> SampleIterator::next_block()
@@ -639,8 +690,10 @@ DecoderErrorOr<Block> SampleIterator::next_block()
         }
 
         m_position = streamer.position();
-        if (block.has_value())
+        if (block.has_value()) {
+            m_last_timestamp = block->timestamp();
             return block.release_value();
+        }
     }
 
     m_current_cluster.clear();
diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h
index e6a607bd00e..b96151d1b49 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h
@@ -75,6 +75,7 @@ class SampleIterator {
 public:
     DecoderErrorOr<Block> next_block();
     Cluster const& current_cluster() { return *m_current_cluster; }
+    Time const& last_timestamp() { return m_last_timestamp; }
 
 private:
     friend class Reader;
@@ -98,6 +99,7 @@ private:
     size_t m_position { 0 };
 
     u64 m_timestamp_scale { 0 };
+    Time m_last_timestamp { Time::min() };
 
     Optional<Cluster> m_current_cluster;
 };