src/math/random/rand.bend

bend-collections-laws-crypto@1.0.0.0 · 324 lines · 52 definitions

1import Base
2import bend-collections-laws-crypto@1.0.0.0/src/math/u64.bend as W
3import bend-collections-laws-crypto@1.0.0.0/src/math/w64.bend as X
4import bend-collections-laws-crypto@1.0.0.0/src/math/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
58def fst(-A: Data, -S: Data, p: A & S) -> A:
59  (a, s) = p
60  a
61
62def snd(-A: Data, -S: Data, p: A & S) -> S:
63  (a, s) = p
64  s
65
66def uint64(~S: Data, ~next: S -> W.U64 & S, s: S) -> W.U64 & S:
67  next(s)
68
69def top32(-S: Data, p: W.U64 & S) -> U32 & S:
70  (x, s) = p
71  (X.hi(x), s)
72
73def uint32(~S: Data, ~next: S -> W.U64 & S, s: S) -> U32 & S:
74  top32(S, next(s))
75
76def 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
80def int64(~S: Data, ~next: S -> W.U64 & S, s: S) -> W.U64 & S:
81  low63(S, next(s))
82
83def top31(-S: Data, p: W.U64 & S) -> U32 & S:
84  (x, s) = p
85  (U32.shr(X.hi(x)), s)
86
87def int32(~S: Data, ~next: S -> W.U64 & S, s: S) -> U32 & S:
88  top31(S, next(s))
89
90# ---- bounded integers: Lemire ----
91
92def 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)
96def is_pow2(+n: W.U64) -> Bool:
97  X.is_zero(and64(n, X.sub(n, W.U64{1, 0})))
98
99def 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)
104def 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
108type Draw<-S: Data> is Data:
109  D{hi: W.U64, lo: W.U64, s: S}
110
111def draw_fin(-S: Data, s: S, p: W.U64 & W.U64) -> Draw<S>:
112  (lo, hi) = p
113  D{hi, lo, s}
114
115def 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
120def 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
127def 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)
133def 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
140def 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
147def 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
154def 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
159def 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)
167def 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
170def uint_below(~S: Data, ~next: S -> W.U64 & S, s: S, +n: W.U64) -> W.U64 & S:
171  uint64n(~S, ~next, s, n)
172
173def 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)
178def 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
182def 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)
190def 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
198def nat64(+n: Nat) -> W.U64:
199  W.U64{w_go(4n, n), w_go(4n, skip(4n, n))}
200
201def to_nat(-S: Data, p: W.U64 & S) -> Nat & S:
202  (x, s) = p
203  (X.n48(x), s)
204
205def 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
213def 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
216def 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
221def 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)
229def 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
232def 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
240def 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
243def float_of(-S: Data, p: W.U64 & S) -> F.F64 & S:
244  (+x, s) = p
245  (to_float(x), s)
246
247def float64(~S: Data, ~next: S -> W.U64 & S, s: S) -> F.F64 & S:
248  float_of(S, next(s))
249
250# ---- permutations ----
251
252def 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
261def 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
270def 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)
278def 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
281def 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
286def 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
294def 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
299def 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
308def 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
311def 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]
319def range(n: Nat) -> List<&2, Nat>:
320  range_go(n, [])
321
322# Go's Perm: a shuffle of 0..n-1
323def perm(~S: Data, ~next: S -> W.U64 & S, s: S, +n: Nat) -> List<&2, Nat> & S:
324  shuffle(~Nat, ~S, ~next, s, range(n))