cbs_av1: Make fake OBU size length field a write option

This is an option to modify the behaviour of the writer, not a syntax
field.

3 files changed, 28 insertions, 12 deletions

diff --git a/libavcodec/cbs_av1.c b/libavcodec/cbs_av1.c
index 4e687ace79..ed9a7b80d4 100644
--- a/libavcodec/cbs_av1.c
+++ b/libavcodec/cbs_av1.c
@@ -138,19 +138,25 @@ static int cbs_av1_read_leb128(CodedBitstreamContext *ctx, GetBitContext *gbc,
     return 0;
 }
 
-/** Minimum byte length will be used to indicate the len128 of value if byte_len is 0. */
 static int cbs_av1_write_leb128(CodedBitstreamContext *ctx, PutBitContext *pbc,
-                                const char *name, uint64_t value, uint8_t byte_len)
+                                const char *name, uint64_t value, int fixed_length)
 {
     int len, i;
     uint8_t byte;
 
     CBS_TRACE_WRITE_START();
 
-    if (byte_len)
-        av_assert0(byte_len >= (av_log2(value) + 7) / 7);
+    len = (av_log2(value) + 7) / 7;
 
-    len = byte_len ? byte_len : (av_log2(value) + 7) / 7;
+    if (fixed_length) {
+        if (fixed_length < len) {
+            av_log(ctx->log_ctx, AV_LOG_ERROR, "OBU is too large for "
+                   "fixed length size field (%d > %d).\n",
+                   len, fixed_length);
+            return AVERROR(EINVAL);
+        }
+        len = fixed_length;
+    }
 
     for (i = 0; i < len; i++) {
         if (put_bits_left(pbc) < 8)
@@ -1006,8 +1012,8 @@ static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
 
     if (obu->header.obu_has_size_field) {
         pbc_tmp = *pbc;
-        if (obu->obu_size_byte_len) {
-            for (int i = 0; i < obu->obu_size_byte_len; i++)
+        if (priv->fixed_obu_size_length) {
+            for (int i = 0; i < priv->fixed_obu_size_length; i++)
                 put_bits(pbc, 8, 0);
         } else {
             // Add space for the size field to fill later.
@@ -1133,7 +1139,8 @@ static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
     end_pos   /= 8;
 
     *pbc = pbc_tmp;
-    err = cbs_av1_write_leb128(ctx, pbc, "obu_size", obu->obu_size, obu->obu_size_byte_len);
+    err = cbs_av1_write_leb128(ctx, pbc, "obu_size", obu->obu_size,
+                               priv->fixed_obu_size_length);
     if (err < 0)
         goto error;
 
@@ -1150,10 +1157,12 @@ static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
     }
 
     if (obu->obu_size > 0) {
-        if (!obu->obu_size_byte_len) {
-            obu->obu_size_byte_len = start_pos - data_pos;
+        if (!priv->fixed_obu_size_length) {
             memmove(pbc->buf + data_pos,
                     pbc->buf + start_pos, header_size);
+        } else {
+            // The size was fixed so the following data was
+            // already written in the correct place.
         }
         skip_put_bytes(pbc, header_size);
 
@@ -1273,6 +1282,8 @@ static const CodedBitstreamUnitTypeDescriptor cbs_av1_unit_types[] = {
 static const AVOption cbs_av1_options[] = {
     { "operating_point",  "Set operating point to select layers to parse from a scalable bitstream",
                           OFFSET(operating_point), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, AV1_MAX_OPERATING_POINTS - 1, 0 },
+    { "fixed_obu_size_length", "Set fixed length of the obu_size field",
+      OFFSET(fixed_obu_size_length), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 8, 0 },
     { NULL }
 };
 
diff --git a/libavcodec/cbs_av1.h b/libavcodec/cbs_av1.h
index a9e2d2284f..7924257164 100644
--- a/libavcodec/cbs_av1.h
+++ b/libavcodec/cbs_av1.h
@@ -401,7 +401,6 @@ typedef struct AV1RawOBU {
     AV1RawOBUHeader header;
 
     size_t obu_size;
-    uint8_t obu_size_byte_len;
 
     union {
         AV1RawSequenceHeader sequence_header;
@@ -468,6 +467,10 @@ typedef struct CodedBitstreamAV1Context {
 
     // AVOptions
     int operating_point;
+    // When writing, fix the length in bytes of the obu_size field.
+    // Writing will fail with an error if an OBU larger than can be
+    // represented by the fixed size is encountered.
+    int fixed_obu_size_length;
 } CodedBitstreamAV1Context;
 
 
diff --git a/libavcodec/vaapi_encode_av1.c b/libavcodec/vaapi_encode_av1.c
index 3ff1c47b53..5a9ff0f798 100644
--- a/libavcodec/vaapi_encode_av1.c
+++ b/libavcodec/vaapi_encode_av1.c
@@ -634,7 +634,6 @@ static int vaapi_encode_av1_init_picture_params(AVCodecContext *avctx,
         }
     }
 
-    fh_obu->obu_size_byte_len = priv->attr_ext2.bits.obu_size_bytes_minus1 + 1;
     ret = vaapi_encode_av1_add_obu(avctx, obu, AV1_OBU_FRAME_HEADER, &priv->fh);
     if (ret < 0)
         goto end;
@@ -839,6 +838,9 @@ static av_cold int vaapi_encode_av1_init(AVCodecContext *avctx)
         priv->attr_ext2.value = attr.value;
     }
 
+    av_opt_set_int(priv->cbc->priv_data, "fixed_obu_size_length",
+                   priv->attr_ext2.bits.obu_size_bytes_minus1 + 1, 0);
+
     ret = vaapi_encode_av1_set_tile(avctx);
     if (ret < 0)
         return ret;