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"); 317