#include "mpz.h"
#include <string.h>

static size_t __ceil(double f)
{
    /* broken for negative values, which are outside of our need here */
    return (size_t)f + (f != (size_t)f);
}

static size_t sizeinitmpz(mpz_t n, size_t size)
{
    size_t limbs = __ceil((double)size / sizeof(mp_limb_t));
    size_t bytes = limbs * sizeof(mp_limb_t);

    mpz_init(n);

    mp_limb_t *ptr = mpz_limbs_write(n, limbs);
    memset(ptr, 0, bytes);
    mpz_limbs_finish(n, 0);

    return limbs;
}

void pszip_mpz_state_init(struct pszip_mpz_state *mpzs, size_t popcnt, size_t size)
{
    sizeinitmpz(mpzs->current, size);
    mpz_inits(mpzs->a, mpzs->b, mpzs->c, mpzs->d, 0);

    mpz_ui_pow_ui(mpzs->current, 2, popcnt);
    mpz_sub_ui(mpzs->current, mpzs->current, 1);
}

void pszip_mpz_state_clear(struct pszip_mpz_state *mpzs)
{
    mpz_clears(mpzs->current, mpzs->a, mpzs->b, mpzs->c, mpzs->d, 0);
}

void pszip_mpz_init_data(mpz_t n, const void *data, size_t size)
{
    size_t limbs = sizeinitmpz(n, size);

    mp_limb_t *ptr = mpz_limbs_modify(n, limbs);
    memcpy(ptr, data, size);
    mpz_limbs_finish(n, limbs);
}

/*
 * Bit hack source:
 * https://graphics.stanford.edu/~seander/bithacks.html#NextBitPermutation
 */

void pszip_mpz_iterate_bithack01(struct pszip_mpz_state *mpzs)
{
    /* Gets current's least significant 0 bits set to 1
     * a = cur | (cur - 1); */
    mpz_sub_ui(mpzs->b, mpzs->current, 1);
    mpz_ior(mpzs->a, mpzs->current, mpzs->b);

    /* Next set to 1 the most significant bit to change,
     * set to 0 the least significant ones, and add the necessary 1 bits.
     * new = (a + 1) | (((~a & -~a) - 1) >> (__builtin_ctz(cur) + 1));
     * __builtin_ctz(cur) -> mpz_scan1(cur, 0) */
    mpz_add_ui(mpzs->b, mpzs->a, 1);    // (a + 1)
    mpz_com(mpzs->c, mpzs->a);          // (~a)
    mpz_neg(mpzs->d, mpzs->c);          // (-~a)
    mpz_and(mpzs->c, mpzs->c, mpzs->d); // (~a & -~a)
    mpz_sub_ui(mpzs->c, mpzs->c, 1);    // ((~a & -~a) - 1)

    /* bitshift requires direct memory access */
    size_t limbs = mpz_size(mpzs->c);
    mp_limb_t *ptr = mpz_limbs_modify(mpzs->c, limbs);
    mpn_rshift(ptr, ptr, limbs, (mpz_scan1(mpzs->current, 0) + 1));
    mpz_limbs_finish(mpzs->c, limbs);

    mpz_ior(mpzs->current, mpzs->b, mpzs->c);
}

void pszip_mpz_iterate_bithack02(struct pszip_mpz_state *mpzs)
{
    /* See comments from bithack01.
     * a = (cur | (cur - 1)) + 1; */
    mpz_sub_ui(mpzs->a, mpzs->current, 1);
    mpz_ior(mpzs->a, mpzs->current, mpzs->a);
    mpz_add_ui(mpzs->a, mpzs->a, 1);

    /* new = a | ((((a & -a) / (cur & -cur)) >> 1) - 1);
     * The right-shift by 1 is implemented as a divide by 2. */
    mpz_neg(mpzs->b, mpzs->a);                // (-a)
    mpz_and(mpzs->b, mpzs->a, mpzs->b);       // (a & -a)
    mpz_neg(mpzs->c, mpzs->current);          // (-cur)
    mpz_and(mpzs->c, mpzs->current, mpzs->c); // (cur & -cur)
    mpz_mul_ui(mpzs->c, mpzs->c, 2);          // (... / (cur & -cur)) >> 1
    mpz_tdiv_q(mpzs->b, mpzs->b, mpzs->c);    // (((a & -a) / (cur & -cur)) >> 1)
    mpz_sub_ui(mpzs->b, mpzs->b, 1);          // (ior) right hand side
    mpz_ior(mpzs->current, mpzs->a, mpzs->b);
}