h2645_parse: Make ff_h2645_packet_split reference-compatible

This is in preparation for a patch for cbs_h2645. Now the packet's
rbsp_buffer can be owned by an AVBuffer.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@googlemail.com>

9 files changed, 70 insertions, 15 deletions

diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 666970ed03..574a53a60a 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -612,7 +612,7 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 
         err = ff_h2645_packet_split(&priv->read_packet,
                                     frag->data + start, end - start,
-                                    ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1);
+                                    ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 0);
         if (err < 0) {
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC SPS array.\n");
             return err;
@@ -636,7 +636,7 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 
         err = ff_h2645_packet_split(&priv->read_packet,
                                     frag->data + start, end - start,
-                                    ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1);
+                                    ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 0);
         if (err < 0) {
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC PPS array.\n");
             return err;
@@ -690,7 +690,7 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 
             err = ff_h2645_packet_split(&priv->read_packet,
                                         frag->data + start, end - start,
-                                        ctx->log_ctx, 1, 2, AV_CODEC_ID_HEVC, 1);
+                                        ctx->log_ctx, 1, 2, AV_CODEC_ID_HEVC, 1, 0);
             if (err < 0) {
                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split "
                        "HVCC array %d (%d NAL units of type %d).\n",
@@ -709,7 +709,7 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
                                     frag->data, frag->data_size,
                                     ctx->log_ctx,
                                     priv->mp4, priv->nal_length_size,
-                                    codec_id, 1);
+                                    codec_id, 1, 0);
         if (err < 0)
             return err;
 
diff --git a/libavcodec/extract_extradata_bsf.c b/libavcodec/extract_extradata_bsf.c
index f37427c7e1..17e5deb96b 100644
--- a/libavcodec/extract_extradata_bsf.c
+++ b/libavcodec/extract_extradata_bsf.c
@@ -157,7 +157,7 @@ static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
     }
 
     ret = ff_h2645_packet_split(&s->h2645_pkt, pkt->data, pkt->size,
-                                ctx, 0, 0, ctx->par_in->codec_id, 1);
+                                ctx, 0, 0, ctx->par_in->codec_id, 1, 0);
     if (ret < 0)
         return ret;
 
diff --git a/libavcodec/h2645_parse.c b/libavcodec/h2645_parse.c
index aaa4b8f443..942f2c5d71 100644
--- a/libavcodec/h2645_parse.c
+++ b/libavcodec/h2645_parse.c
@@ -343,9 +343,51 @@ static int find_next_start_code(const uint8_t *buf, const uint8_t *next_avc)
     return i + 3;
 }
 
+static void alloc_rbsp_buffer(H2645RBSP *rbsp, unsigned int size, int use_ref)
+{
+    if (size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE)
+        goto fail;
+    size += AV_INPUT_BUFFER_PADDING_SIZE;
+
+    if (rbsp->rbsp_buffer_alloc_size >= size &&
+        (!rbsp->rbsp_buffer_ref || av_buffer_is_writable(rbsp->rbsp_buffer_ref)))
+        return;
+
+    size = FFMIN(size + size / 16 + 32, INT_MAX);
+
+    if (rbsp->rbsp_buffer_ref)
+        av_buffer_unref(&rbsp->rbsp_buffer_ref);
+    else
+        av_free(rbsp->rbsp_buffer);
+
+    rbsp->rbsp_buffer = av_malloc(size);
+    if (!rbsp->rbsp_buffer)
+        goto fail;
+    rbsp->rbsp_buffer_alloc_size = size;
+
+    if (use_ref) {
+        rbsp->rbsp_buffer_ref = av_buffer_create(rbsp->rbsp_buffer, size,
+                                                 NULL, NULL, 0);
+        if (!rbsp->rbsp_buffer_ref)
+            goto fail;
+    }
+
+    return;
+
+fail:
+    rbsp->rbsp_buffer_alloc_size = 0;
+    if (rbsp->rbsp_buffer_ref) {
+        av_buffer_unref(&rbsp->rbsp_buffer_ref);
+        rbsp->rbsp_buffer = NULL;
+    } else
+        av_freep(&rbsp->rbsp_buffer);
+
+    return;
+}
+
 int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
                           void *logctx, int is_nalff, int nal_length_size,
-                          enum AVCodecID codec_id, int small_padding)
+                          enum AVCodecID codec_id, int small_padding, int use_ref)
 {
     GetByteContext bc;
     int consumed, ret = 0;
@@ -353,7 +395,8 @@ int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
     int64_t padding = small_padding ? 0 : MAX_MBPAIR_SIZE;
 
     bytestream2_init(&bc, buf, length);
-    av_fast_padded_malloc(&pkt->rbsp.rbsp_buffer, &pkt->rbsp.rbsp_buffer_alloc_size, length + padding);
+    alloc_rbsp_buffer(&pkt->rbsp, length + padding, use_ref);
+
     if (!pkt->rbsp.rbsp_buffer)
         return AVERROR(ENOMEM);
 
@@ -474,6 +517,10 @@ void ff_h2645_packet_uninit(H2645Packet *pkt)
     }
     av_freep(&pkt->nals);
     pkt->nals_allocated = 0;
-    av_freep(&pkt->rbsp.rbsp_buffer);
+    if (pkt->rbsp.rbsp_buffer_ref) {
+        av_buffer_unref(&pkt->rbsp.rbsp_buffer_ref);
+        pkt->rbsp.rbsp_buffer = NULL;
+    } else
+        av_freep(&pkt->rbsp.rbsp_buffer);
     pkt->rbsp.rbsp_buffer_alloc_size = pkt->rbsp.rbsp_buffer_size = 0;
 }
diff --git a/libavcodec/h2645_parse.h b/libavcodec/h2645_parse.h
index 6dbba8fe4a..2c29ca517c 100644
--- a/libavcodec/h2645_parse.h
+++ b/libavcodec/h2645_parse.h
@@ -23,6 +23,7 @@
 
 #include <stdint.h>
 
+#include "libavutil/buffer.h"
 #include "avcodec.h"
 #include "get_bits.h"
 
@@ -66,6 +67,7 @@ typedef struct H2645NAL {
 
 typedef struct H2645RBSP {
     uint8_t *rbsp_buffer;
+    AVBufferRef *rbsp_buffer_ref;
     int rbsp_buffer_alloc_size;
     int rbsp_buffer_size;
 } H2645RBSP;
@@ -92,10 +94,15 @@ int ff_h2645_extract_rbsp(const uint8_t *src, int length, H2645RBSP *rbsp,
  * the data is contained in the input buffer pointed to by buf.
  * Otherwise, the unescaped data is part of the rbsp_buffer described by the
  * packet's H2645RBSP.
+ *
+ * If the packet's rbsp_buffer_ref is not NULL, the underlying AVBuffer must
+ * own rbsp_buffer. If not and rbsp_buffer is not NULL, use_ref must be 0.
+ * If use_ref is set, rbsp_buffer will be reference-counted and owned by
+ * the underlying AVBuffer of rbsp_buffer_ref.
  */
 int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
                           void *logctx, int is_nalff, int nal_length_size,
-                          enum AVCodecID codec_id, int small_padding);
+                          enum AVCodecID codec_id, int small_padding, int use_ref);
 
 /**
  * Free all the allocated memory in the packet.
diff --git a/libavcodec/h264_parse.c b/libavcodec/h264_parse.c
index 34ffe3b1fe..290ab681a9 100644
--- a/libavcodec/h264_parse.c
+++ b/libavcodec/h264_parse.c
@@ -359,7 +359,7 @@ static int decode_extradata_ps(const uint8_t *data, int size, H264ParamSets *ps,
     H2645Packet pkt = { 0 };
     int i, ret = 0;
 
-    ret = ff_h2645_packet_split(&pkt, data, size, logctx, is_avc, 2, AV_CODEC_ID_H264, 1);
+    ret = ff_h2645_packet_split(&pkt, data, size, logctx, is_avc, 2, AV_CODEC_ID_H264, 1, 0);
     if (ret < 0) {
         ret = 0;
         goto fail;
diff --git a/libavcodec/h264dec.c b/libavcodec/h264dec.c
index 00d922fbe9..837c3b7538 100644
--- a/libavcodec/h264dec.c
+++ b/libavcodec/h264dec.c
@@ -622,8 +622,8 @@ static int decode_nal_units(H264Context *h, const uint8_t *buf, int buf_size)
             h->is_avc = 1;
     }
 
-    ret = ff_h2645_packet_split(&h->pkt, buf, buf_size, avctx, h->is_avc,
-                                h->nal_length_size, avctx->codec_id, avctx->flags2 & AV_CODEC_FLAG2_FAST);
+    ret = ff_h2645_packet_split(&h->pkt, buf, buf_size, avctx, h->is_avc, h->nal_length_size,
+                                avctx->codec_id, avctx->flags2 & AV_CODEC_FLAG2_FAST, 0);
     if (ret < 0) {
         av_log(avctx, AV_LOG_ERROR,
                "Error splitting the input into NAL units.\n");
diff --git a/libavcodec/hevc_parse.c b/libavcodec/hevc_parse.c
index b1b27eef09..dddb293df6 100644
--- a/libavcodec/hevc_parse.c
+++ b/libavcodec/hevc_parse.c
@@ -29,7 +29,8 @@ static int hevc_decode_nal_units(const uint8_t *buf, int buf_size, HEVCParamSets
     int ret = 0;
     H2645Packet pkt = { 0 };
 
-    ret = ff_h2645_packet_split(&pkt, buf, buf_size, logctx, is_nalff, nal_length_size, AV_CODEC_ID_HEVC, 1);
+    ret = ff_h2645_packet_split(&pkt, buf, buf_size, logctx, is_nalff,
+                                nal_length_size, AV_CODEC_ID_HEVC, 1, 0);
     if (ret < 0) {
         goto done;
     }
diff --git a/libavcodec/hevc_parser.c b/libavcodec/hevc_parser.c
index f0b0789990..b444b99955 100644
--- a/libavcodec/hevc_parser.c
+++ b/libavcodec/hevc_parser.c
@@ -194,7 +194,7 @@ static int parse_nal_units(AVCodecParserContext *s, const uint8_t *buf,
     ff_hevc_reset_sei(sei);
 
     ret = ff_h2645_packet_split(&ctx->pkt, buf, buf_size, avctx, ctx->is_avc,
-                                ctx->nal_length_size, AV_CODEC_ID_HEVC, 1);
+                                ctx->nal_length_size, AV_CODEC_ID_HEVC, 1, 0);
     if (ret < 0)
         return ret;
 
diff --git a/libavcodec/hevcdec.c b/libavcodec/hevcdec.c
index 10bf2563c0..b2a87d55db 100644
--- a/libavcodec/hevcdec.c
+++ b/libavcodec/hevcdec.c
@@ -3026,7 +3026,7 @@ static int decode_nal_units(HEVCContext *s, const uint8_t *buf, int length)
     /* split the input packet into NAL units, so we know the upper bound on the
      * number of slices in the frame */
     ret = ff_h2645_packet_split(&s->pkt, buf, length, s->avctx, s->is_nalff,
-                                s->nal_length_size, s->avctx->codec_id, 1);
+                                s->nal_length_size, s->avctx->codec_id, 1, 0);
     if (ret < 0) {
         av_log(s->avctx, AV_LOG_ERROR,
                "Error splitting the input into NAL units.\n");