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
 10struct primes {
 11	struct rcu_head rcu;
 12	unsigned long last, sz;
 13	unsigned long primes[];
 14};
 15
 16#if BITS_PER_LONG == 64
 17static 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
 42static 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
 63static DEFINE_MUTEX(lock);
 64static const struct primes __rcu *primes = RCU_INITIALIZER(&small_primes);
 65
 66static unsigned long selftest_max;
 67
 68static 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
 81static 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
 89static 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
108static 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
155unlock:
156	mutex_unlock(&lock);
157	return true;
158}
159
160static 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 */
187unsigned 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}
207EXPORT_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 */
220bool 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}
241EXPORT_SYMBOL(is_prime_number);
242
243static 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
263static 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
294err:
295	dump_primes();
296	return -EINVAL;
297}
298
299static int __init primes_init(void)
300{
301	return selftest(selftest_max);
302}
303
304static void __exit primes_exit(void)
305{
306	free_primes();
307}
308
309module_init(primes_init);
310module_exit(primes_exit);
311
312module_param_named(selftest, selftest_max, ulong, 0400);
313
314MODULE_AUTHOR("Intel Corporation");
315MODULE_LICENSE("GPL");