FFmpeg/libavcodec/vulkan/ffv1_dec_setup.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
 */

uint8_t setup_state[CONTEXT_SIZE];

uint get_usymbol(inout RangeCoder c)
{
    if (get_rac_direct(c, setup_state[0]))
        return 0;

    int e = 0;
    while (get_rac_direct(c, setup_state[1 + min(e, 9)])) { // 1..10
        e++;
        if (e > 31) {
            corrupt = true;
            return 0;
        }
    }

    uint a = 1;
    for (int i = e - 1; i >= 0; i--) {
        a <<= 1;
        a |= uint(get_rac_direct(c, setup_state[22 + min(i, 9)]));  // 22..31
    }

    return a;
}

bool decode_slice_header(inout SliceContext sc)
{
    [[unroll]]
    for (int i = 0; i < CONTEXT_SIZE; i++)
        setup_state[i] = uint8_t(128);

    uint sx = get_usymbol(sc.c);
    uint sy = get_usymbol(sc.c);
    uint sw = get_usymbol(sc.c) + 1;
    uint sh = get_usymbol(sc.c) + 1;

    if (sx < 0 || sy < 0 || sw <= 0 || sh <= 0 ||
        sx > (gl_NumWorkGroups.x - sw) || sy > (gl_NumWorkGroups.y - sh) ||
        corrupt) {
        return true;
    }

    /* Set coordinates */
    uint sxs = slice_coord(img_size.x, sx     , gl_NumWorkGroups.x, chroma_shift.x);
    uint sxe = slice_coord(img_size.x, sx + sw, gl_NumWorkGroups.x, chroma_shift.x);
    uint sys = slice_coord(img_size.y, sy     , gl_NumWorkGroups.y, chroma_shift.y);
    uint sye = slice_coord(img_size.y, sy + sh, gl_NumWorkGroups.y, chroma_shift.y);

    sc.slice_pos = ivec2(sxs, sys);
    sc.slice_dim = ivec2(sxe - sxs, sye - sys);
    sc.slice_rct_coef = ivec2(1, 1);
    sc.slice_coding_mode = int(0);

    for (uint i = 0; i < codec_planes; i++) {
        uint idx = get_usymbol(sc.c);
        if (idx >= quant_table_count)
            return true;
        sc.quant_table_idx[i] = uint8_t(idx);
    }

    get_usymbol(sc.c);
    get_usymbol(sc.c);
    get_usymbol(sc.c);

    if (version >= 4) {
        sc.slice_reset_contexts = get_rac_direct(sc.c, setup_state[0]);
        sc.slice_coding_mode = get_usymbol(sc.c);
        if (sc.slice_coding_mode != 1 && colorspace == 1) {
            sc.slice_rct_coef.x = int(get_usymbol(sc.c));
            sc.slice_rct_coef.y = int(get_usymbol(sc.c));
            if (sc.slice_rct_coef.x + sc.slice_rct_coef.y > 4)
                return true;
        }
    }

    return false;
}

void golomb_init(inout SliceContext sc)
{
    if (version == 3 && micro_version > 1 || version > 3) {
        setup_state[0] = uint8_t(129);
        get_rac_direct(sc.c, setup_state[0]);
    }

    uint64_t ac_byte_count = sc.c.bytestream - sc.c.bytestream_start - 1;
    init_get_bits(sc.gb, u8buf(sc.c.bytestream_start + ac_byte_count),
                  sc.c.bytestream_end - sc.c.bytestream_start - ac_byte_count);
}

void main(void)
{
    const uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x;

    u8buf bs = u8buf(slice_data + slice_offsets[2*slice_idx + 0]);
    uint32_t slice_size = slice_offsets[2*slice_idx + 1];

    rac_init_dec(slice_ctx[slice_idx].c,
                 bs, slice_size);

    if (slice_idx == (gl_NumWorkGroups.x*gl_NumWorkGroups.y - 1))
        get_rac_equi(slice_ctx[slice_idx].c);

    decode_slice_header(slice_ctx[slice_idx]);

    if (golomb == 1)
        golomb_init(slice_ctx[slice_idx]);

    if (ec != 0 && check_crc != 0) {
        uint32_t crc = crcref;
        for (int i = 0; i < slice_size; i++)
            crc = crc_ieee[(crc & 0xFF) ^ uint32_t(bs[i].v)] ^ (crc >> 8);

        slice_status[2*slice_idx + 0] = crc;
    }

    slice_status[2*slice_idx + 1] = corrupt ? uint32_t(corrupt) : overread;
}