random.bend
bend-kit-random@0.1.0.0 · 168 lines · 0 definitions
1# Seeded pseudo-random numbers: xoshiro128**, unbiased ranges, floats and shuffles. Not for cryptography.
2import Base
3
4# A pure generator: each draw returns the value beside the next state.
5# import ./random/random.bend as Rand
6# (x, r) = Rand.next(Rand.seed(42))
7
8# xoshiro128** state (https://prng.di.unimi.it/xoshiro128starstar.c). It is never all zero.
9type Rng is Data:
10 Rng{a: U32, b: U32, c: U32, d: U32}
11
12def rotl(+x: U32, l: Nat, r: Nat) -> U32:
13 ((x << l) .|. (x >> r) : U32)
14
15# splitmix32's output mixer, a bijection on U32.
16def mix(+z: U32) -> U32:
17 +y = ((z .^. (z >> 16n)) * 569420461 : U32)
18 +w = ((y .^. (y >> 15n)) * 1935289751 : U32)
19 (w .^. (w >> 15n) : U32)
20
21# The first four splitmix32 outputs from s. Distinct seeds give distinct states.
22def seed(+s: U32) -> Rng:
23 Rng{mix((s + 2654435769 : U32)), mix((s + 1013904242 : U32)),
24 mix((s + 3668340011 : U32)), mix((s + 2027808484 : U32))}
25
26def from_state.if(zero: Bool, a: U32, b: U32, c: U32, d: U32) -> Rng:
27 match zero:
28 case True{}:
29 seed(0)
30 case False{}:
31 Rng{a, b, c, d}
32
33# The state (a, b, c, d) as is, to replay a reference stream. All zero gives seed(0).
34def from_state(+a: U32, +b: U32, +c: U32, +d: U32) -> Rng:
35 from_state.if(U32.is_zero((a .|. b .|. c .|. d : U32)), a, b, c, d)
36
37# One xoshiro128** step: a uniform U32 and the next state.
38def next(r: Rng) -> U32 & Rng:
39 match r:
40 case Rng{+a, +b, +c, +d}:
41 +c1 = (c .^. a : U32)
42 +d1 = (d .^. b : U32)
43 ((rotl((b * 5 : U32), 7n, 25n) * 9 : U32),
44 Rng{(a .^. d1 : U32), (b .^. c1 : U32), (c1 .^. (b << 9n) : U32),
45 rotl(d1, 11n, 21n)})
46
47def Bool.case(-A: Type, b: Bool, t: Unit -> A, f: Unit -> A) -> A:
48 match b:
49 case True{}:
50 t(Unit{})
51 case False{}:
52 f(Unit{})
53
54# Rejects draws below t = 2^32 mod n, as OpenBSD's arc4random_uniform does.
55def below.go(fuel: Nat, +n: U32, +t: U32, xr: U32 & Rng) -> U32 & Rng:
56 match fuel:
57 case 0n:
58 (x, r) = xr
59 ((x % n : U32), r)
60 case 1n+p:
61 (+x, +r) = xr
62 Bool.case(U32 & Rng, U32.is_ge(x, t),
63 u => ((x % n : U32), r),
64 u => below.go(p, n, t, next(r)))
65
66def below.if(small: Bool, +n: U32, r: Rng) -> U32 & Rng:
67 match small:
68 case True{}:
69 (0, r)
70 case False{}:
71 below.go(64n, n, ((0 - n : U32) % n : U32), next(r))
72
73# A uniform U32 in [0, n). Each draw is rejected with odds under 1/2, so after
74# 64 rejections, odds 2^-64, it takes x mod n. For n < 2 it is 0 and r is unchanged.
75def below(+n: U32, r: Rng) -> U32 & Rng:
76 below.if(U32.is_lt(n, 2), n, r)
77
78def range.fin(lo: U32, xr: U32 & Rng) -> U32 & Rng:
79 (x, r) = xr
80 ((lo + x : U32), r)
81
82# A uniform U32 in [lo, hi). For hi <= lo + 1 it is lo and r is unchanged.
83def range(+lo: U32, hi: U32, r: Rng) -> U32 & Rng:
84 range.fin(lo, below((hi - lo : U32), r))
85
86def unit.fin(xr: U32 & Rng) -> F32 & Rng:
87 (x, r) = xr
88 (F32.div(U32.to_f32((x >> 8n : U32)), 16777216.0), r)
89
90# A uniform F32 in [0, 1), from the draw's top 24 bits.
91def unit(r: Rng) -> F32 & Rng:
92 unit.fin(next(r))
93
94def shuffle.fill(~T: Data, xs: List<&2, T>, +i: U32, a: Array<T>) -> Array<T>:
95 match xs:
96 case Nil{}:
97 a
98 case Con{h, t}:
99 shuffle.fill(~T, t, (i + 1 : U32), Array.set(T, a, i, h))
100
101def shuffle.swap.set(~T: Data, +i: U32, +j: U32, xi: T, ax: Array<T> & T) -> Array<T>:
102 (a, xj) = ax
103 Array.set(T, Array.set(T, a, i, xj), j, xi)
104
105def shuffle.swap(~T: Data, +i: U32, +j: U32, ax: Array<T> & T) -> Array<T>:
106 (a, xi) = ax
107 shuffle.swap.set(~T, i, j, xi, Array.get(T, a, j))
108
109# Fisher-Yates: at slot i, jr is below(i + 1); slot i swaps with slot j.
110# At k = 0, i is 0 and below(1) drew nothing.
111def shuffle.go(~T: Data, k: Nat, +i: U32, a: Array<T>, jr: U32 & Rng) -> Array<T> & Rng:
112 match k:
113 case 0n:
114 (j, r) = jr
115 (a, r)
116 case 1n+p:
117 (+j, r) = jr
118 shuffle.go(~T, p, (i - 1 : U32), shuffle.swap(~T, i, j, Array.get(T, a, i)),
119 below(i, r))
120
121# Slots [0, k] of the array, with ax the read of slot k, prepended to acc.
122def shuffle.read(~T: Data, k: Nat, ax: Array<T> & T, acc: List<&2, T>) -> List<&2, T>:
123 match k:
124 case 0n:
125 (a, x) = ax
126 x <> acc
127 case 1n+(+p):
128 (a, x) = ax
129 shuffle.read(~T, p, Array.get(T, a, U32.from_nat(p)), x <> acc)
130
131def shuffle.fin(~T: Data, +m: U32, ar: Array<T> & Rng) -> List<&2, T> & Rng:
132 (a, r) = ar
133 (shuffle.read(~T, U32.to_nat(m), Array.get(T, a, m), Nil{}), r)
134
135def shuffle.run(~T: Data, +h: T, +xs: List<&2, T>, r: Rng) -> List<&2, T> & Rng:
136 +m = (U32.from_nat(List.length(&2, T, xs)) - 1 : U32)
137 shuffle.fin(~T, m,
138 shuffle.go(~T, U32.to_nat(m), m,
139 shuffle.fill(~T, xs, 0, Array.new(T, Nat.add(1n, U32.log2(m)), h)),
140 below((m + 1 : U32), r)))
141
142# A uniform permutation of xs. Every draw is below(i + 1), so a list of n
143# items uses n - 1 draws.
144def shuffle(~T: Data, xs: List<&2, T>, r: Rng) -> List<&2, T> & Rng:
145 match xs:
146 case Nil{}:
147 (Nil{}, r)
148 case Con{+h, t}:
149 shuffle.run(~T, h, h <> t, r)
150
151# 16 bytes from the OS (getentropy, or crypto.getRandomValues in JS).
152def entropy() -> IO(Result<&1, &1, U32 & String, U32 & U32 & U32 & U32>):
153 import "./effs/random.c"
154 import "./effs/random.js"
155
156def from_os.fin(e: Result<&1, &1, U32 & String, U32 & U32 & U32 & U32>) ->
157 Result<&1, &1, U32 & String, Rng>:
158 match e:
159 case Fail{err}:
160 Fail{err}
161 case Done{(a, b, c, d)}:
162 Done{from_state(a, b, c, d)}
163
164# A generator seeded from OS entropy, for when runs need not repeat.
165def from_os() -> IO(Result<&1, &1, U32 & String, Rng>):
166 do IO<Result<&1, &1, U32 & String, Rng>>:
167 e : Result<&1, &1, U32 & String, U32 & U32 & U32 & U32> <- entropy()
168 return from_os.fin(e)