avcodec/cbs_av1: implement missing set_frame_refs() function

Defined in Section 7.8

This finishes implementing support for frames using
frame_refs_short_signaling.

Signed-off-by: James Almer <jamrial@gmail.com>

1 files changed, 112 insertions, 10 deletions

diff --git a/libavcodec/cbs_av1_syntax_template.c b/libavcodec/cbs_av1_syntax_template.c
index c843cfa02b..796897ab79 100644
--- a/libavcodec/cbs_av1_syntax_template.c
+++ b/libavcodec/cbs_av1_syntax_template.c
@@ -339,6 +339,117 @@ static int FUNC(temporal_delimiter_obu)(CodedBitstreamContext *ctx, RWContext *r
     return 0;
 }
 
+static int FUNC(set_frame_refs)(CodedBitstreamContext *ctx, RWContext *rw,
+                                AV1RawFrameHeader *current)
+{
+    CodedBitstreamAV1Context *priv = ctx->priv_data;
+    const AV1RawSequenceHeader *seq = priv->sequence_header;
+    static const uint8_t ref_frame_list[AV1_NUM_REF_FRAMES - 2] = {
+        AV1_REF_FRAME_LAST2, AV1_REF_FRAME_LAST3, AV1_REF_FRAME_BWDREF,
+        AV1_REF_FRAME_ALTREF2, AV1_REF_FRAME_ALTREF
+    };
+    int8_t ref_frame_idx[AV1_REFS_PER_FRAME], used_frame[AV1_NUM_REF_FRAMES];
+    int8_t shifted_order_hints[AV1_NUM_REF_FRAMES];
+    int cur_frame_hint, latest_order_hint, earliest_order_hint, ref;
+    int i, j;
+
+    for (i = 0; i < AV1_REFS_PER_FRAME; i++)
+        ref_frame_idx[i] = -1;
+    ref_frame_idx[AV1_REF_FRAME_LAST - AV1_REF_FRAME_LAST] = current->last_frame_idx;
+    ref_frame_idx[AV1_REF_FRAME_GOLDEN - AV1_REF_FRAME_LAST] = current->golden_frame_idx;
+
+    for (i = 0; i < AV1_NUM_REF_FRAMES; i++)
+        used_frame[i] = 0;
+    used_frame[current->last_frame_idx] = 1;
+    used_frame[current->golden_frame_idx] = 1;
+
+    cur_frame_hint = 1 << (seq->order_hint_bits_minus_1);
+    for (i = 0; i < AV1_NUM_REF_FRAMES; i++)
+        shifted_order_hints[i] = cur_frame_hint +
+                                 cbs_av1_get_relative_dist(seq, priv->ref[i].order_hint,
+                                                           current->order_hint);
+
+    latest_order_hint = shifted_order_hints[current->last_frame_idx];
+    earliest_order_hint = shifted_order_hints[current->golden_frame_idx];
+
+    ref = -1;
+    for (i = 0; i < AV1_NUM_REF_FRAMES; i++) {
+        int hint = shifted_order_hints[i];
+        if (!used_frame[i] && hint >= cur_frame_hint &&
+            (ref < 0 || hint >= latest_order_hint)) {
+            ref = i;
+            latest_order_hint = hint;
+        }
+    }
+    if (ref >= 0) {
+        ref_frame_idx[AV1_REF_FRAME_ALTREF - AV1_REF_FRAME_LAST] = ref;
+        used_frame[ref] = 1;
+    }
+
+    ref = -1;
+    for (i = 0; i < AV1_NUM_REF_FRAMES; i++) {
+        int hint = shifted_order_hints[i];
+        if (!used_frame[i] && hint >= cur_frame_hint &&
+            (ref < 0 || hint < earliest_order_hint)) {
+            ref = i;
+            earliest_order_hint = hint;
+        }
+    }
+    if (ref >= 0) {
+        ref_frame_idx[AV1_REF_FRAME_BWDREF - AV1_REF_FRAME_LAST] = ref;
+        used_frame[ref] = 1;
+    }
+
+    ref = -1;
+    for (i = 0; i < AV1_NUM_REF_FRAMES; i++) {
+        int hint = shifted_order_hints[i];
+        if (!used_frame[i] && hint >= cur_frame_hint &&
+            (ref < 0 || hint < earliest_order_hint)) {
+            ref = i;
+            earliest_order_hint = hint;
+        }
+    }
+    if (ref >= 0) {
+        ref_frame_idx[AV1_REF_FRAME_ALTREF2 - AV1_REF_FRAME_LAST] = ref;
+        used_frame[ref] = 1;
+    }
+
+    for (i = 0; i < AV1_REFS_PER_FRAME - 2; i++) {
+        int ref_frame = ref_frame_list[i];
+        if (ref_frame_idx[ref_frame - AV1_REF_FRAME_LAST] < 0 ) {
+            ref = -1;
+            for (j = 0; j < AV1_NUM_REF_FRAMES; j++) {
+                int hint = shifted_order_hints[j];
+                if (!used_frame[j] && hint < cur_frame_hint &&
+                    (ref < 0 || hint >= latest_order_hint)) {
+                    ref = j;
+                    latest_order_hint = hint;
+                }
+            }
+            if (ref >= 0) {
+                ref_frame_idx[ref_frame - AV1_REF_FRAME_LAST] = ref;
+                used_frame[ref] = 1;
+            }
+        }
+    }
+
+    ref = -1;
+    for (i = 0; i < AV1_NUM_REF_FRAMES; i++) {
+        int hint = shifted_order_hints[i];
+        if (ref < 0 || hint < earliest_order_hint) {
+            ref = i;
+            earliest_order_hint = hint;
+        }
+    }
+    for (i = 0; i < AV1_REFS_PER_FRAME; i++) {
+        if (ref_frame_idx[i] < 0)
+            ref_frame_idx[i] = ref;
+        infer(ref_frame_idx[i], ref_frame_idx[i]);
+    }
+
+    return 0;
+}
+
 static int FUNC(superres_params)(CodedBitstreamContext *ctx, RWContext *rw,
                                  AV1RawFrameHeader *current)
 {
@@ -1307,16 +1418,7 @@ static int FUNC(uncompressed_header)(CodedBitstreamContext *ctx, RWContext *rw,
             if (current->frame_refs_short_signaling) {
                 fb(3, last_frame_idx);
                 fb(3, golden_frame_idx);
-
-                for (i = 0; i < AV1_REFS_PER_FRAME; i++) {
-                    if (i == 0)
-                        infer(ref_frame_idx[i], current->last_frame_idx);
-                    else if (i == AV1_REF_FRAME_GOLDEN -
-                                  AV1_REF_FRAME_LAST)
-                        infer(ref_frame_idx[i], current->golden_frame_idx);
-                    else
-                        infer(ref_frame_idx[i], -1);
-                }
+                CHECK(FUNC(set_frame_refs)(ctx, rw, current));
             }
         }