Improved Atkin-Bernstein sieve for generating primes

This is a class from my personal code library, and from a package which deals with integer sequences. It implements an interface

package org.cheddarmonk.math.sequence;

public interface IntegerSequence extends Iterable<java.math.BigInteger> {
	public int offset();
	public int size();
	public java.math.BigInteger get(int index);
}

and, as the comments mention, generates primes using an Atkin-Bernstein sieve. The interface is intended to allow operations such as composition of sequences, but in addition to that interface I find it helpful to have methods Primes.prev and Primes.next.

There are a couple of TODOs about thread safety, but I'm bearing in mind YAGNI and commenting them only for future reference if I do turn out to need it.

package org.cheddarmonk.math.sequence;

import java.math.BigInteger;
import java.util.*;
import org.cheddarmonk.math.MathExt;
import org.cheddarmonk.util.LRUCache;

/**
 * The sequence of primes, OEIS A000040.
 *
 * The primes are enumerated using a paged implementation of the Atkin-Bernstein sieve, based on a port of Bernstein's
 * primegen but with some modifications. In particular, the binary quadratic forms used should give about a 6% speed
 * optimisation over those used by Atkin and Bernstein.
 *
 * https://cr.yp.to/papers/primesieves.pdf
 */
public class Primes implements IntegerSequence {
	/**
	 * This is effectively a cache of certain values of the prime number function Pi(n).
	 * It allows us to identify the page containing the nth prime number without recomputing all the pages up to it.
	 */
	private static List<Integer> primesUpToPage = new ArrayList<Integer>();

	// Page size.
	private static final int B = 1001 << 6;
	// Modulus for sharding between binary quadratic forms.
	private static final int M = 60;
	// An upper bound on the number of primes in a range of width M*B.
	private static final int PI_MB = 272895;
	// The number of totients of M.
	private static final int PHI_M = 16;
	// The totients of M.
	private static final int[] totients = new int[] { 1, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 49, 53, 59 };
	// invTotients[totients[i]] == i; invTotients[non-totient] == -1
	private static final int[] invTotients;
	// Primes smaller than M (and null-terminator).
	private static final long[] page0 = new long[] { 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 0 };
	// Parameters for tracing the hyperbolic BQF used for 11+60Z and 59+60Z.
	private static final int[][] hyperbolic = new int[][] {
		{2, 2, 1, 0}, {2, 2, 11, -2}, {2, 2, 19, -6}, {2, 2, 29, -14}, {2, 3, 4, 0}, {2, 3, 14, -3}, {2, 3, 16, -4}, {2, 3, 26, -11},
		{2, 5, 2, 1}, {2, 5, 8, 0}, {2, 5, 22, -7}, {2, 5, 28, -12}, {2, 7, 4, 2}, {2, 7, 14, -1}, {2, 7, 16, -2}, {2, 7, 26, -9},
		{2, 8, 1, 3}, {2, 8, 11, 1}, {2, 8, 19, -3}, {2, 8, 29, -11}, {2, 10, 7, 4}, {2, 10, 13, 2}, {2, 10, 17, 0}, {2, 10, 23, -4},
		{15, 1, 2, -1}, {15, 1, 8, -2}, {15, 1, 22, -9}, {15, 1, 28, -14}, {15, 4, 7, -1}, {15, 4, 13, -3}, {15, 4, 17, -5}, {15, 4, 23, -9},
		{15, 5, 4, 0}, {15, 5, 14, -3}, {15, 5, 16, -4}, {15, 5, 26, -11}, {15, 6, 7, 0}, {15, 6, 13, -2}, {15, 6, 17, -4}, {15, 6, 23, -8},
		{15, 9, 2, 3}, {15, 9, 8, 2}, {15, 9, 22, -5}, {15, 9, 28, -10}, {15, 10, 1, 4}, {15, 10, 11, 2}, {15, 10, 19, -2}, {15, 10, 29, -10}
	};
	// Parameters for tracing the elliptic BQFs used for all totients except 11 and 59.
	private static final int[][] elliptic;

	// Squares of primes 5 < q < 240
	private static final int[] qqtab = new int[] {
		49, 121, 169, 289, 361, 529, 841, 961, 1369, 1681, 1849, 2209, 2809,
		3481, 3721, 4489, 5041, 5329, 6241, 6889, 7921, 9409, 10201, 10609, 11449, 11881,
		12769, 16129, 17161, 18769, 19321, 22201, 22801, 24649, 26569, 27889, 29929, 32041, 32761,
		36481, 37249, 38809, 39601, 44521, 49729, 51529, 52441, 54289, 57121
	};
	// If a_i == q^{-2} (mod 60) is the reciprocal of qq[i], qq60tab[i] = qq[i] + (1 - a_i * qq[i]) / 60
	private static int[] qq60tab = new int[] {
		9, 119, 31, 53, 355, 97, 827, 945, 251, 1653, 339, 405, 515,
		3423, 3659, 823, 4957, 977, 6137, 1263, 7789, 1725, 10031, 1945, 2099, 11683,
		2341, 2957, 16875, 3441, 18999, 21831, 22421, 4519, 4871, 5113, 5487, 31507, 32215,
		35873, 6829, 7115, 38941, 43779, 9117, 9447, 51567, 9953, 56169
	};

	static {
		// Page 0 (primes up to M) has 17 primes.
		primesUpToPage.add(17);
		// Most of the tables will break if M changes, but this pre-calculated table for the first 100M primes will break if B changes.
		if (B == (1001 << 6)) {
			int[] precalc = new int[] {
				272894, 247211, 239110, 234158, 230680, 227918, 225599, 223972, 221931, 220720, 219546, 218432, 217246, 216306, 215542, 214597,
				213847, 213392, 212774, 211825, 211247, 210990, 210292, 209726, 209395, 208811, 208540, 208141, 207692, 207119, 206728, 206442,
				206392, 205811, 205512, 205067, 204929, 204760, 204329, 203832, 203886, 203381, 203252, 203075, 202868, 202677, 202367, 201862,
				201952, 201306, 201689, 200777, 201038, 200984, 200591, 200444, 200088, 199993, 200008, 199641, 199245, 199872, 199207, 199006,
				198863, 198636, 198666, 198331, 198207, 198188, 197667, 198018, 197773, 197434, 197106, 197480, 197001, 197040, 196843, 196897,
				196621, 196508, 196185, 196135, 196227, 195871, 195830, 195704, 195888, 195626, 195296, 195250, 195307, 194867, 195086, 194983,
				194483, 195095, 195068, 194479, 194342, 194126, 194273, 194037, 193991, 193952, 193836, 193725, 193647, 193310, 193835, 193228,
				193169, 193256, 193240, 193084, 193117, 192675, 192930, 192628, 192762, 192728, 192301, 192188, 192727, 192243, 192061, 192021,
				192127, 192023, 191436, 191612, 191361, 191506, 191869, 191398, 191473, 191582, 191163, 191486, 191077, 190823, 191168, 191472,
				190717, 190638, 190822, 190700, 190576, 190558, 190580, 190541, 190135, 190798, 190387, 190045, 190418, 190124, 189936, 190077,
				189523, 189864, 189794, 190018, 189422, 189761, 189566, 189298, 189545, 189376, 189391, 189499, 189096, 189172, 189113, 188917,
				189003, 188712, 189203, 189131, 188911, 188520, 188792, 188681, 188566, 188530, 188383, 188629, 188391, 188029, 188277, 188218,
				188641, 187859, 187754, 188109, 188168, 187812, 187769, 187811, 188150, 187506, 187866, 187908, 187636, 187928, 187484, 187629,
				187392, 187853, 187461, 187137, 187438, 187303, 187018, 187120, 187262, 187121, 187111, 187078, 186904, 186867, 186939, 186548,
				186961, 186781, 186663, 186581, 186660, 186504, 186499, 186203, 186582, 186313, 186324, 186572, 186647, 186485, 186193, 186450,
				186165, 185703, 186248, 186033, 185978, 185888, 186066, 185877, 185894, 186309, 185811, 185731, 185836, 185379, 185625, 185855,
				185736, 185591, 185478, 185828, 185029, 185576, 185436, 185489, 185490, 185229, 185353, 185114, 185030, 185266, 185416, 185002,
				184804, 185166, 184758, 184902, 185073, 185148, 185091, 184550, 184723, 185014, 185037, 184617, 184796, 184771, 184614, 184554,
				184774, 184436, 184317, 184766, 184097, 184276, 184173, 183962, 184263, 183731, 184573, 184463, 184376, 183870, 184538, 184063,
				184313, 184004, 183963, 183690, 183871, 183951, 183887, 184087, 183557, 183771, 183856, 183832, 183738, 183817, 183465, 183465,
				183641, 183785, 183271, 183425, 183842, 183729, 183289, 183405, 183396, 183761, 183430, 183304, 183346, 183209, 183282, 183444,
				182922, 182850, 183230, 183228, 183451, 183053, 182674, 183229, 182761, 182907, 183032, 183283, 182894, 182927, 182974, 182685,
				183150, 182894, 182784, 182689, 182663, 182396, 182523, 182612, 182388, 182822, 182704, 182402, 182362, 182377, 182448, 182401,
				182490, 182471, 182426, 182298, 182284, 182248, 182397, 181913, 182511, 182101, 182220, 182231, 182360, 181811, 181835, 182368,
				182274, 181909, 182299, 182178, 182261, 181941, 181684, 182194, 181636, 181611, 181863, 181931, 181768, 182327, 181416, 181691,
				181545, 181641, 181793, 181863, 181606, 181514, 181167, 181461, 181368, 181521, 181819, 181652, 181238, 181563, 181280, 181284,
				181374, 181661, 181071, 181271, 181136, 181351, 181078, 181311, 181554, 181314, 181226, 180752, 181158, 180933, 180837, 180914,
				181032, 181215, 180984, 180824, 181097, 180932, 181035, 181235, 180690, 180759, 180741, 181135, 180935, 180605, 181009, 180941,
				180423, 180792, 180503, 180408, 180856, 180576, 180430, 180878, 180542, 180575, 180451, 180655, 180453, 180257, 180729, 180485,
				180474, 180320, 180671, 180765, 180256, 180288, 180384, 180112, 180286, 180039, 180257, 180327, 180128, 180386, 180027, 180189,
				179973, 180005, 180244, 179927, 180280, 179653, 179957, 180064, 180192, 179774, 179651, 180189, 180027, 180031, 179893, 180024,
				179720, 179675, 179793, 179764, 179682, 179768, 179999, 179800, 179532, 179688, 179731, 180055, 179640, 179508, 179496, 179520,
				179813, 179785, 179624, 179523, 179594, 179409, 179728, 179210, 179746, 179828, 179282, 179796, 179465, 179304, 179281, 179150,
				179312, 179392, 179302
			};
			int total = 17;
			for (int delta : precalc) primesUpToPage.add(total += delta);
		}

		invTotients = new int[M];
		Arrays.fill(invTotients, -1);
		for (int i = 0; i < totients.length; i++) invTotients[totients[i]] = i;

		// Calculate the parameters for tracing the elliptic BQFs from a table of the BQF used for each totient.
		// E.g. for 17+60Z we use 5x^2 + 3y^2.
		int[][] bqfs = new int[][] {
			{1, 15, 1}, {7, 3, 1}, {13, 4, 1}, {17, 5, 3}, {19, 15, 1}, {23, 5, 3}, {29, 4, 1},
			{31, 15, 1}, {37, 4, 1}, {41, 4, 1}, {43, 3, 1}, {47, 5, 3}, {49, 15, 1}, {53, 5, 3}
		};
		List<int[]> parmSets = new ArrayList<int[]>();
		for (int[] bqf : bqfs) parmSets.addAll(initElliptic(bqf[0], bqf[1], bqf[2]));
		elliptic = parmSets.toArray(new int[0][]);
	}

	/**
	 * Cache of the most recently accessed pages, which is likely to give a significant performance boost to most
	 * use cases not involving the iterator (which has its own cache).
	 */
	private static final LRUCache<Integer, long[]> pageCache = new LRUCache<Integer, long[]>(16);

	@Override
	public int offset() {
		return 0;
	}

	@Override
	public int size() {
		return Integer.MAX_VALUE >> 1;
	}

	@Override
	public BigInteger get(int index) {
		// TODO Thread safety around access to primesUpToPage.
		int known = primesUpToPage.get(primesUpToPage.size() - 1);
		long[] contents = null;
		while (index >= known) {
			contents = getPage(primesUpToPage.size());
			for (long p : contents) {
				if (p == 0) break;
				known++;
			}
			primesUpToPage.add(known);
		}

		if (contents != null) {
			// index < known, but index >= primesUpToPage.get(primesUpToPage.size() - 2);
			return BigInteger.valueOf(contents[index - primesUpToPage.get(primesUpToPage.size() - 2)]);
		}

		// Find the right page number.
		for (int i = 0; true; i++) {
			if (index < primesUpToPage.get(i)) {
				contents = getPage(i);
				return BigInteger.valueOf(contents[index - (i == 0 ? 0 : primesUpToPage.get(i - 1))]);
			}
		}
	}

	@Override
	public Iterator<BigInteger> iterator() {
		return new Iterator<BigInteger>() {
			private int page = -1;
			private long[] contents = new long[1];
			private int contentIdx = 0;

			public boolean hasNext() {
				return true;
			}

			public BigInteger next() {
				while (contents[contentIdx] == 0) {
					page++;
					contents = getPage(page);
					contentIdx = 0;
				}

				return BigInteger.valueOf(contents[contentIdx++]);
			}

			public void remove() {
				throw new UnsupportedOperationException();
			}
		};
	}

	/**
	 * Returns the largest prime less than n.
	 * Note that this is only guaranteed for n up to about 2^50.
	 */
	public static long prev(long n) {
		if (n < 3) throw new IllegalArgumentException("There are no primes less than 2");

		// Page extends up to 60 + page * B * M (not inclusive).
		int page = (int)((n - 60 + (B * M - 1)) / (B * M)); // ceil((n - 60) / (B * M))
		while (true) {
			long[] contents = getPage(page);
			if (contents[0] == 0 || contents[0] >= n) {
				page--;
				continue;
			}

			long prev = contents[0];
			for (long p : contents) {
				if (p >= n) return prev;
				prev = p;
			}
		}
	}

	/**
	 * Returns the smallest prime greater than n.
	 * Note that this is only guaranteed for n up to about 2^50.
	 */
	public static long next(long n) {
		// Special-case negative values.
		if (n < 2) return 2;

		// Page extends up to 60 + page * B * M (not inclusive).
		int page = (int)((n - 58 + (B * M - 1)) / (B * M)); // ceil((n - 58) / (B * M))
		while (true) {
			long[] contents = getPage(page);
			for (long p : contents) {
				if (p == 0) break;
				if (p > n) return p;
			}

			page++;
		}
	}

	@Override
	public String toString() {
		return "Primes (A000040)";
	}

	/**
	 * Produces a set of parameters for traceElliptic to find solutions to ax^2 + cy^2 == d (mod M).
	 * @param d The target residue.
	 * @param a Binary quadratic form parameter.
	 * @param c Binary quadratic form parameter.
	 */
	private static List<int[]> initElliptic(final int d, final int a, final int c) {
		List<int[]> rv = new ArrayList<int[]>();

		// The basic idea is that we maintain an invariant of the form
		//     M k = a x^2 + c y^2 - d
		// Therefore we increment x in steps F such that
		//     a((x + F)^2 - x^2) == 0 (mod M)
		// and similarly for y in steps G.
		int F = computeIncrement(a, M), G = computeIncrement(c, M);
		for (int f = 1; f <= F; f++) {
			for (int g = 1; g <= G; g++) {
				if ((a*f*f + c*g*g - d) % M == 0) {
					rv.add(new int[] { invTotients[d], (2*f + F)*a*F/M, (2*g + G)*c*G/M, (a*f*f + c*g*g - d)/M, 2*a*F*F/M, 2*c*G*G/M });
				}
			}
		}

		return rv;
	}

	private static int computeIncrement(int a, int M) {
		// Find smallest F such that M | 2aF and M | aF^2
		int l = M / MathExt.gcd(M, 2 * a);
		for (int F = l; true; F += l) {
			if (a*F*F % M == 0) return F;
		}
	}

	/**
	 * Gets the primes in the interval [60 + (page - 1) * B * M, 60 + page * B * M].
	 * @param page
	 * @return An array of the primes, in ascending order, which is null-terminated.
	 */
	public static long[] getPage(int page) {
		if (page == 0) return page0;

		if (page > 33520) throw new ArithmeticException("There's an overflow somewhere beginning about page 33521");

		// TODO Not thread-safe.
		long[] rv = pageCache.get(page);
		if (rv != null) return rv;
		long[][] buf = new long[PHI_M][B >> 6];
		long L = 1 + (page - 1) * B;

		int[] Lmodqq = new int[qqtab.length];
		for (int i = 0; i < Lmodqq.length; i++) Lmodqq[i] = (int)(L % qqtab[i]);

		for (long[] arr : buf) Arrays.fill(arr, -1); // TODO Can probably get a minor optimisation by inverting this
		for (int[] parms : elliptic) traceElliptic(buf[parms[0]], parms[1], parms[2], parms[3] - L, parms[4], parms[5], Lmodqq, totients[parms[0]]);
		for (int[] parms : hyperbolic) traceHyperbolic(buf[parms[0]], parms[1], parms[2], parms[3] - L, Lmodqq, totients[parms[0]]);

		// We need to filter down to squarefree numbers.
		long pg_base = L * M;
		squarefreeMid(buf, pg_base, 247, 1, 38);
		squarefreeMid(buf, pg_base, 253, 1, 38);
		squarefreeMid(buf, pg_base, 257, 1, 38);
		squarefreeMid(buf, pg_base, 263, 1, 38);
		squarefreeMid(buf, pg_base, 241, 0, 2);
		squarefreeMid(buf, pg_base, 251, 0, 2);
		squarefreeMid(buf, pg_base, 259, 0, 2);
		squarefreeMid(buf, pg_base, 269, 0, 2);

		// Extract the results.
		rv = new long[PI_MB];
		long[] transpose = new long[PHI_M];
		for (int j = 0, off = 0; j < (B >> 6); j++) {
			// Reduce cache locality problems by transposing.
			for (int k = 0; k < PHI_M; k++) transpose[k] = buf[k][j];
			for (long mask = 1L; mask != 0; mask <<= 1) {
				for (int k = 0; k < PHI_M; k++) {
					if ((transpose[k] & mask) == 0) rv[off++] = pg_base + totients[k];
				}

				pg_base += M;
			}
		}

		// TODO Not thread-safe.
		pageCache.put(page, rv);

		return rv;
	}

	// NB This is generalised somewhat from primegen's implementation.
	private static void traceElliptic(final long[] buf, int x, int y, long start, final int cF2, final int cG2, final int[] Lmodqq, final int d) {
		// Bring the annular segment into the range of ints.
		start += 1000000000;
		while (start < 0) {
			start += x;
			x += cF2;
		}
		start -= 1000000000;
		int i = (int)start;

		while (i < B) {
			i += x;
			x += cF2;
		}

		while (true) {
			x -= cF2;
			if (x <= cF2 >> 1) {
				// It makes no sense that doing this in here should perform well, but empirically it does much better than
				// only eliminating the squares once.
				squarefreeTiny(buf, Lmodqq, d);
				return;
			}
			i -= x;

			while (i < 0) {
				i += y;
				y += cG2;
			}

			int i0 = i, y0 = y;
			while (i < B) {
				buf[i >> 6] ^= 1L << i;
				i += y;
				y += cG2;
			}
			i = i0;
			y = y0;
		}
	}

	// This only handles 3x^2 - y^2, and is closer to a direct port of primegen.
	private static void traceHyperbolic(final long[] a, int x, int y, long start, final int[] Lmodqq, final int d) {
		x += 5;
		y += 15;

		// Bring the segment into the range of ints.
		start += 1000000000;
		while (start < 0) {
			start += x;
			x += 10;
		}
		start -= 1000000000;
		int i = (int)start;

		while (i < 0) {
			i += x;
			x += 10;
		}

		while (true) {
			x += 10;
			while (i >= B) {
				if (x <= y) {
					squarefreeTiny(a, Lmodqq, d);
					return;
				}
				i -= y;
				y += 30;
			}

			int i0 = i, y0 = y;
			while (i >= 0 && y < x) {
				a[i >> 6] ^= 1L << i;
				i -= y;
				y += 30;
			}
			i = i0 + x - 10;
			y = y0;
		}
	}

	private static void squarefreeTiny(final long[] a, final int[] Lmodqq, final int d) {
		for (int j = 0; j < qqtab.length; ++j) {
			int qq = qqtab[j];
			int k = qq - 1 - ((Lmodqq[j] + qq60tab[j] * d - 1) % qq);
			while (k < B) {
				a[k >> 6] |= 1L << k;
				k += qq;
			}
		}
	}

	private static void squarefreeMid(long[][] buf, final long base, int q, int dqq, int di) {
		int qq = q * q;
		q = M * q + (M * M / 4);

		while (qq < M * B) {
			int i = qq - (int)(base % qq);
			if ((i & 1) == 0) i += qq;

			if (i < M * B) {
				int qqhigh = ((qq / M) << 1) + dqq;
				int ilow = i % M;
				int ihigh = i / M;
				while (ihigh < B) {
					int n = invTotients[ilow];
					if (n >= 0) buf[n][ihigh >> 6] |= 1L << ihigh;

					ilow += di;
					ihigh += qqhigh;
					if (ilow >= M) {
						ilow -= M;
						ihigh += 1;
					}
				}
			}

			qq += q;
			q += M * M / 2;
		}

		squarefreebig(buf, base, q, qq);
	}

	private static void squarefreebig(long[][] buf, final long base, int q, long qq) {
		long bound = base + M * B;
		while (qq < bound) {
			long i = qq - (base % qq);
			if ((i & 1) == 0) i += qq;

			if (i < M * B) {
				int pos = (int)i;
				int n = invTotients[pos % M];
				if (n >= 0) {
					int ihigh = pos / M;
					buf[n][ihigh >> 6] |= 1L << ihigh;
				}
			}

			qq += q;
			q += M * M / 2;
		}
	}
}

As I hope is evident, I've tried to optimise for performance, and I would particularly appreciate comments on anything I've overlooked there.

posted about 4 years ago

CC BY-NC-SA 4.0

Peter Taylor‭

1302 reputation 7 34 141 7