Communities

Writing
Writing
Codidact Meta
Codidact Meta
The Great Outdoors
The Great Outdoors
Photography & Video
Photography & Video
Scientific Speculation
Scientific Speculation
Cooking
Cooking
Electrical Engineering
Electrical Engineering
Judaism
Judaism
Languages & Linguistics
Languages & Linguistics
Software Development
Software Development
Mathematics
Mathematics
Christianity
Christianity
Code Golf
Code Golf
Music
Music
Physics
Physics
Linux Systems
Linux Systems
Power Users
Power Users
Tabletop RPGs
Tabletop RPGs

Dashboard
Notifications
Mark all as read
Code Reviews

Improved Atkin-Bernstein sieve for generating primes

+4
−0

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.

Why does this post require moderator attention?
You might want to add some details to your flag.
Why should this post be closed?

1 comment thread

General comments (4 comments)

0 answers

Sign up to answer this question »