From 3bfef8bfe0d6f196c809d50140c31d4f10a6b533 Mon Sep 17 00:00:00 2001
From: Zaggy1024 <zaggy1024@gmail.com>
Date: Mon, 6 Feb 2023 01:25:02 -0600
Subject: [PATCH] LibVideo: Pass the current sample to demuxers to lazily seek
 better

In cases where the PlaybackManager's earliest buffered or displayed
sample is closer to the seek target than the demuxer's chosen keyframe,
we don't want to seek at all. To enable this, demuxers now receive an
optional parameter with the earliest timestamp that the caller can
still access.

The demuxer in turn returns an optional to indicate when a seek was not
needed, which allows PlaybackManager to avoid clearing its queue and
re-decoding frames.
---
 .../Libraries/LibVideo/Containers/Demuxer.h   |  4 +-
 .../Containers/Matroska/MatroskaDemuxer.cpp   | 20 +++++++++-
 .../Containers/Matroska/MatroskaDemuxer.h     |  2 +-
 .../LibVideo/Containers/Matroska/Reader.cpp   | 37 ++++++-------------
 .../LibVideo/Containers/Matroska/Reader.h     |  6 +--
 .../Libraries/LibVideo/PlaybackManager.cpp    | 32 ++++++++++++----
 Userland/Libraries/LibVideo/PlaybackManager.h |  2 +-
 7 files changed, 61 insertions(+), 42 deletions(-)

diff --git a/Userland/Libraries/LibVideo/Containers/Demuxer.h b/Userland/Libraries/LibVideo/Containers/Demuxer.h
index 1a2d2355d2..0915b8685e 100644
--- a/Userland/Libraries/LibVideo/Containers/Demuxer.h
+++ b/Userland/Libraries/LibVideo/Containers/Demuxer.h
@@ -29,7 +29,9 @@ public:
     }
 
     // Returns the timestamp of the keyframe that was seeked to.
-    virtual DecoderErrorOr<Time> seek_to_most_recent_keyframe(Track track, Time timestamp) = 0;
+    // The value is `Optional` to allow the demuxer to decide not to seek so that it can keep its position
+    // in the case that the timestamp is closer to the current time than the nearest keyframe.
+    virtual DecoderErrorOr<Optional<Time>> seek_to_most_recent_keyframe(Track track, Time timestamp, Optional<Time> earliest_available_sample = OptionalNone()) = 0;
 
     virtual DecoderErrorOr<Time> duration() = 0;
 
diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp b/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp
index 62dbb1625b..674a6428ac 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.cpp
@@ -5,6 +5,7 @@
  */
 
 #include "MatroskaDemuxer.h"
+#include "AK/Debug.h"
 
 namespace Video::Matroska {
 
@@ -53,7 +54,7 @@ DecoderErrorOr<MatroskaDemuxer::TrackStatus*> MatroskaDemuxer::get_track_status(
     return &m_track_statuses.get(track).release_value();
 }
 
-DecoderErrorOr<Time> MatroskaDemuxer::seek_to_most_recent_keyframe(Track track, Time timestamp)
+DecoderErrorOr<Optional<Time>> MatroskaDemuxer::seek_to_most_recent_keyframe(Track track, Time timestamp, Optional<Time> earliest_available_sample)
 {
     // Removing the track status will cause us to start from the beginning.
     if (timestamp.is_zero()) {
@@ -62,7 +63,22 @@ DecoderErrorOr<Time> MatroskaDemuxer::seek_to_most_recent_keyframe(Track track,
     }
 
     auto& track_status = *TRY(get_track_status(track));
-    TRY(m_reader.seek_to_random_access_point(track_status.iterator, timestamp));
+    auto seeked_iterator = TRY(m_reader.seek_to_random_access_point(track_status.iterator, timestamp));
+    VERIFY(seeked_iterator.last_timestamp().has_value());
+
+    auto last_sample = earliest_available_sample;
+    if (!last_sample.has_value()) {
+        last_sample = track_status.iterator.last_timestamp();
+    }
+    if (last_sample.has_value()) {
+        bool skip_seek = seeked_iterator.last_timestamp().value() <= last_sample.value() && last_sample.value() <= timestamp;
+        dbgln_if(MATROSKA_DEBUG, "The last available sample at {}ms is {}closer to target timestamp {}ms than the keyframe at {}ms, {}", last_sample->to_milliseconds(), skip_seek ? ""sv : "not "sv, timestamp.to_milliseconds(), seeked_iterator.last_timestamp()->to_milliseconds(), skip_seek ? "skipping seek"sv : "seeking"sv);
+        if (skip_seek) {
+            return OptionalNone();
+        }
+    }
+
+    track_status.iterator = move(seeked_iterator);
     return track_status.iterator.last_timestamp();
 }
 
diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.h b/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.h
index bb55a8911d..660175ca2c 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.h
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/MatroskaDemuxer.h
@@ -27,7 +27,7 @@ public:
 
     DecoderErrorOr<Vector<Track>> get_tracks_for_type(TrackType type) override;
 
-    DecoderErrorOr<Time> seek_to_most_recent_keyframe(Track track, Time timestamp) override;
+    DecoderErrorOr<Optional<Time>> seek_to_most_recent_keyframe(Track track, Time timestamp, Optional<Time> earliest_available_sample = OptionalNone()) override;
 
     DecoderErrorOr<Time> duration() override;
 
diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp
index 62d9daaf94..b634c3bae9 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.cpp
@@ -806,7 +806,7 @@ DecoderErrorOr<void> Reader::seek_to_cue_for_timestamp(SampleIterator& iterator,
     return {};
 }
 
-static DecoderErrorOr<bool> find_keyframe_before_timestamp(SampleIterator& iterator, Time const& timestamp)
+static DecoderErrorOr<void> search_clusters_for_keyframe_before_timestamp(SampleIterator& iterator, Time const& timestamp)
 {
 #if MATROSKA_DEBUG
     size_t inter_frames_count;
@@ -837,10 +837,9 @@ static DecoderErrorOr<bool> find_keyframe_before_timestamp(SampleIterator& itera
         dbgln("Seeked to a keyframe with {} inter frames to skip", inter_frames_count);
 #endif
         iterator = last_keyframe.release_value();
-        return true;
     }
 
-    return false;
+    return {};
 }
 
 DecoderErrorOr<bool> Reader::has_cues_for_track(u64 track_number)
@@ -849,37 +848,23 @@ DecoderErrorOr<bool> Reader::has_cues_for_track(u64 track_number)
     return m_cues.contains(track_number);
 }
 
-DecoderErrorOr<void> Reader::seek_to_random_access_point(SampleIterator& iterator, Time timestamp)
+DecoderErrorOr<SampleIterator> Reader::seek_to_random_access_point(SampleIterator iterator, Time timestamp)
 {
-    if (iterator.m_last_timestamp == timestamp)
-        return {};
-
     if (TRY(has_cues_for_track(iterator.m_track.track_number()))) {
-        auto seeked_iterator = iterator;
-        TRY(seek_to_cue_for_timestamp(seeked_iterator, timestamp));
-        VERIFY(seeked_iterator.m_last_timestamp <= timestamp);
-
-        // We only need to seek to a keyframe if it's not faster to continue from the current position.
-        if (timestamp < iterator.m_last_timestamp || seeked_iterator.m_last_timestamp > iterator.m_last_timestamp)
-            iterator = seeked_iterator;
-        return {};
+        TRY(seek_to_cue_for_timestamp(iterator, timestamp));
+        VERIFY(iterator.last_timestamp().has_value() && iterator.last_timestamp().value() <= timestamp);
+        return iterator;
     }
 
-    // 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 (!iterator.last_timestamp().has_value() || timestamp < iterator.last_timestamp().value()) {
         // 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.track_number()));
-        if (!TRY(find_keyframe_before_timestamp(iterator, timestamp)))
-            return DecoderError::corrupted("No random access points found"sv);
-
-        return {};
+        TRY(search_clusters_for_keyframe_before_timestamp(iterator, timestamp));
+        return iterator;
     }
 
-    auto seeked_iterator = iterator;
-    if (TRY(find_keyframe_before_timestamp(seeked_iterator, timestamp)))
-        iterator = seeked_iterator;
-    VERIFY(iterator.last_timestamp() <= timestamp);
-    return {};
+    TRY(search_clusters_for_keyframe_before_timestamp(iterator, timestamp));
+    return iterator;
 }
 
 DecoderErrorOr<Optional<Vector<CuePoint> const&>> Reader::cue_points_for_track(u64 track_number)
diff --git a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h
index 5df0912d37..8ee14ec951 100644
--- a/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h
+++ b/Userland/Libraries/LibVideo/Containers/Matroska/Reader.h
@@ -38,7 +38,7 @@ public:
     DecoderErrorOr<size_t> track_count();
 
     DecoderErrorOr<SampleIterator> create_sample_iterator(u64 track_number);
-    DecoderErrorOr<void> seek_to_random_access_point(SampleIterator&, Time);
+    DecoderErrorOr<SampleIterator> seek_to_random_access_point(SampleIterator, Time);
     DecoderErrorOr<Optional<Vector<CuePoint> const&>> cue_points_for_track(u64 track_number);
     DecoderErrorOr<bool> has_cues_for_track(u64 track_number);
 
@@ -83,7 +83,7 @@ class SampleIterator {
 public:
     DecoderErrorOr<Block> next_block();
     Cluster const& current_cluster() { return *m_current_cluster; }
-    Time const& last_timestamp() { return m_last_timestamp; }
+    Optional<Time> const& last_timestamp() { return m_last_timestamp; }
 
 private:
     friend class Reader;
@@ -107,7 +107,7 @@ private:
     // Must always point to an element ID or the end of the stream.
     size_t m_position { 0 };
 
-    Time m_last_timestamp { Time::min() };
+    Optional<Time> m_last_timestamp;
 
     Optional<Cluster> m_current_cluster;
 };
diff --git a/Userland/Libraries/LibVideo/PlaybackManager.cpp b/Userland/Libraries/LibVideo/PlaybackManager.cpp
index 84bfddf81d..9a808402b8 100644
--- a/Userland/Libraries/LibVideo/PlaybackManager.cpp
+++ b/Userland/Libraries/LibVideo/PlaybackManager.cpp
@@ -114,11 +114,11 @@ void PlaybackManager::seek_to_timestamp(Time target_timestamp)
     TRY_OR_FATAL_ERROR(m_playback_handler->seek(target_timestamp, m_seek_mode));
 }
 
-Time PlaybackManager::seek_demuxer_to_most_recent_keyframe(Time timestamp)
+Optional<Time> PlaybackManager::seek_demuxer_to_most_recent_keyframe(Time timestamp, Optional<Time> earliest_available_sample)
 {
     // FIXME: When the demuxer is getting samples off the main thread in the future, this needs to
     //        mutex so that seeking can't happen while that thread is getting a sample.
-    auto result = m_demuxer->seek_to_most_recent_keyframe(m_selected_video_track, timestamp);
+    auto result = m_demuxer->seek_to_most_recent_keyframe(m_selected_video_track, timestamp, move(earliest_available_sample));
     if (result.is_error())
         on_decoder_error(result.release_error());
     return result.release_value();
@@ -459,18 +459,30 @@ private:
 
     ErrorOr<void> on_enter() override
     {
-        auto keyframe_timestamp = manager().seek_demuxer_to_most_recent_keyframe(m_target_timestamp);
-        dbgln_if(PLAYBACK_MANAGER_DEBUG, "{} seeking to timestamp target {}ms, selected keyframe at {}ms", m_seek_mode == SeekMode::Accurate ? "Accurate"sv : "Fast"sv, m_target_timestamp.to_milliseconds(), keyframe_timestamp.to_milliseconds());
+        auto earliest_available_sample = manager().m_last_present_in_media_time;
+        if (manager().m_next_frame.has_value() && manager().m_next_frame->is_frame()) {
+            earliest_available_sample = min(earliest_available_sample, manager().m_next_frame->timestamp());
+        }
+        auto keyframe_timestamp = manager().seek_demuxer_to_most_recent_keyframe(m_target_timestamp, earliest_available_sample);
+
+#if PLAYBACK_MANAGER_DEBUG
+        auto seek_mode_name = m_seek_mode == SeekMode::Accurate ? "Accurate"sv : "Fast"sv;
+        if (keyframe_timestamp.has_value()) {
+            dbgln("{} seeking to timestamp target {}ms, selected keyframe at {}ms", seek_mode_name, m_target_timestamp.to_milliseconds(), keyframe_timestamp->to_milliseconds());
+        } else {
+            dbgln("{} seeking to timestamp target {}ms, demuxer kept its iterator position", seek_mode_name, m_target_timestamp.to_milliseconds());
+        }
+#endif
 
         if (m_seek_mode == SeekMode::Fast) {
-            m_target_timestamp = keyframe_timestamp;
+            m_target_timestamp = keyframe_timestamp.value_or(earliest_available_sample);
         }
 
-        if (m_target_timestamp < manager().m_last_present_in_media_time) {
+        if (keyframe_timestamp.has_value()) {
             dbgln_if(PLAYBACK_MANAGER_DEBUG, "Timestamp is earlier than current media time, clearing queue");
             manager().m_frame_queue->clear();
             manager().m_next_frame.clear();
-        } else if (manager().m_next_frame.has_value() && manager().m_next_frame.value().timestamp() > m_target_timestamp) {
+        } else if (m_target_timestamp >= manager().m_last_present_in_media_time && manager().m_next_frame.has_value() && manager().m_next_frame.value().timestamp() > m_target_timestamp) {
             manager().m_last_present_in_media_time = m_target_timestamp;
             return exit_seek();
         }
@@ -571,7 +583,11 @@ private:
 
     ErrorOr<void> play() override
     {
-        manager().m_last_present_in_media_time = manager().seek_demuxer_to_most_recent_keyframe(Time::zero());
+        manager().m_next_frame.clear();
+        manager().m_frame_queue->clear();
+        auto start_timestamp = manager().seek_demuxer_to_most_recent_keyframe(Time::zero());
+        VERIFY(start_timestamp.has_value());
+        manager().m_last_present_in_media_time = start_timestamp.release_value();
         return replace_handler_and_delete_this<PlayingStateHandler>();
     }
     bool is_playing() override { return false; };
diff --git a/Userland/Libraries/LibVideo/PlaybackManager.h b/Userland/Libraries/LibVideo/PlaybackManager.h
index 2006fa6121..3a92b2453d 100644
--- a/Userland/Libraries/LibVideo/PlaybackManager.h
+++ b/Userland/Libraries/LibVideo/PlaybackManager.h
@@ -126,7 +126,7 @@ private:
 
     void start_timer(int milliseconds);
     void timer_callback();
-    Time seek_demuxer_to_most_recent_keyframe(Time timestamp);
+    Optional<Time> seek_demuxer_to_most_recent_keyframe(Time timestamp, Optional<Time> earliest_available_sample = OptionalNone());
 
     bool decode_and_queue_one_sample();
     void on_decode_timer();