src/math/w64.bend source
src/math/w64.bend on the hub · documented module
import Baseimport ./u64.bend as W# Fast unsigned 64-bit arithmetic on the two-limb U64 of u64.bend, for the U64# instance of the generic math. Only native U32 operations and Nat values# below 2^48 (the runtime's Nat bound) are used:## mul32(a, b) the full 64-bit product of two U32, by 16-bit halves# div32(a, d) a / d and a mod d for a 64-bit a and a U32 d > 0, the# low word by two 16-bit digits (partial dividends < 2^48)# divmod(a, b) 64 / 64: for b >= 2^32 the quotient has 32 bits and is# estimated from the top 32 bits of b (never under and# at most 2 over, Knuth's Theorem B), then corrected down# mulmod(a, b, m) a * b mod m for a, b < m: the 128-bit product reduced# twice by the same estimate# isqrt(n) an F32 estimate, one Newton step, then correction## Tested against Python's integers (tools/check_generic.py); not proved.def lo(a: W.U64) -> U32: match a: case W.U64{l, h}: ldef hi(a: W.U64) -> U32: match a: case W.U64{l, h}: hdef b32(b: Bool) -> U32: match b: case True{}: 1 case False{}: 0def mk(+l: U32, +h: U32) -> W.U64: W.U64{l, h}def of32(+x: U32) -> W.U64: W.U64{x, 0}def is_zero(+a: W.U64) -> Bool: Bool.and(U32.is_zero(lo(a)), U32.is_zero(hi(a)))def lt(+a: W.U64, +b: W.U64) -> Bool: Bool.or(U32.is_lt(hi(a), hi(b)), Bool.and(U32.is_eq(hi(a), hi(b)), U32.is_lt(lo(a), lo(b))))def le(+a: W.U64, +b: W.U64) -> Bool: Bool.not(lt(b, a))def eq(+a: W.U64, +b: W.U64) -> Bool: Bool.and(U32.is_eq(lo(a), lo(b)), U32.is_eq(hi(a), hi(b)))# ---- addition and subtraction (wrapping) ----def add_fin(+l: U32, +al: U32, +ah: U32, +bh: U32) -> W.U64: W.U64{l, U32.add(U32.add(ah, bh), b32(U32.is_lt(l, al)))}def add(+a: W.U64, +b: W.U64) -> W.U64: add_fin(U32.add(lo(a), lo(b)), lo(a), hi(a), hi(b))def add_over_fin(+s: U32, +ah: U32, c: Bool) -> Bool: Bool.or(U32.is_lt(s, ah), Bool.and(c, U32.is_eq(s, 4294967295)))# a + b >= 2^64def add_over(+a: W.U64, +b: W.U64) -> Bool: add_over_fin(U32.add(hi(a), hi(b)), hi(a), U32.is_lt(U32.add(lo(a), lo(b)), lo(a)))# a - b mod 2^64def sub(+a: W.U64, +b: W.U64) -> W.U64: W.U64{U32.sub(lo(a), lo(b)), U32.sub(U32.sub(hi(a), hi(b)), b32(U32.is_lt(lo(a), lo(b))))}def half(+a: W.U64) -> W.U64: W.U64{U32.add(U32.shr(lo(a)), U32.mul(U32.and(hi(a), 1), 2147483648)), U32.shr(hi(a))}def odd(+a: W.U64) -> Bool: U32.is_eq(U32.and(lo(a), 1), 1)# ---- products ----def mul32_fin(+ll: U32, +mid: U32, +hh: U32, mc: Bool) -> W.U64: add_fin(U32.add(ll, U32.mul(mid, 65536)), ll, hh, U32.add(U32.shrn(mid, 16n), U32.mul(b32(mc), 65536)))def mul32_mid(+ll: U32, +lh: U32, +hl: U32, +hh: U32) -> W.U64: mul32_fin(ll, U32.add(lh, hl), hh, U32.is_lt(U32.add(lh, hl), lh))def mul32_h(+al: U32, +ah: U32, +bl: U32, +bh: U32) -> W.U64: mul32_mid(U32.mul(al, bl), U32.mul(al, bh), U32.mul(ah, bl), U32.mul(ah, bh))# the full product of two U32def mul32(+a: U32, +b: U32) -> W.U64: mul32_h(U32.and(a, 65535), U32.shrn(a, 16n), U32.and(b, 65535), U32.shrn(b, 16n))# a * b mod 2^64def mul(+a: W.U64, +b: W.U64) -> W.U64: add(mul32(lo(a), lo(b)), W.U64{0, U32.add(U32.mul(lo(a), hi(b)), U32.mul(hi(a), lo(b)))})def mul_over_one(+p: W.U64, +c: W.U64) -> Bool: Bool.or(Bool.not(U32.is_zero(hi(c))), U32.is_lt(U32.add(hi(p), lo(c)), hi(p)))# a * b >= 2^64def mul_over(+a: W.U64, +b: W.U64) -> Bool: Bool.or(Bool.and(Bool.not(U32.is_zero(hi(a))), Bool.not(U32.is_zero(hi(b)))), mul_over_one(mul32(lo(a), lo(b)), add(mul32(hi(a), lo(b)), mul32(lo(a), hi(b)))))# ---- 64 / 32 ----# One digit of the long division by d: (r * base + g) / d and mod d, for a# remainder r < d < 2^32 and a digit g < base = 2^16, so every partial# dividend r * 2^16 + g is below 2^48 (the runtime's Nat bound). The low word# is taken as two 16-bit digits; the base is written 256 * 256 so the proof# checker never expands a 17-bit literal.def dig_t(+r: Nat, +base: Nat, +g: Nat) -> Nat: Nat.add(Nat.mul(r, base), g)def dig1(+lo: U32) -> Nat: U32.to_nat(U32.shrn(lo, 16n))def dig2(+lo: U32) -> Nat: U32.to_nat(U32.and(lo, 65535))def div32_fin(+qh: U32, +q1: Nat, +t2: Nat, +d: Nat) -> W.U64 & U32: (W.U64{U32.from_nat(Nat.add(Nat.mul(q1, 65536n), Nat.div(t2, d))), qh}, U32.from_nat(Nat.mod(t2, d)))def div32_t1(+qh: U32, +lo: U32, +t1: Nat, +d: Nat) -> W.U64 & U32: div32_fin(qh, Nat.div(t1, d), dig_t(Nat.mod(t1, d), 65536n, dig2(lo)), d)# (a / d, a mod d) for d > 0: the high word natively, the low word by digitsdef div32(+a: W.U64, +d: U32) -> W.U64 & U32: div32_t1(U32.div(hi(a), d), lo(a), dig_t(U32.to_nat(U32.mod(hi(a), d)), 65536n, dig1(lo(a))), U32.to_nat(d))def mod_step(+r: Nat, +base: Nat, +g: Nat, +d: Nat) -> Nat: Nat.mod(dig_t(r, base, g), d)# a mod d for d > 0def mod32(+a: W.U64, +d: U32) -> U32: U32.from_nat(mod_step(mod_step(U32.to_nat(U32.mod(hi(a), d)), 65536n, dig1(lo(a)), U32.to_nat(d)), 65536n, dig2(lo(a)), U32.to_nat(d)))# ---- shifts by a variable count ----# 2^k for k < 32def pow2(k: Nat) -> U32: match k: case 0n: 1 case 1n: 2 case 2n: 4 case 3n: 8 case 4n: 16 case 5n: 32 case 6n: 64 case 7n: 128 case 8n: 256 case 9n: 512 case 10n: 1024 case 11n: 2048 case 12n: 4096 case 13n: 8192 case 14n: 16384 case 15n: 32768 case 16n: 65536 case 17n: 131072 case 18n: 262144 case 19n: 524288 case 20n: 1048576 case 21n: 2097152 case 22n: 4194304 case 23n: 8388608 case 24n: 16777216 case 25n: 33554432 case 26n: 67108864 case 27n: 134217728 case 28n: 268435456 case 29n: 536870912 case 30n: 1073741824 case _: 2147483648def bl_pick(+x: U32, +t: U32, +k: Nat, +base: Nat, big: Bool) -> Nat & U32: match big: case True{}: (Nat.add(base, k), U32.shrn(x, k)) case False{}: (base, x)def bl_step(+k: Nat, st: Nat & U32) -> Nat & U32: (+base, +x) = st bl_pick(x, pow2(k), k, base, U32.is_le(pow2(k), x))def bl_fin(st: Nat & U32) -> Nat: (+base, +x) = st Nat.add(base, U32.to_nat(U32.min(x, 1)))# the number of bits of x (0 for 0), by a binary search on the top bitdef bitlen(+x: U32) -> Nat: bl_fin(bl_step(1n, bl_step(2n, bl_step(4n, bl_step(8n, bl_step(16n, (0n, x)))))))def shl_lt(+a: W.U64, +k: Nat) -> W.U64: match k: case 0n: a case 1n+ +j: W.U64{U32.mul(lo(a), pow2(1n+j)), U32.add(U32.mul(hi(a), pow2(1n+j)), U32.shrn(lo(a), Nat.sub(32n, 1n+j)))}def shl_ge(+a: W.U64, +k: Nat) -> W.U64: W.U64{0, U32.mul(lo(a), pow2(Nat.sub(k, 32n)))}def shl_pick(+a: W.U64, +k: Nat, small: Bool) -> W.U64: match small: case True{}: shl_lt(a, k) case False{}: shl_ge(a, k)# a << k mod 2^64 for k < 64def shl(+a: W.U64, +k: Nat) -> W.U64: shl_pick(a, k, Nat.is_lt(k, 32n))def shr_lt(+a: W.U64, +k: Nat) -> W.U64: match k: case 0n: a case 1n+ +j: W.U64{U32.add(U32.shrn(lo(a), 1n+j), U32.mul(hi(a), pow2(Nat.sub(32n, 1n+j)))), U32.shrn(hi(a), 1n+j)}def shr_ge(+a: W.U64, +k: Nat, huge: Bool) -> W.U64: match huge: case True{}: W.U64{0, 0} case False{}: W.U64{U32.shrn(hi(a), Nat.sub(k, 32n)), 0}def shr_pick(+a: W.U64, +k: Nat, small: Bool) -> W.U64: match small: case True{}: shr_lt(a, k) case False{}: shr_ge(a, k, Nat.is_le(64n, k))# a >> k (0 for k >= 64)def shr(+a: W.U64, +k: Nat) -> W.U64: shr_pick(a, k, Nat.is_lt(k, 32n))def fst_q(p: W.U64 & U32) -> W.U64: (q, r) = p qdef snd_r(p: W.U64 & U32) -> U32: (q, r) = p rdef q_clamp_z(+l: U32, z: Bool) -> U32: match z: case True{}: l case False{}: 4294967295# min(q, 2^32 - 1)def q_clamp(+q: W.U64) -> U32: q_clamp_z(lo(q), U32.is_zero(hi(q)))# the product of a U32 and a U64 as (low 64 bits, top 32 bits)def mul_32_64_fin(+p0: W.U64, +p1: W.U64) -> W.U64 & U32: (W.U64{lo(p0), U32.add(hi(p0), lo(p1))}, U32.add(hi(p1), b32(U32.is_lt(U32.add(hi(p0), lo(p1)), hi(p0)))))def mul_32_64(+q: U32, +b: W.U64) -> W.U64 & U32: mul_32_64_fin(mul32(q, lo(b)), mul32(q, hi(b)))# the 96-bit product p exceeds x = xh * 2^64 + xldef over96(+xl: W.U64, +xh: U32, p: W.U64 & U32) -> Bool: (+pl, +ph) = p Bool.or(U32.is_lt(xh, ph), Bool.and(U32.is_eq(xh, ph), lt(xl, pl)))# q - 1 while q * b > x: at most q steps (fuel q + 1)def q_fix(fuel: Nat, +xl: W.U64, +xh: U32, +b: W.U64, +q: U32, over: Bool) -> U32: match fuel over: case 0n _: q case 1n+f False{}: q case 1n+f True{}: q_fix(f, xl, xh, b, U32.sub(q, 1), over96(xl, xh, mul_32_64(U32.sub(q, 1), b)))# floor(x / b) from an estimate q >= floor(x / b): down while q * b > x (at# most q steps, fuel q + 1)def q_start(+xl: W.U64, +xh: U32, +b: W.U64, +q: U32) -> U32: q_fix(Nat.add(U32.to_nat(q), 1n), xl, xh, b, q, over96(xl, xh, mul_32_64(q, b)))# Knuth's estimate for x = xh * 2^64 + xl < b * 2^32 and b >= 2^32, t =# bits(hi(b)): with Y = x >> t (below 2^64) and B = b >> t in [2^31, 2^32),# min(floor(Y / B), 2^32 - 1) is at least floor(x / b) (as B 2^t <= b) and at# most 2 over it (Theorem B), so q_start takes a step or twodef est_lo(+q1: Nat, +t2: Nat, +d: Nat) -> U32: U32.from_nat(Nat.add(Nat.mul(q1, 65536n), Nat.div(t2, d)))def est_mid(+lo: U32, +t1: Nat, +d: Nat) -> U32: est_lo(Nat.div(t1, d), dig_t(Nat.mod(t1, d), 65536n, dig2(lo)), d)# min(floor(y / d), 2^32 - 1) for d > 0: when hi(y) >= d the quotient needs# more than 32 bits; otherwise the high word is its own remainder, and div32's# two 16-bit Nat digit steps finish the division without dividing the high# word.def est_pick(+y: W.U64, +d: U32, big: Bool) -> U32: match big: case True{}: 4294967295 case False{}: est_mid(lo(y), dig_t(U32.to_nat(hi(y)), 65536n, dig1(lo(y))), U32.to_nat(d))def est32(+y: W.U64, +d: U32) -> U32: est_pick(y, d, U32.is_le(d, hi(y)))def q_est(+xl: W.U64, +xh: U32, +b: W.U64, +t: Nat) -> U32: est32(add(shr(xl, t), shl(W.U64{xh, 0}, Nat.sub(64n, t))), lo(shr(b, t)))# floor(x / b) for x = xh * 2^64 + xl < b * 2^32 and b >= 2^32, t = bits(hi(b))def q96(+xl: W.U64, +xh: U32, +b: W.U64, +t: Nat) -> U32: q_start(xl, xh, b, q_est(xl, xh, b, t))# ---- 64 / 64 ----def dm_small(p: W.U64 & U32) -> W.U64 & W.U64: (+q, +r) = p (q, W.U64{r, 0})def dm_big(+a: W.U64, +b: W.U64, +q: U32) -> W.U64 & W.U64: (W.U64{q, 0}, sub(a, fst_q(mul_32_64(q, b))))def dm_pick(+a: W.U64, +b: W.U64, small: Bool) -> W.U64 & W.U64: match small: case True{}: dm_small(div32(a, lo(b))) case False{}: dm_big(a, b, q96(a, 0, b, bitlen(hi(b))))# (a / b, a mod b) for b > 0def divmod(+a: W.U64, +b: W.U64) -> W.U64 & W.U64: dm_pick(a, b, U32.is_zero(hi(b)))def pfst(p: W.U64 & W.U64) -> W.U64: (q, r) = p qdef psnd(p: W.U64 & W.U64) -> W.U64: (q, r) = p rdef quot(+a: W.U64, +b: W.U64) -> W.U64: pfst(divmod(a, b))def rem(+a: W.U64, +b: W.U64) -> W.U64: psnd(divmod(a, b))# ---- a * b mod m ----# x mod m for x = xh * 2^32 + x0 with xh < m and m >= 2^32def red96_z(+xh: W.U64, +x0: U32, +m: W.U64, small: Bool) -> W.U64: match small: case True{}: W.U64{x0, lo(xh)} case False{}: sub(W.U64{x0, lo(xh)}, fst_q(mul_32_64(q96(W.U64{x0, lo(xh)}, hi(xh), m, bitlen(hi(m))), m)))# (the zero test of hi(m) only guards the digit: red96 is called with# hi(m) != 0; it also stops the proof checker from expanding the digit in# every statement that names red96)def red96(+xh: W.U64, +x0: U32, +m: W.U64) -> W.U64: red96_z(xh, x0, m, U32.is_zero(hi(m)))def mul128_fin(+p00: W.U64, +mid: W.U64, c1: Bool, +p11: W.U64) -> W.U64 & W.U64: (W.U64{lo(p00), U32.add(hi(p00), lo(mid))}, add(add(p11, W.U64{hi(mid), b32(c1)}), W.U64{b32(U32.is_lt(U32.add(hi(p00), lo(mid)), hi(p00))), 0}))def mul128_mid(+p00: W.U64, +p01: W.U64, +p10: W.U64, +p11: W.U64) -> W.U64 & W.U64: mul128_fin(p00, add(p01, p10), add_over(p01, p10), p11)# the full product as (low 64 bits, high 64 bits)def mul128(+a: W.U64, +b: W.U64) -> W.U64 & W.U64: mul128_mid(mul32(lo(a), lo(b)), mul32(lo(a), hi(b)), mul32(hi(a), lo(b)), mul32(hi(a), hi(b)))def mm_big(+m: W.U64, p: W.U64 & W.U64) -> W.U64: (+pl, +ph) = p red96(red96(ph, hi(pl), m), lo(pl), m)def mm_pick(+a: W.U64, +b: W.U64, +m: W.U64, small: Bool) -> W.U64: match small: case True{}: W.U64{mod32(mul32(lo(a), lo(b)), lo(m)), 0} case False{}: mm_big(m, mul128(a, b))# a * b mod m for a, b < mdef mulmod(+a: W.U64, +b: W.U64, +m: W.U64) -> W.U64: mm_pick(a, b, m, U32.is_zero(hi(m)))# ---- integer square root ----# r - 1 while r * r > x: from r <= 65535 at most r steps (fuel r + 1), and# r * r never wrapsdef down32(fuel: Nat, +x: U32, +r: U32, over: Bool) -> U32: match fuel over: case 0n _: r case 1n+f False{}: r case 1n+f True{}: down32(f, x, U32.sub(r, 1), U32.is_lt(x, U32.mul(U32.sub(r, 1), U32.sub(r, 1))))def up_ok(+x: U32, +r: U32) -> Bool: Bool.and(U32.is_lt(r, 65535), U32.is_le(U32.mul(U32.add(r, 1), U32.add(r, 1)), x))# r + 1 while r < 65535 and (r + 1)^2 <= x: at most 65535 - r stepsdef up32(fuel: Nat, +x: U32, +r: U32, up: Bool) -> U32: match fuel up: case 0n _: r case 1n+f False{}: r case 1n+f True{}: up32(f, x, U32.add(r, 1), up_ok(x, U32.add(r, 1)))def isqrt32_fix(+x: U32, +r: U32) -> U32: up32(1n+U32.to_nat(U32.sub(65535, r)), x, r, up_ok(x, r))# floor(sqrt(x)) from any estimate r <= 65535: down while r^2 > x, then up# while (r + 1)^2 <= x. The F32 square root lands within 1, so both loops# take at most a step or two; the fuel covers any estimate, so the result# never depends on F32's accuracy.def isqrt32_est(+x: U32, +r: U32) -> U32: isqrt32_fix(x, down32(1n+U32.to_nat(r), x, r, U32.is_lt(x, U32.mul(r, r))))def isqrt32(+x: U32) -> U32: isqrt32_est(x, U32.min(F32.to_u32(F32.sqrt(U32.to_f32(x))), 65535))def sq_over(+n: W.U64, +r: U32) -> Bool: lt(n, mul32(r, r))def down64(fuel: Nat, +n: W.U64, +r: U32, over: Bool) -> U32: match fuel over: case 0n _: r case 1n+f False{}: r case 1n+f True{}: down64(f, n, U32.sub(r, 1), sq_over(n, U32.sub(r, 1)))def up64_ok(+n: W.U64, +r: U32) -> Bool: Bool.and(U32.is_lt(r, 4294967295), Bool.not(sq_over(n, U32.add(r, 1))))def up64(fuel: Nat, +n: W.U64, +r: U32, up: Bool) -> U32: match fuel up: case 0n _: r case 1n+f False{}: r case 1n+f True{}: up64(f, n, U32.add(r, 1), up64_ok(n, U32.add(r, 1)))def isqrt64_fix(+n: W.U64, +r: U32) -> U32: up64(1n+U32.to_nat(U32.sub(4294967295, r)), n, r, up64_ok(n, r))# floor(sqrt(n)) from any estimate r1: down while r^2 > n, then up while# (r + 1)^2 <= n; the fuel covers any estimatedef isqrt_newton(+n: W.U64, +r1: U32) -> U32: isqrt64_fix(n, down64(1n+U32.to_nat(r1), n, r1, sq_over(n, r1)))# one Newton step from the F32 estimate r0 > 0 lands within a few units of# floor(sqrt(n)); the correction makes the result exact for any r0def isqrt_big(+n: W.U64, +r0: U32) -> W.U64: W.U64{isqrt_newton(n, q_clamp(half(add(fst_q(div32(n, r0)), W.U64{r0, 0})))), 0}def f32_clamp32(+f: F32, big: Bool) -> U32: match big: case True{}: 4294967295 case False{}: F32.to_u32(f)def isqrt_est(+f: F32) -> U32: U32.max(f32_clamp32(f, F32.is_le(4294967040.0, f)), 1)def isqrt_pick(+n: W.U64, small: Bool) -> W.U64: match small: case True{}: W.U64{isqrt32(lo(n)), 0} case False{}: isqrt_big(n, isqrt_est(F32.sqrt(F32.add(F32.mul(U32.to_f32(hi(n)), 4294967296.0), U32.to_f32(lo(n))))))def isqrt(+n: W.U64) -> W.U64: isqrt_pick(n, U32.is_zero(hi(n)))def w_pow2(+k: Nat, small: Bool) -> W.U64: match small: case True{}: W.U64{pow2(k), 0} case False{}: W.U64{0, pow2(Nat.sub(k, 32n))}# ---- general shifts and leading zeros (the software F64) ----def or_bit(+a: W.U64, b: Bool) -> W.U64: W.U64{U32.or(lo(a), b32(b)), hi(a)}def jam_pick(+a: W.U64, +k: Nat, huge: Bool) -> W.U64: match huge: case True{}: W.U64{b32(Bool.not(is_zero(a))), 0} case False{}: or_bit(shr(a, k), Bool.not(eq(shl(shr(a, k), k), a)))# a >> k with the bits shifted out OR-ed into bit 0 (SoftFloat's# shiftRightJam64)def shr_jam(+a: W.U64, +k: Nat) -> W.U64: jam_pick(a, k, Nat.is_le(64n, k))def clz_pick(+a: W.U64, high: Bool) -> Nat: match high: case True{}: Nat.sub(32n, bitlen(hi(a))) case False{}: Nat.sub(64n, bitlen(lo(a)))# the number of leading zero bits of a (64 for 0)def clz(+a: W.U64) -> Nat: clz_pick(a, Bool.not(U32.is_zero(hi(a))))# a with bit 0 cleared when bdef clear0(+a: W.U64, b: Bool) -> W.U64: W.U64{U32.sub(lo(a), U32.and(lo(a), b32(b))), hi(a)}def cmp_pick(less: Bool, same: Bool) -> Cmp: match less same: case True{} _: LT{} case False{} True{}: EQ{} case False{} False{}: GT{}def cmp(+a: W.U64, +b: W.U64) -> Cmp: cmp_pick(lt(a, b), eq(a, b))# ---- Montgomery multiplication (odd m, 2^32 <= m < 2^63, R = 2^64) ----# REDC (Montgomery 1985; the word-level algorithm HACL* and Fiat Crypto# verify): for t < m^2 and mp = -1/m mod 2^64, u = t mp mod 2^64 makes# t + u m divisible by 2^64, and r = (t + u m) / 2^64 < 2 m is t / 2^64 mod m# after one conditional subtraction.# x (2 - m x): doubles the correct low bits of an inverse of m mod 2^64def inv_step(+m: W.U64, +x: W.U64) -> W.U64: mul(x, sub(W.U64{2, 0}, mul(m, x)))# -1 / m mod 2^64 for odd m (m is its own inverse mod 8; five Newton steps# reach 96 >= 64 bits)def minv(+m: W.U64) -> W.U64: sub(W.U64{0, 0}, inv_step(m, inv_step(m, inv_step(m, inv_step(m, inv_step(m, m))))))def sub_if(+r: W.U64, +m: W.U64, big: Bool) -> W.U64: match big: case True{}: sub(r, m) case False{}: rdef redc_r(+m: W.U64, +r: W.U64) -> W.U64: sub_if(r, m, le(m, r))def redc_p(+m: W.U64, +tl: W.U64, +th: W.U64, p: W.U64 & W.U64) -> W.U64: (+pl, +ph) = p redc_r(m, add(add(th, ph), W.U64{b32(add_over(tl, pl)), 0}))# t / 2^64 mod m for t = (tl, th) < m^2def redc(+m: W.U64, +mp: W.U64, t: W.U64 & W.U64) -> W.U64: (+tl, +th) = t redc_p(m, tl, th, mul128(mul(tl, mp), m))# a b / 2^64 mod m for a, b < mdef mont(+m: W.U64, +mp: W.U64, +a: W.U64, +b: W.U64) -> W.U64: redc(m, mp, mul128(a, b))def mbit(odd: Bool, +m: W.U64, +mp: W.U64, +b: W.U64, +acc: W.U64) -> W.U64: match odd: case False{}: acc case True{}: mont(m, mp, acc, b)# pow_mod_go of generic.bend on Montgomery forms x 2^64 mod mdef mpow_go(fuel: Nat, +m: W.U64, +mp: W.U64, +b: W.U64, +acc: W.U64, st: W.U64 & Bool) -> W.U64: match fuel st: case 0n _: acc case 1n+f Tuple{e, True{}}: acc case 1n+f Tuple{+e, False{}}: mpow_go(f, m, mp, mont(m, mp, b, b), mbit(odd(e), m, mp, b, acc), (half(e), is_zero(half(e))))# x 2^64 mod m for x < m, m >= 2^32def to_mont(+x: W.U64, +m: W.U64) -> W.U64: red96(red96(x, 0, m), 0, m)# m odd with 2^32 <= m < 2^63 and m minv(m) = -1 mod 2^64 (always so for odd# m; checked, so the proof needs no Hensel lemma)def mont_ok(+m: W.U64) -> Bool: Bool.and(Bool.and(odd(m), eq(mul(m, minv(m)), W.U64{4294967295, 4294967295})), Bool.and(Bool.not(U32.is_zero(hi(m))), U32.is_lt(hi(m), 2147483648)))def mpow_m(+b: W.U64, +e: W.U64, +m: W.U64, +mp: W.U64) -> W.U64: mont(m, mp, mpow_go(140n, m, mp, to_mont(b, m), to_mont(rem(W.U64{1, 0}, m), m), (e, is_zero(e))), W.U64{1, 0})# b^e mod m for b < m and mont_ok(m)def mpow(+b: W.U64, +e: W.U64, +m: W.U64) -> W.U64: mpow_m(b, e, m, minv(m))# ---- 32-bit Montgomery multiplication (odd m, R = 2^32) ----def inv32_step(+m: U32, +x: U32) -> U32: U32.mul(x, U32.sub(2, U32.mul(m, x)))# -1 / m mod 2^32 for odd mdef minv32(+m: U32) -> U32: U32.sub(0, inv32_step(m, inv32_step(m, inv32_step(m, inv32_step(m, m)))))def sub_if32(+r: W.U64, +m: U32, big: Bool) -> U32: match big: case True{}: lo(sub(r, W.U64{m, 0})) case False{}: lo(r)def redc32_r(+m: U32, +r: W.U64) -> U32: sub_if32(r, m, le(W.U64{m, 0}, r))def redc32_s(+m: U32, +s1: U32, +s2: U32, +c1: Bool) -> U32: redc32_r(m, W.U64{s2, U32.add(b32(c1), b32(U32.is_lt(s2, s1)))})def redc32_p(+m: U32, +t: W.U64, +p: W.U64) -> U32: redc32_s(m, U32.add(hi(t), hi(p)), U32.add(U32.add(hi(t), hi(p)), b32(U32.is_lt(U32.add(lo(t), lo(p)), lo(t)))), U32.is_lt(U32.add(hi(t), hi(p)), hi(t)))def redc32(+m: U32, +mp: U32, +t: W.U64) -> U32: redc32_p(m, t, mul32(U32.mul(lo(t), mp), m))def mont32(+m: U32, +mp: U32, +a: U32, +b: U32) -> U32: redc32(m, mp, mul32(a, b))def mbit32(odd: Bool, +m: U32, +mp: U32, +b: U32, +acc: U32) -> U32: match odd: case False{}: acc case True{}: mont32(m, mp, acc, b)def mpow32_go(fuel: Nat, +m: U32, +mp: U32, +b: U32, +acc: U32, st: U32 & Bool) -> U32: match fuel st: case 0n _: acc case 1n+f Tuple{e, True{}}: acc case 1n+f Tuple{+e, False{}}: mpow32_go(f, m, mp, mont32(m, mp, b, b), mbit32(U32.is_eq(U32.and(e, 1), 1), m, mp, b, acc), (U32.shr(e), U32.is_zero(U32.shr(e))))def mont32_ok(+m: U32) -> Bool: Bool.and(U32.is_eq(U32.and(m, 1), 1), U32.is_eq(U32.mul(m, minv32(m)), 4294967295))def mpow32_m(+b: U32, +e: U32, +m: U32, +mp: U32) -> U32: mont32(m, mp, mpow32_go(140n, m, mp, mod32(W.U64{0, b}, m), mod32(W.U64{0, U32.mod(1, m)}, m), (e, U32.is_zero(e))), 1)def mpow32(+b: U32, +e: U32, +m: U32) -> U32: mpow32_m(b, e, m, minv32(m))# a and b fit 48 bits (runtime Nats) and a has more than 16: the binary gcd# then hands them to Euclid on Nat (native division); below 2^16 its few# steps are cheaperdef small2(+a: W.U64, +b: W.U64) -> Bool: Bool.and(Bool.and(U32.is_lt(hi(a), 65536), U32.is_lt(hi(b), 65536)), Bool.or(Bool.not(U32.is_zero(hi(a))), U32.is_lt(65535, lo(a))))def n48(+a: W.U64) -> Nat: Nat.add(U32.to_nat(lo(a)), Nat.mul(Nat.mul(U32.to_nat(hi(a)), 65536n), 65536n))# the words of g < 2^64 from h = g div 2^32 (two divisions by 2^16, so no# closed 2^32 appears: the proof checker would expand it in unary)def of48_h(+g: Nat, +h: Nat) -> W.U64: W.U64{U32.from_nat(Nat.sub(g, Nat.mul(Nat.mul(h, 65536n), 65536n))), U32.from_nat(h)}def of48(+g: Nat) -> W.U64: of48_h(g, Nat.div(Nat.div(g, 65536n), 65536n))