FFmpeg/libavcodec/vulkan/rangecoder.comp

/*
 * FFv1 codec
 *
 * Copyright (c) 2024 Lynne <dev@lynne.ee>
 *
 * This file is part of FFmpeg.
 *
 * 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.
 *
 * 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 FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

struct RangeCoder {
    uint64_t bytestream_start;
    uint64_t bytestream;
    uint64_t bytestream_end;

    int low;
    int range;
    uint16_t outstanding_count;
    uint8_t outstanding_byte;
};

/* Full renorm version that can handle outstanding_byte == 0xFF */
void renorm_encoder_full(inout RangeCoder c)
{
    int bs_cnt = 0;
    u8buf bytestream = u8buf(c.bytestream);

    if (c.outstanding_byte == 0xFF) {
        c.outstanding_byte = uint8_t(c.low >> 8);
    } else if (c.low <= 0xFF00) {
        bytestream[bs_cnt++].v = c.outstanding_byte;
        uint16_t cnt = c.outstanding_count;
        for (; cnt > 0; cnt--)
            bytestream[bs_cnt++].v = uint8_t(0xFF);
        c.outstanding_count = uint16_t(0);
        c.outstanding_byte = uint8_t(c.low >> 8);
    } else if (c.low >= 0x10000) {
        bytestream[bs_cnt++].v = c.outstanding_byte + uint8_t(1);
        uint16_t cnt = c.outstanding_count;
        for (; cnt > 0; cnt--)
            bytestream[bs_cnt++].v = uint8_t(0x00);
        c.outstanding_count = uint16_t(0);
        c.outstanding_byte = uint8_t(bitfieldExtract(c.low, 8, 8));
    } else {
        c.outstanding_count++;
    }

    c.bytestream += bs_cnt;
    c.range <<= 8;
    c.low = bitfieldInsert(0, c.low, 8, 8);
}

/* Cannot deal with outstanding_byte == -1 in the name of speed */
void renorm_encoder(inout RangeCoder c)
{
    uint16_t oc = c.outstanding_count + uint16_t(1);
    int low = c.low;

    c.range <<= 8;
    c.low = bitfieldInsert(0, low, 8, 8);

    if (low > 0xFF00 && low < 0x10000) {
        c.outstanding_count = oc;
        return;
    }

    u8buf bs = u8buf(c.bytestream);
    uint8_t outstanding_byte = c.outstanding_byte;

    c.bytestream        = uint64_t(bs) + oc;
    c.outstanding_count = uint16_t(0);
    c.outstanding_byte  = uint8_t(low >> 8);

    uint8_t obs = uint8_t(low > 0xFF00);
    uint8_t fill = obs - uint8_t(1); /* unsigned underflow */

    bs[0].v = outstanding_byte + obs;
    for (int i = 1; i < oc; i++)
        bs[i].v = fill;
}

void put_rac_norenorm(inout RangeCoder c, uint64_t state, bool bit)
{
    u8buf sb = u8buf(state);
    uint val = uint(sb.v);
    int range1 = uint16_t((c.range * val) >> 8);

#ifdef DEBUG
    if (val == 0)
        debugPrintfEXT("Error: state is zero (addr: 0x%lx)", uint64_t(sb));
    if (range1 >= c.range)
        debugPrintfEXT("Error: range1 >= c.range");
    if (range1 <= 0)
        debugPrintfEXT("Error: range1 <= 0");
#endif

    int diff = c.range - range1;
    if (bit) {
        c.low   += diff;
        c.range  = range1;
    } else {
        c.range  = diff;
    }

    sb.v = zero_one_state[(uint(bit) << 8) + val];

#ifdef DEBUG
    if (sb.v == 0)
        debugPrintfEXT("Error: inserted zero state from tab %i idx %i", bit, val);
#endif
}

/* Equiprobable bit */
void put_rac_equi(inout RangeCoder c, bool bit)
{
    int range1 = c.range >> 1;

#ifdef DEBUG
    if (range1 >= c.range)
        debugPrintfEXT("Error: range1 >= c.range");
    if (range1 <= 0)
        debugPrintfEXT("Error: range1 <= 0");
#endif

    if (bit) {
        c.low   += c.range - range1;
        c.range  = range1;
    } else {
        c.range -= range1;
    }

    if (c.range < 0x100)
        renorm_encoder(c);
}

void put_rac_terminate(inout RangeCoder c)
{
    int range1 = (c.range * 129) >> 8;

#ifdef DEBUG
    if (range1 >= c.range)
        debugPrintfEXT("Error: range1 >= c.range");
    if (range1 <= 0)
        debugPrintfEXT("Error: range1 <= 0");
#endif

    c.range -= range1;
    if (c.range < 0x100)
        renorm_encoder(c);
}

/* Return the number of bytes written. */
uint32_t rac_terminate(inout RangeCoder c)
{
    put_rac_terminate(c);
    c.range = uint16_t(0xFF);
    c.low  += 0xFF;
    renorm_encoder(c);
    c.range = uint16_t(0xFF);
    renorm_encoder(c);

#ifdef DEBUG
    if (c.low != 0)
        debugPrintfEXT("Error: c.low != 0");
    if (c.range < 0x100)
        debugPrintfEXT("Error: range < 0x100");
#endif

    return uint32_t(uint64_t(c.bytestream) - uint64_t(c.bytestream_start));
}

void rac_init(out RangeCoder r, u8buf data, uint buf_size)
{
    r.bytestream_start = uint64_t(data);
    r.bytestream = uint64_t(data);
    r.bytestream_end = uint64_t(data) + buf_size;
    r.low = 0;
    r.range = 0xFF00;
    r.outstanding_count = uint16_t(0);
    r.outstanding_byte = uint8_t(0xFF);
}

/* Decoder */
uint overread;
bool corrupt;

void rac_init_dec(out RangeCoder r, u8buf data, uint buf_size)
{
    overread = 0;
    corrupt = false;

    /* Skip priming bytes */
    rac_init(r, OFFBUF(u8buf, data, 2), buf_size - 2);

    u8vec2 prime = u8vec2buf(data).v;
    /* Switch endianess of the priming bytes */
    r.low = pack16(prime.yx);

    if (r.low >= 0xFF00) {
        r.low = 0xFF00;
        r.bytestream_end = uint64_t(data) + 2;
    }
}

void refill(inout RangeCoder c)
{
    c.range <<= 8;
    c.low   <<= 8;
    if (c.bytestream < c.bytestream_end) {
        c.low += u8buf(c.bytestream).v;
        c.bytestream++;
    } else {
        overread++;
    }
}

bool get_rac(inout RangeCoder c, uint64_t state)
{
    u8buf sb = u8buf(state);
    int val = int(sb.v);
    int range1 = -int(c.range * val >> 8);
    int ranged = c.range + range1;

    bool bit = c.low >= ranged;
    int bv = bit ? 0xFFFFFFFF : 0;
    sb.v = zero_one_state[(bv & 0x100) + val];

    c.low = c.low - (bv & ranged);
    c.range = (ranged & ~bv) - (range1 & bv);

    if (c.range < 0x100)
        refill(c);

    return bit;
}

bool get_rac_equi(inout RangeCoder c)
{
    int range1 = c.range >> 1;

    c.range -= range1;

    bool bit = c.low >= c.range;
    if (bit) {
        c.low -= c.range;
        c.range = range1;
    }

    if (c.range < 0x100)
        refill(c);

    return bit;
}