1 files changed, 190 insertions, 27 deletions

diff --git a/libavformat/assenc.c b/libavformat/assenc.c
index 751485dfb5..52226168c2 100644
--- a/libavformat/assenc.c
+++ b/libavformat/assenc.c
@@ -2,76 +2,238 @@
  * SSA/ASS muxer
  * Copyright (c) 2008 Michael Niedermayer
  *
- * This file is part of Libav.
+ * This file is part of FFmpeg.
  *
- * Libav is free software; you can redistribute it and/or
+ * FFmpeg is free software; you can redistribute it and/or
  * modify it under the terms of the GNU Lesser General Public
  * License as published by the Free Software Foundation; either
  * version 2.1 of the License, or (at your option) any later version.
  *
- * Libav is distributed in the hope that it will be useful,
+ * FFmpeg is distributed in the hope that it will be useful,
  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  * Lesser General Public License for more details.
  *
  * You should have received a copy of the GNU Lesser General Public
- * License along with Libav; if not, write to the Free Software
+ * License along with FFmpeg; if not, write to the Free Software
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */
 
+#include "libavutil/avstring.h"
 #include "avformat.h"
+#include "internal.h"
 
-typedef struct ASSContext{
-    unsigned int extra_index;
-}ASSContext;
+#include "libavutil/opt.h"
+
+typedef struct DialogueLine {
+    int readorder;
+    char *line;
+    struct DialogueLine *prev, *next;
+} DialogueLine;
+
+typedef struct ASSContext {
+    const AVClass *class;
+    int write_ts; // 0: ssa (timing in payload), 1: ass (matroska like)
+    int expected_readorder;
+    DialogueLine *dialogue_cache;
+    DialogueLine *last_added_dialogue;
+    int cache_size;
+    int ssa_mode;
+    int ignore_readorder;
+    uint8_t *trailer;
+    size_t trailer_size;
+} ASSContext;
 
 static int write_header(AVFormatContext *s)
 {
     ASSContext *ass = s->priv_data;
-    AVCodecContext *avctx= s->streams[0]->codec;
-    uint8_t *last= NULL;
+    AVCodecContext *avctx = s->streams[0]->codec;
 
-    if(s->nb_streams != 1 || avctx->codec_id != AV_CODEC_ID_SSA){
+    if (s->nb_streams != 1 || (avctx->codec_id != AV_CODEC_ID_SSA &&
+                               avctx->codec_id != AV_CODEC_ID_ASS)) {
         av_log(s, AV_LOG_ERROR, "Exactly one ASS/SSA stream is needed.\n");
-        return -1;
+        return AVERROR(EINVAL);
     }
+    ass->write_ts = avctx->codec_id == AV_CODEC_ID_ASS;
+    avpriv_set_pts_info(s->streams[0], 64, 1, 100);
+    if (avctx->extradata_size > 0) {
+        size_t header_size = avctx->extradata_size;
+        uint8_t *trailer = strstr(avctx->extradata, "\n[Events]");
 
-    while(ass->extra_index < avctx->extradata_size){
-        uint8_t *p  = avctx->extradata + ass->extra_index;
-        uint8_t *end= strchr(p, '\n');
-        if(!end) end= avctx->extradata + avctx->extradata_size;
-        else     end++;
+        if (trailer)
+            trailer = strstr(trailer, "Format:");
+        if (trailer)
+            trailer = strstr(trailer, "\n");
 
-        avio_write(s->pb, p, end-p);
-        ass->extra_index += end-p;
+        if (trailer++) {
+            header_size = (trailer - avctx->extradata);
+            ass->trailer_size = avctx->extradata_size - header_size;
+            if (ass->trailer_size)
+                ass->trailer = trailer;
+        }
 
-        if(last && !memcmp(last, "[Events]", 8))
-            break;
-        last=p;
+        avio_write(s->pb, avctx->extradata, header_size);
+        if (avctx->extradata[header_size - 1] != '\n')
+            avio_write(s->pb, "\r\n", 2);
+        ass->ssa_mode = !strstr(avctx->extradata, "\n[V4+ Styles]");
+        if (!strstr(avctx->extradata, "\n[Events]"))
+            avio_printf(s->pb, "[Events]\r\nFormat: %s, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\r\n",
+                        ass->ssa_mode ? "Marked" : "Layer");
     }
-
     avio_flush(s->pb);
 
     return 0;
 }
 
+static void purge_dialogues(AVFormatContext *s, int force)
+{
+    int n = 0;
+    ASSContext *ass = s->priv_data;
+    DialogueLine *dialogue = ass->dialogue_cache;
+
+    while (dialogue && (dialogue->readorder == ass->expected_readorder || force)) {
+        DialogueLine *next = dialogue->next;
+        if (dialogue->readorder != ass->expected_readorder) {
+            av_log(s, AV_LOG_WARNING, "ReadOrder gap found between %d and %d\n",
+                   ass->expected_readorder, dialogue->readorder);
+            ass->expected_readorder = dialogue->readorder;
+        }
+        avio_printf(s->pb, "Dialogue: %s\r\n", dialogue->line);
+        if (dialogue == ass->last_added_dialogue)
+            ass->last_added_dialogue = next;
+        av_freep(&dialogue->line);
+        av_free(dialogue);
+        if (next)
+            next->prev = NULL;
+        dialogue = ass->dialogue_cache = next;
+        ass->expected_readorder++;
+        n++;
+    }
+    ass->cache_size -= n;
+    if (n > 1)
+        av_log(s, AV_LOG_DEBUG, "wrote %d ASS lines, cached dialogues: %d, waiting for event id %d\n",
+               n, ass->cache_size, ass->expected_readorder);
+}
+
+static void insert_dialogue(ASSContext *ass, DialogueLine *dialogue)
+{
+    DialogueLine *cur, *next = NULL, *prev = NULL;
+
+    /* from the last added to the end of the list */
+    if (ass->last_added_dialogue) {
+        for (cur = ass->last_added_dialogue; cur; cur = cur->next) {
+            if (cur->readorder > dialogue->readorder)
+                break;
+            prev = cur;
+            next = cur->next;
+        }
+    }
+
+    /* from the beginning to the last one added */
+    if (!prev) {
+        next = ass->dialogue_cache;
+        for (cur = next; cur != ass->last_added_dialogue; cur = cur->next) {
+            if (cur->readorder > dialogue->readorder)
+                break;
+            prev = cur;
+            next = cur->next;
+        }
+    }
+
+    if (prev) {
+        prev->next = dialogue;
+        dialogue->prev = prev;
+    } else {
+        dialogue->prev = ass->dialogue_cache;
+        ass->dialogue_cache = dialogue;
+    }
+    if (next) {
+        next->prev = dialogue;
+        dialogue->next = next;
+    }
+    ass->cache_size++;
+    ass->last_added_dialogue = dialogue;
+}
+
 static int write_packet(AVFormatContext *s, AVPacket *pkt)
 {
-    avio_write(s->pb, pkt->data, pkt->size);
+    ASSContext *ass = s->priv_data;
+
+    if (ass->write_ts) {
+        long int layer;
+        char *p = pkt->data;
+        int64_t start = pkt->pts;
+        int64_t end   = start + pkt->duration;
+        int hh1, mm1, ss1, ms1;
+        int hh2, mm2, ss2, ms2;
+        DialogueLine *dialogue = av_mallocz(sizeof(*dialogue));
+
+        if (!dialogue)
+            return AVERROR(ENOMEM);
+
+        dialogue->readorder = strtol(p, &p, 10);
+        if (dialogue->readorder < ass->expected_readorder)
+            av_log(s, AV_LOG_WARNING, "Unexpected ReadOrder %d\n",
+                   dialogue->readorder);
+        if (*p == ',')
+            p++;
+
+        if (ass->ssa_mode && !strncmp(p, "Marked=", 7))
+            p += 7;
+
+        layer = strtol(p, &p, 10);
+        if (*p == ',')
+            p++;
+        hh1 = (int)(start / 360000);    mm1 = (int)(start / 6000) % 60;
+        hh2 = (int)(end   / 360000);    mm2 = (int)(end   / 6000) % 60;
+        ss1 = (int)(start / 100) % 60;  ms1 = (int)(start % 100);
+        ss2 = (int)(end   / 100) % 60;  ms2 = (int)(end   % 100);
+        if (hh1 > 9) hh1 = 9, mm1 = 59, ss1 = 59, ms1 = 99;
+        if (hh2 > 9) hh2 = 9, mm2 = 59, ss2 = 59, ms2 = 99;
+
+        dialogue->line = av_asprintf("%s%ld,%d:%02d:%02d.%02d,%d:%02d:%02d.%02d,%s",
+                                     ass->ssa_mode ? "Marked=" : "",
+                                     layer, hh1, mm1, ss1, ms1, hh2, mm2, ss2, ms2, p);
+        if (!dialogue->line) {
+            av_free(dialogue);
+            return AVERROR(ENOMEM);
+        }
+        insert_dialogue(ass, dialogue);
+        purge_dialogues(s, ass->ignore_readorder);
+    } else {
+        avio_write(s->pb, pkt->data, pkt->size);
+    }
+
     return 0;
 }
 
 static int write_trailer(AVFormatContext *s)
 {
     ASSContext *ass = s->priv_data;
-    AVCodecContext *avctx= s->streams[0]->codec;
 
-    avio_write(s->pb, avctx->extradata      + ass->extra_index,
-                      avctx->extradata_size - ass->extra_index);
+    purge_dialogues(s, 1);
+
+    if (ass->trailer) {
+        avio_write(s->pb, ass->trailer, ass->trailer_size);
+    }
 
     return 0;
 }
 
+#define OFFSET(x) offsetof(ASSContext, x)
+#define E AV_OPT_FLAG_ENCODING_PARAM
+static const AVOption options[] = {
+    { "ignore_readorder", "write events immediately, even if they're out-of-order", OFFSET(ignore_readorder), AV_OPT_TYPE_INT, {.i64 = 0}, 0, 1, E },
+    { NULL },
+};
+
+static const AVClass ass_class = {
+    .class_name = "ass muxer",
+    .item_name  = av_default_item_name,
+    .option     = options,
+    .version    = LIBAVUTIL_VERSION_INT,
+};
+
 AVOutputFormat ff_ass_muxer = {
     .name           = "ass",
     .long_name      = NULL_IF_CONFIG_SMALL("SSA (SubStation Alpha) subtitle"),
@@ -82,5 +244,6 @@ AVOutputFormat ff_ass_muxer = {
     .write_header   = write_header,
     .write_packet   = write_packet,
     .write_trailer  = write_trailer,
-    .flags          = AVFMT_GLOBALHEADER | AVFMT_NOTIMESTAMPS,
+    .flags          = AVFMT_GLOBALHEADER | AVFMT_NOTIMESTAMPS | AVFMT_TS_NONSTRICT,
+    .priv_class     = &ass_class,
 };