src/math/natural.bend

bend-collections-laws-math@1.0.0.0 · 348 lines · 47 definitions

1import Base
2import bend-collections-laws-math@1.0.0.0/src/math/pow2.bend as P2
3
4# Exact integer functions of a Python-style math library on the natural
5# numbers (Python's non-negative int), after the design reference's
6# section 4.2 (math module integer functions), 3.1 (divmod, three-argument
7# pow) and 10 (ilog, iroot, clamp, mod_inverse).
8#
9#   gcd(a, b), gcd_all(xs)       greatest common divisor; gcd_all([]) == 0
10#   lcm(a, b), lcm_all(xs)       least common multiple;   lcm_all([]) == 1
11#   isqrt(n)                     floor(sqrt(n)), exact for every n
12#   iroot(n, k)                  floor of the k-th root (k >= 1)
13#   ilog(n, b)                   floor(log_b(n)), exact (n >= 1, b >= 2)
14#   factorial(n), perm(n, k)     n!, n! / (n - k)! (0 when k > n)
15#   comb(n, k)                   n choose k (0 when k > n)
16#   prod(xs), sum(xs)            product and sum of a list; prod([]) == 1
17#   pow_mod(b, e, m)             b^e mod m by binary exponentiation
18#   mod_inverse(a, m)            x with a*x == 1 (mod m), x < m
19#   divmod(a, b)                 (a // b, a % b)
20#   bit_length(n)                bits of n without leading zeros
21#   clamp(x, lo, hi)             x limited to [lo, hi]
22#
23# Errors follow the reference's error model (section 2.3) as values, not
24# exceptions: ZeroDivision where Python raises ZeroDivisionError or the
25# "cannot be 0" ValueError of pow, Domain for the other ValueErrors, and
26# NotInvertible for pow(a, -1, m) without an inverse.
27#
28# Nat is exact; Bend's native backend stops the program on a Nat above
29# 2^48 - 1, so results (and the intermediate products of comb, lcm and
30# pow_mod) must stay below that at run time. The proofs are over the
31# mathematical naturals.
32
33type MathError is Data:
34  ZeroDivision{}
35  Domain{}
36  NotInvertible{}
37
38# the last nonzero remainder of the extended Euclid loop and its
39# coefficient: a * coef == gcd (mod m)
40type Bezout is Data:
41  BZ{gcd: Nat, coef: Nat}
42
43# a // b and a % b
44type QuotRem is Data:
45  QR{quot: Nat, rem: Nat}
46
47# ---- greatest common divisor, least common multiple ----
48
49# Euclid's algorithm; b <= fuel throughout (every step lowers b).
50def gcd_go(fuel: Nat, +a: Nat, +b: Nat) -> Nat:
51  match fuel b:
52    case 0n _:
53      a
54    case 1n+f 0n:
55      a
56    case 1n+f 1n+ +bp:
57      gcd_go(f, 1n+bp, Nat.mod(a, 1n+bp))
58
59def gcd(+a: Nat, +b: Nat) -> Nat:
60  gcd_go(b, a, b)
61
62def lcm(+a: Nat, +b: Nat) -> Nat:
63  match a b:
64    case 0n _:
65      0n
66    case 1n+ap 0n:
67      0n
68    case 1n+ +ap 1n+ +bp:
69      Nat.mul(Nat.div(1n+ap, gcd(1n+ap, 1n+bp)), 1n+bp)
70
71def gcd_all_go(xs: List<&2, Nat>, +acc: Nat) -> Nat:
72  match xs:
73    case Nil{}:
74      acc
75    case Con{+x, rest}:
76      gcd_all_go(rest, gcd(acc, x))
77
78def gcd_all(xs: List<&2, Nat>) -> Nat:
79  gcd_all_go(xs, 0n)
80
81def lcm_all_go(xs: List<&2, Nat>, +acc: Nat) -> Nat:
82  match xs:
83    case Nil{}:
84      acc
85    case Con{+x, rest}:
86      lcm_all_go(rest, lcm(acc, x))
87
88def lcm_all(xs: List<&2, Nat>) -> Nat:
89  lcm_all_go(xs, 1n)
90
91# ---- sums and products ----
92
93def sum_go(xs: List<&2, Nat>, +acc: Nat) -> Nat:
94  match xs:
95    case Nil{}:
96      acc
97    case Con{+x, rest}:
98      sum_go(rest, Nat.add(acc, x))
99
100def sum(xs: List<&2, Nat>) -> Nat:
101  sum_go(xs, 0n)
102
103def prod_go(xs: List<&2, Nat>, +acc: Nat) -> Nat:
104  match xs:
105    case Nil{}:
106      acc
107    case Con{+x, rest}:
108      prod_go(rest, Nat.mul(acc, x))
109
110def prod(xs: List<&2, Nat>) -> Nat:
111  prod_go(xs, 1n)
112
113# ---- factorial, permutations, combinations ----
114
115def factorial_go(n: Nat, +acc: Nat) -> Nat:
116  match n:
117    case 0n:
118      acc
119    case 1n+ +p:
120      factorial_go(p, Nat.mul(1n+p, acc))
121
122def factorial(+n: Nat) -> Nat:
123  factorial_go(n, 1n)
124
125# n * (n - 1) * ... * (n - k + 1)
126def perm_go(k: Nat, +m: Nat, +acc: Nat) -> Nat:
127  match k:
128    case 0n:
129      acc
130    case 1n+j:
131      perm_go(j, Nat.sub(m, 1n), Nat.mul(acc, m))
132
133def perm(+n: Nat, +k: Nat) -> Nat:
134  perm_go(k, n, 1n)
135
136# r_{i+1} = r_i * (n - i) / (i + 1): every r_i is C(n, i), so the division
137# is exact (reference section 8.6).
138def comb_go(j: Nat, +n: Nat, +i: Nat, +r: Nat) -> Nat:
139  match j:
140    case 0n:
141      r
142    case 1n+ +jp:
143      comb_go(jp, n, 1n+i, Nat.div(Nat.mul(r, Nat.sub(n, i)), 1n+i))
144
145def comb_small(+n: Nat, +k: Nat, +big: Bool) -> Nat:
146  match big:
147    case True{}:
148      0n
149    case False{}:
150      comb_go(Nat.min(k, Nat.sub(n, k)), n, 0n, 1n)
151
152def comb(+n: Nat, +k: Nat) -> Nat:
153  comb_small(n, k, Nat.is_lt(n, k))
154
155# ---- bits ----
156
157def bit_length_go(fuel: Nat, +n: Nat, +k: Nat) -> Nat:
158  match fuel n:
159    case 0n _:
160      k
161    case 1n+f 0n:
162      k
163    case 1n+f 1n+ +np:
164      bit_length_go(f, Nat.div(1n+np, 2n), 1n+k)
165
166def bit_length(+n: Nat) -> Nat:
167  bit_length_go(n, n, 0n)
168
169# ---- roots and logarithms ----
170
171# Binary search for the last r in [lo, hi) where the test holds, given it
172# holds at lo and fails at hi (reference section 8.5 notes the float-sqrt
173# shortcut is wrong above 2^52; this search is exact). `more` is lo + 1 < hi
174# and `ok` is the test at the midpoint; the fuel bounds the halvings.
175def mid(+lo: Nat, +hi: Nat) -> Nat:
176  Nat.div(Nat.add(lo, hi), 2n)
177
178# r^k <= n without forming a product above n: `ok` is acc * r <= n, asked
179# as acc <= n / r, and acc * r is only formed once it is known to fit.
180def pow_le_go(k: Nat, +r: Nat, +acc: Nat, +n: Nat, +ok: Bool) -> Bool:
181  match k ok:
182    case 0n _:
183      True{}
184    case 1n+j False{}:
185      False{}
186    case 1n+j True{}:
187      pow_le_go(j, r, Nat.mul(acc, r), n, Nat.is_le(Nat.mul(acc, r), Nat.div(n, r)))
188
189def root_le(+n: Nat, +k: Nat, +r: Nat) -> Bool:
190  match k r:
191    case 0n _:
192      Nat.is_le(1n, n)
193    case 1n+kp 0n:
194      True{}
195    case 1n+ +kp 1n+ +rp:
196      pow_le_go(1n+kp, 1n+rp, 1n, n, Nat.is_le(1n, Nat.div(n, 1n+rp)))
197
198# the test is r^k <= n
199def search_go(fuel: Nat, +n: Nat, +k: Nat, +lo: Nat, +hi: Nat, +more: Bool, +ok: Bool) -> Nat:
200  match fuel more ok:
201    case 0n _ _:
202      lo
203    case 1n+f False{} _:
204      lo
205    case 1n+f True{} True{}:
206      search_go(f, n, k, mid(lo, hi), hi, Nat.is_lt(1n+mid(lo, hi), hi), root_le(n, k, mid(mid(lo, hi), hi)))
207    case 1n+f True{} False{}:
208      search_go(f, n, k, lo, mid(lo, hi), Nat.is_lt(1n+lo, mid(lo, hi)), root_le(n, k, mid(lo, mid(lo, hi))))
209
210def search(+n: Nat, +k: Nat, +lo: Nat, +hi: Nat) -> Nat:
211  search_go(Nat.sub(hi, lo), n, k, lo, hi, Nat.is_lt(1n+lo, hi), root_le(n, k, mid(lo, hi)))
212
213# the largest r with r * r <= n, by Heron's (Newton's) iteration as in Lean 4
214# Mathlib's Nat.sqrt: g := (g + n / g) / 2 while that decreases g. Started
215# from 2^(b/2 + 1) >= sqrt(n) (n < 2^b, b = bit_length(n)) it converges in
216# O(log b) steps; the fuel g bounds the (strictly decreasing) steps.
217def sqrt_next(+n: Nat, +g: Nat) -> Nat:
218  Nat.div(Nat.add(g, Nat.div(n, g)), 2n)
219
220def sqrt_iter(fuel: Nat, +n: Nat, +g: Nat, +m: Nat, +down: Bool) -> Nat:
221  match fuel down:
222    case 0n _:
223      g
224    case 1n+f False{}:
225      g
226    case 1n+f True{}:
227      sqrt_iter(f, n, m, sqrt_next(n, m), Nat.is_lt(sqrt_next(n, m), m))
228
229def isqrt_from(+n: Nat, +g: Nat) -> Nat:
230  sqrt_iter(g, n, g, sqrt_next(n, g), Nat.is_lt(sqrt_next(n, g), g))
231
232def isqrt(+n: Nat) -> Nat:
233  isqrt_from(n, P2.pow2t(1n+Nat.div(bit_length(n), 2n)))
234
235# the largest r with r^k <= n (k >= 1)
236def iroot_k(+n: Nat, +k: Nat) -> Nat:
237  match k:
238    case 0n:
239      0n
240    case 1n:
241      n
242    case 2n+ +kp:
243      search(n, 2n+kp, 0n, P2.pow2t(1n+Nat.div(bit_length(n), 2n+kp)))
244
245def iroot(+n: Nat, +k: Nat) -> Result<&2, &2, MathError, Nat>:
246  match k:
247    case 0n:
248      Fail{Domain{}}
249    case 1n+ +kp:
250      Done{iroot_k(n, 1n+kp)}
251
252# the largest k with b^k <= n: p is b^k, `up` is p * b <= n, asked as
253# p <= n / b so no product above n is formed
254def ilog_go(fuel: Nat, +n: Nat, +b: Nat, +k: Nat, +p: Nat, +up: Bool) -> Nat:
255  match fuel up:
256    case 0n _:
257      k
258    case 1n+f False{}:
259      k
260    case 1n+f True{}:
261      ilog_go(f, n, b, 1n+k, Nat.mul(p, b), Nat.is_le(Nat.mul(p, b), Nat.div(n, b)))
262
263def ilog_ok(+n: Nat, +b: Nat, +bad: Bool) -> Result<&2, &2, MathError, Nat>:
264  match bad:
265    case True{}:
266      Fail{Domain{}}
267    case False{}:
268      Done{ilog_go(n, n, b, 0n, 1n, Nat.is_le(1n, Nat.div(n, b)))}
269
270def ilog(+n: Nat, +b: Nat) -> Result<&2, &2, MathError, Nat>:
271  ilog_ok(n, b, Bool.or(Nat.is_eq(n, 0n), Nat.is_lt(b, 2n)))
272
273# ---- modular arithmetic ----
274
275# right-to-left binary exponentiation, reducing mod m after every product:
276# acc * base^e == b^e0 (mod m) throughout
277def pow_mod_odd(+m: Nat, +bit: Nat, +base: Nat, +acc: Nat) -> Nat:
278  match bit:
279    case 0n:
280      acc
281    case 1n+z:
282      Nat.mod(Nat.mul(acc, base), m)
283
284def pow_mod_go(fuel: Nat, +m: Nat, +e: Nat, +base: Nat, +acc: Nat) -> Nat:
285  match fuel e:
286    case 0n _:
287      acc
288    case 1n+f 0n:
289      acc
290    case 1n+f 1n+ +ep:
291      pow_mod_go(f, m, Nat.div(1n+ep, 2n), Nat.mod(Nat.mul(base, base), m), pow_mod_odd(m, Nat.mod(1n+ep, 2n), base, acc))
292
293def pow_mod(+b: Nat, +e: Nat, +m: Nat) -> Result<&2, &2, MathError, Nat>:
294  match m:
295    case 0n:
296      Fail{ZeroDivision{}}
297    case 1n+ +mp:
298      Done{pow_mod_go(e, 1n+mp, e, Nat.mod(b, 1n+mp), Nat.mod(1n, 1n+mp))}
299
300# Extended Euclid on (m, a mod m) keeping only the coefficient of a, mod m:
301# a * s0 == r0 and a * s1 == r1 (mod m). The remainders are gcd's.
302def inv_step(+m: Nat, +q: Nat, +s0: Nat, +s1: Nat) -> Nat:
303  Nat.mod(Nat.add(s0, Nat.sub(m, Nat.mod(Nat.mul(q, s1), m))), m)
304
305def inv_go(fuel: Nat, +m: Nat, +r0: Nat, +s0: Nat, +r1: Nat, +s1: Nat) -> Bezout:
306  match fuel r1:
307    case 0n _:
308      BZ{r0, s0}
309    case 1n+f 0n:
310      BZ{r0, s0}
311    case 1n+f 1n+ +rp:
312      inv_go(f, m, 1n+rp, s1, Nat.mod(r0, 1n+rp), inv_step(m, Nat.div(r0, 1n+rp), s0, s1))
313
314def inv_fin(+m: Nat, bz: Bezout) -> Result<&2, &2, MathError, Nat>:
315  match bz:
316    case BZ{1n, s}:
317      Done{Nat.mod(s, m)}
318    case BZ{0n, s}:
319      Fail{NotInvertible{}}
320    case BZ{2n+g, s}:
321      Fail{NotInvertible{}}
322
323# pow(a, -1, m): the x < m with a * x == 1 (mod m)
324def mod_inverse(+a: Nat, +m: Nat) -> Result<&2, &2, MathError, Nat>:
325  match m:
326    case 0n:
327      Fail{ZeroDivision{}}
328    case 1n+ +mp:
329      inv_fin(1n+mp, inv_go(1n+mp, 1n+mp, 1n+mp, 0n, Nat.mod(a, 1n+mp), 1n))
330
331# ---- division and clamping ----
332
333def divmod(+a: Nat, +b: Nat) -> Result<&2, &2, MathError, QuotRem>:
334  match b:
335    case 0n:
336      Fail{ZeroDivision{}}
337    case 1n+ +bp:
338      Done{QR{Nat.div(a, 1n+bp), Nat.mod(a, 1n+bp)}}
339
340def clamp_ok(+x: Nat, +lo: Nat, +hi: Nat, +bad: Bool) -> Result<&2, &2, MathError, Nat>:
341  match bad:
342    case True{}:
343      Fail{Domain{}}
344    case False{}:
345      Done{Nat.min(Nat.max(x, lo), hi)}
346
347def clamp(+x: Nat, +lo: Nat, +hi: Nat) -> Result<&2, &2, MathError, Nat>:
348  clamp_ok(x, lo, hi, Nat.is_lt(hi, lo))