avcodec/cbs: Factor out common code for writing units

All cbs-functions to write units share a common pattern:
1. They check whether they have a write buffer (that is used to store
the unit's data until the needed size becomes known after writing the
unit when a dedicated buffer will be allocated).
2. They use this buffer for a PutBitContext.
3. The (codec-specific) writing takes place through the PutBitContext.
4. The return value is checked. AVERROR(ENOSPC) here always indicates
that the buffer was too small and leads to a reallocation of said
buffer.
5. The final buffer will be allocated and the data copied.

This commit factors this common code out in a single function in cbs.c.

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

1 files changed, 9 insertions, 52 deletions

diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index bd172100fc..42e4dcf5ac 100644
--- a/libavcodec/cbs_vp9.c
+++ b/libavcodec/cbs_vp9.c
@@ -522,62 +522,28 @@ static int cbs_vp9_read_unit(CodedBitstreamContext *ctx,
 }
 
 static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
-                              CodedBitstreamUnit *unit)
+                              CodedBitstreamUnit *unit,
+                              PutBitContext *pbc)
 {
-    CodedBitstreamVP9Context *priv = ctx->priv_data;
     VP9RawFrame *frame = unit->content;
-    PutBitContext pbc;
     int err;
 
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
-    err = cbs_vp9_write_frame(ctx, &pbc, frame);
-    if (err == AVERROR(ENOSPC)) {
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
+    err = cbs_vp9_write_frame(ctx, pbc, frame);
     if (err < 0)
         return err;
 
     // Frame must be byte-aligned.
-    av_assert0(put_bits_count(&pbc) % 8 == 0);
-
-    unit->data_size        = put_bits_count(&pbc) / 8;
-    unit->data_bit_padding = 0;
-    flush_put_bits(&pbc);
+    av_assert0(put_bits_count(pbc) % 8 == 0);
 
     if (frame->data) {
-        if (unit->data_size + frame->data_size >
-            priv->write_buffer_size) {
-            priv->write_buffer_size *= 2;
-            goto reallocate_and_try_again;
-        }
+        if (frame->data_size > put_bits_left(pbc) / 8)
+            return AVERROR(ENOSPC);
 
-        memcpy(priv->write_buffer + unit->data_size,
-               frame->data, frame->data_size);
-        unit->data_size += frame->data_size;
+        flush_put_bits(pbc);
+        memcpy(put_bits_ptr(pbc), frame->data, frame->data_size);
+        skip_put_bytes(pbc, frame->data_size);
     }
 
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
-
     return 0;
 }
 
@@ -671,13 +637,6 @@ static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static void cbs_vp9_close(CodedBitstreamContext *ctx)
-{
-    CodedBitstreamVP9Context *priv = ctx->priv_data;
-
-    av_freep(&priv->write_buffer);
-}
-
 const CodedBitstreamType ff_cbs_type_vp9 = {
     .codec_id          = AV_CODEC_ID_VP9,
 
@@ -687,6 +646,4 @@ const CodedBitstreamType ff_cbs_type_vp9 = {
     .read_unit         = &cbs_vp9_read_unit,
     .write_unit        = &cbs_vp9_write_unit,
     .assemble_fragment = &cbs_vp9_assemble_fragment,
-
-    .close             = &cbs_vp9_close,
 };