/*
 * Copyright © 2019-2026, VideoLAN and dav2d authors
 * Copyright © 2019-2026, Two Orioles, LLC
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 *    list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "tests/checkasm/internal.h"

#include "src/cpu.h"
#include "src/msac.h"

#include <inttypes.h>
#include <stdio.h>
#include <string.h>

#define BUF_SIZE 128

/* The normal code doesn't use function pointers */
typedef unsigned (*decode_symbol_adapt_fn)(MsacContext *s, uint16_t *cdf,
                                           size_t n_symbols);
typedef unsigned (*decode_adapt_fn)(MsacContext *s, uint16_t *cdf);
typedef unsigned (*decode_bool_bypass_fn)(MsacContext *s);
typedef unsigned (*decode_bools_bypass_fn)(MsacContext *s, unsigned n_bits);
typedef unsigned (*decode_unary_bypass6_fn)(MsacContext *s, unsigned max_bits);
typedef unsigned (*decode_unary_bypass21_fn)(MsacContext *s);

typedef struct {
    decode_symbol_adapt_fn   decode_symbol_adapt4;
    decode_symbol_adapt_fn   decode_symbol_adapt8;
    decode_adapt_fn          decode_bool_adapt;
    decode_bool_bypass_fn    decode_bool_bypass;
    decode_bools_bypass_fn   decode_bools_bypass;
    decode_unary_bypass6_fn  decode_unary_bypass6;
    decode_unary_bypass21_fn decode_unary_bypass21;
} MsacDSPContext;

static void randomize_cdf(uint16_t *const cdf, int n) {
    for (int i = n + 1; i < 8; i++)
        cdf[i] = 0; // padding
    cdf[n] = (rnd() % 125) << 8; // para
    if (rnd() & 1)
        cdf[--n] = 0; // explicitly test the zero edge case
    for (int c = 0; n > 0; n--) {
        c += rnd() % (32768 - c - n) + 1;
        cdf[n - 1] = c;
    }
}

/* memcmp() on structs can have weird behavior due to padding etc. */
static int msac_cmp(const MsacContext *const a, const MsacContext *const b) {
    return a->buf_pos != b->buf_pos || a->buf_end != b->buf_end ||
           a->dif != b->dif || a->rng != b->rng || a->cnt != b->cnt ||
           a->allow_update_cdf != b->allow_update_cdf;
}

static void msac_dump(unsigned c_res, unsigned a_res,
                      const MsacContext *const a, const MsacContext *const b,
                      const uint16_t *const cdf_a, const uint16_t *const cdf_b,
                      const int num_cdf)
{
    if (c_res != a_res)
        fprintf(stderr, "c_res %u a_res %u\n", c_res, a_res);
    if (a->buf_pos != b->buf_pos)
        fprintf(stderr, "buf_pos %p vs %p\n", a->buf_pos, b->buf_pos);
    if (a->buf_end != b->buf_end)
        fprintf(stderr, "buf_end %p vs %p\n", a->buf_end, b->buf_end);
    if (a->dif != b->dif)
        fprintf(stderr, "dif %016"PRIx64" vs %016"PRIx64"\n", a->dif, b->dif);
    if (a->rng != b->rng)
        fprintf(stderr, "rng %u vs %u\n", a->rng, b->rng);
    if (a->cnt != b->cnt)
        fprintf(stderr, "cnt %d vs %d\n", a->cnt, b->cnt);
    if (a->allow_update_cdf != b->allow_update_cdf)
        fprintf(stderr, "allow_update_cdf %d vs %d\n",
                a->allow_update_cdf, b->allow_update_cdf);
    if (num_cdf && memcmp(cdf_a, cdf_b, sizeof(*cdf_a) * (num_cdf + 1))) {
        fprintf(stderr, "cdf:\n");
        for (int i = 0; i <= num_cdf; i++)
            fprintf(stderr, " %5u", cdf_a[i]);
        fprintf(stderr, "\n");
        for (int i = 0; i <= num_cdf; i++)
            fprintf(stderr, " %5u", cdf_b[i]);
        fprintf(stderr, "\n");
        for (int i = 0; i <= num_cdf; i++)
            fprintf(stderr, "     %c", cdf_a[i] != cdf_b[i] ? 'x' : '.');
        fprintf(stderr, "\n");
    }
}

#define CHECK_SYMBOL_ADAPT(n) do {                                         \
    if (check_func(c->decode_symbol_adapt##n,                              \
                   "msac_decode_symbol_adapt%d", n))                       \
    {                                                                      \
        for (int cdf_update = 0; cdf_update <= 1; cdf_update++) {          \
            for (int ns = 1; ns < n; ns++) {                               \
                dav2d_msac_init(&s_c, buf, BUF_SIZE, !cdf_update);         \
                s_a = s_c;                                                 \
                randomize_cdf(cdf[0], ns);                                 \
                memcpy(cdf[1], cdf[0], sizeof(*cdf));                      \
                while (s_c.cnt >= 0) {                                     \
                    unsigned c_res = call_ref(&s_c, cdf[0], ns);           \
                    unsigned a_res = call_new(&s_a, cdf[1], ns);           \
                    if (c_res != a_res || msac_cmp(&s_c, &s_a) ||          \
                        memcmp(cdf[0], cdf[1], sizeof(**cdf) * (ns + 1)))  \
                    {                                                      \
                        if (fail()) {                                      \
                            fprintf(stderr, "n_symbols = %d\n", ns);       \
                            msac_dump(c_res, a_res, &s_c, &s_a,            \
                                      cdf[0], cdf[1], ns);                 \
                        }                                                  \
                        break;                                             \
                    }                                                      \
                }                                                          \
                if (cdf_update && ns == n - 1)                             \
                    bench_new(alternate(&s_c, &s_a),                       \
                              alternate(cdf[0], cdf[1]), ns);              \
            }                                                              \
        }                                                                  \
    }                                                                      \
} while (0)

static void check_decode_symbol_adapt(MsacDSPContext *const c, uint8_t *const buf) {
    ALIGN_STK_16(uint16_t, cdf, 2, [8]);
    MsacContext s_c, s_a;

    declare_func(unsigned, MsacContext *s, uint16_t *cdf, size_t n_symbols);
    CHECK_SYMBOL_ADAPT(4);
    CHECK_SYMBOL_ADAPT(8);
}

static void check_decode_bool_adapt(MsacDSPContext *const c, uint8_t *const buf) {
    MsacContext s_c, s_a;

    declare_func(unsigned, MsacContext *s, uint16_t *cdf);
    if (check_func(c->decode_bool_adapt, "msac_decode_bool_adapt")) {
        uint16_t cdf[2][2];
        for (int cdf_update = 0; cdf_update <= 1; cdf_update++) {
            dav2d_msac_init(&s_c, buf, BUF_SIZE, !cdf_update);
            s_a = s_c;
            cdf[0][0] = cdf[1][0] = rnd() % 32767 + 1;
            cdf[0][1] = cdf[1][1] = (rnd() % 125) << 8;
            while (s_c.cnt >= 0) {
                unsigned c_res = call_ref(&s_c, cdf[0]);
                unsigned a_res = call_new(&s_a, cdf[1]);
                if (c_res != a_res || msac_cmp(&s_c, &s_a) ||
                    memcmp(cdf[0], cdf[1], sizeof(*cdf)))
                {
                    if (fail())
                        msac_dump(c_res, a_res, &s_c, &s_a, cdf[0], cdf[1], 1);
                    break;
                }
            }
            if (cdf_update)
                bench_new(alternate(&s_c, &s_a), alternate(cdf[0], cdf[1]));
        }
    }
}

static unsigned generate_bypass_rng(const MsacContext *const s) {
    const unsigned dif = (unsigned)(s->dif >> 48);
    const unsigned r = rnd();
    assert(dif < 0xfffe);
    if (dif & 0x8000)
        return (dif + (r % (0xfffe - dif)) + 2) & 0xfffe;
    return 0x8000 | (r & 0x7ffe);
}

static void check_decode_bool_bypass(MsacDSPContext *const c, uint8_t *const buf) {
    MsacContext s_c, s_a;

    declare_func(unsigned, MsacContext *s);
    if (check_func(c->decode_bool_bypass, "msac_decode_bool_bypass")) {
        dav2d_msac_init(&s_c, buf, BUF_SIZE, 1);
        s_a = s_c;
        while (s_c.cnt >= 0) {
            s_a.rng = s_c.rng = generate_bypass_rng(&s_c);
            unsigned c_res = call_ref(&s_c);
            unsigned a_res = call_new(&s_a);
            if (c_res != a_res || msac_cmp(&s_c, &s_a)) {
                if (fail())
                    msac_dump(c_res, a_res, &s_c, &s_a, NULL, NULL, 0);
                break;
            }
        }
        s_a.dif = s_c.dif >>= 1;
        s_c.rng = 0xfc92; // Somewhat arbitrarily chosen to produce
        s_a.rng = 0xdb6e; // a reasonably diverse branch pattern.
        bench_new(alternate(&s_c, &s_a));
    }
}

static void check_decode_bools_bypass(MsacDSPContext *const c, uint8_t *const buf) {
    MsacContext s_c, s_a;

    declare_func(unsigned, MsacContext *s, unsigned n_bits);
    if (check_func(c->decode_bools_bypass, "msac_decode_bools_bypass")) {
        dav2d_msac_init(&s_c, buf, BUF_SIZE, 1);
        s_a = s_c;
        while (s_c.cnt >= 0) {
            s_a.rng = s_c.rng = generate_bypass_rng(&s_c);
            const int n_bits = 1 + (rnd() & 31);
            unsigned c_res = call_ref(&s_c, n_bits);
            unsigned a_res = call_new(&s_a, n_bits);
            if (c_res != a_res || msac_cmp(&s_c, &s_a)) {
                if (fail()) {
                    fprintf(stderr, "n_bits = %d\n", n_bits);
                    msac_dump(c_res, a_res, &s_c, &s_a, NULL, NULL, 0);
                }
                break;
            }
        }
        s_a.dif = s_c.dif >>= 1;
        s_c.rng = 0xfc92;
        s_a.rng = 0xdb6e;
        bench_new(alternate(&s_c, &s_a), 8);
    }
}

static void check_decode_unary_bypass6(MsacDSPContext *const c, uint8_t *const buf) {
    MsacContext s_c, s_a;

    declare_func(unsigned, MsacContext *s, unsigned max_bits);
    if (check_func(c->decode_unary_bypass6, "msac_decode_unary_bypass6")) {
        dav2d_msac_init(&s_c, buf, BUF_SIZE, 1);
        s_a = s_c;
        while (s_c.cnt >= 0) {
            s_a.rng = s_c.rng = generate_bypass_rng(&s_c);
            const int max_bits = 5 + (rnd() & 1);
            unsigned c_res = call_ref(&s_c, max_bits);
            unsigned a_res = call_new(&s_a, max_bits);
            if (c_res != a_res || msac_cmp(&s_c, &s_a)) {
                if (fail()) {
                    fprintf(stderr, "max_bits = %d\n", max_bits);
                    msac_dump(c_res, a_res, &s_c, &s_a, NULL, NULL, 0);
                }
                break;
            }
        }
        s_a.dif = s_c.dif >>= 1;
        s_c.rng = 0xfc92;
        s_a.rng = 0xdb6e;
        bench_new(alternate(&s_c, &s_a), 6);
    }
}

static void check_decode_unary_bypass21(MsacDSPContext *const c, uint8_t *const buf) {
    MsacContext s_c, s_a;

    declare_func(unsigned, MsacContext *s);
    if (check_func(c->decode_unary_bypass21, "msac_decode_unary_bypass21")) {
        dav2d_msac_init(&s_c, buf, BUF_SIZE, 1);
        s_a = s_c;
        while (s_c.cnt >= 0) {
            s_a.rng = s_c.rng = generate_bypass_rng(&s_c);
            unsigned c_res = call_ref(&s_c);
            unsigned a_res = call_new(&s_a);
            if (c_res != a_res || msac_cmp(&s_c, &s_a)) {
                if (fail())
                    msac_dump(c_res, a_res, &s_c, &s_a, NULL, NULL, 0);
                break;
            }
        }
        s_a.dif = s_c.dif >>= 1;
        s_c.rng = 0xfc92;
        s_a.rng = 0xdb6e;
        bench_new(alternate(&s_c, &s_a));
    }
}

static void check_decode_adapt(MsacDSPContext *const c, uint8_t *const buf) {
    check_decode_symbol_adapt(c, buf);
    check_decode_bool_adapt(c, buf);
    report("decode_adapt");
}

static void check_decode_bypass(MsacDSPContext *const c, uint8_t *const buf) {
    check_decode_bool_bypass(c, buf);
    check_decode_bools_bypass(c, buf);
    check_decode_unary_bypass6(c, buf);
    check_decode_unary_bypass21(c, buf);
    report("decode_bypass");
}

void checkasm_check_msac(void) {
    /* For performance reasons entropy decoding functions are called directly
     * instead of through function pointers. For testing purposes however we
     * do want to use functions pointers . */
    MsacDSPContext c;
    c.decode_symbol_adapt4  = dav2d_msac_decode_symbol_adapt_c;
    c.decode_symbol_adapt8  = dav2d_msac_decode_symbol_adapt_c;
    c.decode_bool_adapt     = dav2d_msac_decode_bool_adapt_c;
    c.decode_bool_bypass    = dav2d_msac_decode_bool_bypass_c;
    c.decode_bools_bypass   = dav2d_msac_decode_bools_bypass_c;
    c.decode_unary_bypass6  = dav2d_msac_decode_unary_bypass_c;
    c.decode_unary_bypass21 = dav2d_msac_decode_unary_bypass21_c;

#if HAVE_ASM
#if ARCH_AARCH64
    if (dav2d_get_cpu_flags() & DAV2D_ARM_CPU_FLAG_NEON) {
        c.decode_bool_adapt     = dav2d_msac_decode_bool_adapt_neon;
        c.decode_bool_bypass    = dav2d_msac_decode_bool_bypass_neon;
        c.decode_bools_bypass   = dav2d_msac_decode_bools_bypass_neon;
        c.decode_unary_bypass6  = dav2d_msac_decode_unary_bypass_neon;
        c.decode_unary_bypass21  = dav2d_msac_decode_unary_bypass21_neon;
    }
#elif ARCH_X86_64
    if (dav2d_get_cpu_flags() & DAV2D_X86_CPU_FLAG_SSE2) {
        c.decode_symbol_adapt4  = dav2d_msac_decode_symbol_adapt4_sse2;
        c.decode_symbol_adapt8  = dav2d_msac_decode_symbol_adapt8_sse2;
        c.decode_bool_adapt     = dav2d_msac_decode_bool_adapt_sse2;
        c.decode_bool_bypass    = dav2d_msac_decode_bool_bypass_sse2;
        c.decode_bools_bypass   = dav2d_msac_decode_bools_bypass_sse2;
    }
    if (dav2d_get_cpu_flags() & DAV2D_X86_CPU_FLAG_AVX2) {
        c.decode_unary_bypass6  = dav2d_msac_decode_unary_bypass6_avx2;
        c.decode_unary_bypass21 = dav2d_msac_decode_unary_bypass21_avx2;
    }
    if (dav2d_get_cpu_flags() & DAV2D_X86_CPU_FLAG_AVX512ICL) {
        c.decode_unary_bypass21 = dav2d_msac_decode_unary_bypass21_avx512icl;
    }
#endif
#endif

    uint8_t buf[BUF_SIZE];
    for (int i = 0; i < BUF_SIZE; i++)
        buf[i] = rnd();

    check_decode_adapt(&c, buf);
    check_decode_bypass(&c, buf);
}
