1	import Base
2	import ./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	
33	type 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)
40	type Bezout is Data:
41	  BZ{gcd: Nat, coef: Nat}
42	
43	# a // b and a % b
44	type 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).
50	def 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	
59	def gcd(+a: Nat, +b: Nat) -> Nat:
60	  gcd_go(b, a, b)
61	
62	def 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	
71	def 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	
78	def gcd_all(xs: List<&2, Nat>) -> Nat:
79	  gcd_all_go(xs, 0n)
80	
81	def 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	
88	def lcm_all(xs: List<&2, Nat>) -> Nat:
89	  lcm_all_go(xs, 1n)
90	
91	# ---- sums and products ----
92	
93	def 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	
100	def sum(xs: List<&2, Nat>) -> Nat:
101	  sum_go(xs, 0n)
102	
103	def 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	
110	def prod(xs: List<&2, Nat>) -> Nat:
111	  prod_go(xs, 1n)
112	
113	# ---- factorial, permutations, combinations ----
114	
115	def 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	
122	def factorial(+n: Nat) -> Nat:
123	  factorial_go(n, 1n)
124	
125	# n * (n - 1) * ... * (n - k + 1)
126	def 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	
133	def 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).
138	def 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	
145	def 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	
152	def comb(+n: Nat, +k: Nat) -> Nat:
153	  comb_small(n, k, Nat.is_lt(n, k))
154	
155	# ---- bits ----
156	
157	def 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	
166	def 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.
175	def 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.
180	def 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	
189	def 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
199	def 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	
210	def 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.
217	def sqrt_next(+n: Nat, +g: Nat) -> Nat:
218	  Nat.div(Nat.add(g, Nat.div(n, g)), 2n)
219	
220	def 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	
229	def 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	
232	def 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)
236	def 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	
245	def 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
254	def 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	
263	def 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	
270	def 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
277	def 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	
284	def 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	
293	def 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.
302	def 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	
305	def 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	
314	def 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)
324	def 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	
333	def 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	
340	def 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	
347	def clamp(+x: Nat, +lo: Nat, +hi: Nat) -> Result<&2, &2, MathError, Nat>:
348	  clamp_ok(x, lo, hi, Nat.is_lt(hi, lo))