Rework packet finding functions

After seeking, GetFrame() tries to find out where the demuxer ended up
after the seek by trying to recognize the first packet that comes out of
the demuxer afterwards by its PTS (or DTS if Frames.UseDTS is true), or
by its Pos if that fails.
Similarly, DecodePacket() needs to recognize the packet to see if it's a
second field of some interlaced frame.
When the same PTS can appear multiple times in the same file (e.g. in
files with edit lists), this can cause issues.

This commit replaces the FrameFromPTS / FrameFromPos functions and the
fallback logic in GetFrame() by a single FindPacket function that takes
an entire AVPacket.
FindPacket() tries to find the unique packet in the index whose
TS/Pos/Flags match the given packet. If no such packet exists,
it tries a fallback sequence of comparing fewer fields.

The current fallback sequence is mostly chosen based on intuition - for
well-behaved files simply checking all fields should work fine. If
examples come up where this fallback sequence becomes relevant, it can
be adjusted.

Author: arch1t3cht

src/core/track.cpp

@@ -25,6 +25,7 @@
 #include "indexing.h"
 
 #include <algorithm>
+#include <array>
 #include <cmath>
 #include <cstdlib>
 
@@ -157,23 +158,79 @@ static bool PTSComparison(FrameInfo FI1, FrameInfo FI2) {
     return FI1.PTS < FI2.PTS;
 }
 
-int FFMS_Track::FrameFromPTS(int64_t PTS, bool AllowHidden) const {
+enum class AVPacketProp {
+    TS,     // can be PTS or DTS depending on UseDTS
+    Pos,
+    Hidden,
+    Key,
+};
+
+const std::array<std::vector<AVPacketProp>, 5> FindPacketCheckSequence = {{
+    {AVPacketProp::TS, AVPacketProp::Pos, AVPacketProp::Hidden, AVPacketProp::Key},
+    {AVPacketProp::TS, AVPacketProp::Pos, AVPacketProp::Hidden},
+    {AVPacketProp::TS, AVPacketProp::Pos},
+    {AVPacketProp::TS},
+    {AVPacketProp::Pos},
+}};
+
+// Attempts to find the given AVPacket in the track's list of FrameInfos,
+// returning the FrameInfo's index if one was found and -1 if not.
+//
+// The finding logic begins by trying to find a *unique* FrameInfo whose
+// PTS, Pos, and flags match the given packet.
+// If no such packet exists, it tries a fallback sequence of comparing fewer fields.
+//
+// The current fallback sequence is mostly chosen based on intuition - for
+// well-behaved files simply checking all fields should work fine.
+// If examples come up where this fallback sequence becomes relevant, it can
+// be adjusted.
+int FFMS_Track::FindPacket(const AVPacket &packet) const {
     FrameInfo F;
-    F.PTS = PTS;
+    F.PTS = UseDTS ? packet.dts : packet.pts;
 
-    auto Pos = std::lower_bound(begin(), end(), F, PTSComparison);
-    while (Pos != end() && (!AllowHidden && Pos->Skipped()) && Pos->PTS == PTS)
-        Pos++;
+    auto SameTSRange = std::equal_range(begin(), end(), F, PTSComparison);;
 
-    if (Pos == end() || Pos->PTS != PTS)
-        return -1;
-    return std::distance(begin(), Pos);
-}
+    for (auto const& checks : FindPacketCheckSequence) {
+        // If the check sequence includes the timestamp, we can begin with a binary search
+        // since the FrameInfos are sorted by their timestamp, and then do a linear search in the found interval.
+        // If the check sequence does not include the timestamp, we have to do a full linear search.
+        auto Begin = begin();
+        auto End = end();
+
+        if (checks[0] == AVPacketProp::TS) {
+            Begin = SameTSRange.first;
+            End = SameTSRange.second;
+        }
+
+        int found = 0;
+        int result = -1;
+
+        for (auto it = Begin; it < End; it++) {
+            bool match = std::all_of(checks.cbegin(), checks.cend(), [&](AVPacketProp check) {
+                    switch (check) {
+                        case AVPacketProp::TS:
+                            return it->PTS == F.PTS;
+                        case AVPacketProp::Pos:
+                            return it->FilePos == packet.pos;
+                        case AVPacketProp::Hidden:
+                            return it->MarkedHidden == (packet.flags & AV_PKT_FLAG_DISCARD);
+                        case AVPacketProp::Key:
+                            return it->KeyFrame == (packet.flags & AV_PKT_FLAG_KEY);
+                    }
+                    return false;
+                });
+
+            if (match) {
+                found++;
+                result = std::distance(begin(), it);
+            }
+        }
+
+        if (found == 1) {
+            return result;
+        }
+    }
 
-int FFMS_Track::FrameFromPos(int64_t Pos, bool AllowHidden) const {
-    for (size_t i = 0; i < size(); i++)
-        if (Data->Frames[i].FilePos == Pos && (AllowHidden || !Data->Frames[i].Skipped()))
-            return static_cast<int>(i);
     return -1;
 }
 

src/core/track.h

@@ -27,6 +27,7 @@
 #include <vector>
 #include <memory>
 
+struct AVPacket;
 class ZipFile;
 
 struct FrameInfo {
@@ -82,8 +83,7 @@ struct FFMS_Track {
     void FillAudioGaps();
 
     int FindClosestVideoKeyFrame(int Frame) const;
-    int FrameFromPTS(int64_t PTS, bool AllowHidden = false) const;
-    int FrameFromPos(int64_t Pos, bool AllowHidden = false) const;
+    int FindPacket(const AVPacket &packet) const;
     int ClosestFrameFromPTS(int64_t PTS) const;
     int RealFrameNumber(int Frame) const;
     int VisibleFrameCount() const;

src/core/videosource.cpp

@@ -744,7 +744,7 @@ bool FFMS_VideoSource::DecodePacket(const AVPacket &Packet) {
     std::swap(DecodeFrame, LastDecodedFrame);
     ResendPacket = false;
 
-    int PacketNum = Frames.FrameFromPTS(Frames.UseDTS ? Packet.dts : Packet.pts, true);
+    int PacketNum = Frames.FindPacket(Packet);
     bool PacketHidden = !!(Packet.flags & AV_PKT_FLAG_DISCARD) || (PacketNum != -1 && Frames[PacketNum].MarkedHidden);
     bool SecondField = PacketNum != -1 && Frames[PacketNum].SecondField;
 
@@ -862,10 +862,8 @@ SmartAVPacket FFMS_VideoSource::DecodeNextFrame() {
             continue;
         }
 
-        if (!FirstPacket->data || (Frames.UseDTS ? FirstPacket->dts : FirstPacket->pts) < 0) {
-            av_packet_unref(FirstPacket.get());
+        if (!FirstPacket->data)
             av_packet_ref(FirstPacket.get(), Packet.get());
-        }
 
         bool FrameFinished = DecodePacket(*Packet);
         if (ResendPacket)
@@ -966,23 +964,7 @@ FFMS_Frame *FFMS_VideoSource::GetFrame(int n) {
             continue;
 
         int64_t StartTime = FirstPacket->data == nullptr ? AV_NOPTS_VALUE : (Frames.UseDTS ? FirstPacket->dts : FirstPacket->pts);
-
-        if (StartTime == AV_NOPTS_VALUE && !Frames.HasTS) {
-            if (FirstPacket->data) {
-                CurrentFrame = Frames.FrameFromPos(FirstPacket->pos);
-                if (CurrentFrame >= 0)
-                    continue;
-            }
-            // If the track doesn't have timestamps or file positions then
-            // just trust that we got to the right place, since we have no
-            // way to tell where we are
-            else {
-                CurrentFrame = n;
-                continue;
-            }
-        }
-
-        CurrentFrame = Frames.FrameFromPTS(StartTime);
+        CurrentFrame = Frames.FindPacket(*FirstPacket);
 
         // Is the seek destination time known? Does it belong to a frame?
         if (CurrentFrame < 0) {

src/core/videosource.h

@@ -131,7 +131,7 @@ struct FFMS_VideoSource {
     void SetVideoProperties();
     bool DecodePacket(const AVPacket &Packet);
 
-    // Returns the first packet read that has nonzero PTS/DTS (depending on Frames.UseDTS)
+    // Returns the first packet read
     SmartAVPacket DecodeNextFrame();
     bool SeekTo(int n, int SeekOffset);
     int Seek(int n);