xref: /openbmc/linux/lib/math/prime_numbers.c (revision 20e2fc42) 
1 // SPDX-License-Identifier: GPL-2.0-only
2 #define pr_fmt(fmt) "prime numbers: " fmt "\n"
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",
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!",
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",
286 			       last, x, next_prime_number(last));
287 			goto err;
288 		}
289 		last = x;
290 	}
291 
292 	pr_info("selftest(%lu) passed, last prime was %lu", 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