random.bend

bend-kit-random@0.1.0.1 · 168 lines · 26 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)