prime/prime.go

// A bunch of helper functions for dealing with small prime numbers that would fit in an unsigned integer.
//
// This isn't a serious scientifically minded package, or heavily optimized, it's just a quick and dirty tool
// for grabbing prime numbers should you find yourself needing them.
package prime

import (
	"fmt"
	"iter"
	"math"
	"slices"
	"sync"
)

var (
	// initialize nextToCheck to the first odd prime number.
	nextToCheck uint = 3

	// initialize the list to 2 since this value otherwise can't be reached by incrementing by 2 by an odd starting point.
	primes = []uint{2}

	// protect the above two variables, which are modified primarily by the functions [upTo] and [firstN].
	// the test package also has a [resetCache] function to reinitialize them.
	m sync.RWMutex
)

const (
	// this _must_ be larger than 1, because we assume the last item in the list to be odd when we change how switch how we generate these numbers.
	generateLimit = 1 << 16

	// how many new primes to request at a time.
	generateStep = 1 << 8
)

// returns a sub-slice of primes that could potentially be a factor of n; primes must contain all the prime numbers up to sqrt(n) in ascending order.
func possibleFactors(primes []uint, n uint) []uint {
	i, found := slices.BinarySearch(primes, uint(math.Sqrt(float64(n))))
	if found {
		i++
	}
	return primes[:i:i]
}

// returns true if n is prime; primes must contain all the prime numbers up to sqrt(n) in ascending order
func test(primes []uint, n uint) bool {
	if len(primes) > 0 && n <= primes[len(primes)-1] {
		// if the number to test would be inside the list of primes, then check the list using a binary search.
		_, found := slices.BinarySearch(primes, n)
		return found
	}

	if n < 2 {
		// 0 and 1 won't have any prime factors and would otherwise get classified as prime.
		return false
	}

	// otherwise, test if it's prime by checking against its prime factors.
	for _, factor := range possibleFactors(primes, n) {
		if n%factor == 0 {
			return false
		}
	}

	return true
}

// returns a slice containing all the primes up to n inclusive (and possibly more than that).
func upTo(n uint) []uint {
	m.RLock()
	if n < nextToCheck || nextToCheck == 1 {
		primes := primes
		m.RUnlock()
		return primes
	}

	m.RUnlock()
	m.Lock()
	defer m.Unlock()

	for ; n >= nextToCheck && nextToCheck != 1; nextToCheck += 2 {
		if test(primes, nextToCheck) {
			primes = append(primes, nextToCheck)
		}
	}

	return primes
}

// returns a slice containing the first n primes (and possibly more than that),
// and also possibly less if the nth prime can't fit in a uint.
func firstN(n int) []uint {
	m.RLock()
	if n <= len(primes) || nextToCheck == 1 {
		primes := primes
		m.RUnlock()
		return primes
	}

	m.RUnlock()
	m.Lock()
	defer m.Unlock()

	for ; n > len(primes) && nextToCheck != 1; nextToCheck += 2 {
		if test(primes, nextToCheck) {
			primes = append(primes, nextToCheck)
		}
	}

	return primes
}

// FirstN returns a slice containing the first n prime numbers in ascending order.
//
// The returned slice should be considered read-only and must not be modified.
//
// The returned slice might contain less than n elements if those numbers wouldn't fit in a uint.
func FirstN(n int) []uint {
	return firstN(n)[:n:n]
}

// Test returns true if n is prime.
func Test(n uint) bool {
	return test(upTo(uint(math.Sqrt(float64(n)))), n)
}

// Next returns the prime number at some relative offset to n, or 0 if there is no applicable number.
//
// Next(n, 0) will return n if n is prime, otherwise 0.
// Next(n, -1) will return the next prime number before n, if one exists (i.e., n >= 3)
// Next(n, 1) will return the next prime number after n, if representable in a uint.
// Next(n, 100) will return the 100th prime number after n, if representable in a uint.
func Next(n uint, offset int) uint {
	primes := upTo(uint(math.Sqrt(float64(n))))

	// if the list of primes already contains our number (upTo may return more than requested), possibly we can do this based on indexes alone.
	if primes[len(primes)-1] >= n {
		i, found := slices.BinarySearch(primes, n)

		if !found {
			if offset == 0 {
				return 0
			}

			n = primes[i]

			if offset > 0 {
				offset--
			}
		}

		if offset <= 0 {
			if i+offset >= 0 && i+offset <= i {
				return primes[i+offset]
			}

			// there is definitely nothing before this, stop now.
			return 0
		}

		// offset > 0
		if i+offset > i && i+offset < len(primes) {
			// we know exactly which prime is at this offset.
			return primes[i+offset]
		}

		// the list doesn't contain a value for primes[offset+i],
		// so adjust n to the last prime in the list, and decrease offset by the number of numbers we skipped over.
		n = primes[len(primes)-1]
		offset -= len(primes) - 1 - i
	}

	switch {
	// is n even?
	case n%2 == 0:
		switch {
		case offset < 0:
			if n <= 2 {
				// n=0 or n=2, either way there are no primes before them.
				// note: this path isn't normally encountered because the the primes list generally always contains
				// at least 2, and this situation will have been handled via direct index calculation.
				return 0
			}

			// decrease to make n odd.
			n--
		case offset > 0:
			// increase to make n odd.
			n++
		default:
			if n == 2 {
				// the only even prime number.
				// note: this path isn't normally encountered because the the primes list generally always contains
				// at least 2, and this situation will have been handled via direct index calculation.
				return 2
			}

			// some other even number, not prime.
			return 0
		}

		// if n is already prime, we need to skip over it.
	case test(primes, n):
		switch {
		case offset < 0:
			switch {
			case n == 3:
				if offset == -1 {
					return 2
				}

				fallthrough
			case n == 2:
				return 0
			}

			n -= 2
		case offset > 0:
			n += 2
		default:
			return n
		}
	}

	// at this point, offset will be signed, and n will be an odd number greater or equal to 3.
	if offset == 0 || n < 3 || n%2 == 0 {
		panic(fmt.Sprintf("n=%d offset=%d", n, offset))
	}

	if offset < 0 {
		// count backwards

		for ; n != 1; n -= 2 {
			if test(primes, n) {
				offset++

				if offset == 0 {
					return n
				}
			}
		}

		// n=1, meaning n=2 was skipped due to decrementing by 2 each time.
		// if offset = -1, then return the 2 we skipped.
		if offset == -1 {
			return 2
		}

		return 0
	}

	// count forwards
	listLegalTo := uintSqr(primes[len(primes)-1])

	for ; n != 1; n += 2 {
		for n > listLegalTo {
			primes = firstN(len(primes) + generateStep)
			listLegalTo = uintSqr(primes[len(primes)-1])
		}

		if test(primes, n) {
			offset--

			if offset == 0 {
				return n
			}
		}
	}

	// integer overflow happened, the remaining offset primes aren't representable via uint.
	return 0
}

// UpTo returns a slice containing all the up to n inclusive in ascending order.
//
// The returned slice should be considered read-only and must not be modified.
func UpTo(n uint) []uint {
	primes := upTo(n)
	// note that we don't add 1 here like we do in [possibleFactors],
	// as n can be [math.MaxUint] here, which would overflow.
	// instead, we add 1 if the requested number was found in the list.
	i, found := slices.BinarySearch(primes, n)
	if found {
		i++
	}
	return primes[:i:i]
}

// squares a, returning [math.MaxUint] if the operation would overflow.
func uintSqr(a uint) uint {
	if a > 0 && a > math.MaxUint/a {
		// if squaring would overflow, return MaxUint.
		return math.MaxUint
	}

	return a * a
}

// All returns an iterator yielding all the prime numbers up to [math.MaxUint] in ascending order.
func All() iter.Seq[uint] {
	return func(yield func(uint) bool) {
		var primes []uint

		for i := 0; ; i++ {
			if i == len(primes) {
				if i >= generateLimit {
					break
				}

				primes = firstN(min(i+generateStep, generateLimit))
			}

			if !yield(primes[i]) {
				return
			}
		}

		listLegalTo := uintSqr(primes[len(primes)-1])

		// nextCandidate _should_ be odd, so it should be equal to 1 when it overflows.
		for nextCandidate := primes[len(primes)-1] + 2; nextCandidate != 1; nextCandidate += 2 {
			if nextCandidate > listLegalTo {
				// need to grow the list.
				primes = firstN(len(primes) + generateStep)
				listLegalTo = uintSqr(primes[len(primes)-1])
			}

			if test(primes, nextCandidate) && !yield(nextCandidate) {
				return
			}
		}
	}
}

func factorize(primes []uint, n uint) iter.Seq2[uint, int] {
	return func(yield func(uint, int) bool) {
		n := n

		for _, factor := range possibleFactors(primes, n) {
			exponent := 0

			for n%factor == 0 {
				n /= factor
				exponent++
			}

			if exponent > 0 && !yield(factor, exponent) {
				return
			}
		}

		// if n didn't reach zero, then n is prime.
		if n > 1 {
			yield(n, 1)
		}
	}
}

// Factorize returns an iterator that yields the prime factors of n and their exponents.
func Factorize(n uint) iter.Seq2[uint, int] {
	return factorize(upTo(uint(math.Sqrt(float64(n)))), n)
}

// LCM computes the least common multiple. This really doesn't have much to do with prime numbers other than the algorithm heavily relying on them.
//
// Returns 0 if the LCM would overflow, or one of the numbers is 0.
func LCM(numbers ...uint) uint {
	// trivial cases:
	switch {
	case len(numbers) == 0:
		return 1
	case len(numbers) == 1:
		return numbers[0]
	case 0 == slices.Min(numbers):
		return 0
	}

	primes := upTo(uint(math.Sqrt(float64(slices.Max(numbers)))))
	factors := map[uint]int{}

	for _, n := range numbers {
		for factor, exponent := range factorize(primes, n) {
			factors[factor] = max(factors[factor], exponent)
		}
	}

	product := uint(1)

	for factor, exponent := range factors {
		for ; exponent > 0; exponent-- {
			if product > math.MaxUint/factor {
				// would overflow
				return 0
			}

			product *= factor
		}
	}

	return product
}

// GCD computes the greatest common divisor. This really doesn't have much to do with prime numbers other than the algorithm heavily relying on them.
//
// Returns 0 if the list is empty or all of the numbers are 0.
func GCD(numbers ...uint) uint {
	// trivial cases:
trivialCases:
	switch {
	case len(numbers) == 0:
		return 0
	case len(numbers) == 1:
		return numbers[0]
	case 0 == slices.Min(numbers):
		// ugh, why did you have to include a zero?
		// clone the input slice, sort it, remove duplicates, then remove the first element, which will be the zero
		// that it taunting us.
		numbers = slices.Clone(numbers)
		slices.Sort(numbers)
		numbers = slices.Compact(numbers)[1:]

		// start over again.
		goto trivialCases
	}

	primes := upTo(uint(math.Sqrt(float64(slices.Max(numbers)))))
	factors := map[uint]int{}
	var seen []uint

	for i, n := range numbers {
		seen = seen[:0]

		for factor, exponent := range factorize(primes, n) {
			if i == 0 {
				// if this is the first number, store the exponent normally.
				factors[factor] = exponent
			} else {
				seen = append(seen, factor)

				if prevExponent, found := factors[factor]; found {
					// use the minimum of this exponent and the previously known one.
					factors[factor] = min(prevExponent, exponent)
				} else {
					// the previous numbers must have had an exponent of 0 for this,
					// and min(0, exponent) would be 0.
					// we don't need to store factors with an exponent of 0,
					// so we'd delete it in this case. Except it wasn't found in the map,
					// so there's nothing to delete.
				}
			}
		}

		if i != 0 {
			// if this isn't the first number, and we know about factors that weren't visited,
			// then those unvisited factors implicitly had an exponent of 0, and need to be deleted.
			for factor := range factors {
				// note: factorize returns factors in ascending order, so seen is a sorted list and we
				// can use a binary search.
				if _, found := slices.BinarySearch(seen, factor); !found {
					delete(factors, factor)
				}
			}
		}
	}

	product := uint(1)

	for factor, exponent := range factors {
		for ; exponent > 0; exponent-- {
			// note: overflow can't happen here.
			product *= factor
		}
	}

	return product
}