Convert prime number to eisenstein prime using Cornacchia's algorithm

Murage_Eisenstein.c

//Convert prime number to eisenstein prime using Cornacchia's algorithm
//Described in https://leetarxiv.substack.com/p/computation-of-discrete-logarithms
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <stdbool.h>
#include <time.h>
#include <math.h>
#include <gmp.h>
#include <flint/flint.h>
#include <flint/fmpz.h>
#include <flint/fmpzi.h>
#include <flint/fmpq.h>
#include <flint/fmpz_factor.h>
#define STB_DS_IMPLEMENTATION
#include "stb_ds.h"
#define INFINITE_LOOP_CHECK_SUM 100
// clear && gcc Eisenstein2.c -o m.o -lm -lgmp -lmpfr -lflint && ./m.o 

typedef struct discrete_log_system_struct *System;
typedef struct eisenstein_integer_struct *fmpzE;
struct eisenstein_integer_struct
{
	fmpz_t a;
	fmpz_t b;
	fmpz_t norm;
	int exponent;
};


fmpzE fmpzE_init()
{
	fmpzE e = malloc(sizeof(struct eisenstein_integer_struct));
	e->exponent = 0;
	fmpz_init(e->a);
	fmpz_init(e->b);
	fmpz_init(e->norm);
	return e;
}

void fmpzE_clear(fmpzE e)
{
	fmpz_clear(e->a);
	fmpz_clear(e->b);
	fmpz_clear(e->norm);
	free(e);
}

void fmpzE_add(fmpzE r, fmpzE a, fmpzE b)
{
	fmpz_add(r->a, a->a, b->a);
	fmpz_add(r->b, a->b, b->b);
}

void fmpzE_sub(fmpzE r, fmpzE a, fmpzE b)
{
	fmpz_sub(r->a, a->a, b->a);
	fmpz_sub(r->b, a->b, b->b);
}

void fmpzE_mul(fmpzE r, fmpzE a, fmpzE b)
{
	fmpz_t temp0;fmpz_init(temp0);
	fmpz_mul(r->a, a->a, b->a);
	fmpz_mul(temp0, a->b, b->b);
	fmpz_sub(r->a,r->a,temp0);
	
	
	fmpz_set(r->b, temp0);
	fmpz_neg(r->b, r->b);	
	fmpz_mul(temp0, a->a, b->b);
	fmpz_add(r->b,r->b,temp0);
	fmpz_mul(temp0, a->b, b->a);
	fmpz_add(r->b,r->b,temp0);
	fmpz_clear(temp0);
}

void fmpzE_conjugate(fmpzE result, fmpzE e)
{
	fmpz_sub(result->a, e->a, e->b);
	fmpz_neg(result->b, e->b);
}

void fmpzE_norm(fmpz_t norm, fmpz_t a, fmpz_t b)
{
	fmpz_t temp0;
	fmpz_init(temp0);
	fmpz_mul(temp0, a,a);//a^2
	fmpz_mul(norm, b,b);//b^2
	fmpz_add(norm, norm, temp0);
	
	fmpz_mul(temp0, a, b);//ab
	fmpz_sub(norm, norm, temp0);
	fmpz_clear(temp0);
}

void fmpz_set_mpf_round(fmpz_t result, mpf_t x)
{
	mpf_t f_tmp,frac;
	mpf_init(f_tmp);mpf_init(frac);
	
	mpf_floor(f_tmp, x);
	mpf_sub(frac, x, f_tmp);
	if(mpf_cmp_d(frac, 0.5) < 0) 
	{
		fmpz_set_mpf(result, f_tmp); 
	}
	else
	{
		mpf_ceil(f_tmp, x); 
		fmpz_set_mpf(result, f_tmp);
	}
	mpf_clear(f_tmp);mpf_clear(frac);
	
}

void fmpzE_divide(fmpzE r, fmpzE a, fmpzE b)
{
	fmpzE bConjugate = fmpzE_init();
	fmpzE num = fmpzE_init();
	fmpzE_conjugate(bConjugate, b);
	
	fmpzE_mul(num, a, bConjugate);
	fmpzE_norm(b->norm, b->a, b->b);
	
	mpf_t f_NumA, f_NumB, f_norm;
	mpf_init(f_NumA);mpf_init(f_NumB);mpf_init(f_norm);

	fmpz_get_mpf(f_NumA, num->a);
	fmpz_get_mpf(f_NumB, num->b);
	fmpz_get_mpf(f_norm, b->norm);

	mpf_div(f_NumA, f_NumA, f_norm);
	mpf_div(f_NumB, f_NumB, f_norm);

	
	fmpz_set_mpf_round(r->a, f_NumA);
	fmpz_set_mpf_round(r->b, f_NumB);
	
	mpf_clear(f_NumA);mpf_clear(f_NumB);mpf_clear(f_norm);
	fmpzE_clear(num);
	fmpzE_clear(bConjugate);
}

void fmpzE_mod(fmpzE r, fmpzE a, fmpzE b)
{
	fmpzE q = fmpzE_init();
	fmpzE temp = fmpzE_init();
	
	fmpzE_divide(q, a, b);  
	fmpzE_mul(temp, q, b);   
	fmpzE_sub(r, a, temp);  

	fmpzE_clear(q);
	fmpzE_clear(temp);
}

void fmpzE_gcd(fmpzE result, fmpzE alpha, fmpzE beta)
{
	fmpzE a = fmpzE_init();
	fmpzE b = fmpzE_init();
	fmpzE r = fmpzE_init();
	fmpzE zero = fmpzE_init();
	fmpz_set(a->a, alpha->a);fmpz_set(a->b, alpha->b);
	fmpz_set(b->a, beta->a);fmpz_set(b->b, beta->b);
	fmpz_zero(zero->a);fmpz_zero(zero->b);
	int infiniteLoopCheck = 0;
	while(!(fmpz_equal(b->a, zero->a) && fmpz_equal(b->b, zero->b)))
	{
		//r = a % b
		fmpzE_mod(r, a, b);

		//a = b
		fmpz_set(a->a, b->a);fmpz_set(a->b, b->b);

		//b = r
		fmpz_set(b->a, r->a);fmpz_set(b->b, r->b);
    		
    		if(infiniteLoopCheck >= INFINITE_LOOP_CHECK_SUM)
		{
			printf("Infinite loop in fmpzE_gcd\n");
			assert(infiniteLoopCheck < INFINITE_LOOP_CHECK_SUM);
		}
		infiniteLoopCheck += 1;
    	
    	}

	//result = a
	fmpz_set(result->a, a->a);
	fmpz_set(result->b, a->b);

	fmpzE_clear(a);
	fmpzE_clear(b);
	fmpzE_clear(r);
	fmpzE_clear(zero);
}


void fmpzE_print(fmpzE e)
{
	printf("(");fmpz_print(e->a);printf(" + ");fmpz_print(e->b);printf(")\n");
}

void fmpzE_printPrimeFactors(fmpzE *primeFactors)
{
	for(size_t i = 0; i < arrlen(primeFactors); i++)
	{
		printf("(");fmpz_print(primeFactors[i]->a);printf(" + ");fmpz_print(primeFactors[i]->b);printf(") ^ %d, ",primeFactors[i]->exponent);	
	}
}

void fmpzE_freePrimeFactors(fmpzE *primeFactors)
{
	for(size_t i = 0; i < arrlen(primeFactors); i++)
	{
		fmpzE_clear(primeFactors[i]);
	}
	arrfree(primeFactors);
}

struct discrete_log_system_struct
{
	fmpz_t integerPrime;
	fmpz_t integerPrimeMinusOne;
	fmpz_t integerPrimitiveRoot;
	fmpz_t rootOfUnity;
	fmpz_t twoInverse;
	fmpz_t heegnerNumber;
	fmpz_t heegnerNumberRoot;
	fmpzE complexPrime;
	fmpzE complexPrimitiveRoot;	
};


bool FindTV_Cornacchia(fmpz_t T, fmpz_t V, fmpz_t heegner, fmpz_t root, fmpz_t prime)
{
	//Solve T^2 + |D| * V^2 = p
	fmpz_t a0, a1, remainder, rootP,D,rhs;
	fmpz_init(a0);fmpz_init(a1);fmpz_init(remainder);fmpz_init(D);fmpz_init(rootP);fmpz_init(rhs);
	
	//Find square root of prime
	fmpz_root(rootP, prime, 2);
	fmpz_set(a0, prime);
	fmpz_set(a1, root);
	//Ensure we are working with a nageative heegner number
	assert(fmpz_cmp_ui(heegner, 0) < 0);
	fmpz_mul_si(D, heegner, -1);
	
	while(fmpz_cmp(a1, rootP) > 0)
	{
		//Find remainder = a0 mod a1
		//Ensure a1 != 0
		if(fmpz_cmp_ui(a1, 0) == 0){break;}
		fmpz_mod(remainder, a0, a1);
		fmpz_set(a0, a1);
		fmpz_set(a1, remainder);
	}
	//no solution if a1 = 0
	assert(fmpz_cmp_ui(a1, 0) != 0);
	fmpz_set(T, a1);
	//Compute rhs = (p - T^2) / |D|
	fmpz_mul(rhs, T, T);
	fmpz_sub(rhs, prime, rhs);	
	//rhs must be divisible by |D|
	assert(fmpz_divisible(rhs, D));
	fmpz_fdiv_q(rhs, rhs, D);
	//rhs must be a perfect square
	bool isRhsSquare = fmpz_is_square(rhs);
	if(isRhsSquare)fmpz_sqrt(V, rhs);
	fmpz_clear(a0);fmpz_clear(a1);fmpz_clear(remainder);fmpz_clear(D);fmpz_clear(rootP);fmpz_clear(rhs);
	return isRhsSquare;
}

System CreateSystem(fmpz_t prime)
{
	System system = malloc(sizeof(struct discrete_log_system_struct));
	/*Initialize fmpz_t*/			
	fmpz_init(system->integerPrime);
	fmpz_init(system->integerPrimeMinusOne);
	fmpz_init(system->integerPrimitiveRoot);
	fmpz_init(system->heegnerNumber);
	fmpz_init(system->heegnerNumberRoot);
	fmpz_init(system->rootOfUnity);
	fmpz_init(system->twoInverse);
	/*Initialize fmpzE*/			
	system->complexPrime = fmpzE_init();
	system->complexPrimitiveRoot = fmpzE_init();
	
	/*Ensure prime splits*/			
	fmpz_set_si(system->heegnerNumber, -3);
	int legendre = fmpz_jacobi(system->heegnerNumber, prime);
	if(legendre != 1)
	{
		printf("Error: prime does not split root -3\n");
		assert(legendre == 1);			
	}
	/*Set p, p-1 and 2 inverse*/			
	fmpz_set(system->integerPrime, prime);
	fmpz_sub_ui(system->integerPrimeMinusOne, prime,1);
	fmpz_set_ui(system->twoInverse, 2);
	fmpz_invmod(system->twoInverse, system->twoInverse, system->integerPrime);
	
	/*Find a non-trivial root of unity*/	
	int foundSquareRoot = fmpz_sqrtmod(system->heegnerNumberRoot, system->heegnerNumber, system->integerPrime);
	assert(foundSquareRoot == 1);
	
	//Try either heegnerNumberRoots in roots of unity to find T and V
	bool TV_solutionExists = false;
	int solutionSearch = 0;
	while(TV_solutionExists == false && solutionSearch < 2)
	{
		//Set root of unity
		fmpz_add_si(system->rootOfUnity, system->heegnerNumberRoot, -1);
		fmpz_mul(system->rootOfUnity, system->rootOfUnity, system->twoInverse);
		fmpz_mod(system->rootOfUnity, system->rootOfUnity, system->integerPrime);
		/*Find T and V, we use norm as a temp var*/
		fmpz_mul_ui(system->complexPrime->norm,system->integerPrime, 4);
		//Solve for U^2 + 3V^2 = 4p
		TV_solutionExists = FindTV_Cornacchia(system->complexPrime->a, system->complexPrime->b, system->heegnerNumber, system->heegnerNumberRoot, system->complexPrime->norm);
		
		if(TV_solutionExists == false)
		{
			//Try other heegnerNumberRoot
			fmpz_sub(system->heegnerNumberRoot,system->integerPrime, system->heegnerNumberRoot);	
		}
		solutionSearch += 1;
	}
	
	assert(TV_solutionExists == true);
	//Solve for T and norm
	fmpz_add(system->complexPrime->a, system->complexPrime->a, system->complexPrime->b);
	fmpz_divexact_ui(system->complexPrime->a, system->complexPrime->a, 2);
	fmpzE_norm(system->complexPrime->norm, system->complexPrime->a, system->complexPrime->b);
	assert(fmpz_cmp(system->complexPrime->norm, system->integerPrime) == 0);
	return system;
}

void PrintSystem(System system)
{
	printf("Integer Prime: ");fmpz_print(system->integerPrime);printf("\n");
	printf("Integer Prime-1: ");fmpz_print(system->integerPrimeMinusOne);printf("\n");
	printf("Heegner Number: ");fmpz_print(system->heegnerNumber);printf("\n");
	printf("RootOfUnity (S): ");fmpz_print(system->rootOfUnity);printf("\n");
	printf("T: ");fmpz_print(system->complexPrime->a);printf("\n");
	printf("V: ");fmpz_print(system->complexPrime->b);printf("\n");
	printf("Complex prime: (");fmpz_print(system->complexPrime->a);printf(" + ");fmpz_print(system->complexPrime->b);printf("√-3) Norm = ");fmpz_print(system->complexPrime->norm);printf("\n");
}

void DestroySystem(System system)
{
	fmpz_clear(system->heegnerNumberRoot);
	fmpz_clear(system->twoInverse);
	fmpz_clear(system->rootOfUnity);
	fmpz_clear(system->integerPrimeMinusOne);
	fmpz_clear(system->integerPrime);
	fmpz_clear(system->integerPrimitiveRoot);
	fmpz_clear(system->heegnerNumber);
	fmpzE_clear(system->complexPrime);
	fmpzE_clear(system->complexPrimitiveRoot);
	free(system);
}

void TestSmallSystem()
{
	fmpz_t prime;
	fmpz_init(prime);
	fmpz_set_ui(prime, 20959);
	System system = CreateSystem(prime);
	PrintSystem(system);
	fmpz_clear(prime);	
	DestroySystem(system);
}

void TestBitcoinSystem()
{
	char *primeNumberHexadecimal = "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F";	
	fmpz_t prime;
	fmpz_init(prime);
	fmpz_set_str(prime, primeNumberHexadecimal, 16);
	System system = CreateSystem(prime);
	PrintSystem(system);
	fmpz_clear(prime);	
	DestroySystem(system);
}


int main()
{
	TestSmallSystem();
	TestBitcoinSystem();
	flint_cleanup();
	return 0;
}