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: ++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(size: 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(objp: 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(dst: new->primes, nbits: sz); |
145 | bitmap_copy(dst: new->primes, src: p->primes, nbits: p->sz); |
146 | for (y = 2UL; y < sz; y = find_next_bit(addr: new->primes, size: sz, offset: y + 1)) |
147 | new->last = clear_multiples(x: y, p: new->primes, start: p->sz, end: 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: &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: &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(addr: p->primes, size: p->last, offset: 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(list: true, buf, maskp: p->primes, nmaskbits: 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(objp: 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(max: 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 | |