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

[Andreas Rheinhardt]

This commit makes use of the fact that the final unit buffer is now
allocated within write_unit to avoid a memcpy of scan data: It is now no
longer copied into an intermediate buffer before being written into the
final buffer, but directly written into the final buffer.

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 | 33 +++++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/libavcodec/cbs_jpeg.c b/libavcodec/cbs_jpeg.c
index da22bdf720..dacdfa985b 100644
--- a/libavcodec/cbs_jpeg.c
+++ b/libavcodec/cbs_jpeg.c
@@ -330,24 +330,37 @@ static int cbs_jpeg_write_scan(CodedBitstreamContext *ctx,
                                PutBitContext *pbc)
 {
     JPEGRawScan *scan = unit->content;
+    size_t header_size, size;
     int err;
 
     err = cbs_jpeg_write_scan_header(ctx, pbc, &scan->header);
     if (err < 0)
         return err;
 
-    if (scan->data) {
-        if (scan->data_size * 8 > put_bits_left(pbc))
-            return AVERROR(ENOSPC);
+    av_assert0(put_bits_count(pbc) % 8 == 0);
 
-        av_assert0(put_bits_count(pbc) % 8 == 0);
+    flush_put_bits(pbc);
+    header_size = size = put_bits_count(pbc) / 8;
 
-        flush_put_bits(pbc);
+    if (scan->data) {
+        if (scan->data_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                      - header_size)
+            return AVERROR(ENOMEM);
 
-        memcpy(put_bits_ptr(pbc), scan->data, scan->data_size);
-        skip_put_bytes(pbc, scan->data_size);
+        size += scan->data_size;
     }
 
+    err = ff_cbs_alloc_unit_data(ctx, unit, size);
+    if (err < 0)
+        return err;
+
+    unit->data_bit_padding = 0;
+
+    memcpy(unit->data, pbc->buf, header_size);
+
+    if (scan->data)
+        memcpy(unit->data + header_size, scan->data, scan->data_size);
+
     return 0;
 }
 
@@ -387,9 +400,9 @@ static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
     int err;
 
     if (unit->type == JPEG_MARKER_SOS)
-        err = cbs_jpeg_write_scan   (ctx, unit, pbc);
-    else
-        err = cbs_jpeg_write_segment(ctx, unit, pbc);
+        return cbs_jpeg_write_scan(ctx, unit, pbc);
+
+    err = cbs_jpeg_write_segment(ctx, unit, pbc);
     if (err < 0)
         return err;
 
-- 
2.20.1