/* Tests and benchmarks for "Perfect Squares in C, No Math Library Required".
 * gcc -O2 -march=native perfect-squares-bench.c -lm -o bench-gcc && taskset -c 2 ./bench-gcc > gcc.json
 * Same again with clang, then: jq -n --slurpfile g gcc.json --slurpfile c clang.json '{gcc: $g[0], clang: $c[0]}'
 * Prints one JSON object; every number in the article comes from it. */
#include <stdint.h>
#include <stdbool.h>
#include <stdio.h>
#include <math.h>
#include <time.h>

/* ---- the article's function ---------------------------------------------- */

__attribute__((noinline)) bool is_square(uint64_t n)
{
    if (!((0x0213u >> (n & 15)) & 1))
        return false;

    uint64_t x = n, r = 0, bit = 1ULL << 62;
    while (bit > n)
        bit >>= 2;
    while (bit) {
        uint64_t block = r + bit;
        if (x >= block) {
            x -= block;
            r = (r >> 1) + bit;
        } else {
            r >>= 1;
        }
        bit >>= 2;
    }
    return x == 0;
}

/* same loop without the mod-16 filter, to see what the filter buys */
__attribute__((noinline)) bool is_square_nofilter(uint64_t n)
{
    uint64_t x = n, r = 0, bit = 1ULL << 62;
    while (bit > n)
        bit >>= 2;
    while (bit) {
        uint64_t block = r + bit;
        if (x >= block) { x -= block; r = (r >> 1) + bit; }
        else r >>= 1;
        bit >>= 2;
    }
    return x == 0;
}

/* ---- the warm-ups -------------------------------------------------------- */

/* attempt 1: subtract 1, 3, 5, ... until nothing's left. sqrt(n) steps. */
__attribute__((noinline)) bool naive(uint64_t n)
{
    uint64_t odd = 1;
    while (n >= odd) {
        n -= odd;
        odd += 2;
    }
    return n == 0;
}

/* attempt 2: binary search on the root. hi can't be n, or mid * mid overflows. */
__attribute__((noinline)) bool bsearch_sq(uint64_t n)
{
    uint64_t lo = 0, hi = 0xFFFFFFFF;
    while (lo <= hi) {
        uint64_t mid = lo + (hi - lo) / 2;
        uint64_t sq = mid * mid;
        if (sq == n) return true;
        if (sq < n)  lo = mid + 1;
        else if (mid == 0) break;
        else hi = mid - 1;
    }
    return false;
}

/* ---- the float versions it's compared against ----------------------------- */

/* the textbook one: sqrt, round, square it back, all in double. Wrong past 2^53. */
__attribute__((noinline)) bool sqrt_all_double(uint64_t n)
{
    double r = round(sqrt((double)n));
    return r * r == (double)n;
}

/* same float root, but squared back as an integer: survives past 2^53 */
__attribute__((noinline)) bool sqrt_int_check(uint64_t n)
{
    uint64_t r = (uint64_t)llround(sqrt((double)n));
    return r * r == n;
}

/* float estimate, then nudge it with exact integer math: correct, used as the reference */
__attribute__((noinline)) bool sqrt_fixed(uint64_t n)
{
    uint64_t r = (uint64_t)sqrt((double)n);
    if (r > 0xFFFFFFFFu) r = 0xFFFFFFFFu;
    while (r * r > n) r--;
    while (r < 0xFFFFFFFFu && (r + 1) * (r + 1) <= n) r++;
    return r * r == n;
}

/* ---- harness --------------------------------------------------------------- */

static uint64_t seed = 0x9E3779B97F4A7C15ull;
static uint64_t rnd(void) { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; return seed; }
static uint64_t rnd_in(uint64_t lo, uint64_t hi) { return lo + rnd() % (hi - lo + 1); }

typedef bool (*fn)(uint64_t);
#define N (1 << 22)
static uint64_t in[N];

/* best of 7 timed passes over the first `count` inputs, in ns per call */
static double bench(fn f, int count)
{
    double best = 1e18;
    for (int rep = 0; rep < 7; rep++) {
        struct timespec a, b;
        volatile unsigned sink;
        unsigned hits = 0;
        clock_gettime(CLOCK_MONOTONIC, &a);
        for (int i = 0; i < count; i++) hits += f(in[i]);
        clock_gettime(CLOCK_MONOTONIC, &b);
        sink = hits; (void)sink;
        double ns = ((b.tv_sec - a.tv_sec) * 1e9 + (b.tv_nsec - a.tv_nsec)) / count;
        if (ns < best) best = ns;
    }
    return best;
}

int main(void)
{
    /* ---- test 1: is_square agrees with the reference ------------------------ */
    long mismatches = 0, checked = 0;
    for (int i = 0; i < 2000000; i++) {
        uint64_t k = rnd() >> 32;
        uint64_t cases[] = { k * k, k * k - 1, k * k + 1, rnd(), (uint64_t)i };
        for (int j = 0; j < 5; j++, checked++)
            mismatches += is_square(cases[j]) != sqrt_fixed(cases[j])
                        || is_square_nofilter(cases[j]) != sqrt_fixed(cases[j])
                        || bsearch_sq(cases[j]) != sqrt_fixed(cases[j]);
        mismatches += naive((uint64_t)i) != sqrt_fixed((uint64_t)i);  /* naive is too slow for the rest */
    }
    uint64_t edges[] = { 0, 1, 2, 3, 4, UINT64_MAX, 0xFFFFFFFE00000001ull /* (2^32-1)^2 */ };
    for (int j = 0; j < 7; j++, checked++)
        mismatches += is_square(edges[j]) != sqrt_fixed(edges[j])
                    || bsearch_sq(edges[j]) != sqrt_fixed(edges[j]);
    printf("{\"checked\":%ld,\"mismatches\":%ld,\n", checked, mismatches);

    /* ---- attempt 1 on the worst input: one call, timed ---------------------- */
    volatile uint64_t worst = UINT64_MAX;
    struct timespec a, b;
    clock_gettime(CLOCK_MONOTONIC, &a);
    bool naive_hit = naive(worst);
    clock_gettime(CLOCK_MONOTONIC, &b);
    printf("\"naiveMaxSeconds\":%.3f,\"naiveMaxResult\":%d,\n",
           (b.tv_sec - a.tv_sec) + (b.tv_nsec - a.tv_nsec) / 1e9, naive_hit);

    /* ---- test 2: where does the all-double version start lying? ------------- */
    uint64_t k = (uint64_t)sqrt(9007199254740992.0);  /* sqrt(2^53) */
    while (!sqrt_all_double(k * k + 1)) k++;
    printf("\"firstFooledK\":%llu,\n", (unsigned long long)k);

    /* fooled rate for k^2 - 1 and k^2 + 1, by how many bits k^2 has */
    printf("\"fooledByBits\":[\n");
    for (int bits = 48; bits <= 64; bits++) {
        uint64_t lo = (uint64_t)ceil(sqrt(ldexp(1, bits - 1))), hi = (uint64_t)sqrt(ldexp(1, bits));
        if (hi > 0xFFFFFFFFu) hi = 0xFFFFFFFFu;
        long below = 0, above = 0, int_wrong = 0, samples = 200000;
        for (long i = 0; i < samples; i++) {
            uint64_t r = rnd_in(lo, hi - 1);
            below += sqrt_all_double(r * r - 1);
            above += sqrt_all_double(r * r + 1);
            int_wrong += sqrt_int_check(r * r - 1) || sqrt_int_check(r * r + 1) || !sqrt_int_check(r * r);
        }
        printf("  {\"bits\":%d,\"below\":%.4f,\"above\":%.4f,\"intCheckWrong\":%ld}%s\n", bits,
               (double)below / samples, (double)above / samples, int_wrong, bits < 64 ? "," : "");
    }
    printf("],\n");

    /* ---- bench 1: five implementations, two input mixes --------------------- */
    const char *names[] = { "sqrt_all_double", "sqrt_fixed", "bsearch_sq", "is_square", "is_square_nofilter" };
    fn fs[] = { sqrt_all_double, sqrt_fixed, bsearch_sq, is_square, is_square_nofilter };
    printf("\"bench\":[\n");
    for (int mix = 0; mix < 2; mix++) {
        for (int i = 0; i < N; i++) {
            uint64_t r = rnd() >> 32;
            in[i] = mix == 0 ? rnd() : r * r + (rnd() & 1);  /* random, or squares and squares + 1 */
        }
        for (int f = 0; f < 5; f++)
            printf("  {\"mix\":\"%s\",\"fn\":\"%s\",\"ns\":%.3f}%s\n", mix ? "squares" : "random",
                   names[f], bench(fs[f], N), mix == 1 && f == 4 ? "" : ",");
    }
    printf("],\n");

    /* ---- bench 2: cost vs size of n (squares only, so every call runs the loop) */
    printf("\"byBits\":[\n");
    for (int bits = 2; bits <= 64; bits += 2) {
        for (int i = 0; i < N / 16; i++) {
            uint64_t r = rnd_in(1ull << (bits / 2 - 1), (1ull << (bits / 2)) - 1);
            in[i] = r * r;
        }
        printf("  {\"bits\":%d,\"is_square\":%.3f,\"sqrt_fixed\":%.3f}%s\n", bits,
               bench(is_square, N / 16), bench(sqrt_fixed, N / 16), bits < 64 ? "," : "");
    }
    printf("]}\n");
}