Refactor the sequence detection function to detect other sequence elements

5 files changed, 26 insertions, 13 deletions

diff --git a/decoder/LAVVideo/Media.cpp b/decoder/LAVVideo/Media.cpp
index fd4d3d54..c5ccc910 100644
--- a/decoder/LAVVideo/Media.cpp
+++ b/decoder/LAVVideo/Media.cpp
@@ -747,8 +747,9 @@ void fillDXVAExtFormat(DXVA2_ExtendedFormat &fmt, int range, int primaries, int
 
 extern "C" const uint8_t *avpriv_find_start_code(const uint8_t *p, const uint8_t *end, uint32_t *state);
 
-const uint8_t* CheckForEndOfSequence(AVCodecID codec, const uint8_t *buf, long len, uint32_t *state)
+int CheckForSequenceMarkers(AVCodecID codec, const uint8_t *buf, long len, uint32_t *state, const uint8_t **pos)
 {
+  int status = 0;
   if (buf && len > 0) {
     const uint8_t *p = buf, *end = buf + len;
     if (codec == AV_CODEC_ID_MPEG2VIDEO) {
@@ -756,13 +757,18 @@ const uint8_t* CheckForEndOfSequence(AVCodecID codec, const uint8_t *buf, long l
         p = avpriv_find_start_code(p, end, state);
         if (*state == 0x000001b7) {
           DbgLog((LOG_TRACE, 50, L"Found SEQ_END_CODE at %p (end: %p)", p, end));
-          return p;
+          status |= STATE_EOS_FOUND;
+          if (pos)
+            *pos = p;
+          return status;
+        } else if (*state == 0x000001b8) {
+          status |= STATE_GOP_FOUND;
         }
       }
     }
-    return nullptr;
+    return status;
   }
-  return nullptr;
+  return status;
 }
 
 int sws_scale2(struct SwsContext *c, const uint8_t *const srcSlice[], const ptrdiff_t srcStride[], int srcSliceY, int srcSliceH, uint8_t *const dst[], const ptrdiff_t dstStride[])
diff --git a/decoder/LAVVideo/Media.h b/decoder/LAVVideo/Media.h
index fa0b2c8a..ae036583 100644
--- a/decoder/LAVVideo/Media.h
+++ b/decoder/LAVVideo/Media.h
@@ -34,6 +34,10 @@ const codec_config_t *get_codec_config(LAVVideoCodec codec);
 
 int flip_plane(BYTE *buffer, int stride, int height);
 void fillDXVAExtFormat(DXVA2_ExtendedFormat &fmt, int range, int primaries, int matrix, int transfer);
-const uint8_t* CheckForEndOfSequence(AVCodecID codec, const uint8_t *buf, long len, uint32_t *state);
+
+#define STATE_NOT_FOUND 0
+#define STATE_EOS_FOUND 1
+#define STATE_GOP_FOUND 2
+int CheckForSequenceMarkers(AVCodecID codec, const uint8_t *buf, long len, uint32_t *state, const uint8_t **pos = nullptr);
 
 int sws_scale2(struct SwsContext *c, const uint8_t *const srcSlice[], const ptrdiff_t srcStride[], int srcSliceY, int srcSliceH, uint8_t *const dst[], const ptrdiff_t dstStride[]);
diff --git a/decoder/LAVVideo/decoders/avcodec.cpp b/decoder/LAVVideo/decoders/avcodec.cpp
index 1ae33d43..5b5c44de 100644
--- a/decoder/LAVVideo/decoders/avcodec.cpp
+++ b/decoder/LAVVideo/decoders/avcodec.cpp
@@ -774,8 +774,8 @@ STDMETHODIMP CDecAvcodec::Decode(const BYTE *buffer, int buflen, REFERENCE_TIME
           avpkt.pts = rtStart;
           avpkt.duration = 0;
 
-          const uint8_t *eosmarker = CheckForEndOfSequence(m_nCodecId, avpkt.data, avpkt.size, &m_MpegParserState);
-          if (eosmarker) {
+          int state = CheckForSequenceMarkers(m_nCodecId, avpkt.data, avpkt.size, &m_MpegParserState);
+          if (state & STATE_EOS_FOUND) {
             bEndOfSequence = TRUE;
           }
         } else {
diff --git a/decoder/LAVVideo/decoders/cuvid.cpp b/decoder/LAVVideo/decoders/cuvid.cpp
index eb385418..ecc73d5c 100644
--- a/decoder/LAVVideo/decoders/cuvid.cpp
+++ b/decoder/LAVVideo/decoders/cuvid.cpp
@@ -1173,11 +1173,13 @@ STDMETHODIMP CDecCuvid::Decode(const BYTE *buffer, int buflen, REFERENCE_TIME rt
     pCuvidPacket.payload_size = buflen;
 
     if (m_VideoDecoderInfo.CodecType == cudaVideoCodec_MPEG2) {
-      const uint8_t *eosmarker = CheckForEndOfSequence(AV_CODEC_ID_MPEG2VIDEO, buffer, buflen, &m_MpegParserState);
-      const uint8_t *end = buffer+buflen;
+      const uint8_t *eosmarker = nullptr;
+      const uint8_t *end = buffer + buflen;
+      int status = CheckForSequenceMarkers(AV_CODEC_ID_MPEG2VIDEO, buffer, buflen, &m_MpegParserState, &eosmarker);
+
       // If we found a EOS marker, but its not at the end of the packet, then split the packet
       // to be able to individually decode the frame before the EOS, and then decode the remainder
-      if (eosmarker && eosmarker != end) {
+      if (status & STATE_EOS_FOUND && eosmarker && eosmarker != end) {
         Decode(buffer, (int)(eosmarker - buffer), rtStart, rtStop, bSyncPoint, bDiscontinuity);
 
         rtStart = rtStop = AV_NOPTS_VALUE;
diff --git a/decoder/LAVVideo/decoders/quicksync.cpp b/decoder/LAVVideo/decoders/quicksync.cpp
index cfaf141e..6028a75f 100644
--- a/decoder/LAVVideo/decoders/quicksync.cpp
+++ b/decoder/LAVVideo/decoders/quicksync.cpp
@@ -532,11 +532,12 @@ STDMETHODIMP CDecQuickSync::Decode(const BYTE *buffer, int buflen, REFERENCE_TIM
   }
 
   if (m_Codec == FourCC_MPG2) {
-    const uint8_t *eosmarker = CheckForEndOfSequence(AV_CODEC_ID_MPEG2VIDEO, buffer, buflen, &m_MpegParserState);
-    const uint8_t *end = buffer+buflen;
+    const uint8_t *eosmarker = nullptr;
+    const uint8_t *end = buffer + buflen;
+    int status = CheckForSequenceMarkers(AV_CODEC_ID_MPEG2VIDEO, buffer, buflen, &m_MpegParserState, &eosmarker);
     // If we found a EOS marker, but its not at the end of the packet, then split the packet
     // to be able to individually decode the frame before the EOS, and then decode the remainder
-    if (eosmarker && eosmarker != end) {
+    if (status & STATE_EOS_FOUND && eosmarker && eosmarker != end) {
       Decode(buffer, (int)(eosmarker - buffer), rtStart, rtStop, bSyncPoint, bDiscontinuity);
 
       rtStart = rtStop = AV_NOPTS_VALUE;