lavc/hevc_refs: Fix the logic of find_ref_idx()

Currently find_ref_idx() would trigger 2 scans in DPB to find the
requested POC:
1. Firstly, ignore MSB of ref->poc and search for the requested POC;
2. Secondly, compare the entire ref->poc with requested POC;

For long term reference, we are able to only check LSB if MSB is not
presented(e.g. delta_poc_msb_present_flag == 0). However, for short
term reference, we should never ignore poc's MSB and it should be
kind of bit-exact. (Details in 8.3.2)

Otherwise this leads to decoding failures like:
[hevc @ 0x5638f4328600] Error constructing the frame RPS.
[hevc @ 0x5638f4328600] Error parsing NAL unit #2.
[hevc @ 0x5638f4338a80] Could not find ref with POC 21
Error while decoding stream #0:0: Invalid data found when processing input

Search the requested POC based on whether MSB is used, and avoid
the 2-times scan for DPB buffer. This benefits both native HEVC
decoder and integrated HW decoders.

Signed-off-by: Xu Guangxin <guangxin.xu@intel.com>
Signed-off-by: Linjie Fu <linjie.fu@intel.com>

1 files changed, 20 insertions, 18 deletions

diff --git a/libavcodec/hevc_refs.c b/libavcodec/hevc_refs.c
index 7870a72fd6..73aa6d8caf 100644
--- a/libavcodec/hevc_refs.c
+++ b/libavcodec/hevc_refs.c
@@ -358,24 +358,26 @@ int ff_hevc_slice_rpl(HEVCContext *s)
     return 0;
 }
 
-static HEVCFrame *find_ref_idx(HEVCContext *s, int poc)
+static HEVCFrame *find_ref_idx(HEVCContext *s, int poc, uint8_t use_msb)
 {
     int i;
-    int LtMask = (1 << s->ps.sps->log2_max_poc_lsb) - 1;
 
-    for (i = 0; i < FF_ARRAY_ELEMS(s->DPB); i++) {
-        HEVCFrame *ref = &s->DPB[i];
-        if (ref->frame->buf[0] && (ref->sequence == s->seq_decode)) {
-            if ((ref->poc & LtMask) == poc)
-                return ref;
+    if (use_msb) {
+        for (i = 0; i < FF_ARRAY_ELEMS(s->DPB); i++) {
+            HEVCFrame *ref = &s->DPB[i];
+            if (ref->frame->buf[0] && (ref->sequence == s->seq_decode)) {
+                if (ref->poc == poc)
+                    return ref;
+            }
         }
-    }
-
-    for (i = 0; i < FF_ARRAY_ELEMS(s->DPB); i++) {
-        HEVCFrame *ref = &s->DPB[i];
-        if (ref->frame->buf[0] && ref->sequence == s->seq_decode) {
-            if (ref->poc == poc || (ref->poc & LtMask) == poc)
-                return ref;
+    } else {
+        int LtMask = (1 << s->ps.sps->log2_max_poc_lsb) - 1;
+        for (i = 0; i < FF_ARRAY_ELEMS(s->DPB); i++) {
+            HEVCFrame *ref = &s->DPB[i];
+            if (ref->frame->buf[0] && ref->sequence == s->seq_decode) {
+                if ((ref->poc & LtMask) == poc)
+                    return ref;
+            }
         }
     }
 
@@ -427,9 +429,9 @@ static HEVCFrame *generate_missing_ref(HEVCContext *s, int poc)
 
 /* add a reference with the given poc to the list and mark it as used in DPB */
 static int add_candidate_ref(HEVCContext *s, RefPicList *list,
-                             int poc, int ref_flag)
+                             int poc, int ref_flag, uint8_t use_msb)
 {
-    HEVCFrame *ref = find_ref_idx(s, poc);
+    HEVCFrame *ref = find_ref_idx(s, poc, use_msb);
 
     if (ref == s->ref || list->nb_refs >= HEVC_MAX_REFS)
         return AVERROR_INVALIDDATA;
@@ -485,7 +487,7 @@ int ff_hevc_frame_rps(HEVCContext *s)
         else
             list = ST_CURR_AFT;
 
-        ret = add_candidate_ref(s, &rps[list], poc, HEVC_FRAME_FLAG_SHORT_REF);
+        ret = add_candidate_ref(s, &rps[list], poc, HEVC_FRAME_FLAG_SHORT_REF, 1);
         if (ret < 0)
             goto fail;
     }
@@ -495,7 +497,7 @@ int ff_hevc_frame_rps(HEVCContext *s)
         int poc  = long_rps->poc[i];
         int list = long_rps->used[i] ? LT_CURR : LT_FOLL;
 
-        ret = add_candidate_ref(s, &rps[list], poc, HEVC_FRAME_FLAG_LONG_REF);
+        ret = add_candidate_ref(s, &rps[list], poc, HEVC_FRAME_FLAG_LONG_REF, long_rps->poc_msb_present[i]);
         if (ret < 0)
             goto fail;
     }