LibMedia+LibWeb: Speculatively find keyframes near fast seek targets

This is the new way of handling fast seeks. Instead of delegating the
logic all the way down the pipeline to the decoder thread's seek
handler, we can just determine the timestamp we want to seek to ahead
of time.

Libraries/LibMedia/Containers/Matroska/MatroskaDemuxer.cpp

@@ -11,6 +11,7 @@
 #include <LibMedia/Containers/Matroska/Utilities.h>
 #include <LibMedia/DecoderError.h>
 #include <LibMedia/MediaStream.h>
+#include <LibMedia/SeekMode.h>
 
 #include "MatroskaDemuxer.h"
 
@@ -98,6 +99,25 @@ DecoderErrorOr<ReadonlyBytes> MatroskaDemuxer::get_codec_initialization_data_for
     return TRY(m_reader.track_for_track_number(track.identifier()))->codec_private_data();
 }
 
+AK::Duration MatroskaDemuxer::select_fast_seek_target_for_track(Track const& track, AK::Duration target, SeekMode mode)
+{
+    auto cue_points = m_reader.cue_points_for_track(track.identifier());
+    if (!cue_points.has_value() || cue_points->is_empty())
+        return target;
+    auto const& points = cue_points.value();
+    auto at_or_before = Reader::find_cue_point_index_at_or_before(points, m_reader.duration(), target);
+
+    auto const& cue_at_or_before = points[at_or_before];
+    if (mode == SeekMode::FastBefore)
+        return cue_at_or_before.timestamp;
+
+    VERIFY(mode == SeekMode::FastAfter);
+    auto after_index = at_or_before + 1;
+    if (after_index >= points.size())
+        return target;
+    return points[after_index].timestamp;
+}
+
 DecoderErrorOr<DemuxerSeekResult> MatroskaDemuxer::seek_to_most_recent_keyframe(Track const& track, AK::Duration timestamp, DemuxerSeekOptions options)
 {
     auto& track_status = get_track_status(track);

Libraries/LibMedia/Containers/Matroska/MatroskaDemuxer.h

@@ -29,6 +29,7 @@ public:
     virtual DecoderErrorOr<Vector<Track>> get_tracks_for_type(TrackType) override;
     virtual DecoderErrorOr<Optional<Track>> get_preferred_track_for_type(TrackType) override;
 
+    virtual AK::Duration select_fast_seek_target_for_track(Track const&, AK::Duration target, SeekMode) override;
     virtual DecoderErrorOr<DemuxerSeekResult> seek_to_most_recent_keyframe(Track const&, AK::Duration timestamp, DemuxerSeekOptions) override;
 
     virtual DecoderErrorOr<AK::Duration> duration_of_track(Track const&) override;

Libraries/LibMedia/Containers/Matroska/Reader.cpp

@@ -1081,40 +1081,40 @@ DecoderErrorOr<void> Reader::parse_cues(Streamer& streamer)
     return {};
 }
 
+size_t Reader::find_cue_point_index_at_or_before(Vector<TrackCuePoint> const& cue_points, Optional<AK::Duration> total_duration, AK::Duration target)
+{
+    VERIFY(!cue_points.is_empty());
+
+    // Take a guess at where in the cues the target will be and correct from there.
+    size_t index = 0;
+    if (total_duration.has_value() && total_duration->to_nanoseconds() > 0)
+        index = clamp(((target.to_nanoseconds() * cue_points.size()) / total_duration->to_nanoseconds()), 0, cue_points.size() - 1);
+    dbgln_if(MATROSKA_DEBUG, "Finding Matroska cue points for timestamp {}ms starting from cue at {}ms", target.to_milliseconds(), cue_points[index].timestamp.to_milliseconds());
+
+    if (cue_points[index].timestamp > target) {
+        while (index > 0 && cue_points[index].timestamp > target) {
+            --index;
+            dbgln_if(MATROSKA_DEBUG, "Checking previous cue point {}ms", cue_points[index].timestamp.to_milliseconds());
+        }
+        if (cue_points[index].timestamp > target)
+            return 0;
+        return index;
+    }
+
+    while (index + 1 < cue_points.size()) {
+        auto const& next_cue_point = cue_points[index + 1];
+        dbgln_if(MATROSKA_DEBUG, "Checking future cue point {}ms", next_cue_point.timestamp.to_milliseconds());
+        if (next_cue_point.timestamp > target)
+            break;
+        ++index;
+    }
+    return index;
+}
+
 DecoderErrorOr<void> Reader::seek_to_cue_for_timestamp(SampleIterator& iterator, AK::Duration const& timestamp, Vector<TrackCuePoint> const& cue_points, CuePointTarget target)
 {
-    // Take a guess at where in the cues the timestamp will be and correct from there.
-    auto duration = m_segment_information.duration();
-    size_t index = 0;
-    if (duration.has_value())
-        index = clamp(((timestamp.to_nanoseconds() * cue_points.size()) / duration->to_nanoseconds()), 0, cue_points.size() - 1);
-
-    auto const* prev_cue_point = &cue_points[index];
-    dbgln_if(MATROSKA_DEBUG, "Finding Matroska cue points for timestamp {}ms starting from cue at {}ms", timestamp.to_milliseconds(), prev_cue_point->timestamp.to_milliseconds());
-
-    if (prev_cue_point->timestamp == timestamp) {
-        TRY(iterator.seek_to_cue_point(*prev_cue_point, target));
-        return {};
-    }
-
-    if (prev_cue_point->timestamp > timestamp) {
-        while (index > 0 && prev_cue_point->timestamp > timestamp) {
-            prev_cue_point = &cue_points[--index];
-            dbgln_if(MATROSKA_DEBUG, "Checking previous cue point {}ms", prev_cue_point->timestamp.to_milliseconds());
-        }
-        TRY(iterator.seek_to_cue_point(*prev_cue_point, target));
-        return {};
-    }
-
-    while (++index < cue_points.size()) {
-        auto const& cue_point = cue_points[index];
-        dbgln_if(MATROSKA_DEBUG, "Checking future cue point {}ms", cue_point.timestamp.to_milliseconds());
-        if (cue_point.timestamp > timestamp)
-            break;
-        prev_cue_point = &cue_point;
-    }
-
-    TRY(iterator.seek_to_cue_point(*prev_cue_point, target));
+    auto index = find_cue_point_index_at_or_before(cue_points, m_segment_information.duration(), timestamp);
+    TRY(iterator.seek_to_cue_point(cue_points[index], target));
     return {};
 }
 

Libraries/LibMedia/Containers/Matroska/Reader.h

@@ -64,6 +64,10 @@ public:
     DecoderErrorOr<SampleIterator> create_sample_iterator(NonnullRefPtr<MediaStreamCursor> const& stream_consumer, u64 track_number);
     DecoderErrorOr<SampleIterator> seek_to_random_access_point(SampleIterator, AK::Duration);
 
+    Optional<Vector<TrackCuePoint> const&> cue_points_for_track(u64 track_number);
+
+    static size_t find_cue_point_index_at_or_before(Vector<TrackCuePoint> const&, Optional<AK::Duration> total_duration, AK::Duration target);
+
 private:
     Reader() = default;
 
@@ -79,7 +83,6 @@ private:
 
     DecoderErrorOr<void> parse_cues(Streamer&);
 
-    Optional<Vector<TrackCuePoint> const&> cue_points_for_track(u64 track_number);
     bool has_cues_for_track(u64 track_number);
     DecoderErrorOr<void> seek_to_cue_for_timestamp(SampleIterator&, AK::Duration const&, Vector<TrackCuePoint> const&, CuePointTarget);
 

Libraries/LibMedia/Demuxer.h

@@ -16,6 +16,7 @@
 #include "CodecID.h"
 #include "CodedFrame.h"
 #include "DecoderError.h"
+#include "SeekMode.h"
 #include "TimeRanges.h"
 #include "Track.h"
 
@@ -50,6 +51,7 @@ public:
 
     virtual DecoderErrorOr<ReadonlyBytes> get_codec_initialization_data_for_track(Track const&) = 0;
 
+    virtual AK::Duration select_fast_seek_target_for_track(Track const&, AK::Duration target, SeekMode) = 0;
     virtual DecoderErrorOr<DemuxerSeekResult> seek_to_most_recent_keyframe(Track const&, AK::Duration timestamp, DemuxerSeekOptions = DemuxerSeekOptions::None) = 0;
 
     virtual DecoderErrorOr<AK::Duration> duration_of_track(Track const&) = 0;

Libraries/LibMedia/FFmpeg/FFmpegDemuxer.cpp

@@ -12,6 +12,7 @@
 #include <LibMedia/FFmpeg/FFmpegDemuxer.h>
 #include <LibMedia/FFmpeg/FFmpegHelpers.h>
 #include <LibMedia/MediaStream.h>
+#include <LibMedia/SeekMode.h>
 
 extern "C" {
 #include <libavformat/avformat.h>
@@ -260,6 +261,21 @@ DecoderErrorOr<Optional<Track>> FFmpegDemuxer::get_preferred_track_for_type(Trac
     return m_stream_info[preferred_index].track;
 }
 
+AK::Duration FFmpegDemuxer::select_fast_seek_target_for_track(Track const&, AK::Duration target, SeekMode)
+{
+    // FIXME: We can do this using the index getter functions, but unfortunately FFmpeg's seek table is in
+    //        DTS -> byte position, so for files with reordered frames (H.264), seeking to a keyframe will often
+    //        result in the first frame back being at a later PTS than the seek target, so we would display a blank
+    //        frame. To avoid this being especially common, just always accurately seek.
+    //
+    //        Note that we can end up showing a blank frame anyway by accurately seeking very close to a keyframe,
+    //        it's just much less likely to happen under normal usage.
+    //
+    //        This FIXME can be dropped when MP4/MOV is demuxed separately from FFmpeg, and then inclusion of an
+    //        index scan here can be re-evaluated.
+    return target;
+}
+
 DecoderErrorOr<DemuxerSeekResult> FFmpegDemuxer::seek_to_most_recent_keyframe(Track const& track, AK::Duration timestamp, DemuxerSeekOptions)
 {
     auto& track_context = get_track_context(track);

Libraries/LibMedia/FFmpeg/FFmpegDemuxer.h

@@ -30,6 +30,7 @@ public:
     virtual DecoderErrorOr<Vector<Track>> get_tracks_for_type(TrackType) override;
     virtual DecoderErrorOr<Optional<Track>> get_preferred_track_for_type(TrackType) override;
 
+    virtual AK::Duration select_fast_seek_target_for_track(Track const&, AK::Duration target, SeekMode) override;
     virtual DecoderErrorOr<DemuxerSeekResult> seek_to_most_recent_keyframe(Track const&, AK::Duration timestamp, DemuxerSeekOptions) override;
 
     virtual DecoderErrorOr<AK::Duration> duration_of_track(Track const&) override;

Libraries/LibMedia/PlaybackStates/SeekingStateHandler.h

@@ -78,9 +78,24 @@ public:
     }
 
 private:
+    AK::Duration choose_timestamp() const
+    {
+        if (m_mode == SeekMode::Accurate)
+            return m_target_timestamp;
+        Optional<AK::Duration> latest_fast_seek_target;
+        for (auto const& video_track_data : manager().m_video_track_datas) {
+            if (video_track_data.display == nullptr)
+                continue;
+            auto fast_seek_target = video_track_data.producer->select_fast_seek_target(m_target_timestamp, m_mode);
+            if (!latest_fast_seek_target.has_value() || fast_seek_target > latest_fast_seek_target.value())
+                latest_fast_seek_target = fast_seek_target;
+        }
+        return latest_fast_seek_target.value_or(m_target_timestamp);
+    }
+
     void begin_seek()
     {
-        m_chosen_timestamp = m_target_timestamp;
+        m_chosen_timestamp = choose_timestamp();
         m_video_seeks_pending.clear();
         m_audio_seek_pending = false;
 

Libraries/LibMedia/Producers/DecodedVideoProducer.cpp

@@ -140,6 +140,11 @@ void DecodedVideoProducer::ThreadData::dispatch_state_if_changed_while_locked(Pi
     });
 }
 
+AK::Duration DecodedVideoProducer::select_fast_seek_target(AK::Duration timestamp, SeekMode mode)
+{
+    return m_thread_data->select_fast_seek_target(timestamp, mode);
+}
+
 void DecodedVideoProducer::seek(AK::Duration timestamp)
 {
     m_thread_data->seek(timestamp);
@@ -228,6 +233,11 @@ void DecodedVideoProducer::ThreadData::seek(AK::Duration timestamp)
     wake();
 }
 
+AK::Duration DecodedVideoProducer::ThreadData::select_fast_seek_target(AK::Duration target, SeekMode mode) const
+{
+    return m_demuxer->select_fast_seek_target_for_track(m_track, target, mode);
+}
+
 void DecodedVideoProducer::ThreadData::wait_for_start()
 {
     auto locker = take_lock();

Libraries/LibMedia/Producers/DecodedVideoProducer.h

@@ -52,6 +52,7 @@ public:
     virtual PipelineStatus pull(RefPtr<VideoFrame>& into) override;
     virtual void set_state_changed_handler(PipelineStateChangeHandler) override;
 
+    AK::Duration select_fast_seek_target(AK::Duration timestamp, SeekMode);
     virtual void seek(AK::Duration timestamp) override;
 
     TimeRanges buffered_time_ranges() const;
@@ -77,6 +78,7 @@ private:
         PipelineStatus pull(RefPtr<VideoFrame>& into);
 
         void seek(AK::Duration timestamp);
+        AK::Duration select_fast_seek_target(AK::Duration target, SeekMode) const;
 
         void wait_for_start();
         bool should_thread_exit_while_locked() const;

Libraries/LibWeb/MediaSourceExtensions/TrackBufferDemuxer.cpp

@@ -179,6 +179,38 @@ Media::DecoderErrorOr<ReadonlyBytes> TrackBufferDemuxer::get_codec_initializatio
     return m_codec_initialization_data.bytes();
 }
 
+AK::Duration TrackBufferDemuxer::select_fast_seek_target_for_track(Media::Track const&, AK::Duration target, Media::SeekMode mode)
+{
+    Sync::MutexLocker locker { m_mutex };
+    if (m_coded_frames.is_empty())
+        return target;
+
+    size_t nearby_index = 0;
+    binary_search(m_coded_frames, target, &nearby_index, [](AK::Duration needle, Media::CodedFrame const& frame) {
+        return needle <=> frame.timestamp();
+    });
+
+    if (mode == Media::SeekMode::FastBefore) {
+        if (m_coded_frames[nearby_index].timestamp() > target)
+            return target;
+        for (auto i = nearby_index; i-- > 0;) {
+            if (m_coded_frames[i].is_keyframe())
+                return m_coded_frames[i].timestamp();
+        }
+        return target;
+    }
+
+    VERIFY(mode == Media::SeekMode::FastAfter);
+    auto start = nearby_index;
+    if (m_coded_frames[start].timestamp() < target)
+        start++;
+    for (auto i = start; i < m_coded_frames.size(); i++) {
+        if (m_coded_frames[i].is_keyframe())
+            return m_coded_frames[i].timestamp();
+    }
+    return target;
+}
+
 Media::DecoderErrorOr<Media::DemuxerSeekResult> TrackBufferDemuxer::seek_to_most_recent_keyframe(Media::Track const&, AK::Duration timestamp, Media::DemuxerSeekOptions)
 {
     Sync::MutexLocker locker { m_mutex };

Libraries/LibWeb/MediaSourceExtensions/TrackBufferDemuxer.h

@@ -42,6 +42,7 @@ public:
     virtual Media::DecoderErrorOr<Media::CodedFrame> get_next_sample_for_track(Media::Track const&) override;
     virtual Media::DecoderErrorOr<Media::CodecID> get_codec_id_for_track(Media::Track const&) override;
     virtual Media::DecoderErrorOr<ReadonlyBytes> get_codec_initialization_data_for_track(Media::Track const&) override;
+    virtual AK::Duration select_fast_seek_target_for_track(Media::Track const&, AK::Duration target, Media::SeekMode) override;
     virtual Media::DecoderErrorOr<Media::DemuxerSeekResult> seek_to_most_recent_keyframe(Media::Track const&, AK::Duration, Media::DemuxerSeekOptions) override;
     virtual Media::DecoderErrorOr<AK::Duration> duration_of_track(Media::Track const&) override;
     virtual Media::DecoderErrorOr<AK::Duration> total_duration() override;