[FFmpeg-devel] [PATCH v2 4/6] avcodec/cbs_h2645: Write slice data directly

[Andreas Rheinhardt]

Up until now, writing the data of a slice uses an intermediate buffer
into which the unit (both header as well as the rest) is assembled
before being copied into a freshly allocated buffer. But given that one
has a very good upper bound on the size needed before one starts copying
the slice data, one can allocate the buffer in advance, copy the already
written header into it and directly assemble the rest of the unit in
this buffer.

It proved easier to potentially allocate one byte more and decrement
the size afterwards if it needs to be decremented than to calculate the
exact amount in advance. It is even easier than the current way of
determining whether the last byte needs to be written or not.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt at gmail.com>
---
 libavcodec/cbs_h2645.c | 64 ++++++++++++++++++++----------------------
 1 file changed, 31 insertions(+), 33 deletions(-)

diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index d3d5342bad..81d0c478ec 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -1094,24 +1094,33 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
 }
 
 static int cbs_h2645_write_slice_data(CodedBitstreamContext *ctx,
+                                      CodedBitstreamUnit *unit,
                                       PutBitContext *pbc, const uint8_t *data,
                                       size_t data_size, int data_bit_start)
 {
     size_t rest  = data_size - (data_bit_start + 7) / 8;
+    int size = (put_bits_count(pbc) - data_bit_start + 7) / 8;
     const uint8_t *pos = data + data_bit_start / 8;
+    int err;
 
     av_assert0(data_bit_start >= 0 &&
                data_size > data_bit_start / 8);
 
-    if (data_size * 8 + 8 > put_bits_left(pbc))
-        return AVERROR(ENOSPC);
+    if (data_size > INT_MAX / 8 - size)
+        return AVERROR(ENOMEM);
+
+    // size might be one too big if in- and output are misaligned.
+    size += data_size;
+
+    err = ff_cbs_alloc_unit_data(ctx, unit, size);
+    if (err < 0)
+        return err;
 
-    if (!rest)
-        goto rbsp_stop_one_bit;
+    // Rebase pbc onto the new buffer.
+    memcpy(unit->data, pbc->buf, put_bits_ptr(pbc) - pbc->buf);
+    rebase_put_bits(pbc, unit->data, size);
 
-    // First copy the remaining bits of the first byte
-    // The above check ensures that we do not accidentally
-    // copy beyond the rbsp_stop_one_bit.
+    // First copy the remaining bits of the first byte.
     if (data_bit_start % 8)
         put_bits(pbc, 8 - data_bit_start % 8,
                  *pos++ & MAX_UINT_BITS(8 - data_bit_start % 8));
@@ -1122,32 +1131,25 @@ static int cbs_h2645_write_slice_data(CodedBitstreamContext *ctx,
         // This happens normally for CABAC.
         flush_put_bits(pbc);
         memcpy(put_bits_ptr(pbc), pos, rest);
-        skip_put_bytes(pbc, rest);
     } else {
         // If not, we have to copy manually.
-        // rbsp_stop_one_bit forces us to special-case
-        // the last byte.
-        uint8_t temp;
-        int i;
 
-        for (; rest > 4; rest -= 4, pos += 4)
+        for (; rest >= 4; rest -= 4, pos += 4)
             put_bits32(pbc, AV_RB32(pos));
 
-        for (; rest > 1; rest--, pos++)
+        for (; rest > 0; rest--, pos++)
             put_bits(pbc, 8, *pos);
 
-    rbsp_stop_one_bit:
-        temp = rest ? *pos : *pos & MAX_UINT_BITS(8 - data_bit_start % 8);
+        // Pad with zeroes
+        flush_put_bits(pbc);
 
-        av_assert0(temp);
-        i = ff_ctz(*pos);
-        temp = temp >> i;
-        i = rest ? (8 - i) : (8 - i - data_bit_start % 8);
-        put_bits(pbc, i, temp);
-        if (put_bits_count(pbc) % 8)
-            put_bits(pbc, 8 - put_bits_count(pbc) % 8, 0);
+        // Correct size if it is too big because of misalignment.
+        if (!unit->data[size - 1])
+            unit->data_size = size - 1;
     }
 
+    unit->data_bit_padding = 0;
+
     return 0;
 }
 
@@ -1207,11 +1209,9 @@ static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
                 return err;
 
             if (slice->data) {
-                err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
-                                                 slice->data_size,
-                                                 slice->data_bit_start);
-                if (err < 0)
-                    return err;
+                return cbs_h2645_write_slice_data(ctx, unit, pbc, slice->data,
+                                                  slice->data_size,
+                                                  slice->data_bit_start);
             } else {
                 // No slice data - that was just the header.
                 // (Bitstream may be unaligned!)
@@ -1341,11 +1341,9 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
                 return err;
 
             if (slice->data) {
-                err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
-                                                 slice->data_size,
-                                                 slice->data_bit_start);
-                if (err < 0)
-                    return err;
+                return cbs_h2645_write_slice_data(ctx, unit, pbc, slice->data,
+                                                  slice->data_size,
+                                                  slice->data_bit_start);
             } else {
                 // No slice data - that was just the header.
             }
-- 
2.20.1