Useful tools around generating and checking prime numbers and factorising integers

primes.go

package main

import (
	"fmt"
	"math"
	"math/big"
	"math/rand"
	"sort"
	"time"
)

var zero = big.NewInt(int64(0))
var one = big.NewInt(int64(1))
var two = big.NewInt(int64(2))

// Sieve of Eratosthenes
// Generate the primes up to n by going through each number up
// to sqrt(n) and crossing out all their multiples. The numbers
// that are left are all prime.
//		https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes
func SieveEratosthenes(n int) []int {
	var primes []int
	if n < 2 {
		return primes
	}
	// Initialise a boolean slice with all true values
	b := make([]bool, n+1)
	for i, _ := range b {
		b[i] = true
	}
	b[0], b[1] = false, false
	// Set all multiples of numbers up to sqrt(n) to false
	for i := 2; i*i <= n; i++ {
		if !b[i] {
			continue
		}
		for j := i * i; j <= n; j += i {
			b[j] = false
		}
	}
	// Primes are remaining true values
	for i, v := range b {
		if v {
			primes = append(primes, i)
		}
	}
	return primes
}

// Segmented Sieve of Eratosthenes
// Generate the primes up to n, but only portions or the range
// 2-n are sieved at a time.
//		https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes
func SieveEratosthenesSegmented(n int) []int {
	// Use chunks with a maximum size of sqrt(n)
	delta := int(math.Sqrt(float64(n)))
	// Use normal sieve to get first chunk of primes
	primes := SieveEratosthenes(delta)
	// Loop through remaining chunks
	min := delta
	max := 2 * delta
	for min <= n {
		// Shorten chunk if max value is greater than sqrt(n)
		if max >= n {
			max = n
		}
		// Initialise boolean array up to chunk max value
		b := make([]bool, max+1)
		for i, _ := range b {
			b[i] = true
		}
		// Remove all multiples of currently found primes
		for i := 0; i < len(primes); i++ {
			// Find lowest common multiple of current prime in range
			lcm := (min / primes[i]) * primes[i]
			if lcm < min {
				lcm += primes[i]
			}
			// Mark all multiples of current prime as not prime
			for j := lcm; j <= max; j += primes[i] {
				b[j-min] = false
			}
		}
		// Add newly found primes to slice
		for k := min; k <= max; k++ {
			if b[k-min] {
				primes = append(primes, k)
			}
		}
		// Move onto next chunk
		min += delta
		max += delta
	}
	return primes
}

// Sieve of Pritchard / Wheel Sieve
// Generate the primes up to n progressively by building up "wheels"
// which represent the numbers up to n not divisible by any of the
// primes processed so far
//		https://en.wikipedia.org/wiki/Sieve_of_Pritchard
func WheelSieve(n int) []int {
	// Initialise boolean array all false
	W := make([]bool, n+1)
	// Start on W[1]={1}
	W[1] = true
	// Initialise values for W[1]
	step, p, wlimit := 1, 2, 2
	primes := []int{}
	// Expand wheels and gather primes up to n
	for p*p <= n {
		if step <= n {
			// Extend wheel by step sized increments
			for i := 0; i < len(W); i++ {
				if W[i] {
					w := i + step
					for w <= wlimit {
						W[w] = true
						w += step
					}
				}
			}
			// Increase step for next wheel
			step = wlimit
		}
		// Initialise next prime and copy wheel
		np := 5
		wcpy := make([]bool, n+1)
		copy(wcpy, W)
		for i := 0; i < len(W); i++ {
			if wcpy[i] {
				// Initialise next prime
				if np == 5 && i > p {
					np = i
				}
				// Remove multiples of current prime
				w := p * i
				if w > wlimit {
					break
				}
				W[w] = false
			}
		}
		// Next prime cant be lower than current prime
		if np < p {
			break
		}
		// Add current prime to prime list
		primes = append(primes, p)
		// Set prime to next prime found
		if p == 2 {
			p = 3
		} else {
			p = np
		}
		// Increase wheel limit to the lower value of step*p or n
		wlimit = func(a, b int) int {
			if a < b {
				return a
			}
			return b
		}(step*p, n)
	}
	// Add primes left in wheel to found primes
	for i := 2; i < len(W); i++ {
		if W[i] {
			primes = append(primes, i)
		}
	}
	// Sort and return
	sort.Slice(primes, func(i, j int) bool { return i < j })
	return primes
}

// Get all primes up to 2^16
var PRIMES = SieveEratosthenes(2 << 16)

// Return new slice containing only the elements in a that are
// not in b.
//		d = {x : x ∈ a and x ∉ b}
func _difference(a, b []int) []int {
	mb := make(map[int]int, len(b))
	for _, v := range b {
		mb[v] = 1
	}
	var diff []int
	for _, v := range a {
		if mb[v] != 1 {
			diff = append(diff, v)
		}
	}
	return diff
}

// Generate a bitmask for primes from lo->hi
func primemask(lo, hi int) *big.Int {
	low := make([]int, lo)
	high := make([]int, hi)
	// Get primes 2->hi
	high = SieveEratosthenes(hi)
	// Only get lower limit primes is lo > 2
	if lo > 2 {
		low = SieveEratosthenes(lo)
	}
	// Use only difference between high and low slices
	diff := _difference(high, low)
	// Create by doing mask |= 1<<prime
	mask := new(big.Int)
	for i := 0; i < len(diff); i++ {
		prime := new(big.Int)
		prime = prime.Lsh(big.NewInt(int64(1)), uint(diff[i]))
		mask = mask.Or(mask, prime)
	}
	return mask
}

// Compute b^e (mod m) efficiently using the right-to-left
// binary method of modular exponentiation. It makes use of
// the identity (a*b) (mod m) = [(a (mod m) * b (mod m)] (mod m)
//		https://en.wikipedia.org/wiki/Modular_exponentiation
func PowMod(b, e, m uint) (uint, error) {
	// Special case
	if m == 1 {
		return 0, nil
	}
	// Catch overflow
	if m*m >= ^uint(0) {
		return 0, fmt.Errorf("Modulus too large")
	}
	// Initialise remainder
	r := uint(1)
	// Use the binary exponentation of e
	for e > 0 {
		// Calculate product of remainder
		if e%2 == 1 {
			r = (r * b) % m
		}
		// Divide e by two
		e >>= uint(1)
		// Repeatedly square b to achieve b=b^(2^i) (mod m)
		b = (b * b) % m
	}
	return r, nil
}

// Compute b^e (mod m) efficiently using the right-to-left
// binary method of modular exponentiation. It makes use of
// the identity (a*b) (mod m) = [(a (mod m) * b (mod m)] (mod m)
//		https://en.wikipedia.org/wiki/Modular_exponentiation
func BigPowMod(b, e, m *big.Int) *big.Int {
	// Special case
	if m.Cmp(one) == 0 {
		return zero
	}
	// Copy arguments
	base := new(big.Int)
	base = base.Set(b)
	exp := new(big.Int)
	exp = exp.Set(e)
	mod := new(big.Int)
	mod = mod.Set(m)
	// Initialise remainder
	r := big.NewInt(int64(1))
	// Use the binary exponentation of e to successively
	// calculate (b*b) (mod m)
	base = base.Mod(base, mod)
	for exp.Cmp(zero) > 0 {
		em := new(big.Int)
		em = em.Mod(exp, two)
		if em.Cmp(one) == 0 {
			r = r.Mul(r, base).Mod(r, mod)
		}
		// Divide exponent by 2
		exp = exp.Rsh(exp, uint(1))
		// Repeatedly square b to achieve b=b^(2^i) (mod m)
		base = base.Mul(base, base).Mod(base, mod)
	}
	return r
}

// Miller-Rabin Primality Test raises the base witnesses
// to the binary exponentiation of n-1. If this is congruent
// to 1 (mod n) then the number is probably prime. This
// method is 100% accurate up to 2^64-1.
// 25 repitions is recommended
//		https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test
func BigIsPrime(n *big.Int, reps int, force2 bool) bool {
	// n must and odd integer > 2
	even := new(big.Int)
	even = even.And(n, one)
	if n.Cmp(two) <= 0 || even.Cmp(zero) == 0 {
		return false
	}
	// Get prime bitmask for primes between 2-2^16
	if n.Cmp(big.NewInt(int64(65536))) <= 0 {
		primeBitmask := primemask(2, 65536)
		// Check that primeBitmask&(1<<n) != 0 -> n == prime
		shift := new(big.Int)
		shift = shift.Lsh(big.NewInt(int64(1)), uint(n.Uint64()))
		prime := new(big.Int)
		prime = prime.And(primeBitmask, shift)
		return prime.Cmp(zero) != 0
	}
	nm1 := big.NewInt(int64(-1))
	nm1 = nm1.Add(n, nm1)
	nm4 := big.NewInt(int64(-4))
	nm4 = nm4.Add(n, nm4)
	// s > 0 and d > 0 such that n-1 = 2^s(d)
	s := nm1.TrailingZeroBits() // get highest power of 2 from n-1
	d := new(big.Int)
	d = d.Rsh(nm1, s) // get odd multiplier such that n-1/2^s=d
	// Get random source
	src := rand.NewSource(time.Now().UTC().UnixNano())
	rand := rand.New(src)
	y := new(big.Int)
	for i := 0; i < reps; i++ {
		// n is always a probable prime to base 1 and n-1
		a := new(big.Int)
		if i == reps-1 && force2 {
			a = big.NewInt(int64(2)) // Force using base 2 once
		} else {
			a = a.Rand(rand, nm4).Add(a, big.NewInt(int64(2))) // random(2, n-2)
		}
		x := BigPowMod(a, d, n)
		for j := uint(0); j < s; j++ {
			y = BigPowMod(x, big.NewInt(int64(2)), n)
			if y.Cmp(one) == 0 && x.Cmp(one) != 0 && x.Cmp(nm1) != 0 {
				return false // nontrivial square root of 1 (mod n)
			}
			x = x.Set(y)
		}
		if y.Cmp(one) != 0 {
			return false
		}
	}
	return true
}

// Quick sort given slice in place by first finding a correctly
// placed pivot point and then recursively sorting the left and
// right subarrays until all subarrays are of length 1 or 0
func quicksort(array []*big.Int) []*big.Int {
	if len(array) <= 1 {
		return array
	}
	p1 := 0
	p2 := len(array) - 1
	// Correctly place the pivot element
	for p1 != p2 {
		if (array[p1].Cmp(array[p2]) > 0 && p1 < p2) || (array[p1].Cmp(array[p2]) < 0 && p1 > p2) {
			array[p1], array[p2] = array[p2], array[p1]
			p1, p2 = p2, p1
		}
		if p1 < p2 {
			p1 += 1
		} else {
			p1 -= 1
		}
	}
	// Sort left and right subarrays
	left := quicksort(array[:p1])
	right := quicksort(array[p1+1:])
	// Return the sorted slice
	return array
}

// Pollard's Rho Factoring algorithm
// Find a nontrivial factor of input number n using Floyd's
// cycle-finding algorithm (turtle and hare pointers)
//		https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
//		https://en.wikipedia.org/wiki/Cycle_detection#Floyd.27s_Tortoise_and_Hare
//		https://link.springer.com/content/pdf/10.1007/BF01933667.pdf?pdf=inline%20link
func _pollardRhoFloyd(N *big.Int, s int) []*big.Int {
	//fmt.Printf("Factorising: %v (%d digits)\n", N, len(N.String()))
	// Initialise constants
	x := big.NewInt(int64(s)) // TODO: look into optimal starting values (or random)
	y := big.NewInt(int64(2))
	d := big.NewInt(int64(1))
	// Loop through cycle
	for d.Cmp(one) == 0 {
		// x = g(x)
		x = x.Mul(x, x).Add(x, one).Mod(x, N)
		// y = g(g(x))
		y = y.Mul(y, y).Add(y, one).Mod(y, N)
		y = y.Mul(y, y).Add(y, one).Mod(y, N)
		// d = gcd(|x - y|, n)
		xy := new(big.Int)
		xy = xy.Sub(x, y).Abs(xy)
		// If d != 1 then the sequence {x[k] mod p} must have a repetition
		// and the difference between x[i] and x[j] at this point must be a
		// multiple of p
		d = d.GCD(nil, nil, xy, N)
	}
	// Get factors
	factors := []*big.Int{}
	if d.Cmp(N) == 0 {
		//fmt.Printf("No factors found for %v, trying with new parameters\n", N)
		facs := _pollardRhoFloyd(N, s+1)
		factors = append(factors, facs...)
	} else {
		d2 := new(big.Int)
		d2 = d2.Div(N, d)
		factors = append(factors, d)
		factors = append(factors, d2)
	}
	// Recursion step to ensure all are prime
	final := []*big.Int{}
	i := 0
	for i < len(factors) {
		prime := BigIsPrime(factors[i], 25, true)
		if prime {
			final = append(final, factors[i])
			i += 1
			continue
		}
		facs := _pollardRhoFloyd(factors[i], 2)
		factors = append(factors, facs...)
		i += 1
	}
	return final
}

func min(a, b int) int {
	if a >= b {
		return b
	}
	return a
}

// Pollard's Rho Factoring algorithm
// Find a nontrivial factor of input number n using Brent's
// cycle-finding algorithm and optimisations
//		https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
//		https://en.wikipedia.org/wiki/Cycle_detection#Brent.27s_algorithm
//		https://maths-people.anu.edu.au/~brent/pd/rpb051i.pdf
func _pollardRhoBrent(N *big.Int) []*big.Int {
	//fmt.Printf("Factorising: %v (%d digits)\n", N, len(N.String()))
	// Get random source
	src := rand.NewSource(time.Now().UnixNano())
	rng := rand.New(src)
	// Initialise constants
	nm2 := big.NewInt(int64(-2))
	nm2 = nm2.Add(N, nm2)
	//y := new(big.Int)
	//y = y.Rand(rng, nm2).Add(y, one)
	y := big.NewInt(int64(2))
	// Randomise c in f(x)=x^2+c (mod N)
	c := new(big.Int)
	c = c.Rand(rng, nm2).Add(c, one)
	for c.Cmp(nm2) == 0 {
		c = c.Rand(rng, nm2).Add(c, one)
	}
	ys := new(big.Int)
	x := new(big.Int)
	m := 1900 // TODO: look into optimal m value
	k := 0
	r := 2
	q := big.NewInt(int64(1))
	G := big.NewInt(int64(1))
	// Loop through cycle
	for G.Cmp(one) == 0 {
		x.Set(y)
		for i := 0; i <= r; i++ {
			// y = f(y)
			y = y.Mul(y, y).Add(y, c).Mod(y, N)
			// k = 0
			k = 0
		}
		for k <= r && G.Cmp(one) == 0 {
			ys.Set(y)
			min := min(m, r-k)
			for i := 0; i <= min; i++ {
				// y = f(y)
				y = y.Mul(y, y).Add(y, c).Mod(y, N)
				// q = q*|x-y| (mod N)
				xy := new(big.Int)
				xy = xy.Sub(x, y).Abs(xy)
				q = q.Mul(q, xy).Mod(q, N)
			}
			G = G.GCD(nil, nil, q, N)
			k += m
		}
		r *= 2
	}
	// Get factors
	factors := []*big.Int{}
	if G.Cmp(N) == 0 {
		G = big.NewInt(int64(1))
		for G.Cmp(one) == 0 {
			// ys = f(ys)
			ys = ys.Mul(ys, ys).Add(ys, c).Mod(ys, N)
			// G = GCD(|x-ys|, N)
			xys := new(big.Int)
			xys = xys.Sub(x, ys).Abs(xys)
			G = G.GCD(nil, nil, xys, N)
		}
	}
	// Failure case
	if G.Cmp(N) == 0 {
		//fmt.Printf("No factors found for %v, trying with new parameters\n", N)
		facs := _pollardRhoBrent(N)
		factors = append(factors, facs...)
	} else {
		// Success
		G2 := new(big.Int)
		G2 = G2.Div(N, G)
		factors = append(factors, G)
		factors = append(factors, G2)
	}
	// Recursion step to ensure all are prime
	final := []*big.Int{}
	i := 0
	for i < len(factors) {
		prime := BigIsPrime(factors[i], 25, true)
		if prime {
			final = append(final, factors[i])
			i += 1
			continue
		}
		facs := _pollardRhoBrent(factors[i])
		factors = append(factors, facs...)
		i += 1
	}
	return final
}

// Struct to store results of factorisation
type Results struct {
	Number  *big.Int
	Digits  int
	Factors []*big.Int
	Seconds float64
	Method  string
}

// Pollard's Rho factorisation method using Brent's cycle
// finding algorithm and optimisations. Works on any integer
// and will take out factors of 2 prior to undertaking the rho
// method. Returns a Result struct containg the details of the
// factorisation
func PollardRho(N *big.Int) Results {
	// Take out any factors of 2
	n := new(big.Int)
	n.Set(N)
	even := new(big.Int)
	even = even.And(n, one)
	factors := []*big.Int{}
	final := []*big.Int{}
	for even.Cmp(zero) == 0 && n.Cmp(two) != 0 {
		n = n.Div(n, two)
		even = even.And(n, one)
		factors = append(factors, two)
	}
	// Prepare results struct
	var res Results
	res.Number = N
	res.Digits = len(n.String())
	start := time.Now()
	if len(n.String()) < 30 {
		// Use Floyd's cycle detection for smaller numbers
		final = _pollardRhoFloyd(n, 2)
		res.Method = "Pollard's Rho (Floyd)"
	} else {
		// Use Brent's optimisations for larger numbers
		final = _pollardRhoBrent(n)
		res.Method = "Pollard's Rho (Brent)"
	}
	// Add time taken in seconds to results struct
	elapsed := time.Since(start)
	res.Seconds = elapsed.Seconds()
	// Sort factors and add to struct
	final = append(final, factors...)
	final = quicksort(final)
	res.Factors = final
	return res
}

func main() {
	str := []string{
		"13290059",
		"77374242946",
		"3227374242948",
		"91283710282317",
		"8931998123874123",
		"1267819267256172189",
		"537823842718237418924961",
		"3281462346123741923846187637",
		"912873612917217374028365434587",
		"4269668178126580590972042829635793",
		"3764382387926236426798765213827368557",
		"340282366920938463463374607431768211459",
		"340282366920938463463374607431768211457", // F7
		"9128361928152783912761243528812738741235645129",
		"37656535734427677898969945243647233538976467987652557",
	}
	// Loop through each string and factor them with both
	// Rho methods in parrallel
	c := make(chan Results, 2)
	for i := 0; i < len(str)-1; i += 2 {
		for j := 0; j < 2; j++ {
			n := new(big.Int)
			n.SetString(str[i+j], 10)
			go func(ch chan<- Results, n *big.Int) {
				res := PollardRho(n)
				ch <- res
			}(c, n)
		}
		for k := 0; k < 2; k++ {
			factors := <-c
			fmt.Printf("%+v\n", factors)
		}
	}
}