src/math/w64.bend

bend-collections-laws-math@1.0.0.0 · 715 lines · 131 definitions

1import Base
2import bend-collections-laws-math@1.0.0.0/src/math/u64.bend as W
3
4# Fast unsigned 64-bit arithmetic on the two-limb U64 of u64.bend, for the U64
5# instance of the generic math. Only native U32 operations and Nat values
6# below 2^48 (the runtime's Nat bound) are used:
7#
8#   mul32(a, b)        the full 64-bit product of two U32, by 16-bit halves
9#   div32(a, d)        a / d and a mod d for a 64-bit a and a U32 d > 0, the
10#                      low word by two 16-bit digits (partial dividends < 2^48)
11#   divmod(a, b)       64 / 64: for b >= 2^32 the quotient has 32 bits and is
12#                      estimated from the top 32 bits of b (never under and
13#                      at most 2 over, Knuth's Theorem B), then corrected down
14#   mulmod(a, b, m)    a * b mod m for a, b < m: the 128-bit product reduced
15#                      twice by the same estimate
16#   isqrt(n)           an F32 estimate, one Newton step, then correction
17#
18# Tested against Python's integers (tools/check_generic.py); not proved.
19
20def lo(a: W.U64) -> U32:
21  match a:
22    case W.U64{l, h}:
23      l
24
25def hi(a: W.U64) -> U32:
26  match a:
27    case W.U64{l, h}:
28      h
29
30def b32(b: Bool) -> U32:
31  match b:
32    case True{}:
33      1
34    case False{}:
35      0
36
37def mk(+l: U32, +h: U32) -> W.U64:
38  W.U64{l, h}
39
40def of32(+x: U32) -> W.U64:
41  W.U64{x, 0}
42
43def is_zero(+a: W.U64) -> Bool:
44  Bool.and(U32.is_zero(lo(a)), U32.is_zero(hi(a)))
45
46def lt(+a: W.U64, +b: W.U64) -> Bool:
47  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))))
48
49def le(+a: W.U64, +b: W.U64) -> Bool:
50  Bool.not(lt(b, a))
51
52def eq(+a: W.U64, +b: W.U64) -> Bool:
53  Bool.and(U32.is_eq(lo(a), lo(b)), U32.is_eq(hi(a), hi(b)))
54
55# ---- addition and subtraction (wrapping) ----
56
57def add_fin(+l: U32, +al: U32, +ah: U32, +bh: U32) -> W.U64:
58  W.U64{l, U32.add(U32.add(ah, bh), b32(U32.is_lt(l, al)))}
59
60def add(+a: W.U64, +b: W.U64) -> W.U64:
61  add_fin(U32.add(lo(a), lo(b)), lo(a), hi(a), hi(b))
62
63def add_over_fin(+s: U32, +ah: U32, c: Bool) -> Bool:
64  Bool.or(U32.is_lt(s, ah), Bool.and(c, U32.is_eq(s, 4294967295)))
65
66# a + b >= 2^64
67def add_over(+a: W.U64, +b: W.U64) -> Bool:
68  add_over_fin(U32.add(hi(a), hi(b)), hi(a), U32.is_lt(U32.add(lo(a), lo(b)), lo(a)))
69
70# a - b mod 2^64
71def sub(+a: W.U64, +b: W.U64) -> W.U64:
72  W.U64{U32.sub(lo(a), lo(b)), U32.sub(U32.sub(hi(a), hi(b)), b32(U32.is_lt(lo(a), lo(b))))}
73
74def half(+a: W.U64) -> W.U64:
75  W.U64{U32.add(U32.shr(lo(a)), U32.mul(U32.and(hi(a), 1), 2147483648)), U32.shr(hi(a))}
76
77def odd(+a: W.U64) -> Bool:
78  U32.is_eq(U32.and(lo(a), 1), 1)
79
80# ---- products ----
81
82def mul32_fin(+ll: U32, +mid: U32, +hh: U32, mc: Bool) -> W.U64:
83  add_fin(U32.add(ll, U32.mul(mid, 65536)), ll, hh, U32.add(U32.shrn(mid, 16n), U32.mul(b32(mc), 65536)))
84
85def mul32_mid(+ll: U32, +lh: U32, +hl: U32, +hh: U32) -> W.U64:
86  mul32_fin(ll, U32.add(lh, hl), hh, U32.is_lt(U32.add(lh, hl), lh))
87
88def mul32_h(+al: U32, +ah: U32, +bl: U32, +bh: U32) -> W.U64:
89  mul32_mid(U32.mul(al, bl), U32.mul(al, bh), U32.mul(ah, bl), U32.mul(ah, bh))
90
91# the full product of two U32
92def mul32(+a: U32, +b: U32) -> W.U64:
93  mul32_h(U32.and(a, 65535), U32.shrn(a, 16n), U32.and(b, 65535), U32.shrn(b, 16n))
94
95# a * b mod 2^64
96def mul(+a: W.U64, +b: W.U64) -> W.U64:
97  add(mul32(lo(a), lo(b)), W.U64{0, U32.add(U32.mul(lo(a), hi(b)), U32.mul(hi(a), lo(b)))})
98
99def mul_over_one(+p: W.U64, +c: W.U64) -> Bool:
100  Bool.or(Bool.not(U32.is_zero(hi(c))), U32.is_lt(U32.add(hi(p), lo(c)), hi(p)))
101
102# a * b >= 2^64
103def mul_over(+a: W.U64, +b: W.U64) -> Bool:
104  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)))))
105
106# ---- 64 / 32 ----
107
108# One digit of the long division by d: (r * base + g) / d and mod d, for a
109# remainder r < d < 2^32 and a digit g < base = 2^16, so every partial
110# dividend r * 2^16 + g is below 2^48 (the runtime's Nat bound). The low word
111# is taken as two 16-bit digits; the base is written 256 * 256 so the proof
112# checker never expands a 17-bit literal.
113def dig_t(+r: Nat, +base: Nat, +g: Nat) -> Nat:
114  Nat.add(Nat.mul(r, base), g)
115
116def dig1(+lo: U32) -> Nat:
117  U32.to_nat(U32.shrn(lo, 16n))
118
119def dig2(+lo: U32) -> Nat:
120  U32.to_nat(U32.and(lo, 65535))
121
122def div32_fin(+qh: U32, +q1: Nat, +t2: Nat, +d: Nat) -> W.U64 & U32:
123  (W.U64{U32.from_nat(Nat.add(Nat.mul(q1, 65536n), Nat.div(t2, d))), qh}, U32.from_nat(Nat.mod(t2, d)))
124
125def div32_t1(+qh: U32, +lo: U32, +t1: Nat, +d: Nat) -> W.U64 & U32:
126  div32_fin(qh, Nat.div(t1, d), dig_t(Nat.mod(t1, d), 65536n, dig2(lo)), d)
127
128# (a / d, a mod d) for d > 0: the high word natively, the low word by digits
129def div32(+a: W.U64, +d: U32) -> W.U64 & U32:
130  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))
131
132def mod_step(+r: Nat, +base: Nat, +g: Nat, +d: Nat) -> Nat:
133  Nat.mod(dig_t(r, base, g), d)
134
135# a mod d for d > 0
136def mod32(+a: W.U64, +d: U32) -> U32:
137  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)))
138
139# ---- shifts by a variable count ----
140
141# 2^k for k < 32
142def pow2(k: Nat) -> U32:
143  match k:
144    case 0n:
145      1
146    case 1n:
147      2
148    case 2n:
149      4
150    case 3n:
151      8
152    case 4n:
153      16
154    case 5n:
155      32
156    case 6n:
157      64
158    case 7n:
159      128
160    case 8n:
161      256
162    case 9n:
163      512
164    case 10n:
165      1024
166    case 11n:
167      2048
168    case 12n:
169      4096
170    case 13n:
171      8192
172    case 14n:
173      16384
174    case 15n:
175      32768
176    case 16n:
177      65536
178    case 17n:
179      131072
180    case 18n:
181      262144
182    case 19n:
183      524288
184    case 20n:
185      1048576
186    case 21n:
187      2097152
188    case 22n:
189      4194304
190    case 23n:
191      8388608
192    case 24n:
193      16777216
194    case 25n:
195      33554432
196    case 26n:
197      67108864
198    case 27n:
199      134217728
200    case 28n:
201      268435456
202    case 29n:
203      536870912
204    case 30n:
205      1073741824
206    case _:
207      2147483648
208
209def bl_pick(+x: U32, +t: U32, +k: Nat, +base: Nat, big: Bool) -> Nat & U32:
210  match big:
211    case True{}:
212      (Nat.add(base, k), U32.shrn(x, k))
213    case False{}:
214      (base, x)
215
216def bl_step(+k: Nat, st: Nat & U32) -> Nat & U32:
217  (+base, +x) = st
218  bl_pick(x, pow2(k), k, base, U32.is_le(pow2(k), x))
219
220def bl_fin(st: Nat & U32) -> Nat:
221  (+base, +x) = st
222  Nat.add(base, U32.to_nat(U32.min(x, 1)))
223
224# the number of bits of x (0 for 0), by a binary search on the top bit
225def bitlen(+x: U32) -> Nat:
226  bl_fin(bl_step(1n, bl_step(2n, bl_step(4n, bl_step(8n, bl_step(16n, (0n, x)))))))
227
228def shl_lt(+a: W.U64, +k: Nat) -> W.U64:
229  match k:
230    case 0n:
231      a
232    case 1n+ +j:
233      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)))}
234
235def shl_ge(+a: W.U64, +k: Nat) -> W.U64:
236  W.U64{0, U32.mul(lo(a), pow2(Nat.sub(k, 32n)))}
237
238def shl_pick(+a: W.U64, +k: Nat, small: Bool) -> W.U64:
239  match small:
240    case True{}:
241      shl_lt(a, k)
242    case False{}:
243      shl_ge(a, k)
244
245# a << k mod 2^64 for k < 64
246def shl(+a: W.U64, +k: Nat) -> W.U64:
247  shl_pick(a, k, Nat.is_lt(k, 32n))
248
249def shr_lt(+a: W.U64, +k: Nat) -> W.U64:
250  match k:
251    case 0n:
252      a
253    case 1n+ +j:
254      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)}
255
256def shr_ge(+a: W.U64, +k: Nat, huge: Bool) -> W.U64:
257  match huge:
258    case True{}:
259      W.U64{0, 0}
260    case False{}:
261      W.U64{U32.shrn(hi(a), Nat.sub(k, 32n)), 0}
262
263def shr_pick(+a: W.U64, +k: Nat, small: Bool) -> W.U64:
264  match small:
265    case True{}:
266      shr_lt(a, k)
267    case False{}:
268      shr_ge(a, k, Nat.is_le(64n, k))
269
270# a >> k (0 for k >= 64)
271def shr(+a: W.U64, +k: Nat) -> W.U64:
272  shr_pick(a, k, Nat.is_lt(k, 32n))
273
274def fst_q(p: W.U64 & U32) -> W.U64:
275  (q, r) = p
276  q
277
278def snd_r(p: W.U64 & U32) -> U32:
279  (q, r) = p
280  r
281
282def q_clamp_z(+l: U32, z: Bool) -> U32:
283  match z:
284    case True{}:
285      l
286    case False{}:
287      4294967295
288
289# min(q, 2^32 - 1)
290def q_clamp(+q: W.U64) -> U32:
291  q_clamp_z(lo(q), U32.is_zero(hi(q)))
292
293# the product of a U32 and a U64 as (low 64 bits, top 32 bits)
294def mul_32_64_fin(+p0: W.U64, +p1: W.U64) -> W.U64 & U32:
295  (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)))))
296
297def mul_32_64(+q: U32, +b: W.U64) -> W.U64 & U32:
298  mul_32_64_fin(mul32(q, lo(b)), mul32(q, hi(b)))
299
300# the 96-bit product p exceeds x = xh * 2^64 + xl
301def over96(+xl: W.U64, +xh: U32, p: W.U64 & U32) -> Bool:
302  (+pl, +ph) = p
303  Bool.or(U32.is_lt(xh, ph), Bool.and(U32.is_eq(xh, ph), lt(xl, pl)))
304
305# q - 1 while q * b > x: at most q steps (fuel q + 1)
306def q_fix(fuel: Nat, +xl: W.U64, +xh: U32, +b: W.U64, +q: U32, over: Bool) -> U32:
307  match fuel over:
308    case 0n _:
309      q
310    case 1n+f False{}:
311      q
312    case 1n+f True{}:
313      q_fix(f, xl, xh, b, U32.sub(q, 1), over96(xl, xh, mul_32_64(U32.sub(q, 1), b)))
314
315# floor(x / b) from an estimate q >= floor(x / b): down while q * b > x (at
316# most q steps, fuel q + 1)
317def q_start(+xl: W.U64, +xh: U32, +b: W.U64, +q: U32) -> U32:
318  q_fix(Nat.add(U32.to_nat(q), 1n), xl, xh, b, q, over96(xl, xh, mul_32_64(q, b)))
319
320# Knuth's estimate for x = xh * 2^64 + xl < b * 2^32 and b >= 2^32, t =
321# bits(hi(b)): with Y = x >> t (below 2^64) and B = b >> t in [2^31, 2^32),
322# min(floor(Y / B), 2^32 - 1) is at least floor(x / b) (as B 2^t <= b) and at
323# most 2 over it (Theorem B), so q_start takes a step or two
324def est_lo(+q1: Nat, +t2: Nat, +d: Nat) -> U32:
325  U32.from_nat(Nat.add(Nat.mul(q1, 65536n), Nat.div(t2, d)))
326
327def est_mid(+lo: U32, +t1: Nat, +d: Nat) -> U32:
328  est_lo(Nat.div(t1, d), dig_t(Nat.mod(t1, d), 65536n, dig2(lo)), d)
329
330# min(floor(y / d), 2^32 - 1) for d > 0: when hi(y) >= d the quotient needs
331# more than 32 bits; otherwise the high word is its own remainder, and div32's
332# two 16-bit Nat digit steps finish the division without dividing the high
333# word.
334def est_pick(+y: W.U64, +d: U32, big: Bool) -> U32:
335  match big:
336    case True{}:
337      4294967295
338    case False{}:
339      est_mid(lo(y), dig_t(U32.to_nat(hi(y)), 65536n, dig1(lo(y))), U32.to_nat(d))
340
341def est32(+y: W.U64, +d: U32) -> U32:
342  est_pick(y, d, U32.is_le(d, hi(y)))
343
344def q_est(+xl: W.U64, +xh: U32, +b: W.U64, +t: Nat) -> U32:
345  est32(add(shr(xl, t), shl(W.U64{xh, 0}, Nat.sub(64n, t))), lo(shr(b, t)))
346
347# floor(x / b) for x = xh * 2^64 + xl < b * 2^32 and b >= 2^32, t = bits(hi(b))
348def q96(+xl: W.U64, +xh: U32, +b: W.U64, +t: Nat) -> U32:
349  q_start(xl, xh, b, q_est(xl, xh, b, t))
350
351# ---- 64 / 64 ----
352
353def dm_small(p: W.U64 & U32) -> W.U64 & W.U64:
354  (+q, +r) = p
355  (q, W.U64{r, 0})
356
357def dm_big(+a: W.U64, +b: W.U64, +q: U32) -> W.U64 & W.U64:
358  (W.U64{q, 0}, sub(a, fst_q(mul_32_64(q, b))))
359
360def dm_pick(+a: W.U64, +b: W.U64, small: Bool) -> W.U64 & W.U64:
361  match small:
362    case True{}:
363      dm_small(div32(a, lo(b)))
364    case False{}:
365      dm_big(a, b, q96(a, 0, b, bitlen(hi(b))))
366
367# (a / b, a mod b) for b > 0
368def divmod(+a: W.U64, +b: W.U64) -> W.U64 & W.U64:
369  dm_pick(a, b, U32.is_zero(hi(b)))
370
371def pfst(p: W.U64 & W.U64) -> W.U64:
372  (q, r) = p
373  q
374
375def psnd(p: W.U64 & W.U64) -> W.U64:
376  (q, r) = p
377  r
378
379def quot(+a: W.U64, +b: W.U64) -> W.U64:
380  pfst(divmod(a, b))
381
382def rem(+a: W.U64, +b: W.U64) -> W.U64:
383  psnd(divmod(a, b))
384
385# ---- a * b mod m ----
386
387# x mod m for x = xh * 2^32 + x0 with xh < m and m >= 2^32
388def red96_z(+xh: W.U64, +x0: U32, +m: W.U64, small: Bool) -> W.U64:
389  match small:
390    case True{}:
391      W.U64{x0, lo(xh)}
392    case False{}:
393      sub(W.U64{x0, lo(xh)}, fst_q(mul_32_64(q96(W.U64{x0, lo(xh)}, hi(xh), m, bitlen(hi(m))), m)))
394
395# (the zero test of hi(m) only guards the digit: red96 is called with
396# hi(m) != 0; it also stops the proof checker from expanding the digit in
397# every statement that names red96)
398def red96(+xh: W.U64, +x0: U32, +m: W.U64) -> W.U64:
399  red96_z(xh, x0, m, U32.is_zero(hi(m)))
400
401def mul128_fin(+p00: W.U64, +mid: W.U64, c1: Bool, +p11: W.U64) -> W.U64 & W.U64:
402  (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}))
403
404def mul128_mid(+p00: W.U64, +p01: W.U64, +p10: W.U64, +p11: W.U64) -> W.U64 & W.U64:
405  mul128_fin(p00, add(p01, p10), add_over(p01, p10), p11)
406
407# the full product as (low 64 bits, high 64 bits)
408def mul128(+a: W.U64, +b: W.U64) -> W.U64 & W.U64:
409  mul128_mid(mul32(lo(a), lo(b)), mul32(lo(a), hi(b)), mul32(hi(a), lo(b)), mul32(hi(a), hi(b)))
410
411def mm_big(+m: W.U64, p: W.U64 & W.U64) -> W.U64:
412  (+pl, +ph) = p
413  red96(red96(ph, hi(pl), m), lo(pl), m)
414
415def mm_pick(+a: W.U64, +b: W.U64, +m: W.U64, small: Bool) -> W.U64:
416  match small:
417    case True{}:
418      W.U64{mod32(mul32(lo(a), lo(b)), lo(m)), 0}
419    case False{}:
420      mm_big(m, mul128(a, b))
421
422# a * b mod m for a, b < m
423def mulmod(+a: W.U64, +b: W.U64, +m: W.U64) -> W.U64:
424  mm_pick(a, b, m, U32.is_zero(hi(m)))
425
426# ---- integer square root ----
427
428# r - 1 while r * r > x: from r <= 65535 at most r steps (fuel r + 1), and
429# r * r never wraps
430def down32(fuel: Nat, +x: U32, +r: U32, over: Bool) -> U32:
431  match fuel over:
432    case 0n _:
433      r
434    case 1n+f False{}:
435      r
436    case 1n+f True{}:
437      down32(f, x, U32.sub(r, 1), U32.is_lt(x, U32.mul(U32.sub(r, 1), U32.sub(r, 1))))
438
439def up_ok(+x: U32, +r: U32) -> Bool:
440  Bool.and(U32.is_lt(r, 65535), U32.is_le(U32.mul(U32.add(r, 1), U32.add(r, 1)), x))
441
442# r + 1 while r < 65535 and (r + 1)^2 <= x: at most 65535 - r steps
443def up32(fuel: Nat, +x: U32, +r: U32, up: Bool) -> U32:
444  match fuel up:
445    case 0n _:
446      r
447    case 1n+f False{}:
448      r
449    case 1n+f True{}:
450      up32(f, x, U32.add(r, 1), up_ok(x, U32.add(r, 1)))
451
452def isqrt32_fix(+x: U32, +r: U32) -> U32:
453  up32(1n+U32.to_nat(U32.sub(65535, r)), x, r, up_ok(x, r))
454
455# floor(sqrt(x)) from any estimate r <= 65535: down while r^2 > x, then up
456# while (r + 1)^2 <= x. The F32 square root lands within 1, so both loops
457# take at most a step or two; the fuel covers any estimate, so the result
458# never depends on F32's accuracy.
459def isqrt32_est(+x: U32, +r: U32) -> U32:
460  isqrt32_fix(x, down32(1n+U32.to_nat(r), x, r, U32.is_lt(x, U32.mul(r, r))))
461
462def isqrt32(+x: U32) -> U32:
463  isqrt32_est(x, U32.min(F32.to_u32(F32.sqrt(U32.to_f32(x))), 65535))
464
465def sq_over(+n: W.U64, +r: U32) -> Bool:
466  lt(n, mul32(r, r))
467
468def down64(fuel: Nat, +n: W.U64, +r: U32, over: Bool) -> U32:
469  match fuel over:
470    case 0n _:
471      r
472    case 1n+f False{}:
473      r
474    case 1n+f True{}:
475      down64(f, n, U32.sub(r, 1), sq_over(n, U32.sub(r, 1)))
476
477def up64_ok(+n: W.U64, +r: U32) -> Bool:
478  Bool.and(U32.is_lt(r, 4294967295), Bool.not(sq_over(n, U32.add(r, 1))))
479
480def up64(fuel: Nat, +n: W.U64, +r: U32, up: Bool) -> U32:
481  match fuel up:
482    case 0n _:
483      r
484    case 1n+f False{}:
485      r
486    case 1n+f True{}:
487      up64(f, n, U32.add(r, 1), up64_ok(n, U32.add(r, 1)))
488
489def isqrt64_fix(+n: W.U64, +r: U32) -> U32:
490  up64(1n+U32.to_nat(U32.sub(4294967295, r)), n, r, up64_ok(n, r))
491
492# floor(sqrt(n)) from any estimate r1: down while r^2 > n, then up while
493# (r + 1)^2 <= n; the fuel covers any estimate
494def isqrt_newton(+n: W.U64, +r1: U32) -> U32:
495  isqrt64_fix(n, down64(1n+U32.to_nat(r1), n, r1, sq_over(n, r1)))
496
497# one Newton step from the F32 estimate r0 > 0 lands within a few units of
498# floor(sqrt(n)); the correction makes the result exact for any r0
499def isqrt_big(+n: W.U64, +r0: U32) -> W.U64:
500  W.U64{isqrt_newton(n, q_clamp(half(add(fst_q(div32(n, r0)), W.U64{r0, 0})))), 0}
501
502def f32_clamp32(+f: F32, big: Bool) -> U32:
503  match big:
504    case True{}:
505      4294967295
506    case False{}:
507      F32.to_u32(f)
508
509def isqrt_est(+f: F32) -> U32:
510  U32.max(f32_clamp32(f, F32.is_le(4294967040.0, f)), 1)
511
512def isqrt_pick(+n: W.U64, small: Bool) -> W.U64:
513  match small:
514    case True{}:
515      W.U64{isqrt32(lo(n)), 0}
516    case False{}:
517      isqrt_big(n, isqrt_est(F32.sqrt(F32.add(F32.mul(U32.to_f32(hi(n)), 4294967296.0), U32.to_f32(lo(n))))))
518
519def isqrt(+n: W.U64) -> W.U64:
520  isqrt_pick(n, U32.is_zero(hi(n)))
521
522def w_pow2(+k: Nat, small: Bool) -> W.U64:
523  match small:
524    case True{}:
525      W.U64{pow2(k), 0}
526    case False{}:
527      W.U64{0, pow2(Nat.sub(k, 32n))}
528
529# ---- general shifts and leading zeros (the software F64) ----
530
531def or_bit(+a: W.U64, b: Bool) -> W.U64:
532  W.U64{U32.or(lo(a), b32(b)), hi(a)}
533
534def jam_pick(+a: W.U64, +k: Nat, huge: Bool) -> W.U64:
535  match huge:
536    case True{}:
537      W.U64{b32(Bool.not(is_zero(a))), 0}
538    case False{}:
539      or_bit(shr(a, k), Bool.not(eq(shl(shr(a, k), k), a)))
540
541# a >> k with the bits shifted out OR-ed into bit 0 (SoftFloat's
542# shiftRightJam64)
543def shr_jam(+a: W.U64, +k: Nat) -> W.U64:
544  jam_pick(a, k, Nat.is_le(64n, k))
545
546def clz_pick(+a: W.U64, high: Bool) -> Nat:
547  match high:
548    case True{}:
549      Nat.sub(32n, bitlen(hi(a)))
550    case False{}:
551      Nat.sub(64n, bitlen(lo(a)))
552
553# the number of leading zero bits of a (64 for 0)
554def clz(+a: W.U64) -> Nat:
555  clz_pick(a, Bool.not(U32.is_zero(hi(a))))
556
557# a with bit 0 cleared when b
558def clear0(+a: W.U64, b: Bool) -> W.U64:
559  W.U64{U32.sub(lo(a), U32.and(lo(a), b32(b))), hi(a)}
560
561def cmp_pick(less: Bool, same: Bool) -> Cmp:
562  match less same:
563    case True{} _:
564      LT{}
565    case False{} True{}:
566      EQ{}
567    case False{} False{}:
568      GT{}
569
570def cmp(+a: W.U64, +b: W.U64) -> Cmp:
571  cmp_pick(lt(a, b), eq(a, b))
572
573# ---- Montgomery multiplication (odd m, 2^32 <= m < 2^63, R = 2^64) ----
574# REDC (Montgomery 1985; the word-level algorithm HACL* and Fiat Crypto
575# verify): for t < m^2 and mp = -1/m mod 2^64, u = t mp mod 2^64 makes
576# t + u m divisible by 2^64, and r = (t + u m) / 2^64 < 2 m is t / 2^64 mod m
577# after one conditional subtraction.
578
579# x (2 - m x): doubles the correct low bits of an inverse of m mod 2^64
580def inv_step(+m: W.U64, +x: W.U64) -> W.U64:
581  mul(x, sub(W.U64{2, 0}, mul(m, x)))
582
583# -1 / m mod 2^64 for odd m (m is its own inverse mod 8; five Newton steps
584# reach 96 >= 64 bits)
585def minv(+m: W.U64) -> W.U64:
586  sub(W.U64{0, 0}, inv_step(m, inv_step(m, inv_step(m, inv_step(m, inv_step(m, m))))))
587
588def sub_if(+r: W.U64, +m: W.U64, big: Bool) -> W.U64:
589  match big:
590    case True{}:
591      sub(r, m)
592    case False{}:
593      r
594
595def redc_r(+m: W.U64, +r: W.U64) -> W.U64:
596  sub_if(r, m, le(m, r))
597
598def redc_p(+m: W.U64, +tl: W.U64, +th: W.U64, p: W.U64 & W.U64) -> W.U64:
599  (+pl, +ph) = p
600  redc_r(m, add(add(th, ph), W.U64{b32(add_over(tl, pl)), 0}))
601
602# t / 2^64 mod m for t = (tl, th) < m^2
603def redc(+m: W.U64, +mp: W.U64, t: W.U64 & W.U64) -> W.U64:
604  (+tl, +th) = t
605  redc_p(m, tl, th, mul128(mul(tl, mp), m))
606
607# a b / 2^64 mod m for a, b < m
608def mont(+m: W.U64, +mp: W.U64, +a: W.U64, +b: W.U64) -> W.U64:
609  redc(m, mp, mul128(a, b))
610
611def mbit(odd: Bool, +m: W.U64, +mp: W.U64, +b: W.U64, +acc: W.U64) -> W.U64:
612  match odd:
613    case False{}:
614      acc
615    case True{}:
616      mont(m, mp, acc, b)
617
618# pow_mod_go of generic.bend on Montgomery forms x 2^64 mod m
619def mpow_go(fuel: Nat, +m: W.U64, +mp: W.U64, +b: W.U64, +acc: W.U64, st: W.U64 & Bool) -> W.U64:
620  match fuel st:
621    case 0n _:
622      acc
623    case 1n+f Tuple{e, True{}}:
624      acc
625    case 1n+f Tuple{+e, False{}}:
626      mpow_go(f, m, mp, mont(m, mp, b, b), mbit(odd(e), m, mp, b, acc), (half(e), is_zero(half(e))))
627
628# x 2^64 mod m for x < m, m >= 2^32
629def to_mont(+x: W.U64, +m: W.U64) -> W.U64:
630  red96(red96(x, 0, m), 0, m)
631
632# m odd with 2^32 <= m < 2^63 and m minv(m) = -1 mod 2^64 (always so for odd
633# m; checked, so the proof needs no Hensel lemma)
634def mont_ok(+m: W.U64) -> Bool:
635  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)))
636
637def mpow_m(+b: W.U64, +e: W.U64, +m: W.U64, +mp: W.U64) -> W.U64:
638  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})
639
640# b^e mod m for b < m and mont_ok(m)
641def mpow(+b: W.U64, +e: W.U64, +m: W.U64) -> W.U64:
642  mpow_m(b, e, m, minv(m))
643
644# ---- 32-bit Montgomery multiplication (odd m, R = 2^32) ----
645
646def inv32_step(+m: U32, +x: U32) -> U32:
647  U32.mul(x, U32.sub(2, U32.mul(m, x)))
648
649# -1 / m mod 2^32 for odd m
650def minv32(+m: U32) -> U32:
651  U32.sub(0, inv32_step(m, inv32_step(m, inv32_step(m, inv32_step(m, m)))))
652
653def sub_if32(+r: W.U64, +m: U32, big: Bool) -> U32:
654  match big:
655    case True{}:
656      lo(sub(r, W.U64{m, 0}))
657    case False{}:
658      lo(r)
659
660def redc32_r(+m: U32, +r: W.U64) -> U32:
661  sub_if32(r, m, le(W.U64{m, 0}, r))
662
663def redc32_s(+m: U32, +s1: U32, +s2: U32, +c1: Bool) -> U32:
664  redc32_r(m, W.U64{s2, U32.add(b32(c1), b32(U32.is_lt(s2, s1)))})
665
666def redc32_p(+m: U32, +t: W.U64, +p: W.U64) -> U32:
667  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)))
668
669def redc32(+m: U32, +mp: U32, +t: W.U64) -> U32:
670  redc32_p(m, t, mul32(U32.mul(lo(t), mp), m))
671
672def mont32(+m: U32, +mp: U32, +a: U32, +b: U32) -> U32:
673  redc32(m, mp, mul32(a, b))
674
675def mbit32(odd: Bool, +m: U32, +mp: U32, +b: U32, +acc: U32) -> U32:
676  match odd:
677    case False{}:
678      acc
679    case True{}:
680      mont32(m, mp, acc, b)
681
682def mpow32_go(fuel: Nat, +m: U32, +mp: U32, +b: U32, +acc: U32, st: U32 & Bool) -> U32:
683  match fuel st:
684    case 0n _:
685      acc
686    case 1n+f Tuple{e, True{}}:
687      acc
688    case 1n+f Tuple{+e, False{}}:
689      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))))
690
691def mont32_ok(+m: U32) -> Bool:
692  Bool.and(U32.is_eq(U32.and(m, 1), 1), U32.is_eq(U32.mul(m, minv32(m)), 4294967295))
693
694def mpow32_m(+b: U32, +e: U32, +m: U32, +mp: U32) -> U32:
695  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)
696
697def mpow32(+b: U32, +e: U32, +m: U32) -> U32:
698  mpow32_m(b, e, m, minv32(m))
699
700# a and b fit 48 bits (runtime Nats) and a has more than 16: the binary gcd
701# then hands them to Euclid on Nat (native division); below 2^16 its few
702# steps are cheaper
703def small2(+a: W.U64, +b: W.U64) -> Bool:
704  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))))
705
706def n48(+a: W.U64) -> Nat:
707  Nat.add(U32.to_nat(lo(a)), Nat.mul(Nat.mul(U32.to_nat(hi(a)), 65536n), 65536n))
708
709# the words of g < 2^64 from h = g div 2^32 (two divisions by 2^16, so no
710# closed 2^32 appears: the proof checker would expand it in unary)
711def of48_h(+g: Nat, +h: Nat) -> W.U64:
712  W.U64{U32.from_nat(Nat.sub(g, Nat.mul(Nat.mul(h, 65536n), 65536n))), U32.from_nat(h)}
713
714def of48(+g: Nat) -> W.U64:
715  of48_h(g, Nat.div(Nat.div(g, 65536n), 65536n))