“even faster” maze counter (which is actually slower)

efmc.c

/* efmc.c - Even Faster Maze Counter

(which actually turns out to be slower, lol)
*/

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <gmp.h>
#include "fmc.h"

/** Small helper functions **/

/* The largest power of two less than or equal to n

(A bit annoying, because finding the MSB is a single instruction
on x86 – BSR – but we only do this once so it doesn't matter.)
*/
inline static int msb(int n)
{
    int p = 1;
    while (p <= n) p <<= 1;
    return p >> 1;
}

inline static void swap(mpz_t **x, mpz_t **y)
{
    mpz_t *t = *x; *x = *y; *y = t;
}

/** Vector operations **/

static mpz_t *vec_init(int n) {
    mpz_t *vec = malloc(n * sizeof(mpz_t));

    for (int i = 0; i < n; ++i)
    {
        mpz_init(vec[i]);
    }

    return vec;
}

static void vec_clear(int n, mpz_t *vec)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_clear(vec[i]);
    }

    free(vec);
}

static void vec_set(int n, mpz_t *dest, mpz_t *src)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_set(dest[i], src[i]);
    }
}

static void vec_set_ui(int n, mpz_t *dest, unsigned int c)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_set_ui(dest[i], c);
    }
}

static void vec_set_si(int n, mpz_t *dest, int c)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_set_si(dest[i], c);
    }
}

static void vec_add(int n, mpz_t *dest, mpz_t *a, mpz_t *b, mpz_t p)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_add(dest[i], a[i], b[i]);
        mpz_fdiv_r(dest[i], dest[i], p);
    }
}

static void vec_sub(int n, mpz_t *dest, mpz_t *a, mpz_t *b, mpz_t p)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_sub(dest[i], a[i], b[i]);
        mpz_fdiv_r(dest[i], dest[i], p);
    }
}

static void vec_mul(int n, mpz_t *dest, mpz_t *a, mpz_t *b, mpz_t p)
{
    for (int i = 0; i < n; ++i)
    {
        mpz_mul(dest[i], a[i], b[i]);
        mpz_fdiv_r(dest[i], dest[i], p);
    }
}

static void vec_product(mpz_t result, int n, mpz_t *vec, mpz_t p)
{
    mpz_set_ui(result, 1);
    for (int i = 0; i < n; ++i)
    {
        mpz_mul(result, result, vec[i]);
        mpz_fdiv_r(result, result, p);
    }
}

/*
static void vec_print(char *label, int n, mpz_t *vec)
{
    return;
    for (int i = 0; i < n; ++i)
    {
        gmp_printf("%s[%d] = %Zd\n", label, i, vec[i]);
    }
}
*/

/** Main **/

/* We want a finite field ℤ/pℤ that contains a (2 width)'th
   root of unity, which means the order of the multiplicative
   group, p-1, needs to be a multiple of (2 width).

   We also want p to be greater than the final answer; obviously
   we don’t yet know the answer, but we know it will be
   ≤ 4^((width - 1) (height - 1)).
*/
inline static void find_modulus(mpz_t p, int width, int height)
{
    mpz_ui_pow_ui(p, 4, (width - 1) * (height - 1));
    mpz_cdiv_q_ui(p, p, 2*width);
    mpz_mul_ui(p, p, 2*width);
    mpz_add_ui(p, p, 1);
    while (!mpz_probab_prime_p(p, 15)) {
        mpz_add_ui(p, p, 2*width);
    }
}

/* Find a primitive 'n'th root of unity in ℤ/pℤ, assuming
that (p - 1) is a multiple of n.
*/
inline static void find_primitive_root(mpz_t u, int n, mpz_t p)
{
    mpz_t pmo, power, x, y;
    bool primitive_root_found;

    mpz_init(pmo);
    mpz_init(power);
    mpz_init(x);

    // pmo := p - 1
    mpz_sub_ui(pmo, p, 1);

    // power := pmo / n
    mpz_divexact_ui(power, pmo, n);

    // initialise the RNG
    gmp_randstate_t randstate;
    gmp_randinit_default(randstate);

    do {
        // Let u be random in the range 1 .. p-1
        mpz_urandomm(u, randstate, pmo);
        mpz_add_ui(u, u, 1);

        // u := u^power
        mpz_powm(u, u, power, p);

        /* Now u is an n'th root of unity, but it may not be
        a primitive root. Check by brute force, since we may
        assume that n is small. */
        primitive_root_found = true;
        for (int i = 2; i < n; ++i)
        {
            mpz_powm_ui(x, u, i, p);
            if (mpz_cmp_ui(x, 1) == 0)
            {
                primitive_root_found = false;
                break;
            }
        }
    }
    while (!primitive_root_found);

    mpz_clear(pmo);
    mpz_clear(power);
    mpz_clear(x);
}

inline static void find_initial_eigenvalues(int n, mpz_t v[], mpz_t u, mpz_t p)
{
    mpz_t temp;

    mpz_init(temp);

    for (int i = 0; i < n; ++i)
    {
        // v[i] := u^i + u^-i + 4
        mpz_powm_ui(v[i], u, i+1, p);
        mpz_invert(temp, v[i], p);
        mpz_add(v[i], v[i], temp);
        mpz_add_ui(v[i], v[i], 4);
        mpz_fdiv_r(v[i], v[i], p);
    }

    mpz_clear(temp);
}

static void find_eigenvalues(mpz_t *ret, int width, int height, mpz_t *m, mpz_t p)
{
    int n = width - 1;
    mpz_t *a, *b, *c, *temp, *new_a, *new_b;

    a = vec_init(n);
    b = vec_init(n);
    c = vec_init(n);
    temp = vec_init(n);
    new_a = vec_init(n);
    new_b = vec_init(n);

    vec_set_si(n, a, -1);
    vec_set_ui(n, b, 0);
    vec_set_ui(n, c, +1);

    for (int bit = msb(height); bit > 0; bit >>= 1)
    {
        /* a, b := b^2 - a^2, bc - ab */
        vec_mul(n, new_a, b, b, p);
        vec_mul(n, temp, a, a, p);
        vec_sub(n, new_a, new_a, temp, p);
        
        vec_mul(n, new_b, b, c, p);
        vec_mul(n, temp, a, b, p);
        vec_sub(n, new_b, new_b, temp, p);
        
        swap(&a, &new_a);
        swap(&b, &new_b);

        if ((height & bit) > 0)
        {
            /* a, b := b, bM - a */
            vec_set(n, new_a, b);
            vec_mul(n, new_b, b, m, p);
            vec_sub(n, new_b, new_b, a, p);
            swap(&a, &new_a);
            swap(&b, &new_b);
        }
        
        /* c := bM - a */
        vec_mul(n, c, b, m, p);
        vec_sub(n, c, c, a, p);
    }

    vec_set(n, ret, b);

    vec_clear(n, a);
    vec_clear(n, b);
    vec_clear(n, c);
    vec_clear(n, temp);
    vec_clear(n, new_a);
    vec_clear(n, new_b);
}

/* Fast Maze Counter.

Count the number of mazes on a 'width'x'height grid,
and store the result in 'out'.
*/
void fmc(mpz_t *out, int width, int height)
{
    int n = width - 1;

    mpz_t p, u;
    mpz_t *initial_eigenvalues, *eigenvalues;

    mpz_init(p);
    mpz_init(u);
    initial_eigenvalues = vec_init(n);
    eigenvalues = vec_init(n);

    find_modulus(p, width, height);
    gmp_printf("p = %Zd\n", p);

    find_primitive_root(u, 2*width, p);
    gmp_printf("u = %Zd\n", u);

    find_initial_eigenvalues(n, initial_eigenvalues, u, p);
    find_eigenvalues(eigenvalues, width, height, initial_eigenvalues, p);

    // vec_print("initial_eigenvalues", n, initial_eigenvalues);
    // vec_print("eigenvalues", n, eigenvalues);

    vec_product(*out, n, eigenvalues, p);

    mpz_clear(p);
    mpz_clear(u);

    vec_clear(n, initial_eigenvalues);
    vec_clear(n, eigenvalues);
}