lib/math/prime_numbers.c

   1 // SPDX-License-Identifier: GPL-2.0-only
   2 #define pr_fmt(fmt) "prime numbers: " fmt
   3
   4 #include <linux/module.h>
   5 #include <linux/mutex.h>
   6 #include <linux/prime_numbers.h>
   7 #include <linux/slab.h>
   8
   9 #define bitmap_size(nbits) (BITS_TO_LONGS(nbits) * sizeof(unsigned long))
  10
  11 struct primes {
  12         struct rcu_head rcu;
  13         unsigned long last, sz;
  14         unsigned long primes[];
  15 };
  16
  17 #if BITS_PER_LONG == 64
  18 static const struct primes small_primes = {
  19         .last = 61,
  20         .sz = 64,
  21         .primes = {
  22                 BIT(2) |
  23                 BIT(3) |
  24                 BIT(5) |
  25                 BIT(7) |
  26                 BIT(11) |
  27                 BIT(13) |
  28                 BIT(17) |
  29                 BIT(19) |
  30                 BIT(23) |
  31                 BIT(29) |
  32                 BIT(31) |
  33                 BIT(37) |
  34                 BIT(41) |
  35                 BIT(43) |
  36                 BIT(47) |
  37                 BIT(53) |
  38                 BIT(59) |
  39                 BIT(61)
  40         }
  41 };
  42 #elif BITS_PER_LONG == 32
  43 static const struct primes small_primes = {
  44         .last = 31,
  45         .sz = 32,
  46         .primes = {
  47                 BIT(2) |
  48                 BIT(3) |
  49                 BIT(5) |
  50                 BIT(7) |
  51                 BIT(11) |
  52                 BIT(13) |
  53                 BIT(17) |
  54                 BIT(19) |
  55                 BIT(23) |
  56                 BIT(29) |
  57                 BIT(31)
  58         }
  59 };
  60 #else
  61 #error "unhandled BITS_PER_LONG"
  62 #endif
  63
  64 static DEFINE_MUTEX(lock);
  65 static const struct primes __rcu *primes = RCU_INITIALIZER(&small_primes);
  66
  67 static unsigned long selftest_max;
  68
  69 static bool slow_is_prime_number(unsigned long x)
  70 {
  71         unsigned long y = int_sqrt(x);
  72
  73         while (y > 1) {
  74                 if ((x % y) == 0)
  75                         break;
  76                 y--;
  77         }
  78
  79         return y == 1;
  80 }
  81
  82 static unsigned long slow_next_prime_number(unsigned long x)
  83 {
  84         while (x < ULONG_MAX && !slow_is_prime_number(++x))
  85                 ;
  86
  87         return x;
  88 }
  89
  90 static unsigned long clear_multiples(unsigned long x,
  91                                      unsigned long *p,
  92                                      unsigned long start,
  93                                      unsigned long end)
  94 {
  95         unsigned long m;
  96
  97         m = 2 * x;
  98         if (m < start)
  99                 m = roundup(start, x);
 100
 101         while (m < end) {
 102                 __clear_bit(m, p);
 103                 m += x;
 104         }
 105
 106         return x;
 107 }
 108
 109 static bool expand_to_next_prime(unsigned long x)
 110 {
 111         const struct primes *p;
 112         struct primes *new;
 113         unsigned long sz, y;
 114
 115         /* Betrand's Postulate (or Chebyshev's theorem) states that if n > 3,
 116          * there is always at least one prime p between n and 2n - 2.
 117          * Equivalently, if n > 1, then there is always at least one prime p
 118          * such that n < p < 2n.
 119          *
 120          * http://mathworld.wolfram.com/BertrandsPostulate.html
 121          * https://en.wikipedia.org/wiki/Bertrand's_postulate
 122          */
 123         sz = 2 * x;
 124         if (sz < x)
 125                 return false;
 126
 127         sz = round_up(sz, BITS_PER_LONG);
 128         new = kmalloc(sizeof(*new) + bitmap_size(sz),
 129                       GFP_KERNEL | __GFP_NOWARN);
 130         if (!new)
 131                 return false;
 132
 133         mutex_lock(&lock);
 134         p = rcu_dereference_protected(primes, lockdep_is_held(&lock));
 135         if (x < p->last) {
 136                 kfree(new);
 137                 goto unlock;
 138         }
 139
 140         /* Where memory permits, track the primes using the
 141          * Sieve of Eratosthenes. The sieve is to remove all multiples of known
 142          * primes from the set, what remains in the set is therefore prime.
 143          */
 144         bitmap_fill(new->primes, sz);
 145         bitmap_copy(new->primes, p->primes, p->sz);
 146         for (y = 2UL; y < sz; y = find_next_bit(new->primes, sz, y + 1))
 147                 new->last = clear_multiples(y, new->primes, p->sz, sz);
 148         new->sz = sz;
 149
 150         BUG_ON(new->last <= x);
 151
 152         rcu_assign_pointer(primes, new);
 153         if (p != &small_primes)
 154                 kfree_rcu((struct primes *)p, rcu);
 155
 156 unlock:
 157         mutex_unlock(&lock);
 158         return true;
 159 }
 160
 161 static void free_primes(void)
 162 {
 163         const struct primes *p;
 164
 165         mutex_lock(&lock);
 166         p = rcu_dereference_protected(primes, lockdep_is_held(&lock));
 167         if (p != &small_primes) {
 168                 rcu_assign_pointer(primes, &small_primes);
 169                 kfree_rcu((struct primes *)p, rcu);
 170         }
 171         mutex_unlock(&lock);
 172 }
 173
 174 /**
 175  * next_prime_number - return the next prime number
 176  * @x: the starting point for searching to test
 177  *
 178  * A prime number is an integer greater than 1 that is only divisible by
 179  * itself and 1.  The set of prime numbers is computed using the Sieve of
 180  * Eratoshenes (on finding a prime, all multiples of that prime are removed
 181  * from the set) enabling a fast lookup of the next prime number larger than
 182  * @x. If the sieve fails (memory limitation), the search falls back to using
 183  * slow trial-divison, up to the value of ULONG_MAX (which is reported as the
 184  * final prime as a sentinel).
 185  *
 186  * Returns: the next prime number larger than @x
 187  */
 188 unsigned long next_prime_number(unsigned long x)
 189 {
 190         const struct primes *p;
 191
 192         rcu_read_lock();
 193         p = rcu_dereference(primes);
 194         while (x >= p->last) {
 195                 rcu_read_unlock();
 196
 197                 if (!expand_to_next_prime(x))
 198                         return slow_next_prime_number(x);
 199
 200                 rcu_read_lock();
 201                 p = rcu_dereference(primes);
 202         }
 203         x = find_next_bit(p->primes, p->last, x + 1);
 204         rcu_read_unlock();
 205
 206         return x;
 207 }
 208 EXPORT_SYMBOL(next_prime_number);
 209
 210 /**
 211  * is_prime_number - test whether the given number is prime
 212  * @x: the number to test
 213  *
 214  * A prime number is an integer greater than 1 that is only divisible by
 215  * itself and 1. Internally a cache of prime numbers is kept (to speed up
 216  * searching for sequential primes, see next_prime_number()), but if the number
 217  * falls outside of that cache, its primality is tested using trial-divison.
 218  *
 219  * Returns: true if @x is prime, false for composite numbers.
 220  */
 221 bool is_prime_number(unsigned long x)
 222 {
 223         const struct primes *p;
 224         bool result;
 225
 226         rcu_read_lock();
 227         p = rcu_dereference(primes);
 228         while (x >= p->sz) {
 229                 rcu_read_unlock();
 230
 231                 if (!expand_to_next_prime(x))
 232                         return slow_is_prime_number(x);
 233
 234                 rcu_read_lock();
 235                 p = rcu_dereference(primes);
 236         }
 237         result = test_bit(x, p->primes);
 238         rcu_read_unlock();
 239
 240         return result;
 241 }
 242 EXPORT_SYMBOL(is_prime_number);
 243
 244 static void dump_primes(void)
 245 {
 246         const struct primes *p;
 247         char *buf;
 248
 249         buf = kmalloc(PAGE_SIZE, GFP_KERNEL);
 250
 251         rcu_read_lock();
 252         p = rcu_dereference(primes);
 253
 254         if (buf)
 255                 bitmap_print_to_pagebuf(true, buf, p->primes, p->sz);
 256         pr_info("primes.{last=%lu, .sz=%lu, .primes[]=...x%lx} = %s\n",
 257                 p->last, p->sz, p->primes[BITS_TO_LONGS(p->sz) - 1], buf);
 258
 259         rcu_read_unlock();
 260
 261         kfree(buf);
 262 }
 263
 264 static int selftest(unsigned long max)
 265 {
 266         unsigned long x, last;
 267
 268         if (!max)
 269                 return 0;
 270
 271         for (last = 0, x = 2; x < max; x++) {
 272                 bool slow = slow_is_prime_number(x);
 273                 bool fast = is_prime_number(x);
 274
 275                 if (slow != fast) {
 276                         pr_err("inconsistent result for is-prime(%lu): slow=%s, fast=%s!\n",
 277                                x, slow ? "yes" : "no", fast ? "yes" : "no");
 278                         goto err;
 279                 }
 280
 281                 if (!slow)
 282                         continue;
 283
 284                 if (next_prime_number(last) != x) {
 285                         pr_err("incorrect result for next-prime(%lu): expected %lu, got %lu\n",
 286                                last, x, next_prime_number(last));
 287                         goto err;
 288                 }
 289                 last = x;
 290         }
 291
 292         pr_info("%s(%lu) passed, last prime was %lu\n", __func__, x, last);
 293         return 0;
 294
 295 err:
 296         dump_primes();
 297         return -EINVAL;
 298 }
 299
 300 static int __init primes_init(void)
 301 {
 302         return selftest(selftest_max);
 303 }
 304
 305 static void __exit primes_exit(void)
 306 {
 307         free_primes();
 308 }
 309
 310 module_init(primes_init);
 311 module_exit(primes_exit);
 312
 313 module_param_named(selftest, selftest_max, ulong, 0400);
 314
 315 MODULE_AUTHOR("Intel Corporation");
 316 MODULE_LICENSE("GPL");