1	import Base
2	
3	# Integer extras on the natural numbers, after the design reference's
4	# section 3.2 (int.bit_count), 8.7 (the extended Euclidean algorithm) and 10
5	# (egcd, is_prime), next to natural.bend's section 4.2 functions.
6	#
7	#   bit_count(n)       the number of 1 bits of n (Python's int.bit_count)
8	#   egcd(a, b)         gcd(a, b) with its Bezout coefficients (see EGcd)
9	#   is_prime(n)        n is prime, by trial division up to sqrt(n)
10	#
11	# Nat is exact; Bend's native backend stops the program on a Nat above
12	# 2^48 - 1 (natural.bend), so is_prime is meant for n below that. The proofs
13	# are over the mathematical naturals.
14	
15	# gcd(a, b) and signed Bezout coefficients split into magnitudes and a sign:
16	#   neg == False:  a * x == gcd + b * y      (a*x - b*y == gcd)
17	#   neg == True:   b * y == gcd + a * x      (b*y - a*x == gcd)
18	# Python's egcd would return one negative coefficient; the unsigned library
19	# returns both magnitudes and says which side carries the minus sign.
20	type EGcd is Data:
21	  EG{gcd: Nat, x: Nat, y: Nat, neg: Bool}
22	
23	# ---- bit count ----
24	
25	# n halves every step, so n steps are enough
26	def bit_count_go(fuel: Nat, +n: Nat, +acc: Nat) -> Nat:
27	  match fuel n:
28	    case 0n _:
29	      acc
30	    case 1n+f 0n:
31	      acc
32	    case 1n+f 1n+ +np:
33	      bit_count_go(f, Nat.div(1n+np, 2n), Nat.add(acc, Nat.mod(1n+np, 2n)))
34	
35	def bit_count(+n: Nat) -> Nat:
36	  bit_count_go(n, n, 0n)
37	
38	# ---- extended Euclid ----
39	
40	# From the child (b, a mod b) to the parent (a, b): with q = a div b the
41	# coefficients become (y', x' + q*y') and the sign flips (Knuth 4.5.2,
42	# algorithm X, in magnitudes).
43	def eg_up(+q: Nat, r: EGcd) -> EGcd:
44	  match r:
45	    case EG{+g, +x, +y, neg}:
46	      EG{g, y, Nat.add(x, Nat.mul(q, y)), Bool.not(neg)}
47	
48	# b <= fuel throughout, as in natural.bend's gcd_go
49	def egcd_go(fuel: Nat, +a: Nat, +b: Nat) -> EGcd:
50	  match fuel b:
51	    case 0n _:
52	      EG{a, 1n, 0n, False{}}
53	    case 1n+f 0n:
54	      EG{a, 1n, 0n, False{}}
55	    case 1n+ +f 1n+ +bp:
56	      eg_up(Nat.div(a, 1n+bp), egcd_go(f, 1n+bp, Nat.mod(a, 1n+bp)))
57	
58	def egcd(+a: Nat, +b: Nat) -> EGcd:
59	  egcd_go(b, a, b)
60	
61	# ---- primality ----
62	
63	# trial division: no d in [2, n) divides n iff none in [2, isqrt(n)] does.
64	# `done` is n < d*d, `divides` is d | n, both for the current d.
65	def prime_go(fuel: Nat, +n: Nat, +d: Nat, done: Bool, divides: Bool) -> Bool:
66	  match fuel done divides:
67	    case 0n _ _:
68	      True{}
69	    case 1n+f True{} _:
70	      True{}
71	    case 1n+f False{} True{}:
72	      False{}
73	    case 1n+ +f False{} False{}:
74	      prime_go(f, n, 1n+d, Nat.is_lt(n, Nat.mul(1n+d, 1n+d)), Nat.is_eq(Nat.mod(n, 1n+d), 0n))
75	
76	def is_prime(+n: Nat) -> Bool:
77	  match n:
78	    case 0n:
79	      False{}
80	    case 1n:
81	      False{}
82	    case 2n+ +p:
83	      prime_go(2n+p, 2n+p, 2n, Nat.is_lt(2n+p, 4n), Nat.is_eq(Nat.mod(2n+p, 2n), 0n))