encode.c

#include <stdbool.h>
#include <byteswap.h>

#include <libavformat/avformat.h>
#include <libswscale/swscale.h>

// TODO: VFR maybe?
#define FRAMERATE 60

struct bound_info {
    uint32_t width;
    uint32_t height;
    uint64_t start_time;
    size_t frames;
};

struct packet_switch {
    uint32_t width;
    uint32_t height;
    enum AVPixelFormat format;
    uint8_t bpp;
};

struct packet_update {
    uint32_t x;
    uint32_t y;
    uint32_t w;
    uint32_t h;
    uint64_t timestamp;
    size_t datalen;
};

static bool maybe_open_avio(AVFormatContext *context, const char *uri)
{
    int ret;

    if (context->oformat->flags & AVFMT_NOFILE)
        return true;

    if ((ret = avio_open(&context->pb, uri, AVIO_FLAG_WRITE)) < 0) {
        fprintf(stderr, "Unable to open '%s': %s\n", uri, av_err2str(ret));
        return false;
    }

    return true;
}

static void maybe_close_avio(AVFormatContext *context)
{
    if (context->oformat->flags & AVFMT_NOFILE)
        return;

    avio_closep(&context->pb);
}

static void *alloc_read(FILE *ifile, size_t len)
{
    void *buf;

    if ((buf = malloc(len)) == NULL) {
        fprintf(stderr, "Unable to allocate buffer for "
                "intermediate video packet: %s\n", strerror(errno));
        return NULL;
    }
    if (fread(buf, 1, len, ifile) != len) {
        fprintf(stderr, "Unable to read intermediate video packet with "
                "size %zu\n", len);
        return NULL;
    }

    return buf;
}

static bool convert_endian(void *data, size_t len, uint8_t bpp)
{
    const uint16_t endian_test = 1;

    if (*(uint8_t*)&endian_test == 0)
        return true;

    switch (bpp) {
        case 2:
            for (int i = 0; i < len; ++i, data += bpp)
                *(uint16_t*)data = bswap_16(*(uint16_t*)data);
            break;
        case 4:
            for (int i = 0; i < len; ++i, data += bpp)
                *(uint32_t*)data = bswap_32(*(uint32_t*)data);
            break;
        default:
            fprintf(stderr, "Unable to handle pixel byte size of %d.\n", bpp);
            return false;
    }
    return true;
}

static AVFrame *alloc_frame(enum AVPixelFormat pix_fmt, uint32_t width,
                            uint32_t height)
{
    AVFrame *frame;

    if ((frame = av_frame_alloc()) == NULL)
        return NULL;

    frame->format = pix_fmt;
    frame->width  = width;
    frame->height = height;

    if (av_frame_get_buffer(frame, 0) < 0) {
        fprintf(stderr, "Could not allocate frame data.\n");
        return NULL;
    }

    return frame;
}

static struct packet_switch *parse_switch(FILE *ifile)
{
    void *buf;
    struct packet_switch *out;
    uint8_t tmp_format;

    if ((buf = alloc_read(ifile, 10)) == NULL)
        return NULL;

    if ((out = malloc(sizeof(struct packet_switch))) == NULL) {
        fprintf(stderr, "Unable to allocate packet_switch: %s\n",
                strerror(errno));
        free(buf);
        return NULL;
    }

    memcpy(&out->width,  buf,     4);
    memcpy(&out->height, buf + 4, 4);
    memcpy(&tmp_format,  buf + 8, 1);
    memcpy(&out->bpp,    buf + 9, 1);

    free(buf);

    switch (tmp_format) {
        case 1: out->format = AV_PIX_FMT_RGB555BE; break;
        case 2: out->format = AV_PIX_FMT_RGB565BE; break;
        case 3: out->format = AV_PIX_FMT_0RGB;     break;
        case 4: out->format = AV_PIX_FMT_RGB0;     break;
        case 5: out->format = AV_PIX_FMT_BGR0;     break;
        default:
            fprintf(stderr, "Unknown pixel format %d in switch directive.\n",
                    tmp_format);
            free(out);
            return NULL;
    }

    return out;
}

static struct packet_update *parse_update(FILE *ifile, uint8_t bpp)
{
    void *buf;
    struct packet_update *out;

    if ((buf = alloc_read(ifile, 24)) == NULL)
        return NULL;

    if ((out = malloc(sizeof(struct packet_update))) == NULL) {
        fprintf(stderr, "Unable to allocate packet_update: %s\n",
                strerror(errno));
        free(buf);
        return NULL;
    }

    memcpy(&out->x,         buf,      4);
    memcpy(&out->y,         buf + 4,  4);
    memcpy(&out->w,         buf + 8,  4);
    memcpy(&out->h,         buf + 12, 4);
    memcpy(&out->timestamp, buf + 16, 8);

    free(buf);

    out->datalen = out->w * bpp * out->h;

    return out;
}

static struct bound_info *get_bounds(FILE *ifile)
{
    struct bound_info *out;
    uint32_t width = 0, height = 0;
    uint64_t start_time = 0;
    uint8_t bpp = 0;
    size_t frames = 0;
    char opcode;

    struct packet_switch *sw = NULL;
    struct packet_update *up;

    while (!feof(ifile)) {
        opcode = getc(ifile);
        switch (opcode) {
            case 'S':
                if ((sw = parse_switch(ifile)) == NULL)
                    return NULL;
                bpp = sw->bpp;
                if (sw->width > width)
                    width = sw->width;
                if (sw->height > height)
                    height = sw->height;
                free(sw);
                break;
            case 'U':
                if ((up = parse_update(ifile, bpp)) == NULL)
                    return NULL;
                if (start_time == 0)
                    start_time = up->timestamp;
                if (fseek(ifile, up->datalen, SEEK_CUR) == -1) {
                    fprintf(stderr, "Unable to seek past framedata: %s\n",
                            strerror(errno));
                    free(up);
                    return NULL;
                }
                free(up);
                frames++;
                break;
            case EOF:
                break;
            default:
                fprintf(stderr, "Unknown opcode 0x%02x when parsing "
                        "intermediate format.\n", opcode);
                return NULL;
        }
    }

    if (width == 0 || height == 0) {
        fprintf(stderr, "Couldn't get size after processing %zu frames.\n",
                frames);
        return NULL;
    }

    if ((out = malloc(sizeof(struct bound_info))) == NULL) {
        fprintf(stderr, "Unable to allocate bound_info: %s\n",
                strerror(errno));
        return NULL;
    }

    out->width = width;
    out->height = height;
    out->start_time = start_time;
    out->frames = frames;

    rewind(ifile);
    return out;
}

static bool encode_frame(AVCodecContext *context, AVFormatContext *fcontext,
                         AVFrame *frame, AVPacket *packet)
{
    int ret;

    ret = avcodec_send_frame(context, frame);

    if (ret == AVERROR(EAGAIN)) {
        ret = avcodec_receive_packet(context, packet);
        if (ret == AVERROR_EOF)
            return true;
        else if (ret < 0) {
            fprintf(stderr, "Error encoding frame: %s\n",
                    av_err2str(ret));
            return false;
        }

        av_interleaved_write_frame(fcontext, packet);
        av_packet_unref(packet);

        return encode_frame(context, fcontext, frame, packet);
    }

    return true;
}

static bool encode_frames(FILE *ifile, AVCodecContext *context,
                          AVFormatContext *fcontext, AVFrame *oframe,
                          uint64_t start_time, size_t frames)
{
    bool status = true, first_frame = true;
    char opcode;
    size_t offset;
    void *data;
    size_t frameno = 0;
    uint64_t newpts;

    struct packet_switch *sw = NULL;
    struct packet_update *up;

    struct SwsContext *swcontext = NULL;
    AVFrame *frame;
    AVPacket *packet;

    packet = av_packet_alloc();

    while (!feof(ifile)) {
        opcode = getc(ifile);
        switch (opcode) {
            case 'S':
                if (sw != NULL) {
                    sws_freeContext(swcontext);
                    av_frame_free(&frame);
                    free(sw);
                }

                if ((sw = parse_switch(ifile)) == NULL)
                    goto out_err;

                swcontext = sws_getContext(sw->width, sw->height, sw->format,
                                           context->width, context->height,
                                           context->pix_fmt, SWS_BICUBIC,
                                           NULL, NULL, NULL);
                if (swcontext == NULL) {
                    fputs("Couldn't initialize conversion context!\n", stderr);
                    goto out_err;
                }

                frame = alloc_frame(sw->format, sw->width, sw->height);
                av_frame_make_writable(frame);
                memset(frame->data[0], 0, sw->height * sw->width * sw->bpp);

                break;
            case 'U':
                if (sw == NULL)
                    continue;

                if ((up = parse_update(ifile, sw->bpp)) == NULL)
                    goto out_err;

                if ((data = alloc_read(ifile, up->datalen)) == NULL) {
                    free(up);
                    goto out_err;
                }

                if (!convert_endian(data, up->w * up->h, sw->bpp)) {
                    free(up);
                    free(data);
                    goto out_err;
                }

                av_frame_make_writable(frame);
                offset = up->x * sw->bpp + up->y * sw->width * sw->bpp;
                memcpy(frame->data[0] + offset, data, up->datalen);

                newpts = (up->timestamp - start_time)
                       / (1000000000 / FRAMERATE);

                free(up);
                free(data);

                if (oframe->pts >= newpts && !first_frame)
                    continue;
                else
                    first_frame = false;

                oframe->pts = newpts;

                // TODO: Only scale the slice!
                sws_scale(swcontext, (const uint8_t * const *)frame->data,
                          frame->linesize, 0, sw->height, oframe->data,
                          oframe->linesize);

                fprintf(stderr, "\rEncoding frame %zu of %zu... ", ++frameno,
                        frames);
                fflush(stderr);

                if (!encode_frame(context, fcontext, oframe, packet))
                    goto out_err;

                break;
            case EOF:
                goto out;
            default:
                fprintf(stderr, "Unknown opcode 0x%02x when parsing "
                        "intermediate format.\n", opcode);
                goto out_err;
        }
    }

    goto out;

out_err:
    status = false;

out:
    fprintf(stderr, "\rEncoded %zu frames out of %zu.\n", frameno, frames);

    if (sw != NULL) {
        sws_freeContext(swcontext);
        av_frame_free(&frame);
        free(sw);
    }

    av_packet_free(&packet);

    return status;
}

int main(int argc, char **argv)
{
    FILE *ifile = NULL;
    int ret, ecode = EXIT_SUCCESS;
    struct bound_info *bounds;
    uint64_t start_time;
    size_t frames;

    const AVCodec *codec;
    AVFormatContext *fcontext = NULL;
    AVCodecContext *context = NULL;
    AVStream *stream;
    AVFrame *oframe = NULL;
    AVDictionary *opt = NULL;

    if (argc != 3) {
        fprintf(stderr, "Usage: %s <intermediate_format_file> <output_file>\n",
                argv[0]);
        return EXIT_FAILURE;
    }

    if ((ifile = fopen(argv[1], "rb")) == NULL) {
        fprintf(stderr, "Unable to open input file '%s': %s\n",
                argv[1], strerror(errno));
        return EXIT_FAILURE;
    }

    av_register_all();

    if (avformat_alloc_output_context2(&fcontext, NULL, NULL, argv[2]) < 0) {
        fprintf(stderr, "Couldn't deduce format for output file '%s'.\n",
                argv[2]);
        goto out_err;
    }

    if (fcontext->oformat->video_codec == AV_CODEC_ID_NONE) {
        fprintf(stderr, "Unable to determine video codec for '%s'.\n",
                argv[2]);
        goto out_err;
    }

    codec = avcodec_find_encoder(fcontext->oformat->video_codec);
    if (codec == NULL) {
        fprintf(stderr, "Could not find video encoder for '%s'.\n",
                avcodec_get_name(fcontext->oformat->video_codec));
        goto out_err;
    }

    if ((stream = avformat_new_stream(fcontext, NULL)) == NULL) {
        fputs("Unable to allocate stream.\n", stderr);
        goto out_err;
    }

    if ((context = avcodec_alloc_context3(codec)) == NULL) {
        fputs("Unable to allocate context for video codec.\n", stderr);
        goto out_err;
    }

    if ((bounds = get_bounds(ifile)) == NULL)
        goto out_err;

    context->width = bounds->width;
    context->height = bounds->height;
    start_time = bounds->start_time;
    frames = bounds->frames;

    free(bounds);

    context->codec_id = fcontext->oformat->video_codec;
    // FIXME: Determine automatically or at least with some profile.
    context->bit_rate = 4000000;

    // FIXME: Clean up!
    context->time_base = (AVRational){1, FRAMERATE};
    context->framerate = (AVRational){FRAMERATE, 1};
    stream->time_base = context->time_base;
    stream->avg_frame_rate = context->framerate;

    // XXX: Get rid of this!
    context->gop_size = 12;

    if (codec->pix_fmts == NULL) {
        fprintf(stderr, "Unable to determine pixel format for codec '%s'.\n",
                avcodec_get_name(fcontext->oformat->video_codec));
        goto out_err;
    }

    context->pix_fmt = codec->pix_fmts[0];

    av_dict_copy(&opt, NULL, 0);

    if (avcodec_open2(context, codec, &opt) != 0) {
        fprintf(stderr, "Unable to open context with codec '%s'.\n",
                avcodec_get_name(fcontext->oformat->video_codec));
        goto out_err;
    }

    if (avcodec_parameters_from_context(stream->codecpar, context) < 0) {
        fputs("Unable to copy stream parameters.\n", stderr);
        goto out_err;
    }

    av_dump_format(fcontext, 0, argv[2], 1);

    oframe = alloc_frame(context->pix_fmt, context->width, context->height);
    if (oframe == NULL)
        goto out_err;

    if (!maybe_open_avio(fcontext, argv[2]))
        goto out_err;

    if ((ret = avformat_write_header(fcontext, &opt)) < 0) {
        fprintf(stderr, "Unable to write stream header to '%s': %s\n",
                argv[2], av_err2str(ret));
        goto out_err;
    }

    if (!encode_frames(ifile, context, fcontext, oframe, start_time, frames))
        goto out_err;

    if ((ret = av_write_trailer(fcontext)) < 0) {
        fprintf(stderr, "Unable to write stream trailer to '%s': %s\n",
                argv[2], av_err2str(ret));
        goto out_err;
    }

    goto out;

out_err:
    ecode = EXIT_FAILURE;

out:
    if (oframe != NULL)
        av_frame_free(&oframe);
    if (context != NULL)
        avcodec_free_context(&context);
    if (fcontext != NULL) {
        if (fcontext->pb != NULL)
            maybe_close_avio(fcontext);
        avformat_free_context(fcontext);
    }
    if (ifile != NULL)
        fclose(ifile);

    return ecode;
}