From f901d75bf18c866933f90c052fb990736aee7c47 Mon Sep 17 00:00:00 2001
From: James Almer <jamrial@gmail.com>
Date: Thu, 3 Sep 2020 15:39:32 -0300
Subject: [PATCH] avformat/yuv4mpegenc: simplify writing the header

Actually write it in yuv4_write_header() instead of with the first
packet.

Signed-off-by: James Almer <jamrial@gmail.com>
---
 libavformat/yuv4mpegenc.c | 35 ++++++++++++++---------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/libavformat/yuv4mpegenc.c b/libavformat/yuv4mpegenc.c
index e84dbf9568..c4781042bd 100644
--- a/libavformat/yuv4mpegenc.c
+++ b/libavformat/yuv4mpegenc.c
@@ -26,14 +26,16 @@
 
 #define Y4M_LINE_MAX 256
 
-static int yuv4_generate_header(AVFormatContext *s, char* buf)
+static int yuv4_write_header(AVFormatContext *s)
 {
     AVStream *st;
+    AVIOContext *pb = s->pb;
     int width, height;
     int raten, rated, aspectn, aspectd, n;
     char inter;
     const char *colorspace = "";
     const char *colorrange = "";
+    char buf[Y4M_LINE_MAX + 1];
     int field_order;
 
     st     = s->streams[0];
@@ -173,7 +175,15 @@ static int yuv4_generate_header(AVFormatContext *s, char* buf)
                  Y4M_MAGIC, width, height, raten, rated, inter,
                  aspectn, aspectd, colorspace, colorrange);
 
-    return n;
+    if (n < 0) {
+        av_log(s, AV_LOG_ERROR,
+               "Error. YUV4MPEG stream header write failed.\n");
+        return AVERROR(EIO);
+    }
+
+    avio_write(pb, buf, strlen(buf));
+
+    return 0;
 }
 
 
@@ -182,26 +192,12 @@ static int yuv4_write_packet(AVFormatContext *s, AVPacket *pkt)
     AVStream *st = s->streams[pkt->stream_index];
     AVIOContext *pb = s->pb;
     AVFrame *frame;
-    int* first_pkt = s->priv_data;
     int width, height, h_chroma_shift, v_chroma_shift;
     int i;
-    char buf2[Y4M_LINE_MAX + 1];
     uint8_t *ptr, *ptr1, *ptr2;
 
     frame = (AVFrame *)pkt->data;
 
-    /* for the first packet we have to output the header as well */
-    if (*first_pkt) {
-        *first_pkt = 0;
-        if (yuv4_generate_header(s, buf2) < 0) {
-            av_log(s, AV_LOG_ERROR,
-                   "Error. YUV4MPEG stream header write failed.\n");
-            return AVERROR(EIO);
-        } else {
-            avio_write(pb, buf2, strlen(buf2));
-        }
-    }
-
     /* construct frame header */
 
     avio_printf(s->pb, "%s\n", Y4M_FRAME_MAGIC);
@@ -279,10 +275,8 @@ static int yuv4_write_packet(AVFormatContext *s, AVPacket *pkt)
     return 0;
 }
 
-static int yuv4_write_header(AVFormatContext *s)
+static int yuv4_init(AVFormatContext *s)
 {
-    int *first_pkt = s->priv_data;
-
     if (s->nb_streams != 1)
         return AVERROR(EIO);
 
@@ -347,7 +341,6 @@ static int yuv4_write_header(AVFormatContext *s)
         return AVERROR(EIO);
     }
 
-    *first_pkt = 1;
     return 0;
 }
 
@@ -355,9 +348,9 @@ AVOutputFormat ff_yuv4mpegpipe_muxer = {
     .name              = "yuv4mpegpipe",
     .long_name         = NULL_IF_CONFIG_SMALL("YUV4MPEG pipe"),
     .extensions        = "y4m",
-    .priv_data_size    = sizeof(int),
     .audio_codec       = AV_CODEC_ID_NONE,
     .video_codec       = AV_CODEC_ID_WRAPPED_AVFRAME,
+    .init              = yuv4_init,
     .write_header      = yuv4_write_header,
     .write_packet      = yuv4_write_packet,
 };