src/math/natural.bend
bend-collections-laws-containers@1.0.0.0 · 348 lines · 0 definitions
1import Base
2import bend-collections-laws-containers@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))