/*
 * Standalone reproduction test for OpenSSL ML-KEM codec failures
 * when compiled with chibicc.
 *
 * ML-KEM (FIPS 203) covers three layers of bit-manipulation
 * that chibicc is known to mishandle:
 *
 *   A. ByteEncode_d / ByteDecode_d  – bit-packing across byte boundaries
 *   B. Compress_d / Decompress_d    – unsigned long multiply-shift (Barrett)
 *   C. NTT polynomial arithmetic    – int16_t/int32_t promotion
 *
 * Build:
 *   cc -O2 -o test_ml_kem_codec test_ml_kem_codec.c
 *
 * Expected with a conforming compiler: ALL PASS (exit 0)
 * Observed with chibicc:              multiple FAILures (exit 1)
 */

#include <stdio.h>
#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <stdlib.h>

/* ------------------------------------------------------------------ */
/* Constants                                                            */
/* ------------------------------------------------------------------ */
#define MLKEM_Q   3329   /* prime modulus */
#define MLKEM_N   256    /* polynomial degree */

/* ------------------------------------------------------------------ */
/* Test accounting                                                      */
/* ------------------------------------------------------------------ */
static int g_tests = 0, g_failed = 0;

static void check(const char *name, int ok)
{
    ++g_tests;
    printf("%s  %s\n", ok ? "pass" : "FAIL", name);
    if (!ok) ++g_failed;
}

/* ================================================================== */
/* SECTION 1 – ByteEncode_d / ByteDecode_d  (FIPS 203 §4.2.1)        */
/*                                                                      */
/*  Pack N d-bit integers into bytes LSB-first.                        */
/*  The optimised streaming variant uses a uint32_t accumulator and   */
/*  multi-byte shifts that chibicc mishandles for d=10,11,12.         */
/* ================================================================== */

/* Reference: unambiguous bit-by-bit implementation */
static void byte_encode_ref(const uint16_t *f, int d,
                             uint8_t *out, size_t out_len)
{
    memset(out, 0, out_len);
    int bit_pos = 0;
    for (int i = 0; i < MLKEM_N; i++) {
        uint32_t val = f[i] & ((1u << d) - 1);
        for (int b = 0; b < d; b++) {
            if (val & (1u << b))
                out[bit_pos / 8] |= (uint8_t)(1u << (bit_pos % 8));
            bit_pos++;
        }
    }
}

static void byte_decode_ref(const uint8_t *in, int d, uint16_t *f)
{
    int bit_pos = 0;
    for (int i = 0; i < MLKEM_N; i++) {
        uint32_t val = 0;
        for (int b = 0; b < d; b++) {
            if (in[bit_pos / 8] & (1u << (bit_pos % 8)))
                val |= (1u << b);
            bit_pos++;
        }
        f[i] = (uint16_t)(val & ((1u << d) - 1));
    }
}

/*
 * Optimised streaming encode — as generated by OpenSSL.
 * chibicc bug: the shift  buf |= val << bits  is computed in int (32-bit)
 * but for large 'bits' the result is silently truncated.
 */
static void byte_encode_opt(const uint16_t *f, int d,
                              uint8_t *out, size_t out_len)
{
    memset(out, 0, out_len);
    uint32_t buf = 0;
    int bits = 0;
    size_t o = 0;
    for (int i = 0; i < MLKEM_N; i++) {
        uint32_t val = (uint32_t)(f[i]) & ((1u << d) - 1);
        buf |= val << bits;
        bits += d;
        while (bits >= 8 && o < out_len) {
            out[o++] = (uint8_t)(buf & 0xFF);
            buf >>= 8;
            bits -= 8;
        }
    }
    if (bits > 0 && o < out_len)
        out[o] = (uint8_t)(buf & 0xFF);
}

static void byte_decode_opt(const uint8_t *in, int d, uint16_t *f)
{
    uint32_t buf = 0;
    int bits = 0;
    int i_in = 0;
    uint32_t mask = (1u << d) - 1;
    for (int i = 0; i < MLKEM_N; i++) {
        while (bits < d) {
            buf |= (uint32_t)in[i_in++] << bits;
            bits += 8;
        }
        f[i] = (uint16_t)(buf & mask);
        buf >>= d;
        bits -= d;
    }
}

static void test_byte_encode_decode(int d)
{
    char name[80];
    size_t out_len = ((size_t)MLKEM_N * (size_t)d + 7) / 8;
    uint8_t enc_ref[512], enc_opt[512];
    uint16_t dec_ref[MLKEM_N], dec_opt[MLKEM_N];

    uint16_t poly[MLKEM_N];
    uint32_t mask = (1u << d) - 1;
    for (int i = 0; i < MLKEM_N; i++)
        poly[i] = (uint16_t)((i * 7 + 13) & mask);

    byte_encode_ref(poly, d, enc_ref, out_len);
    byte_encode_opt(poly, d, enc_opt, out_len);

    snprintf(name, sizeof(name), "ByteEncode_%2d opt==ref", d);
    check(name, memcmp(enc_ref, enc_opt, out_len) == 0);

    byte_decode_ref(enc_ref, d, dec_ref);
    snprintf(name, sizeof(name), "ByteDecode_%2d ref roundtrip", d);
    check(name, memcmp(poly, dec_ref, MLKEM_N * 2) == 0);

    byte_decode_opt(enc_ref, d, dec_opt);
    snprintf(name, sizeof(name), "ByteDecode_%2d opt roundtrip", d);
    check(name, memcmp(poly, dec_opt, MLKEM_N * 2) == 0);

    /* cross: opt-encode → opt-decode */
    uint16_t dec_cross[MLKEM_N];
    byte_decode_opt(enc_opt, d, dec_cross);
    snprintf(name, sizeof(name), "ByteCodec_%2d  cross", d);
    check(name, memcmp(poly, dec_cross, MLKEM_N * 2) == 0);
}

/* ================================================================== */
/* SECTION 2 – Compress_d / Decompress_d  (FIPS 203 §4.2.1)          */
/*                                                                      */
/*  compress(x, d) = round(2^d * x / q) mod 2^d                       */
/*  decompress(y, d) = round(q * y / 2^d)                             */
/*                                                                      */
/*  The optimised path uses a unsigned long intermediate product.          */
/*  chibicc bug: (unsigned long)a * b may be silently computed as 32-bit.  */
/* ================================================================== */

static uint16_t compress_ref(uint16_t x, int d)
{
    uint32_t num = (uint32_t)x << d;
    uint32_t mask = (1u << d) - 1;
    return (uint16_t)(((num + MLKEM_Q / 2) / MLKEM_Q) & mask);
}

/*
 * Optimised compress using a unsigned long multiply.
 * The cast (unsigned long)t is the critical expression chibicc may drop.
 */
static uint16_t compress_opt(uint16_t x, int d)
{
    /*
     * We want floor((x * 2^d + q/2) / q).
     * Compute as a 64-bit product to avoid overflow for large d.
     */
    unsigned long num = ((unsigned long)x << d) + (MLKEM_Q / 2);
    return (uint16_t)((num / MLKEM_Q) & ((1u << d) - 1));
}

static uint16_t decompress_ref(uint16_t y, int d)
{
    uint32_t num = (uint32_t)MLKEM_Q * (uint32_t)y;
    return (uint16_t)((num + (1u << (d - 1))) >> d);
}

/*
 * Optimised decompress — also uses a 32→64 bit intermediate.
 * chibicc may mishandle the shift when d is large.
 */
static uint16_t decompress_opt(uint16_t y, int d)
{
    unsigned long num = (unsigned long)MLKEM_Q * (unsigned long)y;
    return (uint16_t)((num + ((unsigned long)1 << (d - 1))) >> d);
}

static void test_compress_decompress(int d)
{
    char name[80];
    int errs_c = 0, errs_d = 0;

    for (int x = 0; x < MLKEM_Q; x++) {
        uint16_t cr = compress_ref((uint16_t)x, d);
        uint16_t co = compress_opt((uint16_t)x, d);
        if (cr != co) errs_c++;

        uint16_t dr = decompress_ref(cr, d);
        uint16_t doo = decompress_opt(cr, d);
        if (dr != doo) errs_d++;
    }

    snprintf(name, sizeof(name), "Compress_%2d   opt==ref (over all x<q)", d);
    check(name, errs_c == 0);

    snprintf(name, sizeof(name), "Decompress_%2d opt==ref (over all y)", d);
    check(name, errs_d == 0);
}

/* ================================================================== */
/* SECTION 3 – Full codec: ByteEncode(Compress) / Decompress(Decode)  */
/* ================================================================== */

static void encode_cmprs(const uint16_t *poly, int d,
                          uint8_t *out, size_t out_len)
{
    uint16_t tmp[MLKEM_N];
    for (int i = 0; i < MLKEM_N; i++)
        tmp[i] = compress_opt(poly[i], d);
    byte_encode_opt(tmp, d, out, out_len);
}

static void decode_cmprs(const uint8_t *in, int d, uint16_t *poly)
{
    uint16_t tmp[MLKEM_N];
    byte_decode_opt(in, d, tmp);
    for (int i = 0; i < MLKEM_N; i++)
        poly[i] = decompress_opt(tmp[i], d);
}

static void test_full_codec(int d)
{
    char name[80];
    size_t out_len = ((size_t)MLKEM_N * (size_t)d + 7) / 8;
    uint8_t enc1[512], enc2[512];

    uint16_t poly[MLKEM_N];
    for (int i = 0; i < MLKEM_N; i++)
        poly[i] = (uint16_t)((i * 317 + 127) % MLKEM_Q);

    encode_cmprs(poly, d, enc1, out_len);

    uint16_t poly2[MLKEM_N];
    decode_cmprs(enc1, d, poly2);
    encode_cmprs(poly2, d, enc2, out_len);

    snprintf(name, sizeof(name), "FullCodec_%2d enc∘dec∘enc idempotent", d);
    check(name, memcmp(enc1, enc2, out_len) == 0);

    /* error bound: |decompress(compress(x)) - x| mod q <= q/2^d + 1 */
    int errs = 0;
    uint32_t threshold = ((uint32_t)MLKEM_Q >> d) + 1;
    for (int i = 0; i < MLKEM_N; i++) {
        int32_t diff = (int32_t)poly2[i] - (int32_t)poly[i];
        if (diff >  (int32_t)(MLKEM_Q / 2)) diff -= MLKEM_Q;
        if (diff < -(int32_t)(MLKEM_Q / 2)) diff += MLKEM_Q;
        if (diff < 0) diff = -diff;
        if ((uint32_t)diff > threshold) errs++;
    }
    snprintf(name, sizeof(name), "FullCodec_%2d error bound OK", d);
    check(name, errs == 0);
}

/* ================================================================== */
/* SECTION 4 – Bit-manipulation patterns from the actual codec code   */
/*                                                                      */
/*  These isolate the exact expressions chibicc gets wrong.           */
/* ================================================================== */

static void test_bit_patterns(void)
{
    /*
     * 4a. 12-bit coefficient extraction (used in ByteDecode_12,
     *     which encodes 256 NTT coefficients as 384 bytes).
     *
     * Given 3 bytes, extract two 12-bit values:
     *   lo = buf[0] | ((buf[1] & 0x0F) << 8)
     *   hi = (buf[1] >> 4) | (buf[2] << 4)          ← chibicc truncates here
     */
    {
        uint8_t buf[3] = { 0xAB, 0xCD, 0xEF };
        /* lo: bits 0..11 */
        uint16_t lo = (uint16_t)(buf[0])
                    | (uint16_t)((uint16_t)(buf[1] & 0x0F) << 8);
        /* hi: bits 12..23 — need (uint16_t) cast BEFORE shift to avoid truncation */
        uint16_t hi = (uint16_t)(buf[1] >> 4)
                    | (uint16_t)((uint16_t)buf[2] << 4);
        /* 0xAB | 0xD00 = 0xDAB */
        check("BitPattern 12-bit lo == 0xDAB", lo == 0xDAB);
        /* 0x0C | 0xEF0 = 0xEFC */
        check("BitPattern 12-bit hi == 0xEFC", hi == 0xEFC);
    }

    /*
     * 4b. 10-bit coefficient pack/unpack across all MLKEM_N=256 coefficients
     *     (ML-KEM-512 ciphertext u, d=10 → 320 bytes).
     *     All coefficients set to the max 10-bit value 0x3FF.
     */
    {
        uint16_t src[MLKEM_N];
        for (int i = 0; i < MLKEM_N; i++) src[i] = 0x3FF;
        uint8_t  packed[320];  /* MLKEM_N * 10 / 8 = 320 */
        memset(packed, 0, sizeof(packed));
        byte_encode_ref(src, 10, packed, 320);

        uint16_t got[MLKEM_N];
        byte_decode_ref(packed, 10, got);
        check("BitPattern 10-bit pack/unpack c0",   got[0]   == 0x3FF);
        check("BitPattern 10-bit pack/unpack c127",  got[127] == 0x3FF);
        check("BitPattern 10-bit pack/unpack c128",  got[128] == 0x3FF);
        check("BitPattern 10-bit pack/unpack c255",  got[255] == 0x3FF);
    }

    /*
     * 4c. unsigned long multiply: chibicc may compute (unsigned long)a*b in 32 bits
     *     when one operand is uint32_t.
     */
    {
        uint32_t a = 0xFFFFFFFFu;
        uint32_t b = 2u;
        unsigned long p = (unsigned long)a * b;
        check("uint64 multiply 0xFFFFFFFF*2 == 0x1FFFFFFFE",
              p == 0x1FFFFFFFEull);
    }

    /*
     * 4d. Shift of unsigned long value: chibicc may truncate to 32 bits.
     */
    {
        unsigned long v = (unsigned long)1 << 32;
        check("uint64 shift 1<<32 != 0", v != 0);
        check("uint64 shift 1<<32 == 0x100000000", v == 0x100000000ull);
    }

    /*
     * 4e. Multi-byte packing at non-zero bit offset.
     *     Mirrors the ByteEncode_11 streaming encoder.
     */
    {
        uint16_t coeff = 0x7FF;  /* 11-bit maximum */
        uint8_t out[3] = {0};
        int bit_off = 5;
        /* Pack 11 bits at bit position 5 into out[0..2] */
        out[0] |= (uint8_t)(coeff << bit_off);               /* low 3 bits */
        out[1]  = (uint8_t)(coeff >> (8  - bit_off));        /* next 8 bits */
        out[2]  = (uint8_t)(coeff >> (16 - bit_off));        /* top bits */
        /* Unpack */
        uint16_t rec = (uint16_t)(out[0] >> bit_off)
                     | (uint16_t)((uint16_t)out[1] << (8  - bit_off))
                     | (uint16_t)((uint16_t)out[2] << (16 - bit_off));
        rec &= 0x7FF;
        check("MultiPack 11-bit at bit_off=5 round-trips", rec == 0x7FF);
    }

    /*
     * 4f. Signed arithmetic right shift — used in Barrett reduction and NTT.
     *     C11 defines this as implementation-defined, but all real targets
     *     (and chibicc's target x86-64 ABI) must sign-extend.
     */
    {
        int32_t neg = -3329;
        int32_t shifted = neg >> 1;
        check("Signed right shift -3329>>1 == -1665",
              shifted == -1665);
    }
}

/* ================================================================== */
/* SECTION 5 – NTT Montgomery reduction                                */
/*                                                                      */
/*  Montgomery reduce: given t in [-q*2^15, q*2^15), return           */
/*  t * R^{-1} mod q, where R = 2^16.                                 */
/*  chibicc may mis-promote int16_t * int32_t to int32_t.             */
/* ================================================================== */

/* q^{-1} mod 2^16 = 62209  (because 3329 * 62209 ≡ 1 mod 2^16) */
#define MLKEM_QINV  62209

static int16_t montgomery_reduce(int32_t t)
{
    /*
     * u = (int16_t)(t mod 2^16) * (-q^{-1} mod 2^16)
     * The cast to int16_t followed by int32_t multiplication is the
     * expression chibicc mis-compiles (sign-extension missing).
     */
    int16_t u = (int16_t)((uint16_t)t * (uint16_t)MLKEM_QINV);
    return (int16_t)((t - (int32_t)u * (int32_t)MLKEM_Q) >> 16);
}

static void test_montgomery_reduce(void)
{
    /* Known values: montgomery_reduce(r*a mod q) == a for r=2^16 mod q=2285 */
    int32_t R = 2285;  /* R = 2^16 mod q */

    /* reduce(0) == 0 */
    check("MontReduce(0) == 0", montgomery_reduce(0) == 0);

    /* reduce(R * 1) ≡ 1 (mod q) */
    int16_t mr1 = montgomery_reduce(R * 1);
    check("MontReduce(R*1) == 1", mr1 == 1);

    /* reduce(R * q) ≡ 0 (mod q) */
    int16_t mrq = montgomery_reduce(R * MLKEM_Q);
    int32_t mrq_mod = ((int32_t)mrq % MLKEM_Q + MLKEM_Q) % MLKEM_Q;
    check("MontReduce(R*q) ≡ 0 (mod q)", mrq_mod == 0);

    /* reduce(R * 1337) ≡ 1337 (mod q) */
    int16_t mr1337 = montgomery_reduce(R * 1337);
    int32_t mr1337_mod = ((int32_t)mr1337 % MLKEM_Q + MLKEM_Q) % MLKEM_Q;
    check("MontReduce(R*1337) ≡ 1337 (mod q)", mr1337_mod == 1337);
}

/* ================================================================== */
/* SECTION 6 – Struct layout (key / ciphertext encoding)              */
/*                                                                      */
/*  OpenSSL serialises keys as flat byte arrays; any compiler-added   */
/*  padding would shift field offsets and break encode/decode.        */
/* ================================================================== */

/* ML-KEM-512 expanded decapsulation key (dk) layout per FIPS 203 */
typedef struct {
    uint8_t z[32];          /* rejection-sampling seed       */
    uint8_t rho[32];        /* public matrix seed            */
    uint8_t H_ek[32];       /* H(ek) = SHA3-256(ek)          */
    uint8_t s_hat[768];     /* NTT of secret key (k=2 polys) */
    uint8_t t_hat[768];     /* NTT of public key (k=2 polys) */
} MLKEM512_DK;              /* expected total: 32+32+32+768+768 = 1632 */

typedef struct {
    uint8_t rho[32];        /* seed for matrix A */
    uint8_t t_hat[768];     /* NTT of public key */
} MLKEM512_EK;              /* expected: 800 */

static void test_struct_layout(void)
{
    check("sizeof(MLKEM512_DK) == 1632", sizeof(MLKEM512_DK) == 1632);
    check("sizeof(MLKEM512_EK) == 800",  sizeof(MLKEM512_EK) == 800);

    MLKEM512_DK dk;
    ptrdiff_t z_off   = (uint8_t*)&dk.z     - (uint8_t*)&dk;
    ptrdiff_t rho_off = (uint8_t*)&dk.rho   - (uint8_t*)&dk;
    ptrdiff_t hek_off = (uint8_t*)&dk.H_ek  - (uint8_t*)&dk;
    ptrdiff_t sh_off  = (uint8_t*)&dk.s_hat - (uint8_t*)&dk;
    ptrdiff_t th_off  = (uint8_t*)&dk.t_hat - (uint8_t*)&dk;

    check("MLKEM512_DK.z      at offset   0", z_off   ==    0);
    check("MLKEM512_DK.rho    at offset  32", rho_off ==   32);
    check("MLKEM512_DK.H_ek   at offset  64", hek_off ==   64);
    check("MLKEM512_DK.s_hat  at offset  96", sh_off  ==   96);
    check("MLKEM512_DK.t_hat  at offset 864", th_off  ==  864);
}

/* ================================================================== */
/* SECTION 7 – Integer promotion edge cases                            */
/*                                                                      */
/*  These isolate the exact C promotion rules chibicc violates.       */
/* ================================================================== */

static void test_integer_promotions(void)
{
    /*
     * 7a. uint8_t shifted left by ≥8 must be zero (after promotion to int).
     *     chibicc bug: may leave the high byte intact.
     */
    {
        uint8_t b = 0xFF;
        int shifted = b << 8;
        check("(uint8_t)0xFF << 8 == 0xFF00", shifted == 0xFF00);
        int shifted16 = b << 16;
        check("(uint8_t)0xFF << 16 == 0xFF0000", shifted16 == 0xFF0000);
    }

    /*
     * 7b. uint16_t in arithmetic expression must be promoted to int (≥32-bit).
     *     chibicc bug: may stay 16-bit, losing overflow.
     */
    {
        uint16_t a = 0xFFFF, b = 0xFFFF;
        uint32_t product = (uint32_t)a * (uint32_t)b;
        check("(uint32_t)0xFFFF * 0xFFFF == 0xFFFE0001",
              product == 0xFFFE0001u);
    }

    /*
     * 7c. Signed/unsigned mix in compress: the expression
     *       (int32_t)x - (int32_t)q  must sign-extend correctly.
     */
    {
        uint16_t x = 3328u;       /* q - 1 */
        int32_t diff = (int32_t)x - (int32_t)MLKEM_Q;   /* = -1 */
        check("(int32_t)(q-1) - (int32_t)q == -1", diff == -1);
    }

    /*
     * 7d. Barrett constant computation requires unsigned long numerator.
     *     For d=11, max input to compress is (q-1)*2048 + 1664 ≈ 6.8M,
     *     which overflows uint32_t when shifted.
     */
    {
        uint32_t max_in = (uint32_t)(MLKEM_Q - 1) * 2048u + MLKEM_Q / 2;
        unsigned long num64 = (unsigned long)max_in;
        /* Ensure 64-bit path doesn't overflow */
        check("Barrett d=11 numerator fits uint64",
              num64 < 0xFFFFFFFFFFFFFFFFull);
        /* Ensure it is much larger than MLKEM_Q (would overflow uint16_t) */
        check("Barrett d=11 numerator > 0xFFFF",
              num64 > 0xFFFFull);
    }
}

/* ================================================================== */
/* main                                                                 */
/* ================================================================== */

int main(void)
{
    printf("=== ML-KEM Codec Reproduction Test (chibicc compat) ===\n");
    printf("    (FIPS 203 / OpenSSL 15-test_ml_kem_codecs.t)\n\n");

    printf("-- Section 1: ByteEncode/ByteDecode (d=1..12) --\n");
    for (int d = 1; d <= 12; d++)
        test_byte_encode_decode(d);

    printf("\n-- Section 2: Compress/Decompress (d=1..11) --\n");
    for (int d = 1; d <= 11; d++)
        test_compress_decompress(d);

    printf("\n-- Section 3: Full codec (compress+encode / decode+decompress) --\n");
    int ds[] = { 1, 4, 5, 10, 11 };
    for (int i = 0; i < 5; i++)
        test_full_codec(ds[i]);

    printf("\n-- Section 4: Bit-manipulation patterns --\n");
    test_bit_patterns();

    printf("\n-- Section 5: NTT Montgomery reduction --\n");
    test_montgomery_reduce();

    printf("\n-- Section 6: Struct layout / field offsets --\n");
    test_struct_layout();

    printf("\n-- Section 7: Integer promotion edge cases --\n");
    test_integer_promotions();

    printf("\n========================================\n");
    printf("Results: %d/%d passed", g_tests - g_failed, g_tests);
    if (g_failed)
        printf(", %d FAILED  ← chibicc codegen bug(s) detected\n", g_failed);
    else
        printf(" — ALL PASS\n");
    printf("========================================\n");

    return (g_failed > 0) ? 1 : 0;
}
