qs.cpp

#include <gmp.h>
#include <gmpxx.h>
#include <iostream>
#include <cstdlib>
#include <stdint.h>
#include <cmath>
#include <vector>
#include <stack>
#include <algorithm>
#include <sstream>
#include <iostream>
#include <time.h>
#include "utils.hpp"

using namespace std;

// How large chunks should we sieve at a time?
const uint32_t SIEVE_STEP = 65536;

// Max amount prime numbers generated
const uint32_t PRIME_BOUND = 200000;

// Linear relations, how many extra smooth numbers to be generated
const uint32_t LINEAR_RELATIONS = 10;

// Node containing smooth numbers
struct smooth_node {
        uint32_t X;
        vector<uint32_t> smooth_numbers;
        
        smooth_node(uint32_t X, vector<uint32_t> smooth_numbers) {
                this->X = X;
                this->smooth_numbers = smooth_numbers;
        }
};

vector<mpz_class> primes;
stack<mpz_class> factors;

void handle_perfect_power(mpz_class& factor);
bool find_small_factors(mpz_class& factor); // Trial division
void era_sieve();   //Sieve Of Eratosthenes
mpz_class quadratic_sieve(mpz_class &N);

int main() {
        srand(4711);

        // Generate primes with Sieve of Eratosthenes
        era_sieve();

        mpz_class N;
        // Read numbers from stdin until end of file.
        while(cin >> N) {
                factors.push(N);

                while(!factors.empty()) {
                        mpz_class factor = factors.top();
                        factors.pop();

                        // Print if factor is prime
                        if(is_prime(factor)) {
                                cout << factor << endl;
                                continue;
                        } else {
                                // Check for small factors, before running quadratic sieve.
                                if(find_small_factors(factor))
                                    continue;

                                // Quadratic sieve doesn't handle perferct powers very well, handle those separately.
                                if(mpz_perfect_power_p(factor.get_mpz_t())) {
                                    handle_perfect_power(factor);

                                } else {
                                    mpz_class f = quadratic_sieve(factor);

                                    factors.push(f);
                                    factors.push(factor / f);
                                }
                        }
                }

                cout << endl;
        }

        return 0;
}




// A quadratic sieve implementation for integers up to 100 bits. N must be composite.
mpz_class quadratic_sieve(mpz_class &N) {
        vector<uint32_t> factor_base;
        vector<uint32_t> log_primes;

        mpz_class sqrt_N = sqrt(N);

        /* Setting up a Factor Base and a Sieving Interval */       
        uint32_t B;
        {   // Self initialize smoothsness bound
                // Approximate natural logarithm of N.
                float log_N = mpz_sizeinbase(N.get_mpz_t(), 2) * ln2;

                // The optimal smoothsness bound is exp((0.5 + o(1)) * sqrt(log(n)*log(log(n)))).
                B = (uint32_t)ceil(exp((0.5+0.068) * sqrt(log_N * log(log_N)))+550); // 2^(-1/2)
                // B = (B<300) ? 300 : B;
        }

        // Criteria 1:  The primes in factor base must be primes such that 
        //              the Legendre symbol (n/p) = 1.
        // 
        // Criteria 2:  These primes must be less than a bound B.
        //
        // Ref: http://www.cs.virginia.edu/crab/QFS_Simple.pdf 

        {   // Make factor base by sieving.
                char *sieve = new char[B + 1];
                memset(sieve, 1, B + 1);
                for(unsigned long p = 2; p <= B; ++p) {
                        if(!sieve[p])
                                continue;

                        if(mpz_kronecker_ui(N.get_mpz_t(), p) == 1){
                                factor_base.push_back(p);
                                log_primes.push_back(log(p)/log(2));
                        }

                        for(unsigned long i = p; i <= B; i += p)
                                sieve[i] = 0;
                }
                delete[] sieve;
        }


        /* Sieving */
        float *log_Q = new float[SIEVE_STEP];

        // The sieve boundary.
        uint32_t lower_x = 0;
        uint32_t upper_x = SIEVE_STEP;

        // Calculate sieve index (where to start the sieve) for each factor base number.
        uint32_t **sieve_start_indexes = new uint32_t*[2];
        sieve_start_indexes[0] = new uint32_t[factor_base.size()];
        sieve_start_indexes[1] = new uint32_t[factor_base.size()];
        for(uint32_t p = 0; p < factor_base.size(); ++p) {
                // Uses Tonelli Shanks to solve the x^2 = n (mod p) and get the start and end indexes for sieving
                // of the factor base. Start index (x) and end index (N-x).
                uint32_t ssidx[2];
                mpz_class N_mod_P = N % mpz_class(factor_base[p]);
                shanks_tonelli(N_mod_P.get_ui(), factor_base[p], ssidx);

                //start index
                mpz_class a = (ssidx[0] - sqrt_N) % factor_base[p];
                a += factor_base[p];
                sieve_start_indexes[0][p] = a.get_ui();

                //end index
                mpz_class b = (ssidx[1] - sqrt_N) % factor_base[p];
                b += factor_base[p];
                sieve_start_indexes[1][p] = b.get_ui();
        }

        float cheat_log = 0;
        uint32_t estimate = 1;

        // We have a smooth number when all of its prime factors are small. Specifally, 
        // we say m is B-smooth if all of its prime factors are <= B
        vector<smooth_node> smooths;

        // Once we have a set of primes for our factor base, we begin to take numbers
        // x from our sieving interval and calculate Q(x), and check to see if it factors
        // completely over our factor base. If it factors, it is said to have smoothness. If
        // it does not, we throw it out, and we go on to the next element of our sieving
        // interval.

        // Sieve in chunk size to save memory. Sieves until there is enough smooths numbers.
        // NOTE TO SELF: If quadratic sieving gets stuck try use more linear relations.
        while(smooths.size() < (factor_base.size() + LINEAR_RELATIONS)) {
                // Generate our log_Q vector for the sieve, containing log approximations that fit in machine words.
                for(uint32_t t = 1; t < SIEVE_STEP; ++t) {
                        // Log estimate is expensive so dont do it for every Q(x).
                        if(estimate <= (t + lower_x)) {
                                mpz_class y = (sqrt_N + t + lower_x) * (sqrt_N + t + lower_x) - N;

                                // Estimation of 2 logarithm
                                cheat_log = mpz_sizeinbase(y.get_mpz_t(), 2);

                                // The higher t gets, the less the logarithm of log_Q[t] changes.
                                estimate = estimate * 1.8 + 1;
                        }
                        log_Q[t] = cheat_log;
                }

                // Perform log sieve.
                for(uint32_t p = 0; p < factor_base.size(); ++p) {
                        float lg = log_primes[p];
                        for(uint32_t t = 0; t < 2; ++t) {
                                while(sieve_start_indexes[t][p] < upper_x) {
                                        log_Q[sieve_start_indexes[t][p] - lower_x] -= lg;
                                        sieve_start_indexes[t][p] += factor_base[p];
                                }
                                // prime number 2 has one modular root so,
                                // do an early break.
                                if(factor_base[p] == 2)
                                        break;
                        }
                }


                {   // Factor using trail division all values of belonging to null space.
                    // I.e. all values whose logarithms were reduced to approximately zero.
                        float threshold = log(factor_base.back()) / ln2;
                        for(uint32_t i = 0; i < SIEVE_STEP; ++i) {
                                if(fabs(log_Q[i]) < threshold) {
                                        mpz_class y = (sqrt_N + i + lower_x) * (sqrt_N + i + lower_x) - N;
                                        vector<uint32_t> smooth_vector;
                                        for(uint32_t p = 0; p < factor_base.size(); ++p) {
                                                while(mpz_divisible_ui_p(y.get_mpz_t(), factor_base[p])) {
                                                        mpz_divexact_ui(y.get_mpz_t(), y.get_mpz_t(), factor_base[p]);
                                                        smooth_vector.push_back(p);
                                                }
                                        }
 
                                        if(y == 1) {
                                                // Found a V with B-smooth numbers.
                                                smooths.push_back(smooth_node(i + lower_x, smooth_vector));
                                                // Break free from loop if enough smooths numbers has been found.
                                                if(smooths.size() >= (factor_base.size() + LINEAR_RELATIONS))
                                                        break;
                                        } else {
                                                // Not B-smooth numbers, skip.
                                        }
                                }
                        }
                }

                lower_x += SIEVE_STEP;
                upper_x += SIEVE_STEP;
        }

        /* Initial matrix A for gauss elimination */
        uint64_t **A = new uint64_t*[factor_base.size()];

        // The amount of words needed to accomodate a row in the augmented A.
        int row_words = (smooths.size() + sizeof(uint64_t)) / sizeof(uint64_t);

        for(uint32_t i = 0; i < factor_base.size(); ++i) {
                A[i] = new uint64_t[row_words];
                memset(A[i], 0, row_words * sizeof(uint64_t));
        }

        for(uint32_t s = 0; s < smooths.size(); ++s) {
                // For each factor in the smooths number, add the factor to the corresponding element in the A.
                for(uint32_t p = 0; p < smooths[s].smooth_numbers.size(); ++p)
                        flip_bit(s, A[smooths[s].smooth_numbers[p]]);
        }

        // Perform gauss elimination. 
        // Augmented A will have dimension factor_base.size() x (smooths.size() + 1).
        {
                uint32_t i = 0, j = 0;
                while(i < factor_base.size() && j < (smooths.size() + 1)) {
                        uint32_t max_i = i;

                        // Find pivot element.
                        for(uint32_t k = i + 1; k < factor_base.size(); ++k) {
                                if(check_bit(j, A[k]) == 1) {
                                        max_i = k;
                                        break;
                                }
                        }
                        if(check_bit(j, A[max_i]) == 1) {
                                swap(A[i], A[max_i]);
                                
                                for(uint32_t u = i + 1; u < factor_base.size(); ++u) {
                                        if(check_bit(j, A[u]) == 1) {
                                                for(int32_t w = 0; w < row_words; ++w)
                                                        A[u][w] ^= A[i][w];
                                        }
                                }
                                ++i;
                        }
                        ++j;
                }
        }

        mpz_class a;
        mpz_class b;

        // A initialize array to hold a copy of A.
        // Back-substitution on will be performed on the copy.
        uint64_t **back_A = new uint64_t*[factor_base.size()];
        for(uint32_t i = 0; i < factor_base.size(); ++i)
                back_A[i] = new uint64_t[row_words];

        uint32_t *x = new uint32_t[smooths.size()];
        uint32_t *power = new uint32_t[factor_base.size()];

        do {    // Search for a non-trivial factor
                // Deep-copy the gauss eliminated A.
                for(uint32_t i = 0; i < factor_base.size(); ++i)
                        memcpy(back_A[i], A[i], row_words * sizeof(uint64_t));

                // Empty x-vector.
                memset(x, 0, smooths.size() * sizeof(uint32_t));

                {  // Get x by back-substituting A that is in row echelon form.
                        int32_t i = factor_base.size() - 1;

                        while(i >= 0) {
                                // Count non-zero elements in current row.
                                int32_t count = 0;
                                int32_t current = -1;
                                for(uint32_t c = 0; c < smooths.size(); ++c) {
                                        count += check_bit(c, back_A[i]);
                                        current = check_bit(c, back_A[i]) ? c : current;
                                }

                                // Empty row, advance to next.
                                if(count == 0) {
                                        --i;
                                        continue;
                                }

                                // If row has more than 1 non-zero element the system is undetermined.
                                // current x can be choosen freely. Avoid trivial solution by not setting x to 0. 
                                x[current] = count > 1 ? rand() % 2 : check_bit(smooths.size(), back_A[i]);

                                for(int32_t u = 0; u <= i; ++u) {
                                        if(check_bit(current, back_A[u]) == 1) {
                                                if(x[current] == 1)
                                                        flip_bit(smooths.size(), back_A[u]);
                                                clear_bit(current, back_A[u]);
                                        }
                                }

                                if(count == 1)
                                        --i;
                        }
                }
                a = 1;
                b = 1;

                // The way to combine the factor base to get our square.
                memset(power, 0, sizeof(uint32_t) * factor_base.size());
                for(uint32_t i = 0; i < smooths.size(); ++i) {
                        if(x[i] == 1) {
                                for(uint32_t p = 0; p < smooths[i].smooth_numbers.size(); ++p)
                                        ++power[smooths[i].smooth_numbers[p]];
                                b *= (smooths[i].X + sqrt_N);
                        }
                }

                for(uint32_t p = 0; p < factor_base.size(); ++p) {
                        for(uint32_t i = 0; i < (power[p] / 2); ++i)
                                a *= factor_base[p];               
                }

        // If   a =  b (mod N) 
        // or   a = -b (mod N) 
        // then solution is trivial, run the loop again to find a new a and b.
        } while(a % N == b % N || a % N == (-b) % N + N);

        //We have the desired identity (a+b)(a-b) = 0 (mod N)
        //Compute GCD of n with the difference of a and b.
        b -= a;

        mpz_class factor;
        gcd(factor, b, N);

        for(uint32_t i = 0; i < factor_base.size(); ++i) {
                delete[] A[i];
                delete[] back_A[i];
        }
        delete[] A;
        delete[] back_A;
        delete[] sieve_start_indexes[0];
        delete[] sieve_start_indexes[1];
        delete[] sieve_start_indexes;
        delete[] power;
        delete[] log_Q;
        delete[] x;

        return factor;
}

void era_sieve(){
        char *sieve = new char[PRIME_BOUND + 1];
        memset(sieve, 0, PRIME_BOUND + 1);
        for(uint32_t i = 2 ; i < ceil(sqrt(PRIME_BOUND)) ; ++i){ 
                if(sieve[i]) 
                    continue;
                for(int j = i*i; j < PRIME_BOUND; j += i){
                        sieve[j] = 1;
                }
        }
        for(uint32_t i = 2 ; i < PRIME_BOUND ; ++i){
                if(!sieve[i]){
                        primes.push_back((mpz_class)i);
                }
        }
}


void handle_perfect_power(mpz_class& factor){
    {
        mpz_class root, rem;

        // Check root remainders up half of the amount of bits in factor.
        uint32_t max = mpz_sizeinbase(factor.get_mpz_t(), 2) / 2;
        for(uint32_t n = 2; n < max; ++n) {
            mpz_rootrem(root.get_mpz_t(), rem.get_mpz_t(), factor.get_mpz_t(), n);
            if(rem == 0) {
                    // Push the n root factors.
                    for(uint32_t i = 0; i < n; ++i)
                            factors.push(root);
                    break;
            }
        }
    }
}

bool find_small_factors(mpz_class& factor){
    {
        for(uint32_t p = 0; p < primes.size(); ++p) {
                if(mpz_divisible_ui_p(factor.get_mpz_t(), primes[p].get_ui())) {
                        factors.push(primes[p]);
                        factors.push(factor / primes[p]);
                        return true;
                }
        }
        return false;
    }
}