1	import Base
2	import ./u64.bend as W
3	import ./w64.bend as X
4	import ./generic.bend as G
5	import ./instances.bend as I
6	import ./number.bend as NB
7	import ./num.bend as N
8	
9	# Fixed-width integer operations for U32 and the two-limb U64, as Rust
10	# defines them for u32/u64 (design reference section 2.6: "provide explicit
11	# families: checked_add, wrapping_add, saturating_add, overflowing_add"),
12	# plus bit counts, primality and byte serialisation (sections 3.2, 10).
13	#
14	#   checked_{add,sub,mul,div,rem,pow,shl,shr}    Maybe: None on overflow,
15	#                                                a zero divisor, or a
16	#                                                shift >= the width
17	#   wrapping_{add,sub,mul,pow,shl,shr}           modulo 2^width (shifts use
18	#                                                the amount mod the width)
19	#   saturating_{add,sub,mul,pow}                 clamped to [0, 2^width - 1]
20	#   overflowing_{add,sub,mul,pow,shl,shr}        (wrapping result, overflowed)
21	#   bit_count(x)                                 number of 1 bits
22	#   is_prime(x), next_prime(x)                   U32: trial division (exact)
23	#   to_bytes_le/be, from_bytes_le/be             bytes as U32s in [0, 256)
24	#
25	# Unsigned division never overflows, so Rust's wrapping_/saturating_/
26	# overflowing_ div and rem are plain division; checked_div/rem cover the
27	# zero divisor (Rust panics there; here it is None). Shift amounts and pow
28	# exponents are U32, as in Rust. Every operation is the Nat operation on
29	# the values (spec/math/fixed.bend).
30	
31	type Over<-T: Data> is Data:
32	  OV{value: T, overflow: Bool}
33	
34	# ---- shared helpers ----
35	
36	def opt(-T: Data, bad: Bool, +r: T) -> Maybe<&2, T>:
37	  match bad:
38	    case True{}:
39	      None{}
40	    case False{}:
41	      Some{r}
42	
43	def pick(-T: Data, c: Bool, +a: T, +b: T) -> T:
44	  match c:
45	    case True{}:
46	      a
47	    case False{}:
48	      b
49	
50	def res_opt(-T: Data, r: Result<&2, &2, N.NumError, T>) -> Maybe<&2, T>:
51	  match r:
52	    case Done{x}:
53	      Some{x}
54	    case Fail{e}:
55	      None{}
56	
57	def or_top(-T: Data, +top: T, r: Maybe<&2, T>) -> T:
58	  match r:
59	    case Some{x}:
60	      x
61	    case None{}:
62	      top
63	
64	def res_bad(-T: Data, r: Result<&2, &2, N.NumError, T>) -> Bool:
65	  match r:
66	    case Done{x}:
67	      False{}
68	    case Fail{e}:
69	      True{}
70	
71	# ---- U32 ----
72	
73	def u32_max() -> U32:
74	  4294967295
75	
76	def u32_add_over(+a: U32, +b: U32) -> Bool:
77	  I.u32_add_over(a, b)
78	
79	def u32_mul_over(+a: U32, +b: U32) -> Bool:
80	  I.u32_mul_over(a, b)
81	
82	def u32_checked_add(+a: U32, +b: U32) -> Maybe<&2, U32>:
83	  opt(U32, u32_add_over(a, b), U32.add(a, b))
84	
85	def u32_wrapping_add(+a: U32, +b: U32) -> U32:
86	  U32.add(a, b)
87	
88	def u32_saturating_add(+a: U32, +b: U32) -> U32:
89	  pick(U32, u32_add_over(a, b), u32_max(), U32.add(a, b))
90	
91	def u32_overflowing_add(+a: U32, +b: U32) -> Over<U32>:
92	  OV{U32.add(a, b), u32_add_over(a, b)}
93	
94	def u32_checked_sub(+a: U32, +b: U32) -> Maybe<&2, U32>:
95	  opt(U32, U32.is_lt(a, b), U32.sub(a, b))
96	
97	def u32_wrapping_sub(+a: U32, +b: U32) -> U32:
98	  U32.sub(a, b)
99	
100	def u32_saturating_sub(+a: U32, +b: U32) -> U32:
101	  pick(U32, U32.is_lt(a, b), 0, U32.sub(a, b))
102	
103	def u32_overflowing_sub(+a: U32, +b: U32) -> Over<U32>:
104	  OV{U32.sub(a, b), U32.is_lt(a, b)}
105	
106	def u32_checked_mul(+a: U32, +b: U32) -> Maybe<&2, U32>:
107	  opt(U32, u32_mul_over(a, b), U32.mul(a, b))
108	
109	def u32_wrapping_mul(+a: U32, +b: U32) -> U32:
110	  U32.mul(a, b)
111	
112	def u32_saturating_mul(+a: U32, +b: U32) -> U32:
113	  pick(U32, u32_mul_over(a, b), u32_max(), U32.mul(a, b))
114	
115	def u32_overflowing_mul(+a: U32, +b: U32) -> Over<U32>:
116	  OV{U32.mul(a, b), u32_mul_over(a, b)}
117	
118	def u32_checked_div(+a: U32, +b: U32) -> Maybe<&2, U32>:
119	  opt(U32, U32.is_zero(b), U32.div(a, b))
120	
121	def u32_checked_rem(+a: U32, +b: U32) -> Maybe<&2, U32>:
122	  opt(U32, U32.is_zero(b), U32.mod(a, b))
123	
124	def u32_checked_pow(+a: U32, +e: U32) -> Maybe<&2, U32>:
125	  res_opt(U32, G.pow(~U32, ~I.u32_op, ~I.u32_is, a, U32.to_nat(e)))
126	
127	# binary exponentiation modulo 2^32; e halves every step (32 steps)
128	def u32_wpow(fuel: Nat, +e: U32, +b: U32, +acc: U32, z: Bool) -> U32:
129	  match fuel z:
130	    case 0n _:
131	      acc
132	    case 1n+f True{}:
133	      acc
134	    case 1n+ +f False{}:
135	      u32_wpow(f, U32.shr(e), U32.mul(b, b), pick(U32, I.u32_odd(e), U32.mul(acc, b), acc), U32.is_zero(U32.shr(e)))
136	
137	def u32_wrapping_pow(+a: U32, +e: U32) -> U32:
138	  u32_wpow(32n, e, a, 1, U32.is_zero(e))
139	
140	def u32_saturating_pow(+a: U32, +e: U32) -> U32:
141	  or_top(U32, u32_max(), u32_checked_pow(a, e))
142	
143	def u32_overflowing_pow(+a: U32, +e: U32) -> Over<U32>:
144	  OV{u32_wrapping_pow(a, e), res_bad(U32, G.pow(~U32, ~I.u32_op, ~I.u32_is, a, U32.to_nat(e)))}
145	
146	# shifts by s mod 32 (Rust masks the amount the same way); a left shift is
147	# the product with 2^k modulo 2^32
148	def u32_wrapping_shl(+a: U32, +s: U32) -> U32:
149	  U32.mul(a, X.pow2(U32.to_nat(U32.and(s, 31))))
150	
151	def u32_checked_shl(+a: U32, +s: U32) -> Maybe<&2, U32>:
152	  opt(U32, Bool.not(U32.is_lt(s, 32)), u32_wrapping_shl(a, s))
153	
154	def u32_overflowing_shl(+a: U32, +s: U32) -> Over<U32>:
155	  OV{u32_wrapping_shl(a, s), Bool.not(U32.is_lt(s, 32))}
156	
157	def u32_wrapping_shr(+a: U32, +s: U32) -> U32:
158	  U32.shrn(a, U32.to_nat(U32.and(s, 31)))
159	
160	def u32_checked_shr(+a: U32, +s: U32) -> Maybe<&2, U32>:
161	  opt(U32, Bool.not(U32.is_lt(s, 32)), u32_wrapping_shr(a, s))
162	
163	def u32_overflowing_shr(+a: U32, +s: U32) -> Over<U32>:
164	  OV{u32_wrapping_shr(a, s), Bool.not(U32.is_lt(s, 32))}
165	
166	# ---- bit count ----
167	
168	# the 1 bits of a value below 16
169	def ones4(+n: Nat) -> Nat:
170	  match n:
171	    case 0n:
172	      0n
173	    case 1n:
174	      1n
175	    case 2n:
176	      1n
177	    case 3n:
178	      2n
179	    case 4n:
180	      1n
181	    case 5n:
182	      2n
183	    case 6n:
184	      2n
185	    case 7n:
186	      3n
187	    case 8n:
188	      1n
189	    case 9n:
190	      2n
191	    case 10n:
192	      2n
193	    case 11n:
194	      3n
195	    case 12n:
196	      2n
197	    case 13n:
198	      3n
199	    case 14n:
200	      3n
201	    case 15n:
202	      4n
203	    case 16n+m:
204	      0n
205	
206	# the 1 bits of the k low nibbles of x, a nibble at a time
207	def u32_bc_go(fuel: Nat, +x: U32) -> Nat:
208	  match fuel:
209	    case 0n:
210	      0n
211	    case 1n+k:
212	      Nat.add(ones4(U32.to_nat(U32.mod(x, 16))), u32_bc_go(k, U32.shrn(x, 4n)))
213	
214	def u32_bit_count(+x: U32) -> Nat:
215	  u32_bc_go(8n, x)
216	
217	# ---- primality (U32: trial division on the value) ----
218	
219	def u32_is_prime(+x: U32) -> Bool:
220	  NB.is_prime(U32.to_nat(x))
221	
222	# the least prime above n below 2^32: the candidates n+1, n+2, ... stop at
223	# the first prime or when the next candidate wraps to 0
224	def u32_np(fuel: Nat, +m: U32, wrapped: Bool, prime: Bool) -> Maybe<&2, U32>:
225	  match fuel wrapped prime:
226	    case 0n _ _:
227	      None{}
228	    case 1n+f True{} _:
229	      None{}
230	    case 1n+f False{} True{}:
231	      Some{m}
232	    case 1n+ +f False{} False{}:
233	      u32_np(f, U32.add(m, 1), U32.is_zero(U32.add(m, 1)), NB.is_prime(U32.to_nat(U32.add(m, 1))))
234	
235	def u32_next_prime(+n: U32) -> Maybe<&2, U32>:
236	  u32_np(Nat.add(U32.to_nat(U32.not(n)), 1n), U32.add(n, 1), U32.is_zero(U32.add(n, 1)), NB.is_prime(U32.to_nat(U32.add(n, 1))))
237	
238	# ---- bytes (U32 values in [0, 256)) ----
239	
240	def byte32(+x: U32, +k: Nat) -> U32:
241	  U32.mod(U32.shrn(x, Nat.mul(8n, k)), 256)
242	
243	def u32_to_bytes_le(+x: U32) -> List<&2, U32>:
244	  [byte32(x, 0n), byte32(x, 1n), byte32(x, 2n), byte32(x, 3n)]
245	
246	def u32_to_bytes_be(+x: U32) -> List<&2, U32>:
247	  [byte32(x, 3n), byte32(x, 2n), byte32(x, 1n), byte32(x, 0n)]
248	
249	def is_byte(+b: U32) -> Bool:
250	  U32.is_lt(b, 256)
251	
252	# one more base-256 digit below the value so far: b + 256 z
253	def dig32(r: Maybe<&2, U32>, ok: Bool, +b: U32) -> Maybe<&2, U32>:
254	  match r:
255	    case Some{+z}:
256	      match ok:
257	        case True{}:
258	          Some{U32.add(b, U32.mul(z, 256))}
259	        case False{}:
260	          None{}
261	    case None{}:
262	      None{}
263	
264	def from4(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> Maybe<&2, U32>:
265	  dig32(dig32(dig32(dig32(Some{0}, is_byte(b3), b3), is_byte(b2), b2), is_byte(b1), b1), is_byte(b0), b0)
266	
267	def u32_from_bytes_le(bs: List<&2, U32>) -> Maybe<&2, U32>:
268	  match bs:
269	    case [+b0, +b1, +b2, +b3]:
270	      from4(b0, b1, b2, b3)
271	    case _:
272	      None{}
273	
274	def u32_from_bytes_be(bs: List<&2, U32>) -> Maybe<&2, U32>:
275	  match bs:
276	    case [+b3, +b2, +b1, +b0]:
277	      from4(b0, b1, b2, b3)
278	    case _:
279	      None{}
280	
281	# ---- U64 ----
282	
283	def u64_max() -> W.U64:
284	  W.U64{4294967295, 4294967295}
285	
286	def u64_zero() -> W.U64:
287	  W.U64{0, 0}
288	
289	def u64_checked_add(+a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
290	  opt(W.U64, X.add_over(a, b), X.add(a, b))
291	
292	def u64_wrapping_add(+a: W.U64, +b: W.U64) -> W.U64:
293	  X.add(a, b)
294	
295	def u64_saturating_add(+a: W.U64, +b: W.U64) -> W.U64:
296	  pick(W.U64, X.add_over(a, b), u64_max(), X.add(a, b))
297	
298	def u64_overflowing_add(+a: W.U64, +b: W.U64) -> Over<W.U64>:
299	  OV{X.add(a, b), X.add_over(a, b)}
300	
301	# a - b modulo 2^64 for any a, b: a + (2^64 - 1 - b) + 1
302	def u64_wsub(+a: W.U64, +b: W.U64) -> W.U64:
303	  X.add(X.add(a, X.sub(u64_max(), b)), W.U64{1, 0})
304	
305	def u64_checked_sub(+a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
306	  opt(W.U64, X.lt(a, b), X.sub(a, b))
307	
308	def u64_wrapping_sub(+a: W.U64, +b: W.U64) -> W.U64:
309	  u64_wsub(a, b)
310	
311	def u64_saturating_sub(+a: W.U64, +b: W.U64) -> W.U64:
312	  pick(W.U64, X.lt(a, b), u64_zero(), X.sub(a, b))
313	
314	def u64_overflowing_sub(+a: W.U64, +b: W.U64) -> Over<W.U64>:
315	  OV{u64_wsub(a, b), X.lt(a, b)}
316	
317	def u64_checked_mul(+a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
318	  opt(W.U64, X.mul_over(a, b), X.mul(a, b))
319	
320	def u64_wrapping_mul(+a: W.U64, +b: W.U64) -> W.U64:
321	  X.mul(a, b)
322	
323	def u64_saturating_mul(+a: W.U64, +b: W.U64) -> W.U64:
324	  pick(W.U64, X.mul_over(a, b), u64_max(), X.mul(a, b))
325	
326	def u64_overflowing_mul(+a: W.U64, +b: W.U64) -> Over<W.U64>:
327	  OV{X.mul(a, b), X.mul_over(a, b)}
328	
329	# the quotient is only computed for a nonzero divisor (the two-limb
330	# division by zero would build a Nat past the runtime's 2^48 bound)
331	def u64_cdiv(z: Bool, +a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
332	  match z:
333	    case True{}:
334	      None{}
335	    case False{}:
336	      Some{X.quot(a, b)}
337	
338	def u64_crem(z: Bool, +a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
339	  match z:
340	    case True{}:
341	      None{}
342	    case False{}:
343	      Some{X.rem(a, b)}
344	
345	def u64_checked_div(+a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
346	  u64_cdiv(X.is_zero(b), a, b)
347	
348	def u64_checked_rem(+a: W.U64, +b: W.U64) -> Maybe<&2, W.U64>:
349	  u64_crem(X.is_zero(b), a, b)
350	
351	def u64_checked_pow(+a: W.U64, +e: U32) -> Maybe<&2, W.U64>:
352	  res_opt(W.U64, G.pow(~W.U64, ~I.u64_op, ~I.u64_is, a, U32.to_nat(e)))
353	
354	def u64_wpow(fuel: Nat, +e: U32, +b: W.U64, +acc: W.U64, z: Bool) -> W.U64:
355	  match fuel z:
356	    case 0n _:
357	      acc
358	    case 1n+f True{}:
359	      acc
360	    case 1n+ +f False{}:
361	      u64_wpow(f, U32.shr(e), X.mul(b, b), pick(W.U64, I.u32_odd(e), X.mul(acc, b), acc), U32.is_zero(U32.shr(e)))
362	
363	def u64_wrapping_pow(+a: W.U64, +e: U32) -> W.U64:
364	  u64_wpow(32n, e, a, W.U64{1, 0}, U32.is_zero(e))
365	
366	def u64_saturating_pow(+a: W.U64, +e: U32) -> W.U64:
367	  or_top(W.U64, u64_max(), u64_checked_pow(a, e))
368	
369	def u64_overflowing_pow(+a: W.U64, +e: U32) -> Over<W.U64>:
370	  OV{u64_wrapping_pow(a, e), res_bad(W.U64, G.pow(~W.U64, ~I.u64_op, ~I.u64_is, a, U32.to_nat(e)))}
371	
372	def u64_checked_shl(+a: W.U64, +s: U32) -> Maybe<&2, W.U64>:
373	  opt(W.U64, Bool.not(U32.is_lt(s, 64)), X.shl(a, U32.to_nat(U32.and(s, 63))))
374	
375	def u64_wrapping_shl(+a: W.U64, +s: U32) -> W.U64:
376	  X.shl(a, U32.to_nat(U32.and(s, 63)))
377	
378	def u64_overflowing_shl(+a: W.U64, +s: U32) -> Over<W.U64>:
379	  OV{u64_wrapping_shl(a, s), Bool.not(U32.is_lt(s, 64))}
380	
381	def u64_checked_shr(+a: W.U64, +s: U32) -> Maybe<&2, W.U64>:
382	  opt(W.U64, Bool.not(U32.is_lt(s, 64)), X.shr(a, U32.to_nat(U32.and(s, 63))))
383	
384	def u64_wrapping_shr(+a: W.U64, +s: U32) -> W.U64:
385	  X.shr(a, U32.to_nat(U32.and(s, 63)))
386	
387	def u64_overflowing_shr(+a: W.U64, +s: U32) -> Over<W.U64>:
388	  OV{u64_wrapping_shr(a, s), Bool.not(U32.is_lt(s, 64))}
389	
390	def u64_bit_count(+x: W.U64) -> Nat:
391	  Nat.add(u32_bit_count(X.lo(x)), u32_bit_count(X.hi(x)))
392	
393	def u64_to_bytes_le(+x: W.U64) -> List<&2, U32>:
394	  [byte32(X.lo(x), 0n), byte32(X.lo(x), 1n), byte32(X.lo(x), 2n), byte32(X.lo(x), 3n), byte32(X.hi(x), 0n), byte32(X.hi(x), 1n), byte32(X.hi(x), 2n), byte32(X.hi(x), 3n)]
395	
396	def u64_to_bytes_be(+x: W.U64) -> List<&2, U32>:
397	  [byte32(X.hi(x), 3n), byte32(X.hi(x), 2n), byte32(X.hi(x), 1n), byte32(X.hi(x), 0n), byte32(X.lo(x), 3n), byte32(X.lo(x), 2n), byte32(X.lo(x), 1n), byte32(X.lo(x), 0n)]
398	
399	def dig64(r: Maybe<&2, W.U64>, ok: Bool, +b: U32) -> Maybe<&2, W.U64>:
400	  match r:
401	    case Some{+z}:
402	      match ok:
403	        case True{}:
404	          Some{X.add(W.U64{b, 0}, X.mul(z, W.U64{256, 0}))}
405	        case False{}:
406	          None{}
407	    case None{}:
408	      None{}
409	
410	def u64_from_bytes_le(bs: List<&2, U32>) -> Maybe<&2, W.U64>:
411	  match bs:
412	    case [+b0, +b1, +b2, +b3, +b4, +b5, +b6, +b7]:
413	      dig64(dig64(dig64(dig64(dig64(dig64(dig64(dig64(Some{W.U64{0, 0}}, is_byte(b7), b7), is_byte(b6), b6), is_byte(b5), b5), is_byte(b4), b4), is_byte(b3), b3), is_byte(b2), b2), is_byte(b1), b1), is_byte(b0), b0)
414	    case _:
415	      None{}
416	
417	def u64_from_bytes_be(bs: List<&2, U32>) -> Maybe<&2, W.U64>:
418	  match bs:
419	    case [+b7, +b6, +b5, +b4, +b3, +b2, +b1, +b0]:
420	      dig64(dig64(dig64(dig64(dig64(dig64(dig64(dig64(Some{W.U64{0, 0}}, is_byte(b7), b7), is_byte(b6), b6), is_byte(b5), b5), is_byte(b4), b4), is_byte(b3), b3), is_byte(b2), b2), is_byte(b1), b1), is_byte(b0), b0)
421	    case _:
422	      None{}