2
\$\begingroup\$

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, but at the moment they're a long way down my priority list.

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;
        }
    }
}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.