[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_vp9.c | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index a74c4f4cec..c58589f07f 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);
 
+    flush_put_bits(pbc);
+    data_size = header_size = put_bits_count(pbc) / 8;
+    unit->data_bit_padding = 0;
+
     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 ff_cbs_default_write_unit_data(ctx, unit, pbc);
+    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 0;
 }
 
 static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
-- 
2.20.1