geom/primes/cache.go

package primes

import "opslag.de/schobers/geom/ints"

var primes = &cache{}

const (
	primesPerPage int = 8192
)

// Cache represents the cached state of some pre-calculated prime numbers.
type Cache interface {
	IsPrime(n int64) bool
	OffsetOfPrime(n int64) int
	Prime(n int) int64
}

// NewCache creates a new prime number cache.
func NewCache() Cache {
	return &cache{}
}

type cache struct {
	pages []page
}

type page struct {
	primes []int64
	last   int64
}

func (c *cache) OffsetOfPrime(n int64) int {
	p := 0
	for {
		if len(c.pages) == p {
			c.generatePage()
		}
		if c.pages[p].last >= n {
			for i := 0; i < primesPerPage; i++ {
				if c.pages[p].primes[i] >= n {
					return i + p*primesPerPage
				}
			}
		}
		p++
	}
}

func (c *cache) Prime(n int) int64 {
	i := n - 1
	p := i / primesPerPage
	for p >= len(c.pages) {
		c.generatePage()
	}
	offset := i % primesPerPage
	return c.pages[p].primes[offset]
}

func (c *cache) IsPrime(n int64) bool {
	if inRange, isPrime := c.isKnownPrime(n); inRange {
		return isPrime
	}
	upper := ints.Sqrt64(n)
	it := NewIteratorFromCache(c)
	for it.Next() {
		p := it.Get()
		if p > upper {
			return true
		}
		if n%p == 0 {
			return false
		}
	}
	panic("out of primes")
}

func (c *cache) generatePage() {
	pageI := len(c.pages)
	c.pages = append(c.pages, page{})
	var p = &c.pages[pageI]
	p.primes = make([]int64, primesPerPage)
	var i int
	var last int64
	if 0 == pageI {
		copy(p.primes, seed)
		i = 20
		last = p.primes[19]
	} else {
		last = c.pages[pageI-1].last
	}

	wheel := newWheel6(c)
	wheel.resetTo(last)

	for ; i < primesPerPage; i++ {
		prime := wheel.next()
		for !c.IsPrime(prime) {
			prime = wheel.next()
		}
		p.primes[i] = prime
	}
	p.last = p.primes[primesPerPage-1]
}

func (c *cache) isKnownPrime(n int64) (bool, bool) {
	if 0 == len(c.pages) {
		c.generatePage()
	}
	for _, p := range c.pages {
		if n > p.last {
			continue
		}
		l := len(p.primes)
		offset := 0
		shift := l >> 1
		for shift > 0 {
			pp := p.primes[offset+shift-1]
			if n == pp {
				return true, true
			}
			if n > pp {
				offset += shift
			}
			shift >>= 1
		}
		if n == p.last {
			return true, true
		}
		return true, false
	}
	return false, false
}

// Generate n pre-calculates the first n prime numbers for the default cache.
func Generate(n int) {
	primes.Prime(n)
}

// N returns the first n prime numbers.
func N(n int) []int64 {
	Generate(n)

	ps := make([]int64, n)
	var i = 0
	for _, p := range primes.pages {
		var toCopy = primesPerPage
		if n-i < primesPerPage {
			toCopy = n - i
		}
		i += copy(ps[i:], p.primes[:toCopy])
	}
	return ps
}

// Nth returns the nth prime number.
func Nth(n int) int64 {
	return primes.Prime(n)
}