generate.py

from pathlib import Path


def HELPER(u64: bool):
    if u64:
        return """\
static inline uint64_t rotl64(uint64_t x, unsigned int N) {
    N &= 63; return x << N | x >> (64 - N);
}

static inline uint64_t xr_step_inv64(uint64_t x, uint64_t r1, uint64_t r2) {
    x ^= rotl64(x, r1) ^ rotl64(x, r2); r1 += r1; r2 += r2;
    x ^= rotl64(x, r1) ^ rotl64(x, r2); r1 += r1; r2 += r2;
    x ^= rotl64(x, r1) ^ rotl64(x, r2); r1 += r1; r2 += r2;
    x ^= rotl64(x, r1) ^ rotl64(x, r2); r1 += r1; r2 += r2;
    x ^= rotl64(x, r1) ^ rotl64(x, r2); r1 += r1; r2 += r2;
    x ^= rotl64(x, r1) ^ rotl64(x, r2); return x;
}

static inline uint64_t moremur(uint64_t x) {
    x ^= x >> 27;
    x *= 0x3c79ac492ba7b653ULL;
    x ^= x >> 33;
    x *= 0x1c69b3f74ac4ae35ULL;
    x ^= x >> 27;
    return x;
}"""
    else:
        return """\
static inline uint32_t rotl32(uint32_t x, unsigned int N) {
    N &= 31; return x << N | x >> (32 - N);
}

static inline uint32_t xr_step_inv32(uint32_t x, uint32_t r1, uint32_t r2) {
    x ^= rotl32(x, r1) ^ rotl32(x, r2); r1 += r1; r2 += r2;
    x ^= rotl32(x, r1) ^ rotl32(x, r2); r1 += r1; r2 += r2;
    x ^= rotl32(x, r1) ^ rotl32(x, r2); r1 += r1; r2 += r2;
    x ^= rotl32(x, r1) ^ rotl32(x, r2); r1 += r1; r2 += r2;
    x ^= rotl32(x, r1) ^ rotl32(x, r2); return x;
}

static inline uint32_t tibhash(uint32_t x) {
    x ^= x >> 16;
    x *= 0x21f0aaadUL;
    x ^= x >> 15;
    x *= 0x735a2d97UL;
    x ^= x >> 15;
    return x;
}"""


def STATE(u64: bool, quad: bool):
    return f"""static {"uint64_t" if u64 else "uint32_t"} {", ".join("wax" if not quad else "wabx")}, domain_id;"""


def GOLDEN(u64: bool):
    return "0x9e3779b97f4a7c15ULL" if u64 else "0x9e3779b9UL"


def FN_NEXT(prng_name: str, u64: bool, quad: bool):
    ty = "uint64_t" if u64 else "uint32_t"
    rotl = "rotl64" if u64 else "rotl32"
    r1 = "17" if u64 else "10"
    r2 = "42" if u64 else "25"

    return f"""\
inline {ty} {prng_name}_next() {{
{
        f'''\
    {ty} out = x + a;
    {ty} tmp = w ^ x ^ domain_id;
    x  = a ^ {rotl}(a, {r1}) ^ {rotl}(a, {r2});
    a += tmp;
    w += {GOLDEN(u64)};
    return out;'''
        if not quad
        else f'''\
    {ty} out = x + a;
    {ty} tmp = w ^ b ^ domain_id;
    {ty} tmq = a ^ x;
    x  = b ^ {rotl}(b, {r1}) ^ {rotl}(b, {r2});
    a += tmp;
    b += tmq;
    w += {GOLDEN(u64)};
    return out;'''
    }
}}

{ty} {prng_name}_next_back() {{
{
        f'''\
    {ty} a_ = a;
    w -= {GOLDEN(u64)};
    a  = xr_step_inv{64 if u64 else 32}(x, {r1}, {r2});
    x  = (a_ - a) ^ w ^ domain_id;
    return x + a;'''
        if not quad
        else f'''\
    {ty} a_ = a, b_ = b;
    w -= {GOLDEN(u64)};
    b  = xr_step_inv{64 if u64 else 32}(x, {r1}, {r2});
    a  = a_ - (w ^ b ^ domain_id);
    x  =(b_ - b) ^ a;
    return x + a;'''
    }
}}"""


def FN_SET(prng_name: str, u64: bool, quad: bool):
    ty = "uint64_t" if u64 else "uint32_t"

    return f"""\
void {prng_name}_set_domain({ty} domain_id_) {{
    domain_id = domain_id_;
}}

void {prng_name}_set_state({", ".join(f"{ty} {x}_" for x in ("wax" if not quad else "wabx"))}) {{
    {", ".join(f"{x} = {x}_" for x in ("wax" if not quad else "wabx"))};
}}"""


def FN_MIX(prng_name: str, u64: bool, quad: bool):
    return f"""\
void {prng_name}_mix() {{
{
        (
            f'''\
    {prng_name}_next();
    {prng_name}_next();
    uint64_t w_ = moremur(w ^ ~a & x);
    uint64_t a_ = moremur(a ^ ~x & w);
    uint64_t x_ = moremur(x ^ ~w & a);
    w = w_, a = a_, x = x_;'''
            if u64
            else f'''\
    {prng_name}_next();
    {prng_name}_next();
    uint64_t w_ = tibhash(w ^ ~a & x);
    uint64_t a_ = tibhash(a ^ ~x & w);
    uint64_t x_ = tibhash(x ^ ~w & a);
    w = w_, a = a_, x = x_;'''
        )
        if not quad
        else (
            f'''\
    {prng_name}_next();
    w += a; x ^= w; x = rotl64(x, 32);
    b += x; a ^= b; a = rotl64(a, 24);
    w += a; x ^= w; x = rotl64(x, 16);
    b += x; a ^= b; a = rotl64(a, 13);
    w = moremur(w); a = moremur(a);
    b = moremur(b); x = moremur(x);'''
            if u64
            else f'''\
    {prng_name}_next();
    w += a; x ^= w; x = rotl32(x, 16);
    b += x; a ^= b; a = rotl32(a, 12);
    w += a; x ^= w; x = rotl32(x, 8);
    b += x; a ^= b; a = rotl32(a, 7);
    w = tibhash(w); a = tibhash(a);
    b = tibhash(b); x = tibhash(x);'''
        )
    }
}}"""


def C_HEADER(prng_name: str, u64: bool, quad: bool) -> tuple[str, str]:
    ty = "uint64_t" if u64 else "uint32_t"

    return (
        f"{prng_name}.h",
        f"""\
/*
 * {prng_name} PRNG
 * By K--Aethiax
 *
 * Released into the public domain under the CC0 1.0 license. To view a copy
 * of this license, visit: https://creativecommons.org/publicdomain/zero/1.0/
 */

#include <stdint.h>

{ty} {prng_name}_next();

{ty} {prng_name}_next_back();

void {prng_name}_set_domain({ty} domain_id);

void {prng_name}_set_state({", ".join(f"{ty} {x}" for x in ("abc" if not quad else "abcd"))});

void {prng_name}_mix();
""",
    )


def C_SOURCE(prng_name: str, u64: bool, quad: bool):
    return (
        f"{prng_name}.c",
        f"""\
/*
 * {prng_name} PRNG
 * By K--Aethiax
 *
 * Released into the public domain under the CC0 1.0 license. To view a copy
 * of this license, visit: https://creativecommons.org/publicdomain/zero/1.0/
 */

#include <stdint.h>
#include "{prng_name}.h"

{HELPER(u64)}

{STATE(u64, quad)}

{FN_NEXT(prng_name, u64, quad)}

{FN_SET(prng_name, u64, quad)}

{FN_MIX(prng_name, u64, quad)}
""",
    )


def generate(prng_name: str, u64: bool, quad: bool):
    name, content = C_HEADER(prng_name, u64, quad)
    with open(Path(__file__).parent / name, "w", encoding="utf-8") as f:
        f.write(content)

    name, content = C_SOURCE(prng_name, u64, quad)
    with open(Path(__file__).parent / name, "w", encoding="utf-8") as f:
        f.write(content)


if __name__ == "__main__":
    generate("Arata", True, False)
    generate("Arata32", False, False)
    generate("ArataQuad", True, True)
    generate("ArataQuad32", False, True)