1	import Base
2	import ../u64.bend as W
3	import ../w64.bend as X
4	import ../f64.bend as F
5	
6	# Random values from any source, written once: Go's math/rand/v2 Rand
7	# methods, bit for bit (the 64-bit code paths).
8	#
9	# The Source interface. A source is any state type S with a step
10	#
11	#   ~next: S -> W.U64 & S       the next 64-bit output and the new state
12	#
13	# passed as templates (~S, ~next), as src/math/num.bend passes a numeric
14	# type's ~op and ~test: each source compiles to its own copy with next
15	# inlined, and the state is threaded explicitly (Bend is pure: every
16	# function returns its value next to the advanced state). The sources here
17	# are chacha8.bend's ChaCha8 and pcg.bend's PCG:
18	#
19	#   import ./src/math/random.bend as R
20	#   import ./src/math/random/chacha8.bend as C8
21	#   (k, g) = R.uint64n(~C8.ChaCha8, ~R.chacha8_next, g, n)
22	#
23	# A later module (math/statistics: normal, exponential, ...) is written the
24	# same way, once for every source, on top of these functions:
25	#
26	#   def normal(~S: Data, ~next: S -> W.U64 & S, s: S) -> F.F64 & S:
27	#     ... R.float64(~S, ~next, s) ...
28	#
29	# and its laws can be stated and proved for an arbitrary ~next (a template
30	# proof is checked once, for every source; proofs/math/random/rand.bend).
31	#
32	#   uint64(s)            next(s)                           Go Uint64
33	#   uint32(s)            the top 32 bits                   Go Uint32
34	#   int64(s)             the low 63 bits                   Go Int64
35	#   int32(s)             the top 31 bits                   Go Int32
36	#   uint64n(s, n)        uniform in [0, n) for n > 0       Go Uint64N
37	#                        (uint_below is the same function; n == 0 is read
38	#                        as 2^64, Go's internal uint64n(0): a full uint64)
39	#   uint32n(s, n)        uniform in [0, n), n: U32 > 0      Go Uint32N
40	#   intn(s, n)           uniform in [0, n), n: Nat,        Go IntN
41	#                        0 < n < 2^48 (a Nat's run-time bound); 0 for n == 0
42	#   int_range(s, lo, hi) uniform in [lo, hi), lo < hi     lo + IntN(hi - lo)
43	#   float64(s)           uniform multiple of 2^-53 in [0, 1)  Go Float64
44	#   shuffle(s, xs)       a Fisher-Yates permutation of xs  Go Shuffle
45	#   perm(s, n)           a random permutation of 0..n-1    Go Perm
46	#
47	# uint64n is Lemire's nearly divisionless method ("Fast Random Integer
48	# Generation in an Interval", ACM TOMACS 2019), exactly as Go's uint64n:
49	# a power of two n masks the low bits; otherwise the 128-bit product x * n
50	# is split into hi:lo, and x is rejected while lo < 2^64 mod n (computed,
51	# with a division, only when lo < n). The accepted x map onto every k < n
52	# equally often (proved: proofs/math/random/lemire.bend), so the output is
53	# exactly uniform for a uniform source. Go retries forever; here at most
54	# 128 draws are made (Bend requires termination): a uniform source rejects
55	# with probability below 1/2 per draw, so the 128th rejection has
56	# probability below 2^-128, and even then the result is below n.
57	
58	def fst(-A: Data, -S: Data, p: A & S) -> A:
59	  (a, s) = p
60	  a
61	
62	def snd(-A: Data, -S: Data, p: A & S) -> S:
63	  (a, s) = p
64	  s
65	
66	def uint64(~S: Data, ~next: S -> W.U64 & S, s: S) -> W.U64 & S:
67	  next(s)
68	
69	def top32(-S: Data, p: W.U64 & S) -> U32 & S:
70	  (x, s) = p
71	  (X.hi(x), s)
72	
73	def uint32(~S: Data, ~next: S -> W.U64 & S, s: S) -> U32 & S:
74	  top32(S, next(s))
75	
76	def low63(-S: Data, p: W.U64 & S) -> W.U64 & S:
77	  (+x, s) = p
78	  (W.U64{X.lo(x), U32.and(X.hi(x), 2147483647)}, s)
79	
80	def int64(~S: Data, ~next: S -> W.U64 & S, s: S) -> W.U64 & S:
81	  low63(S, next(s))
82	
83	def top31(-S: Data, p: W.U64 & S) -> U32 & S:
84	  (x, s) = p
85	  (U32.shr(X.hi(x)), s)
86	
87	def int32(~S: Data, ~next: S -> W.U64 & S, s: S) -> U32 & S:
88	  top31(S, next(s))
89	
90	# ---- bounded integers: Lemire ----
91	
92	def and64(+a: W.U64, +b: W.U64) -> W.U64:
93	  W.U64{U32.and(X.lo(a), X.lo(b)), U32.and(X.hi(a), X.hi(b))}
94	
95	# n & (n - 1) == 0: n is a power of two (or zero)
96	def is_pow2(+n: W.U64) -> Bool:
97	  X.is_zero(and64(n, X.sub(n, W.U64{1, 0})))
98	
99	def mask(-S: Data, +n: W.U64, p: W.U64 & S) -> W.U64 & S:
100	  (+x, s) = p
101	  (and64(x, X.sub(n, W.U64{1, 0})), s)
102	
103	# 2^64 mod n = (2^64 - n) mod n, for n > 0 (Go: -n % n)
104	def thresh(+n: W.U64) -> W.U64:
105	  X.rem(X.sub(W.U64{0, 0}, n), n)
106	
107	# the draw x as the product x * n = hi:lo, kept with the state
108	type Draw<-S: Data> is Data:
109	  D{hi: W.U64, lo: W.U64, s: S}
110	
111	def draw_fin(-S: Data, s: S, p: W.U64 & W.U64) -> Draw<S>:
112	  (lo, hi) = p
113	  D{hi, lo, s}
114	
115	def draw(-S: Data, +n: W.U64, p: W.U64 & S) -> Draw<S>:
116	  (+x, s) = p
117	  draw_fin(S, s, X.mul128(x, n))
118	
119	# one rejection step: while lo < t, draw again; an accepted draw is kept
120	def again(~S: Data, ~next: S -> W.U64 & S, +n: W.U64, +hi: W.U64, +lo: W.U64, s: S, reject: Bool) -> Draw<S>:
121	  match reject:
122	    case False{}:
123	      D{hi, lo, s}
124	    case True{}:
125	      draw(S, n, next(s))
126	
127	def step(~S: Data, ~next: S -> W.U64 & S, +n: W.U64, +t: W.U64, d: Draw<S>) -> Draw<S>:
128	  match d:
129	    case D{+hi, +lo, s}:
130	      again(~S, ~next, n, hi, lo, s, X.lt(lo, t))
131	
132	# fuel rejection steps (an accepted draw is a fixed point of step)
133	def retry(~S: Data, ~next: S -> W.U64 & S, fuel: Nat, +n: W.U64, +t: W.U64, d: Draw<S>) -> Draw<S>:
134	  match fuel:
135	    case 0n:
136	      d
137	    case 1n+f:
138	      retry(~S, ~next, f, n, t, step(~S, ~next, n, t, d))
139	
140	def result(-S: Data, d: Draw<S>) -> W.U64 & S:
141	  match d:
142	    case D{hi, lo, s}:
143	      (hi, s)
144	
145	# lo >= n >= t accepts at once; otherwise t = 2^64 mod n is computed and at
146	# most 127 more draws are made
147	def lemire_small(~S: Data, ~next: S -> W.U64 & S, +n: W.U64, +hi: W.U64, +lo: W.U64, s: S, small: Bool) -> W.U64 & S:
148	  match small:
149	    case False{}:
150	      (hi, s)
151	    case True{}:
152	      result(S, retry(~S, ~next, 127n, n, thresh(n), D{hi, lo, s}))
153	
154	def lemire_first(~S: Data, ~next: S -> W.U64 & S, +n: W.U64, d: Draw<S>) -> W.U64 & S:
155	  match d:
156	    case D{+hi, +lo, s}:
157	      lemire_small(~S, ~next, n, hi, lo, s, X.lt(lo, n))
158	
159	def uint64n_pick(~S: Data, ~next: S -> W.U64 & S, s: S, +n: W.U64, pow2: Bool) -> W.U64 & S:
160	  match pow2:
161	    case True{}:
162	      mask(S, n, next(s))
163	    case False{}:
164	      lemire_first(~S, ~next, n, draw(S, n, next(s)))
165	
166	# uniform in [0, n) for n > 0 (Go's Uint64N)
167	def uint64n(~S: Data, ~next: S -> W.U64 & S, s: S, +n: W.U64) -> W.U64 & S:
168	  uint64n_pick(~S, ~next, s, n, is_pow2(n))
169	
170	def uint_below(~S: Data, ~next: S -> W.U64 & S, s: S, +n: W.U64) -> W.U64 & S:
171	  uint64n(~S, ~next, s, n)
172	
173	def lo32(-S: Data, p: W.U64 & S) -> U32 & S:
174	  (x, s) = p
175	  (X.lo(x), s)
176	
177	# uniform in [0, n) for n > 0 (Go's Uint32N)
178	def uint32n(~S: Data, ~next: S -> W.U64 & S, s: S, +n: U32) -> U32 & S:
179	  lo32(S, uint64n(~S, ~next, s, W.U64{n, 0}))
180	
181	# the 32-bit word of the low 8k bits of n, a byte at a time
182	def w_go(k: Nat, +n: Nat) -> U32:
183	  match k:
184	    case 0n:
185	      0
186	    case 1n+j:
187	      U32.add(U32.from_nat(Nat.mod(n, 256n)), U32.mul(w_go(j, Nat.div(n, 256n)), X.pow2(8n)))
188	
189	# n div 2^(8k)
190	def skip(k: Nat, +n: Nat) -> Nat:
191	  match k:
192	    case 0n:
193	      n
194	    case 1n+j:
195	      skip(j, Nat.div(n, 256n))
196	
197	# the 64-bit word of n < 2^64
198	def nat64(+n: Nat) -> W.U64:
199	  W.U64{w_go(4n, n), w_go(4n, skip(4n, n))}
200	
201	def to_nat(-S: Data, p: W.U64 & S) -> Nat & S:
202	  (x, s) = p
203	  (X.n48(x), s)
204	
205	def intn_z(~S: Data, ~next: S -> W.U64 & S, s: S, +n: Nat, z: Bool) -> Nat & S:
206	  match z:
207	    case True{}:
208	      (0n, s)
209	    case False{}:
210	      to_nat(S, uint64n(~S, ~next, s, nat64(n)))
211	
212	# uniform in [0, n) for 0 < n < 2^48 (Go's IntN); (0, s) for n == 0
213	def intn(~S: Data, ~next: S -> W.U64 & S, s: S, +n: Nat) -> Nat & S:
214	  intn_z(~S, ~next, s, n, Nat.is_eq(n, 0n))
215	
216	def plus(-S: Data, +lo: Nat, p: Nat & S) -> Nat & S:
217	  (k, s) = p
218	  (Nat.add(lo, k), s)
219	
220	# uniform in [lo, hi) for lo < hi, hi - lo < 2^48; (lo, s) when hi <= lo
221	def int_range(~S: Data, ~next: S -> W.U64 & S, s: S, +lo: Nat, +hi: Nat) -> Nat & S:
222	  plus(S, lo, intn(~S, ~next, s, Nat.sub(hi, lo)))
223	
224	# ---- floats ----
225	
226	# x >> 11 as a double divided by 2^53: for m = x >> 11 > 0 with c leading
227	# zeros (11 <= c <= 63), m << (c - 11) has its top bit at 52, and the value
228	# m 2^-53 is that significand with biased exponent 1033 - c (below 1023)
229	def f53(+m: W.U64, +c: Nat) -> F.F64:
230	  F.Bits{X.lo(X.shl(m, Nat.sub(c, 11n))), U32.add(U32.mul(U32.from_nat(Nat.sub(1033n, c)), 1048576), U32.and(X.hi(X.shl(m, Nat.sub(c, 11n))), 1048575))}
231	
232	def f53_z(+m: W.U64, z: Bool) -> F.F64:
233	  match z:
234	    case True{}:
235	      F.zero(False{})
236	    case False{}:
237	      f53(m, X.clz(m))
238	
239	# Go's Float64: float64(x << 11 >> 11) / (1 << 53), the low 53 bits
240	def to_float(+x: W.U64) -> F.F64:
241	  f53_z(W.U64{X.lo(x), U32.and(X.hi(x), 2097151)}, X.is_zero(W.U64{X.lo(x), U32.and(X.hi(x), 2097151)}))
242	
243	def float_of(-S: Data, p: W.U64 & S) -> F.F64 & S:
244	  (+x, s) = p
245	  (to_float(x), s)
246	
247	def float64(~S: Data, ~next: S -> W.U64 & S, s: S) -> F.F64 & S:
248	  float_of(S, next(s))
249	
250	# ---- permutations ----
251	
252	def nth(-A: Data, xs: List<&2, A>, i: Nat) -> Maybe<&2, A>:
253	  match xs i:
254	    case Nil{} _:
255	      None{}
256	    case Con{h, t} 0n:
257	      Some{h}
258	    case Con{h, t} 1n+p:
259	      nth(A, t, p)
260	
261	def put(-A: Data, xs: List<&2, A>, i: Nat, x: A) -> List<&2, A>:
262	  match xs i:
263	    case Nil{} _:
264	      Nil{}
265	    case Con{h, t} 0n:
266	      Con{x, t}
267	    case Con{h, t} 1n+p:
268	      Con{h, put(A, t, p, x)}
269	
270	def swap_m(-A: Data, xs: List<&2, A>, +i: Nat, +j: Nat, a: Maybe<&2, A>, b: Maybe<&2, A>) -> List<&2, A>:
271	  match a b:
272	    case Some{x} Some{y}:
273	      put(A, put(A, xs, i, y), j, x)
274	    case _ _:
275	      xs
276	
277	# exchange elements i and j (unchanged when either is out of range)
278	def swap(-A: Data, +xs: List<&2, A>, +i: Nat, +j: Nat) -> List<&2, A>:
279	  swap_m(A, xs, i, j, nth(A, xs, i), nth(A, xs, j))
280	
281	def swap_at(-A: Data, -S: Data, xs: List<&2, A>, +i: Nat, p: W.U64 & S) -> List<&2, A> & S:
282	  (+j, s) = p
283	  (swap(A, xs, i, X.n48(j)), s)
284	
285	# the length as a 64-bit word
286	def len64(-A: Data, xs: List<&2, A>) -> W.U64:
287	  match xs:
288	    case Nil{}:
289	      W.U64{0, 0}
290	    case Con{x, rest}:
291	      X.add(len64(A, rest), W.U64{1, 0})
292	
293	# swap element i with a uniform j < n = i + 1
294	def shuffle_step(~A: Data, ~S: Data, ~next: S -> W.U64 & S, +i: Nat, +n: W.U64, st: List<&2, A> & S) -> List<&2, A> & S:
295	  (xs, s) = st
296	  swap_at(A, S, xs, i, uint64n(~S, ~next, s, n))
297	
298	# i = k, k - 1, ..., 1 with the bound n = i + 1 as a 64-bit word
299	def shuffle_go(~A: Data, ~S: Data, ~next: S -> W.U64 & S, k: Nat, +n: W.U64, st: List<&2, A> & S) -> List<&2, A> & S:
300	  match k:
301	    case 0n:
302	      st
303	    case 1n+ +p:
304	      shuffle_go(~A, ~S, ~next, p, X.sub(n, W.U64{1, 0}), shuffle_step(~A, ~S, ~next, 1n+p, n, st))
305	
306	# Go's Shuffle (Fisher-Yates, from the last index down), for fewer than 2^48
307	# elements
308	def shuffle(~A: Data, ~S: Data, ~next: S -> W.U64 & S, s: S, +xs: List<&2, A>) -> List<&2, A> & S:
309	  shuffle_go(~A, ~S, ~next, Nat.sub(List.length(&2, A, xs), 1n), len64(A, xs), (xs, s))
310	
311	def range_go(n: Nat, acc: List<&2, Nat>) -> List<&2, Nat>:
312	  match n:
313	    case 0n:
314	      acc
315	    case 1n+ +p:
316	      range_go(p, Con{p, acc})
317	
318	# [0, 1, ..., n - 1]
319	def range(n: Nat) -> List<&2, Nat>:
320	  range_go(n, [])
321	
322	# Go's Perm: a shuffle of 0..n-1
323	def perm(~S: Data, ~next: S -> W.U64 & S, s: S, +n: Nat) -> List<&2, Nat> & S:
324	  shuffle(~Nat, ~S, ~next, s, range(n))