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