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