prime_numbers.c 6.5 KB

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