lib/math/prime_numbers.c

0001 // SPDX-License-Identifier: GPL-2.0-only
0002 #define pr_fmt(fmt) "prime numbers: " fmt
0003
0004 #include <linux/module.h>
0005 #include <linux/mutex.h>
0006 #include <linux/prime_numbers.h>
0007 #include <linux/slab.h>
0008
0009 #define bitmap_size(nbits) (BITS_TO_LONGS(nbits) * sizeof(unsigned long))
0010
0011 struct primes {
0012     struct rcu_head rcu;
0013     unsigned long last, sz;
0014     unsigned long primes[];
0015 };
0016
0017 #if BITS_PER_LONG == 64
0018 static const struct primes small_primes = {
0019     .last = 61,
0020     .sz = 64,
0021     .primes = {
0022         BIT(2) |
0023         BIT(3) |
0024         BIT(5) |
0025         BIT(7) |
0026         BIT(11) |
0027         BIT(13) |
0028         BIT(17) |
0029         BIT(19) |
0030         BIT(23) |
0031         BIT(29) |
0032         BIT(31) |
0033         BIT(37) |
0034         BIT(41) |
0035         BIT(43) |
0036         BIT(47) |
0037         BIT(53) |
0038         BIT(59) |
0039         BIT(61)
0040     }
0041 };
0042 #elif BITS_PER_LONG == 32
0043 static const struct primes small_primes = {
0044     .last = 31,
0045     .sz = 32,
0046     .primes = {
0047         BIT(2) |
0048         BIT(3) |
0049         BIT(5) |
0050         BIT(7) |
0051         BIT(11) |
0052         BIT(13) |
0053         BIT(17) |
0054         BIT(19) |
0055         BIT(23) |
0056         BIT(29) |
0057         BIT(31)
0058     }
0059 };
0060 #else
0061 #error "unhandled BITS_PER_LONG"
0062 #endif
0063
0064 static DEFINE_MUTEX(lock);
0065 static const struct primes __rcu *primes = RCU_INITIALIZER(&small_primes);
0066
0067 static unsigned long selftest_max;
0068
0069 static bool slow_is_prime_number(unsigned long x)
0070 {
0071     unsigned long y = int_sqrt(x);
0072
0073     while (y > 1) {
0074         if ((x % y) == 0)
0075             break;
0076         y--;
0077     }
0078
0079     return y == 1;
0080 }
0081
0082 static unsigned long slow_next_prime_number(unsigned long x)
0083 {
0084     while (x < ULONG_MAX && !slow_is_prime_number(++x))
0085         ;
0086
0087     return x;
0088 }
0089
0090 static unsigned long clear_multiples(unsigned long x,
0091                      unsigned long *p,
0092                      unsigned long start,
0093                      unsigned long end)
0094 {
0095     unsigned long m;
0096
0097     m = 2 * x;
0098     if (m < start)
0099         m = roundup(start, x);
0100
0101     while (m < end) {
0102         __clear_bit(m, p);
0103         m += x;
0104     }
0105
0106     return x;
0107 }
0108
0109 static bool expand_to_next_prime(unsigned long x)
0110 {
0111     const struct primes *p;
0112     struct primes *new;
0113     unsigned long sz, y;
0114
0115     /* Betrand's Postulate (or Chebyshev's theorem) states that if n > 3,
0116      * there is always at least one prime p between n and 2n - 2.
0117      * Equivalently, if n > 1, then there is always at least one prime p
0118      * such that n < p < 2n.
0119      *
0120      * http://mathworld.wolfram.com/BertrandsPostulate.html
0121      * https://en.wikipedia.org/wiki/Bertrand's_postulate
0122      */
0123     sz = 2 * x;
0124     if (sz < x)
0125         return false;
0126
0127     sz = round_up(sz, BITS_PER_LONG);
0128     new = kmalloc(sizeof(*new) + bitmap_size(sz),
0129               GFP_KERNEL | __GFP_NOWARN);
0130     if (!new)
0131         return false;
0132
0133     mutex_lock(&lock);
0134     p = rcu_dereference_protected(primes, lockdep_is_held(&lock));
0135     if (x < p->last) {
0136         kfree(new);
0137         goto unlock;
0138     }
0139
0140     /* Where memory permits, track the primes using the
0141      * Sieve of Eratosthenes. The sieve is to remove all multiples of known
0142      * primes from the set, what remains in the set is therefore prime.
0143      */
0144     bitmap_fill(new->primes, sz);
0145     bitmap_copy(new->primes, p->primes, p->sz);
0146     for (y = 2UL; y < sz; y = find_next_bit(new->primes, sz, y + 1))
0147         new->last = clear_multiples(y, new->primes, p->sz, sz);
0148     new->sz = sz;
0149
0150     BUG_ON(new->last <= x);
0151
0152     rcu_assign_pointer(primes, new);
0153     if (p != &small_primes)
0154         kfree_rcu((struct primes *)p, rcu);
0155
0156 unlock:
0157     mutex_unlock(&lock);
0158     return true;
0159 }
0160
0161 static void free_primes(void)
0162 {
0163     const struct primes *p;
0164
0165     mutex_lock(&lock);
0166     p = rcu_dereference_protected(primes, lockdep_is_held(&lock));
0167     if (p != &small_primes) {
0168         rcu_assign_pointer(primes, &small_primes);
0169         kfree_rcu((struct primes *)p, rcu);
0170     }
0171     mutex_unlock(&lock);
0172 }
0173
0174 /**
0175  * next_prime_number - return the next prime number
0176  * @x: the starting point for searching to test
0177  *
0178  * A prime number is an integer greater than 1 that is only divisible by
0179  * itself and 1.  The set of prime numbers is computed using the Sieve of
0180  * Eratoshenes (on finding a prime, all multiples of that prime are removed
0181  * from the set) enabling a fast lookup of the next prime number larger than
0182  * @x. If the sieve fails (memory limitation), the search falls back to using
0183  * slow trial-divison, up to the value of ULONG_MAX (which is reported as the
0184  * final prime as a sentinel).
0185  *
0186  * Returns: the next prime number larger than @x
0187  */
0188 unsigned long next_prime_number(unsigned long x)
0189 {
0190     const struct primes *p;
0191
0192     rcu_read_lock();
0193     p = rcu_dereference(primes);
0194     while (x >= p->last) {
0195         rcu_read_unlock();
0196
0197         if (!expand_to_next_prime(x))
0198             return slow_next_prime_number(x);
0199
0200         rcu_read_lock();
0201         p = rcu_dereference(primes);
0202     }
0203     x = find_next_bit(p->primes, p->last, x + 1);
0204     rcu_read_unlock();
0205
0206     return x;
0207 }
0208 EXPORT_SYMBOL(next_prime_number);
0209
0210 /**
0211  * is_prime_number - test whether the given number is prime
0212  * @x: the number to test
0213  *
0214  * A prime number is an integer greater than 1 that is only divisible by
0215  * itself and 1. Internally a cache of prime numbers is kept (to speed up
0216  * searching for sequential primes, see next_prime_number()), but if the number
0217  * falls outside of that cache, its primality is tested using trial-divison.
0218  *
0219  * Returns: true if @x is prime, false for composite numbers.
0220  */
0221 bool is_prime_number(unsigned long x)
0222 {
0223     const struct primes *p;
0224     bool result;
0225
0226     rcu_read_lock();
0227     p = rcu_dereference(primes);
0228     while (x >= p->sz) {
0229         rcu_read_unlock();
0230
0231         if (!expand_to_next_prime(x))
0232             return slow_is_prime_number(x);
0233
0234         rcu_read_lock();
0235         p = rcu_dereference(primes);
0236     }
0237     result = test_bit(x, p->primes);
0238     rcu_read_unlock();
0239
0240     return result;
0241 }
0242 EXPORT_SYMBOL(is_prime_number);
0243
0244 static void dump_primes(void)
0245 {
0246     const struct primes *p;
0247     char *buf;
0248
0249     buf = kmalloc(PAGE_SIZE, GFP_KERNEL);
0250
0251     rcu_read_lock();
0252     p = rcu_dereference(primes);
0253
0254     if (buf)
0255         bitmap_print_to_pagebuf(true, buf, p->primes, p->sz);
0256     pr_info("primes.{last=%lu, .sz=%lu, .primes[]=...x%lx} = %s\n",
0257         p->last, p->sz, p->primes[BITS_TO_LONGS(p->sz) - 1], buf);
0258
0259     rcu_read_unlock();
0260
0261     kfree(buf);
0262 }
0263
0264 static int selftest(unsigned long max)
0265 {
0266     unsigned long x, last;
0267
0268     if (!max)
0269         return 0;
0270
0271     for (last = 0, x = 2; x < max; x++) {
0272         bool slow = slow_is_prime_number(x);
0273         bool fast = is_prime_number(x);
0274
0275         if (slow != fast) {
0276             pr_err("inconsistent result for is-prime(%lu): slow=%s, fast=%s!\n",
0277                    x, slow ? "yes" : "no", fast ? "yes" : "no");
0278             goto err;
0279         }
0280
0281         if (!slow)
0282             continue;
0283
0284         if (next_prime_number(last) != x) {
0285             pr_err("incorrect result for next-prime(%lu): expected %lu, got %lu\n",
0286                    last, x, next_prime_number(last));
0287             goto err;
0288         }
0289         last = x;
0290     }
0291
0292     pr_info("%s(%lu) passed, last prime was %lu\n", __func__, x, last);
0293     return 0;
0294
0295 err:
0296     dump_primes();
0297     return -EINVAL;
0298 }
0299
0300 static int __init primes_init(void)
0301 {
0302     return selftest(selftest_max);
0303 }
0304
0305 static void __exit primes_exit(void)
0306 {
0307     free_primes();
0308 }
0309
0310 module_init(primes_init);
0311 module_exit(primes_exit);
0312
0313 module_param_named(selftest, selftest_max, ulong, 0400);
0314
0315 MODULE_AUTHOR("Intel Corporation");
0316 MODULE_LICENSE("GPL");