Add functions for primality testing

project.janet

@@ -52,3 +52,7 @@
   :source @["src/zip.c" "deps/miniz/miniz.c"]
   :defines @{"_LARGEFILE64_SOURCE" true}
   :headers @["deps/miniz/miniz.h"])
+
+(declare-native
+  :name "spork/cmath"
+  :source @["src/cmath.c"])

spork/math.janet

@@ -1,4 +1,5 @@
 (use ./misc)
+(import spork/cmath :prefix "" :export true)
 
 # Statistics
 
@@ -309,7 +310,7 @@
 
 (defn z-score
   ```
-  Gets the standard score for number `x` from mean `m` 
+  Gets the standard score for number `x` from mean `m`
   and standard deviation `d`.
   ```
   [x m d]
@@ -1168,3 +1169,129 @@
             (for i 0 r
               (+= res (* (get-in m [0 i])
                          (perm (minor m 0 i))))))))
+
+# Number Theory
+
+(def- small-primes
+  "primes less than 212"
+  [2 3 5 7 11 13 17 19 23 29 31 37
+   41 43 47 53 59 61 67 71 73 79 83 89
+   97 101 103 107 109 113 127 131 137 139 149 151
+   157 163 167 173 179 181 191 193 197 199 211])
+
+(def- soe-indices
+  "pre-calced sieve of eratosthenes for 2 3 5 7"
+  [1 11 13 17 19 23 29 31 37 41 43 47
+   53 59 61 67 71 73 79 83 89 97 101 103
+   107 109 113 121 127 131 137 139 143 149 151 157
+   163 167 169 173 179 181 187 191 193 197 199 209])
+
+(def- soe-offsets
+  "distances between sieve values"
+  [10 2 4 2 4 6 2 6 4 2 4 6
+   6 2 6 4 2 6 4 6 8 4 2 4
+   2 4 8 6 4 6 2 4 6 2 6 6
+   4 2 4 6 2 6 4 2 4 2 10 2])
+
+(def- soe-table
+  "tabulated offsets mod 105"
+  [0 10 2 0 4 0 0 0 8 0 0 2 0 4 0
+   0 6 2 0 4 0 0 4 6 0 0 6 0 0 2
+   0 6 2 0 4 0 0 4 6 0 0 2 0 4 2
+   0 6 6 0 0 0 0 6 6 0 0 0 0 4 2
+   0 6 2 0 4 0 0 4 6 0 0 2 0 6 2
+   0 6 0 0 4 0 0 4 6 0 0 2 0 4 8
+   0 0 2 0 10 0 0 4 0 0 0 2 0 4 2])
+
+(def- mr-bases-32
+  "bases for miller-rabin determinism to 2^32"
+  [2 7 61])
+
+(def- mr-bases-64
+  "bases for miller-rabin determinism to 2^64"
+  [2 325 9375 28178 450775 9780504 1795265022])
+
+(def- prime-prod
+  "product of primes from 3 to 43"
+  6541380665835015)
+
+(defn- miller-rabin-prp?
+  ``Performs a Miller-Rabin probable prime test on `n` against all given bases.
+  If no bases are given, sufficient bases are used to ensure that the check
+  is deterministic for all `n` less than 2^63.``
+  [n & bases]
+  (def ps
+    (if (empty? bases)
+      (if (int? n) mr-bases-32 mr-bases-64)
+      bases))
+  (var d (- n 1))
+  (var s 0)
+  (while (zero? (mod d 2))
+    (++ s)
+    (/= d 2))
+  (label result
+    (each p ps
+      (var x (powmod p d n))
+      (when (compare< 1 x (- n 1))
+        (for r 1 s
+          (set x (mulmod x x n))
+          (if-not (compare< 1 x (- n 1)) (break)))
+        (if-not (compare= x (- n 1))
+          (return result false))))
+    true))
+
+(defn prime?
+  "A primality test, deterministic for all `n` less than 2^63."
+  [n]
+  (cond
+    (compare<= n 211) (do
+                        (def m (binary-search n small-primes compare<))
+                        (compare= n (in small-primes m)))
+    (zero? (mod n 2)) false
+    (= 0 (jacobi n prime-prod)) false
+    (miller-rabin-prp? n)))
+
+(defn next-prime
+  "Returns the next prime number strictly larger than `n`."
+  [n]
+  (var x (+ n 1))
+  (if (compare<= x 211)
+    (in small-primes (binary-search x small-primes compare<))
+    (label result
+      (def i (mod x 210))
+      (def m (binary-search i soe-indices compare<))
+      (+= x (- (in soe-indices m) i))
+      (def offs [;(tuple/slice soe-offsets m) ;(tuple/slice soe-offsets 0 m)])
+      (forever
+        (each o offs
+          (if (not= 0 (jacobi x prime-prod))
+            (if (miller-rabin-prp? x) (return result x)))
+          (+= x o))))))
+
+(defn primes
+  "A boundless prime number generator."
+  []
+  (coro
+    (each n small-primes (yield n))
+    (def sieve @{221 13 253 11})
+    (def pg (drop 6 (primes)))
+    (var p (resume pg))
+    (var q (* p p))
+    (var n 211)
+    (forever
+      (each offset soe-offsets
+        (+= n offset)
+        (def stp (get sieve n))
+        (cond
+          stp (do
+                (put sieve n nil)
+                (var nxt (/ n stp))
+                (+= nxt (in soe-table (mod nxt 105)))
+                (while (get sieve (* nxt stp))
+                  (+= nxt (in soe-table (mod nxt 105))))
+                (put sieve (* nxt stp) stp))
+          (< n q) (yield n)
+          (do
+            (put sieve (+ q (* p (in soe-table (mod p 105)))) p)
+            (set p (resume pg))
+            (set q (* p p))))))))

src/cmath.c

@@ -0,0 +1,197 @@
+/*
+* Copyright (c) 2023 Calvin Rose
+*
+* Permission is hereby granted, free of charge, to any person obtaining a copy
+* of this software and associated documentation files (the "Software"), to
+* deal in the Software without restriction, including without limitation the
+* rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+* sell copies of the Software, and to permit persons to whom the Software is
+* furnished to do so, subject to the following conditions:
+*
+* The above copyright notice and this permission notice shall be included in
+* all copies or substantial portions of the Software.
+*
+* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+* IN THE SOFTWARE.
+*/
+
+#include <janet.h>
+
+int64_t _mod_impl(int64_t a, int64_t m) {
+    if (m == 0) return a;
+
+    int64_t x = a % m;
+    return (((a ^ m) < 0) && (x != 0)) ? x + m : x;
+}
+
+int64_t _jacobi_impl(int64_t a, int64_t m) {
+    int64_t res = 2;
+    a = _mod_impl(a, m);
+    while (a != 0) {
+        int64_t l = a & -a;
+        a /= l;
+        res ^= (a & m) ^ ((l % 3) & (m ^ (m >> 1)));
+        int64_t tmpa = a;
+        a = _mod_impl(m, a);
+        m = tmpa;
+    }
+    return (m == 1) ? (res & 2) - 1 : 0;
+}
+
+int64_t _invmod_impl(int64_t a, int64_t m) {
+    int64_t x = 0;
+    int64_t u = 1;
+    int64_t n = (m < 0) ? -m : m;
+    a = _mod_impl(a, n);
+    while (a != 0) {
+        int64_t tmpx = x;
+        x = u;
+        u = tmpx - (n / a) * u;
+        int64_t tmpa = a;
+        a = _mod_impl(n, a);
+        n = tmpa;
+    }
+    return (n == 1) ? _mod_impl(x, m) : 0;
+}
+
+#if defined(__SIZEOF_INT128__)
+
+int64_t _mulmod_impl(int64_t a, int64_t b, int64_t m) {
+    if (m == 0) return a * b;
+
+    int64_t x = (((signed __int128)a) * b) % m;
+    return (((a ^ b ^ m) < 0) && (x != 0)) ? x + m : x;
+}
+
+#elif defined(_WIN64) && defined(_MSC_VER)
+
+#include <intrin.h>
+int64_t _mulmod_impl(int64_t a, int64_t b, int64_t m) {
+    if (m == 0) return a * b;
+
+    int64_t r;
+    int64_t s = _mul128(a, b, &r);
+    (void)_div128(r, s, m, &r);
+    return (((a ^ b ^ m) < 0) && (r != 0)) ? r + m : r;
+}
+
+#else
+
+int64_t _mulmod_impl(int64_t a, int64_t b, int64_t m) {
+    int64_t res = 0;
+    a = _mod_impl(a, m);
+    b = _mod_impl(b, m);
+    if (a < b) {
+        int64_t tmpa = a;
+        a = b;
+        b = tmpa;
+    }
+    if (b < 0) b -= m;
+    while (b > 0) {
+        if ((b & 1) == 1)
+            res += (((m - res - a) ^ m) < 0) ? a - m : a;
+        a += (((m - a - a) ^ m) < 0) ? a - m : a;
+        b >>= 1;
+    }
+    return res;
+}
+
+#endif
+
+Janet wrap_nan() {
+#ifdef NAN
+    return janet_wrap_number(NAN);
+#else
+    return janet_wrap_number(0.0 / 0.0);
+#endif
+}
+
+Janet wrap_result(int64_t a, Janet m) {
+    if (!janet_checktype(m, JANET_ABSTRACT))
+        return janet_wrap_number(a);
+
+    const JanetAbstractType *at = janet_abstract_type(janet_unwrap_abstract(m));
+    int64_t *box = janet_abstract(at, sizeof(int64_t));
+    *box = a;
+    return janet_wrap_abstract(box);
+}
+
+JANET_FN(cfun_cmath_jacobi,
+        "(math/jacobi a m)",
+        "Computes the Jacobi Symbol (a|m).") {
+    janet_fixarity(argc, 2);
+    int64_t a = janet_getinteger64(argv, 0);
+    int64_t m = janet_getinteger64(argv, 1);
+
+    return janet_wrap_number(_jacobi_impl(a, m));
+}
+
+JANET_FN(cfun_cmath_invmod,
+        "(math/invmod a m)",
+        "Modular multiplicative inverse of `a` mod `m`. "
+        "Both arguments must be integer. The return value has the same type as `m`. "
+        "If no inverse exists, returns `math/nan` instead.") {
+    janet_fixarity(argc, 2);
+    int64_t a = janet_getinteger64(argv, 0);
+    int64_t m = janet_getinteger64(argv, 1);
+
+    int64_t res = _invmod_impl(a, m);
+    if (res == 0)
+        return wrap_nan();
+
+    return wrap_result(res, argv[1]);
+}
+
+JANET_FN(cfun_cmath_mulmod,
+        "(math/mulmod a b m)",
+        "Modular multiplication of `a` and `b` mod `m`. "
+        "All arguments must be integer. The return value has the same type as `m`.") {
+    janet_fixarity(argc, 3);
+    int64_t a = janet_getinteger64(argv, 0);
+    int64_t b = janet_getinteger64(argv, 1);
+    int64_t m = janet_getinteger64(argv, 2);
+
+    return wrap_result(_mulmod_impl(a, b, m), argv[2]);
+}
+
+JANET_FN(cfun_cmath_powmod,
+        "(math/powmod a b m)",
+        "Modular exponentiation of `a` to the power of `b` mod `m`. "
+        "All arguments must be integer. The return value has the same type as `m`.") {
+    janet_fixarity(argc, 3);
+    int64_t a = janet_getinteger64(argv, 0);
+    int64_t b = janet_getinteger64(argv, 1);
+    int64_t m = janet_getinteger64(argv, 2);
+
+    int64_t res = 1;
+    if (b < 0) {
+        a = _invmod_impl(a, m);
+        if (a == 0)
+            return wrap_nan();
+        b = -b;
+    }
+    while (b > 0) {
+        if ((b & 1) == 1)
+            res = _mulmod_impl(res, a, m);
+        a = _mulmod_impl(a, a, m);
+        b >>= 1;
+    }
+
+    return wrap_result(res, argv[2]);
+}
+
+JANET_MODULE_ENTRY(JanetTable *env) {
+    JanetRegExt cfuns[] = {
+        JANET_REG("jacobi", cfun_cmath_jacobi),
+        JANET_REG("invmod", cfun_cmath_invmod),
+        JANET_REG("mulmod", cfun_cmath_mulmod),
+        JANET_REG("powmod", cfun_cmath_powmod),
+        JANET_REG_END
+    };
+    janet_cfuns_ext(env, "cmath", cfuns);
+}