[FFmpeg-devel] [PATCH v2 6/6] avcodec/cbs_jpeg: Write scan data directly

[Andreas Rheinhardt]

Writing the scan element in cbs_jpeg used an intermediate buffer
to write the scan header followed by the scan data that was copied
into said buffer. Afterwards, the final buffer for the element 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
scan data directly into it, avoiding a memcpy.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt at gmail.com>
---
 libavcodec/cbs_jpeg.c | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/libavcodec/cbs_jpeg.c b/libavcodec/cbs_jpeg.c
index 64fe70beab..25321bce35 100644
--- a/libavcodec/cbs_jpeg.c
+++ b/libavcodec/cbs_jpeg.c
@@ -337,15 +337,25 @@ static int cbs_jpeg_write_scan(CodedBitstreamContext *ctx,
         return err;
 
     if (scan->data) {
-        if (scan->data_size * 8 > put_bits_left(pbc))
-            return AVERROR(ENOSPC);
+        size_t header_size = put_bits_count(pbc) / 8;
 
         av_assert0(put_bits_count(pbc) % 8 == 0);
 
+        if (scan->data_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                      - header_size)
+            return AVERROR(ENOMEM);
+
+        err = ff_cbs_alloc_unit_data(ctx, unit, header_size + scan->data_size);
+        if (err < 0)
+            return err;
+
         flush_put_bits(pbc);
+        unit->data_bit_padding = 0;
+
+        memcpy(unit->data, pbc->buf, header_size);
+        memcpy(unit->data + header_size, scan->data, scan->data_size);
 
-        memcpy(put_bits_ptr(pbc), scan->data, scan->data_size);
-        skip_put_bytes(pbc, scan->data_size);
+        return 1;
     }
 
     return 0;
-- 
2.20.1