[FFmpeg-devel] [PATCH 1/2] avformat/nutenc: fix duration estimation by writing EOR packets

[Paul B Mahol]

Signed-off-by: Paul B Mahol <onemda at gmail.com>
---
 libavformat/nut.h    |   3 +
 libavformat/nutenc.c | 130 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 129 insertions(+), 4 deletions(-)

diff --git a/libavformat/nut.h b/libavformat/nut.h
index a4409ee23d..4b15dca1ea 100644
--- a/libavformat/nut.h
+++ b/libavformat/nut.h
@@ -76,6 +76,9 @@ typedef struct StreamContext {
     int last_flags;
     int skip_until_key_frame;
     int64_t last_pts;
+    int64_t eor_pts;
+    int64_t eor_dts;
+    int eor_flushed;
     int time_base_id;
     AVRational *time_base;
     int msb_pts_shift;
diff --git a/libavformat/nutenc.c b/libavformat/nutenc.c
index e9a3bb49db..9d1c540a9d 100644
--- a/libavformat/nutenc.c
+++ b/libavformat/nutenc.c
@@ -940,6 +940,80 @@ fail:
     return ret;
 }
 
+static void nut_update_max_pts(AVFormatContext *s, int stream_index, int64_t pts)
+{
+    NUTContext *nut = s->priv_data;
+    StreamContext *nus = &nut->stream[stream_index];
+
+    if (!nut->max_pts_tb || av_compare_ts(nut->max_pts, *nut->max_pts_tb, pts, *nus->time_base) < 0) {
+        nut->max_pts = pts;
+        nut->max_pts_tb = nus->time_base;
+    }
+}
+
+static int nut_write_eor(AVFormatContext *s, AVPacket *pkt)
+{
+    AVIOContext *bc = s->pb, *dyn_bc;
+    NUTContext *nut = s->priv_data;
+    StreamContext *nus = &nut->stream[pkt->stream_index];
+    int ret, frame_code, flags, needed_flags;
+    int64_t pts = pkt->pts;
+    int64_t coded_pts;
+    FrameCode *fc;
+
+    coded_pts = pts & ((1 << nus->msb_pts_shift) - 1);
+    if (ff_lsb2full(nus, coded_pts) != pts)
+        coded_pts = pts + (1 << nus->msb_pts_shift);
+
+    nut->last_syncpoint_pos = avio_tell(bc);
+    ret                     = avio_open_dyn_buf(&dyn_bc);
+    if (ret < 0)
+        return ret;
+    put_tt(nut, nus->time_base, dyn_bc, pkt->dts);
+    ff_put_v(dyn_bc, 0);
+
+    if (nut->flags & NUT_BROADCAST) {
+        put_tt(nut, nus->time_base, dyn_bc,
+               av_rescale_q(av_gettime(), AV_TIME_BASE_Q, *nus->time_base));
+    }
+    put_packet(nut, bc, dyn_bc, 1, SYNCPOINT_STARTCODE);
+
+    frame_code  = -1;
+    for (int i = 0; i < 256; i++) {
+        FrameCode *fc = &nut->frame_code[i];
+        int flags     = fc->flags;
+
+        if (flags & FLAG_INVALID)
+            continue;
+
+        if (!(flags & FLAG_CODED)) {
+            continue;
+        }
+
+        frame_code = i;
+        break;
+    }
+    av_assert0(frame_code != -1);
+
+    fc           = &nut->frame_code[frame_code];
+    flags        = fc->flags;
+    needed_flags = FLAG_KEY | FLAG_EOR | FLAG_CODED_PTS | FLAG_STREAM_ID;
+
+    ffio_init_checksum(bc, ff_crc04C11DB7_update, 0);
+    avio_w8(bc, frame_code);
+    if (flags & FLAG_CODED) {
+        ff_put_v(bc, (flags ^ needed_flags) & ~(FLAG_CODED));
+        flags = needed_flags;
+    }
+    if (flags & FLAG_STREAM_ID)  ff_put_v(bc, pkt->stream_index);
+    if (flags & FLAG_CODED_PTS)  ff_put_v(bc, coded_pts);
+    ffio_get_checksum(bc);
+
+    nut_update_max_pts(s, pkt->stream_index, pkt->pts);
+
+    return 0;
+}
+
 static int nut_write_packet(AVFormatContext *s, AVPacket *pkt)
 {
     NUTContext *nut    = s->priv_data;
@@ -956,6 +1030,9 @@ static int nut_write_packet(AVFormatContext *s, AVPacket *pkt)
     int data_size = pkt->size;
     uint8_t *sm_buf = NULL;
 
+    if (!data_size)
+        return nut_write_eor(s, pkt);
+
     if (pkt->pts < 0) {
         av_log(s, AV_LOG_ERROR,
                "Negative pts not supported stream %d, pts %"PRId64"\n",
@@ -1150,10 +1227,7 @@ static int nut_write_packet(AVFormatContext *s, AVPacket *pkt)
             nus->keyframe_pts[nut->sp_count] = pkt->pts;
     }
 
-    if (!nut->max_pts_tb || av_compare_ts(nut->max_pts, *nut->max_pts_tb, pkt->pts, *nus->time_base) < 0) {
-        nut->max_pts = pkt->pts;
-        nut->max_pts_tb = nus->time_base;
-    }
+    nut_update_max_pts(s, pkt->stream_index, pkt->pts);
 
 fail:
     av_freep(&sm_buf);
@@ -1161,6 +1235,53 @@ fail:
     return ret;
 }
 
+static int nut_interleave_packet(AVFormatContext *s, AVPacket *out,
+                                 AVPacket *pkt, int flush)
+{
+    NUTContext *nut = s->priv_data;
+    int ret;
+
+    if (pkt) {
+        StreamContext *nus = &nut->stream[pkt->stream_index];
+
+        nus->eor_pts = FFMAX(pkt->pts + pkt->duration, nus->eor_pts);
+        nus->eor_dts = FFMAX(pkt->dts + pkt->duration, nus->eor_dts);
+    } else if (flush) {
+        StreamContext *nus;
+        int i;
+
+        for (i = 0; i < s->nb_streams; i++) {
+            nus = &nut->stream[i];
+
+            if (nus->eor_flushed)
+                continue;
+
+            if (s->streams[i]->last_in_packet_buffer)
+                continue;
+
+            break;
+        }
+
+        if (i < s->nb_streams) {
+            pkt = av_packet_alloc();
+            if (!pkt)
+                return AVERROR(ENOMEM);
+
+            ret = av_new_packet(pkt, 0);
+            if (ret < 0)
+                return ret;
+            pkt->stream_index = i;
+            pkt->pts = nus->eor_pts;
+            pkt->dts = nus->eor_dts;
+            pkt->duration = 0;
+            nus->eor_flushed = 1;
+            ret = ff_interleave_packet_per_dts(s, out, pkt, 0);
+            pkt = NULL;
+        }
+    }
+    return ff_interleave_packet_per_dts(s, out, pkt, flush);
+}
+
 static int nut_write_trailer(AVFormatContext *s)
 {
     NUTContext *nut = s->priv_data;
@@ -1225,6 +1346,7 @@ AVOutputFormat ff_nut_muxer = {
     .write_header   = nut_write_header,
     .write_packet   = nut_write_packet,
     .write_trailer  = nut_write_trailer,
+    .interleave_packet = nut_interleave_packet,
     .deinit         = nut_write_deinit,
     .flags          = AVFMT_GLOBALHEADER | AVFMT_VARIABLE_FPS,
     .codec_tag      = ff_nut_codec_tags,
-- 
2.17.1