vp9: Add bsf to fix reordering in raw streams

Takes a raw input stream containing frames with correct timestamps but
possibly out of order and inserts additional show-existing-frame
packets to correct the ordering.

3 files changed, 413 insertions, 0 deletions

diff --git a/libavcodec/Makefile b/libavcodec/Makefile
index 80c572c7c4..674c8d0445 100644
--- a/libavcodec/Makefile
+++ b/libavcodec/Makefile
@@ -772,6 +772,7 @@ OBJS-$(CONFIG_NOISE_BSF)                  += noise_bsf.o
 OBJS-$(CONFIG_NULL_BSF)                   += null_bsf.o
 OBJS-$(CONFIG_REMOVE_EXTRADATA_BSF)       += remove_extradata_bsf.o
 OBJS-$(CONFIG_TEXT2MOVSUB_BSF)            += movsub_bsf.o
+OBJS-$(CONFIG_VP9_RAW_REORDER_BSF)        += vp9_raw_reorder_bsf.o
 OBJS-$(CONFIG_VP9_SUPERFRAME_BSF)         += vp9_superframe_bsf.o
 OBJS-$(CONFIG_VP9_SUPERFRAME_SPLIT_BSF)   += vp9_superframe_split_bsf.o
 
diff --git a/libavcodec/bitstream_filters.c b/libavcodec/bitstream_filters.c
index d46fdad81b..79ce40f9ed 100644
--- a/libavcodec/bitstream_filters.c
+++ b/libavcodec/bitstream_filters.c
@@ -38,6 +38,7 @@ extern const AVBitStreamFilter ff_null_bsf;
 extern const AVBitStreamFilter ff_text2movsub_bsf;
 extern const AVBitStreamFilter ff_noise_bsf;
 extern const AVBitStreamFilter ff_remove_extradata_bsf;
+extern const AVBitStreamFilter ff_vp9_raw_reorder_bsf;
 extern const AVBitStreamFilter ff_vp9_superframe_bsf;
 extern const AVBitStreamFilter ff_vp9_superframe_split_bsf;
 
diff --git a/libavcodec/vp9_raw_reorder_bsf.c b/libavcodec/vp9_raw_reorder_bsf.c
new file mode 100644
index 0000000000..7f00f5f103
--- /dev/null
+++ b/libavcodec/vp9_raw_reorder_bsf.c
@@ -0,0 +1,411 @@
+/*
+ * This file is part of Libav.
+ *
+ * Libav 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,
+ * 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
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#include "libavutil/avassert.h"
+#include "libavutil/intmath.h"
+#include "libavutil/log.h"
+#include "libavutil/mem.h"
+#include "libavutil/opt.h"
+
+#include "bitstream.h"
+#include "bsf.h"
+#include "put_bits.h"
+
+#define FRAME_SLOTS 8
+
+typedef struct VP9RawReorderFrame {
+    AVPacket    *packet;
+    int          needs_output;
+    int          needs_display;
+
+    int64_t      pts;
+    int64_t      sequence;
+    unsigned int slots;
+
+    unsigned int profile;
+
+    unsigned int show_existing_frame;
+    unsigned int frame_to_show;
+
+    unsigned int frame_type;
+    unsigned int show_frame;
+    unsigned int refresh_frame_flags;
+} VP9RawReorderFrame;
+
+typedef struct VP9RawReorderContext {
+    int64_t sequence;
+    VP9RawReorderFrame *slot[FRAME_SLOTS];
+    VP9RawReorderFrame *next_frame;
+} VP9RawReorderContext;
+
+static void vp9_raw_reorder_frame_free(VP9RawReorderFrame **frame)
+{
+    if (*frame)
+        av_packet_free(&(*frame)->packet);
+    av_freep(frame);
+}
+
+static void vp9_raw_reorder_clear_slot(VP9RawReorderContext *ctx, int s)
+{
+    if (ctx->slot[s]) {
+        ctx->slot[s]->slots &= ~(1 << s);
+        if (ctx->slot[s]->slots == 0)
+            vp9_raw_reorder_frame_free(&ctx->slot[s]);
+        else
+            ctx->slot[s] = NULL;
+    }
+}
+
+static int vp9_raw_reorder_frame_parse(AVBSFContext *bsf, VP9RawReorderFrame *frame)
+{
+    BitstreamContext bc;
+    int err;
+
+    unsigned int frame_marker;
+    unsigned int profile_low_bit, profile_high_bit, reserved_zero;
+    unsigned int error_resilient_mode;
+    unsigned int frame_sync_code;
+
+    err = bitstream_init8(&bc, frame->packet->data, frame->packet->size);
+    if (err)
+        return err;
+
+    frame_marker = bitstream_read(&bc, 2);
+    if (frame_marker != 2) {
+        av_log(bsf, AV_LOG_ERROR, "Invalid frame marker: %u.\n",
+               frame_marker);
+        return AVERROR_INVALIDDATA;
+    }
+
+    profile_low_bit  = bitstream_read_bit(&bc);
+    profile_high_bit = bitstream_read_bit(&bc);
+    frame->profile = (profile_high_bit << 1) | profile_low_bit;
+    if (frame->profile == 3) {
+        reserved_zero = bitstream_read_bit(&bc);
+        if (reserved_zero != 0) {
+            av_log(bsf, AV_LOG_ERROR, "Profile reserved_zero bit set: "
+                   "unsupported profile or invalid bitstream.\n");
+            return AVERROR_INVALIDDATA;
+        }
+    }
+
+    frame->show_existing_frame = bitstream_read_bit(&bc);
+    if (frame->show_existing_frame) {
+        frame->frame_to_show = bitstream_read(&bc, 3);
+        return 0;
+    }
+
+    frame->frame_type = bitstream_read_bit(&bc);
+    frame->show_frame = bitstream_read_bit(&bc);
+    error_resilient_mode = bitstream_read_bit(&bc);
+
+    if (frame->frame_type == 0) {
+        frame_sync_code = bitstream_read(&bc, 24);
+        if (frame_sync_code != 0x498342) {
+            av_log(bsf, AV_LOG_ERROR, "Invalid frame sync code: %06x.\n",
+                   frame_sync_code);
+            return AVERROR_INVALIDDATA;
+        }
+        frame->refresh_frame_flags = 0xff;
+    } else {
+        unsigned int intra_only;
+
+        if (frame->show_frame == 0)
+            intra_only = bitstream_read_bit(&bc);
+        else
+            intra_only = 0;
+        if (error_resilient_mode == 0) {
+            // reset_frame_context
+            bitstream_skip(&bc, 2);
+        }
+        if (intra_only) {
+            frame_sync_code = bitstream_read(&bc, 24);
+            if (frame_sync_code != 0x498342) {
+                av_log(bsf, AV_LOG_ERROR, "Invalid frame sync code: "
+                       "%06x.\n", frame_sync_code);
+                return AVERROR_INVALIDDATA;
+            }
+            if (frame->profile > 0) {
+                unsigned int color_space;
+                if (frame->profile >= 2) {
+                    // ten_or_twelve_bit
+                    bitstream_skip(&bc, 1);
+                }
+                color_space = bitstream_read(&bc, 3);
+                if (color_space != 7 /* CS_RGB */) {
+                    // color_range
+                    bitstream_skip(&bc, 1);
+                    if (frame->profile == 1 || frame->profile == 3) {
+                        // subsampling
+                        bitstream_skip(&bc, 3);
+                    }
+                } else {
+                    if (frame->profile == 1 || frame->profile == 3)
+                        bitstream_skip(&bc, 1);
+                }
+            }
+            frame->refresh_frame_flags = bitstream_read(&bc, 8);
+        } else {
+            frame->refresh_frame_flags = bitstream_read(&bc, 8);
+        }
+    }
+
+    return 0;
+}
+
+static int vp9_raw_reorder_make_output(AVBSFContext *bsf,
+                                   AVPacket *out,
+                                   VP9RawReorderFrame *last_frame)
+{
+    VP9RawReorderContext *ctx = bsf->priv_data;
+    VP9RawReorderFrame *next_output = last_frame,
+                      *next_display = last_frame, *frame;
+    int s, err;
+
+    for (s = 0; s < FRAME_SLOTS; s++) {
+        frame = ctx->slot[s];
+        if (!frame)
+            continue;
+        if (frame->needs_output && (!next_output ||
+            frame->sequence < next_output->sequence))
+            next_output = frame;
+        if (frame->needs_display && (!next_display ||
+            frame->pts < next_display->pts))
+            next_display = frame;
+    }
+
+    if (!next_output && !next_display)
+        return AVERROR_EOF;
+
+    if (!next_display || (next_output &&
+        next_output->sequence < next_display->sequence))
+        frame = next_output;
+    else
+        frame = next_display;
+
+    if (frame->needs_output && frame->needs_display &&
+        next_output == next_display) {
+        av_log(bsf, AV_LOG_DEBUG, "Output and display frame "
+               "%"PRId64" (%"PRId64") in order.\n",
+               frame->sequence, frame->pts);
+
+        av_packet_move_ref(out, frame->packet);
+
+        frame->needs_output = frame->needs_display = 0;
+    } else if (frame->needs_output) {
+        if (frame->needs_display) {
+            av_log(bsf, AV_LOG_DEBUG, "Output frame %"PRId64" "
+                   "(%"PRId64") for later display.\n",
+                   frame->sequence, frame->pts);
+        } else {
+            av_log(bsf, AV_LOG_DEBUG, "Output unshown frame "
+                   "%"PRId64" (%"PRId64") to keep order.\n",
+                   frame->sequence, frame->pts);
+        }
+
+        av_packet_move_ref(out, frame->packet);
+        out->pts = out->dts;
+
+        frame->needs_output = 0;
+    } else {
+        PutBitContext pb;
+
+        av_assert0(!frame->needs_output && frame->needs_display);
+
+        if (frame->slots == 0) {
+            av_log(bsf, AV_LOG_ERROR, "Attempting to display frame "
+                   "which is no longer available?\n");
+            frame->needs_display = 0;
+            return AVERROR_INVALIDDATA;
+        }
+
+        s = ff_ctz(frame->slots);
+        av_assert0(s < FRAME_SLOTS);
+
+        av_log(bsf, AV_LOG_DEBUG, "Display frame %"PRId64" "
+               "(%"PRId64") from slot %d.\n",
+               frame->sequence, frame->pts, s);
+
+        frame->packet = av_packet_alloc();
+        if (!frame->packet)
+            return AVERROR(ENOMEM);
+
+        err = av_new_packet(out, 2);
+        if (err < 0)
+            return err;
+
+        init_put_bits(&pb, out->data, 2);
+
+        // frame_marker
+        put_bits(&pb, 2, 2);
+        // profile_low_bit
+        put_bits(&pb, 1, frame->profile & 1);
+        // profile_high_bit
+        put_bits(&pb, 1, (frame->profile >> 1) & 1);
+        if (frame->profile == 3) {
+            // reserved_zero
+            put_bits(&pb, 1, 0);
+        }
+        // show_existing_frame
+        put_bits(&pb, 1, 1);
+        // frame_to_show_map_idx
+        put_bits(&pb, 3, s);
+
+        while (put_bits_count(&pb) < 16)
+            put_bits(&pb, 1, 0);
+
+        flush_put_bits(&pb);
+        out->pts = out->dts = frame->pts;
+
+        frame->needs_display = 0;
+    }
+
+    return 0;
+}
+
+static int vp9_raw_reorder_filter(AVBSFContext *bsf, AVPacket *out)
+{
+    VP9RawReorderContext *ctx = bsf->priv_data;
+    VP9RawReorderFrame *frame;
+    AVPacket *in;
+    int err, s;
+
+    if (ctx->next_frame) {
+        frame = ctx->next_frame;
+
+    } else {
+        err = ff_bsf_get_packet(bsf, &in);
+        if (err < 0) {
+            if (err == AVERROR_EOF)
+                return vp9_raw_reorder_make_output(bsf, out, NULL);
+            return err;
+        }
+
+        if (in->data[in->size - 1] & 0xe0 == 0xc0) {
+            av_log(bsf, AV_LOG_ERROR, "Input in superframes is not "
+                   "supported.\n");
+            av_packet_free(&in);
+            return AVERROR(ENOSYS);
+        }
+
+        frame = av_mallocz(sizeof(*frame));
+        if (!frame) {
+            av_packet_free(&in);
+            return AVERROR(ENOMEM);
+        }
+
+        frame->packet   = in;
+        frame->pts      = in->pts;
+        frame->sequence = ++ctx->sequence;
+        err = vp9_raw_reorder_frame_parse(bsf, frame);
+        if (err) {
+            av_log(bsf, AV_LOG_ERROR, "Failed to parse input "
+                   "frame: %d.\n", err);
+            goto fail;
+        }
+
+        frame->needs_output  = 1;
+        frame->needs_display = frame->pts != AV_NOPTS_VALUE;
+
+        if (frame->show_existing_frame)
+            av_log(bsf, AV_LOG_DEBUG, "Show frame %"PRId64" "
+                   "(%"PRId64"): show %u.\n", frame->sequence,
+                   frame->pts, frame->frame_to_show);
+        else
+            av_log(bsf, AV_LOG_DEBUG, "New frame %"PRId64" "
+                   "(%"PRId64"): type %u show %u refresh %02x.\n",
+                   frame->sequence, frame->pts, frame->frame_type,
+                   frame->show_frame, frame->refresh_frame_flags);
+
+        ctx->next_frame = frame;
+    }
+
+    for (s = 0; s < FRAME_SLOTS; s++) {
+        if (!(frame->refresh_frame_flags & (1 << s)))
+            continue;
+        if (ctx->slot[s] && ctx->slot[s]->needs_display &&
+            ctx->slot[s]->slots == (1 << s)) {
+            // We are overwriting this slot, which is last reference
+            // to the frame previously present in it.  In order to be
+            // a valid stream, that frame must already have been
+            // displayed before the pts of the current frame.
+            err = vp9_raw_reorder_make_output(bsf, out, ctx->slot[s]);
+            if (err < 0) {
+                av_log(bsf, AV_LOG_ERROR, "Failed to create "
+                       "output overwriting slot %d: %d.\n",
+                       s, err);
+                // Clear the slot anyway, so we don't end up
+                // in an infinite loop.
+                vp9_raw_reorder_clear_slot(ctx, s);
+                return AVERROR_INVALIDDATA;
+            }
+            return 0;
+        }
+        vp9_raw_reorder_clear_slot(ctx, s);
+    }
+
+    for (s = 0; s < FRAME_SLOTS; s++) {
+        if (!(frame->refresh_frame_flags & (1 << s)))
+            continue;
+        ctx->slot[s] = frame;
+    }
+    frame->slots = frame->refresh_frame_flags;
+
+    if (!frame->refresh_frame_flags) {
+        err = vp9_raw_reorder_make_output(bsf, out, frame);
+        if (err < 0) {
+            av_log(bsf, AV_LOG_ERROR, "Failed to create output "
+                   "for transient frame.\n");
+            ctx->next_frame = NULL;
+            return AVERROR_INVALIDDATA;
+        }
+        if (!frame->needs_display) {
+            vp9_raw_reorder_frame_free(&frame);
+            ctx->next_frame = NULL;
+        }
+        return 0;
+    }
+
+    ctx->next_frame = NULL;
+    return AVERROR(EAGAIN);
+
+fail:
+    vp9_raw_reorder_frame_free(&frame);
+    return err;
+}
+
+static void vp9_raw_reorder_close(AVBSFContext *bsf)
+{
+    VP9RawReorderContext *ctx = bsf->priv_data;
+    int s;
+
+    for (s = 0; s < FRAME_SLOTS; s++)
+        vp9_raw_reorder_clear_slot(ctx, s);
+}
+
+static const enum AVCodecID vp9_raw_reorder_codec_ids[] = {
+    AV_CODEC_ID_VP9, AV_CODEC_ID_NONE,
+};
+
+const AVBitStreamFilter ff_vp9_raw_reorder_bsf = {
+    .name           = "vp9_raw_reorder",
+    .priv_data_size = sizeof(VP9RawReorderContext),
+    .close          = &vp9_raw_reorder_close,
+    .filter         = &vp9_raw_reorder_filter,
+    .codec_ids      = vp9_raw_reorder_codec_ids,
+};