Skip to content

Commit

Permalink
Merge pull request #145 from primo-ppcg/modular-arithmetic
Browse files Browse the repository at this point in the history
Add functions for primality testing
  • Loading branch information
bakpakin committed Aug 10, 2023
2 parents 1d31c68 + 2587b2f commit ce3e437
Show file tree
Hide file tree
Showing 4 changed files with 505 additions and 1 deletion.
4 changes: 4 additions & 0 deletions project.janet
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
129 changes: 128 additions & 1 deletion spork/math.janet
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
(use ./misc)
(import spork/cmath :prefix "" :export true)

# Statistics

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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))))))))
197 changes: 197 additions & 0 deletions src/cmath.c
Original file line number Diff line number Diff line change
@@ -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);
}
Loading

0 comments on commit ce3e437

Please sign in to comment.