[FFmpeg-devel] [PATCH] lavf/matroskaenc: Write Block Keyframe correctly

[Vignesh Venkatasubramanian]

Per matroska Block Structure [1], for keyframes 0th bit of the flag
should not be set (unlike SimpleBlocks). For Blocks, keyframes is
inferred by the absence of ReferenceBlock element (as done by
matroskadec). This CL writes the flag correctly and inserts the
ReferenceBlock element for non-keyframes. The timestamp inserted is
that of the immediately preceding frame (which is true for VP8 and VP9
- the only 2 codecs using the matroska block element as of now). It
also considers all non-video frames (audio, subtitles, metadata) to
be keyframes.

[1] http://www.matroska.org/technical/specs/index.html#block_structure

Signed-off-by: Vignesh Venkatasubramanian <vigneshv at google.com>
---
 libavformat/matroskaenc.c | 33 ++++++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/libavformat/matroskaenc.c b/libavformat/matroskaenc.c
index aa9d119..70d4375 100644
--- a/libavformat/matroskaenc.c
+++ b/libavformat/matroskaenc.c
@@ -95,6 +95,10 @@ typedef struct mkv_track {
 #define MODE_MATROSKAv2 0x01
 #define MODE_WEBM       0x02
 
+/** Maximum number of tracks allowed in a Matroska file (with track numbers in
+ * range 1 to 126 (inclusive) */
+#define MAX_TRACKS 126
+
 typedef struct MatroskaMuxContext {
     const AVClass  *class;
     int             mode;
@@ -125,6 +129,8 @@ typedef struct MatroskaMuxContext {
     uint32_t chapter_id_offset;
     int wrote_chapters;
 
+    int64_t last_track_timestamp[MAX_TRACKS];
+
     int allow_raw_vfw;
 } MatroskaMuxContext;
 
@@ -1354,6 +1360,10 @@ static int mkv_write_header(AVFormatContext *s)
             version = 4;
     }
 
+    for (i = 0; i < MAX_TRACKS; i++) {
+        mkv->last_track_timestamp[i] = 0;
+    }
+
     mkv->tracks = av_mallocz_array(s->nb_streams, sizeof(*mkv->tracks));
     if (!mkv->tracks)
         return AVERROR(ENOMEM);
@@ -1545,7 +1555,7 @@ fail:
 }
 
 static void mkv_write_block(AVFormatContext *s, AVIOContext *pb,
-                            unsigned int blockid, AVPacket *pkt, int flags)
+                            unsigned int blockid, AVPacket *pkt, int keyframe)
 {
     MatroskaMuxContext *mkv = s->priv_data;
     AVCodecContext *codec = s->streams[pkt->stream_index]->codec;
@@ -1554,11 +1564,13 @@ static void mkv_write_block(AVFormatContext *s, AVIOContext *pb,
     int64_t ts = mkv->tracks[pkt->stream_index].write_dts ? pkt->dts : pkt->pts;
     uint64_t additional_id = 0;
     int64_t discard_padding = 0;
+    uint8_t track_number = (mkv->is_dash ? mkv->dash_track_number : (pkt->stream_index + 1));
     ebml_master block_group, block_additions, block_more;
 
     av_log(s, AV_LOG_DEBUG, "Writing block at offset %" PRIu64 ", size %d, "
-           "pts %" PRId64 ", dts %" PRId64 ", duration %d, flags %d\n",
-           avio_tell(pb), pkt->size, pkt->pts, pkt->dts, pkt->duration, flags);
+           "pts %" PRId64 ", dts %" PRId64 ", duration %d, keyframe %d\n",
+           avio_tell(pb), pkt->size, pkt->pts, pkt->dts, pkt->duration,
+           keyframe != 0);
     if (codec->codec_id == AV_CODEC_ID_H264 && codec->extradata_size > 0 &&
         (AV_RB24(codec->extradata) == 1 || AV_RB32(codec->extradata) == 1))
         ff_avc_parse_nal_units_buf(pkt->data, &data, &size);
@@ -1609,13 +1621,19 @@ static void mkv_write_block(AVFormatContext *s, AVIOContext *pb,
     put_ebml_id(pb, blockid);
     put_ebml_num(pb, size + 4, 0);
     // this assumes stream_index is less than 126
-    avio_w8(pb, 0x80 | (mkv->is_dash ? mkv->dash_track_number : (pkt->stream_index + 1)));
+    avio_w8(pb, 0x80 | track_number);
     avio_wb16(pb, ts - mkv->cluster_pts);
-    avio_w8(pb, flags);
+    avio_w8(pb, (blockid == MATROSKA_ID_SIMPLEBLOCK && keyframe) ? (1 << 7) : 0);
     avio_write(pb, data + offset, size);
     if (data != pkt->data)
         av_free(data);
 
+    if (blockid == MATROSKA_ID_BLOCK && !keyframe) {
+        put_ebml_sint(pb, MATROSKA_ID_BLOCKREFERENCE,
+                      mkv->last_track_timestamp[track_number - 1]);
+    }
+    mkv->last_track_timestamp[track_number - 1] = ts - mkv->cluster_pts;
+
     if (discard_padding) {
         put_ebml_sint(pb, MATROSKA_ID_DISCARDPADDING, discard_padding);
     }
@@ -1757,7 +1775,7 @@ static int mkv_write_packet_internal(AVFormatContext *s, AVPacket *pkt, int add_
     relative_packet_pos = avio_tell(s->pb) - mkv->cluster.pos;
 
     if (codec->codec_type != AVMEDIA_TYPE_SUBTITLE) {
-        mkv_write_block(s, pb, MATROSKA_ID_SIMPLEBLOCK, pkt, keyframe << 7);
+        mkv_write_block(s, pb, MATROSKA_ID_SIMPLEBLOCK, pkt, keyframe);
         if (codec->codec_type == AVMEDIA_TYPE_VIDEO && keyframe || add_cue) {
             ret = mkv_add_cuepoint(mkv->cues, pkt->stream_index, dash_tracknum, ts, mkv->cluster_pos, relative_packet_pos, -1);
             if (ret < 0) return ret;
@@ -1772,7 +1790,8 @@ static int mkv_write_packet_internal(AVFormatContext *s, AVPacket *pkt, int add_
         if (pkt->convergence_duration > 0) {
             duration = pkt->convergence_duration;
         }
-        mkv_write_block(s, pb, MATROSKA_ID_BLOCK, pkt, 0);
+        /* All subtitle blocks are considered to be keyframes. */
+        mkv_write_block(s, pb, MATROSKA_ID_BLOCK, pkt, 1);
         put_ebml_uint(pb, MATROSKA_ID_BLOCKDURATION, duration);
         end_ebml_master(pb, blockgroup);
     }
-- 
2.2.0.rc0.207.ga3a616c