lavc/extract_extradata: Use bytestream api

Signed-off-by: Andriy Gelman <andriy.gelman@gmail.com>
Signed-off-by: James Almer <jamrial@gmail.com>

libavcodec/extract_extradata_bsf.c

@@ -27,6 +27,7 @@
 #include "av1.h"
 #include "av1_parse.h"
 #include "bsf.h"
+#include "bytestream.h"
 #include "h2645_parse.h"
 #include "h264.h"
 #include "hevc.h"
@@ -86,7 +87,8 @@ static int extract_extradata_av1(AVBSFContext *ctx, AVPacket *pkt,
 
     if (extradata_size && has_seq) {
         AVBufferRef *filtered_buf = NULL;
-        uint8_t *extradata, *filtered_data;
+        PutByteContext pb_filtered_data, pb_extradata;
+        uint8_t *extradata;
 
         if (s->remove) {
             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
@@ -94,8 +96,6 @@ static int extract_extradata_av1(AVBSFContext *ctx, AVPacket *pkt,
                 return AVERROR(ENOMEM);
             }
             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
-
-            filtered_data = filtered_buf->data;
         }
 
         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
@@ -108,15 +108,17 @@ static int extract_extradata_av1(AVBSFContext *ctx, AVPacket *pkt,
         *data = extradata;
         *size = extradata_size;
 
+        bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
+        if (s->remove)
+            bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
+
         for (i = 0; i < s->av1_pkt.nb_obus; i++) {
             AV1OBU *obu = &s->av1_pkt.obus[i];
             if (val_in_array(extradata_obu_types, nb_extradata_obu_types,
                              obu->type)) {
-                memcpy(extradata, obu->raw_data, obu->raw_size);
-                extradata += obu->raw_size;
+                bytestream2_put_bufferu(&pb_extradata, obu->raw_data, obu->raw_size);
             } else if (s->remove) {
-                memcpy(filtered_data, obu->raw_data, obu->raw_size);
-                filtered_data += obu->raw_size;
+                bytestream2_put_bufferu(&pb_filtered_data, obu->raw_data, obu->raw_size);
             }
         }
 
@@ -180,7 +182,8 @@ static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
         ((ctx->par_in->codec_id == AV_CODEC_ID_HEVC && has_sps && has_vps) ||
          (ctx->par_in->codec_id == AV_CODEC_ID_H264 && has_sps))) {
         AVBufferRef *filtered_buf = NULL;
-        uint8_t *extradata, *filtered_data;
+        PutByteContext pb_filtered_data, pb_extradata;
+        uint8_t *extradata;
 
         if (s->remove) {
             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
@@ -188,8 +191,6 @@ static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
                 return AVERROR(ENOMEM);
             }
             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
-
-            filtered_data = filtered_buf->data;
         }
 
         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
@@ -202,17 +203,19 @@ static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
         *data = extradata;
         *size = extradata_size;
 
+        bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
+        if (s->remove)
+            bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
+
         for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
             H2645NAL *nal = &s->h2645_pkt.nals[i];
             if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
                              nal->type)) {
-                AV_WB24(extradata, 1); // startcode
-                memcpy(extradata + 3, nal->raw_data, nal->raw_size);
-                extradata += 3 + nal->raw_size;
+                bytestream2_put_be24u(&pb_extradata, 1); //startcode
+                bytestream2_put_bufferu(&pb_extradata, nal->raw_data, nal->raw_size);
             } else if (s->remove) {
-                AV_WB24(filtered_data, 1); // startcode
-                memcpy(filtered_data + 3, nal->raw_data, nal->raw_size);
-                filtered_data += 3 + nal->raw_size;
+                bytestream2_put_be24u(&pb_filtered_data, 1); // startcode
+                bytestream2_put_bufferu(&pb_filtered_data, nal->raw_data, nal->raw_size);
             }
         }