[FFmpeg-devel] [PATCH v2 3/6] avcodec/cbs_vp9: Write frame data directly

[Andreas Rheinhardt]

Writing a unit (always a frame) in cbs_vp9 used an intermediate buffer
to write the frame header followed by the frame data that was copied
into said buffer. Afterwards, the final buffer for the frame was
allocated and everything copied into this buffer. But it is trivial to
compute the needed size of the final buffer after having written the
header, so one can allocate the final buffer immediately and copy the
frame data directly into it.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt at gmail.com>
---
 libavcodec/cbs.c          |  4 ++++
 libavcodec/cbs_internal.h |  6 ++++--
 libavcodec/cbs_vp9.c      | 25 +++++++++++++++++++------
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 0badb192d9..9ad2641f6d 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -306,6 +306,10 @@ static int cbs_write_unit_data(CodedBitstreamContext *ctx,
     init_put_bits(&pbc, ctx->write_buffer, ctx->write_buffer_size);
 
     ret = ctx->codec->write_unit(ctx, unit, &pbc);
+    if (ret == 1) {
+        // write_unit has already finished the unit.
+        return 0;
+    }
     if (ret < 0) {
         if (ret == AVERROR(ENOSPC)) {
             // Overflow.
diff --git a/libavcodec/cbs_internal.h b/libavcodec/cbs_internal.h
index 4c5a535ca6..5768baa9ca 100644
--- a/libavcodec/cbs_internal.h
+++ b/libavcodec/cbs_internal.h
@@ -44,8 +44,10 @@ typedef struct CodedBitstreamType {
     int (*read_unit)(CodedBitstreamContext *ctx,
                      CodedBitstreamUnit *unit);
 
-    // Write the data bitstream from unit->content into pbc.
-    // Return value AVERROR(ENOSPC) indicates that pbc was too small.
+    // Write the data bitstream from unit->content into pbc or into unit->data.
+    // Return value AVERROR(ENOSPC) indicates that pbc was too small;
+    // 1 indicates that the unit has already been finished by write_unit
+    // (i.e. unit->data and unit->data_ref have been allocated and filled).
     int (*write_unit)(CodedBitstreamContext *ctx,
                       CodedBitstreamUnit *unit,
                       PutBitContext *pbc);
diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index 42e4dcf5ac..bc074c4631 100644
--- a/libavcodec/cbs_vp9.c
+++ b/libavcodec/cbs_vp9.c
@@ -526,6 +526,7 @@ static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
                               PutBitContext *pbc)
 {
     VP9RawFrame *frame = unit->content;
+    size_t data_size, header_size;
     int err;
 
     err = cbs_vp9_write_frame(ctx, pbc, frame);
@@ -535,16 +536,28 @@ static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
     // Frame must be byte-aligned.
     av_assert0(put_bits_count(pbc) % 8 == 0);
 
+    data_size = header_size = put_bits_count(pbc) / 8;
+    unit->data_bit_padding = 0;
+    flush_put_bits(pbc);
+
     if (frame->data) {
-        if (frame->data_size > put_bits_left(pbc) / 8)
-            return AVERROR(ENOSPC);
+        if (frame->data_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                       - header_size)
+            return AVERROR(ENOMEM);
 
-        flush_put_bits(pbc);
-        memcpy(put_bits_ptr(pbc), frame->data, frame->data_size);
-        skip_put_bytes(pbc, frame->data_size);
+        data_size += frame->data_size;
     }
 
-    return 0;
+    err = ff_cbs_alloc_unit_data(ctx, unit, data_size);
+    if (err < 0)
+        return err;
+
+    memcpy(unit->data, pbc->buf, header_size);
+
+    if (frame->data)
+        memcpy(unit->data + header_size, frame->data, frame->data_size);
+
+    return 1;
 }
 
 static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
-- 
2.20.1