GNU Linux-libre 4.19.211-gnu1
[releases.git] / lib / prime_numbers.c
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");