spec/math/random/rand.bend

bend-collections-laws-crypto@1.0.0.0 · 116 lines · 12 definitions

1import Base
2import bend-collections-laws-crypto@1.0.0.0/spec/lib/common.bend as C
3import bend-collections-laws-crypto@1.0.0.0/src/math/u64.bend as W
4import bend-collections-laws-crypto@1.0.0.0/spec/math/w64.bend as SW
5import bend-collections-laws-crypto@1.0.0.0/spec/math/random/source.bend as SRC
6
7# Executable specification of the bounded draws, shuffles and permutations
8# of src/math/random/rand.bend, on natural numbers, transcribed from Go's
9# math/rand/v2 rand.go (uint64n, Shuffle, Perm) and Lemire, "Fast Random
10# Integer Generation in an Interval" (ACM TOMACS 29(1), 2019, Algorithm 5).
11# Widths are parameters: w = 64 for Uint64N (a 32-bit variant is w = 32);
12# no constant 2^w is ever formed (C.low, C.high, C.fits and pow2mod work on
13# a symbolic w), as spec/lib/common.bend explains.
14#
15#   and_bits(w, a, b)        the bitwise and of the low w bits
16#   pow2mod(w, n)            2^w mod n (Go's -n % n = (2^w - n) mod n)
17#   draw(w, x, n)            Go's uint64n decision for one source output x:
18#                            n == 0 (read as 2^w): Some{x}; n a power of two
19#                            (n & (n - 1) == 0): Some{x & (n - 1)}; otherwise
20#                            x * n = hi * 2^w + lo is accepted with hi unless
21#                            lo < n and lo < 2^w mod n (rejected: None)
22#   count(w, n, k, N)        #{x < N : draw(w, x, n) == Some{k}}
23#   below(~S, ~next, f, n, s)  the bounded draw from a source: the first
24#                            accepted draw among 1 + f draws (Go's loop,
25#                            which is unbounded; the last candidate if every
26#                            draw is rejected)
27#   occurrences(~rel, v, xs) #{i : rel(xs[i], v)}: two lists are permutations
28#                            of each other when these agree for an equality
29#                            rel and every v
30
31def b2n(b: Bool) -> Nat:
32  match b:
33    case True{}:
34      1n
35    case False{}:
36      0n
37
38def and_bits(w: Nat, +a: Nat, +b: Nat) -> Nat:
39  match w:
40    case 0n:
41      0n
42    case 1n+k:
43      Nat.add(Nat.mul(C.bit(a), C.bit(b)), Nat.double(and_bits(k, C.half(a), C.half(b))))
44
45def pow2mod(w: Nat, +n: Nat) -> Nat:
46  match w:
47    case 0n:
48      Nat.mod(1n, n)
49    case 1n+k:
50      Nat.mod(Nat.double(pow2mod(k, n)), n)
51
52def accept(+hi: Nat, ok: Bool) -> Maybe<&2, Nat>:
53  match ok:
54    case True{}:
55      Some{hi}
56    case False{}:
57      None{}
58
59# Lemire: the product x * n = hi * 2^w + lo
60def lemire(+w: Nat, +n: Nat, +hi: Nat, +lo: Nat) -> Maybe<&2, Nat>:
61  accept(hi, Bool.not(Bool.and(Nat.is_lt(lo, n), Nat.is_lt(lo, pow2mod(w, n)))))
62
63def draw_pos(+w: Nat, +x: Nat, +n: Nat, +m: Nat, pow2: Bool) -> Maybe<&2, Nat>:
64  match pow2:
65    case True{}:
66      Some{and_bits(w, x, m)}
67    case False{}:
68      lemire(w, n, C.high(w, Nat.mul(x, n)), C.low(w, Nat.mul(x, n)))
69
70def draw(+w: Nat, +x: Nat, n: Nat) -> Maybe<&2, Nat>:
71  match n:
72    case 0n:
73      Some{x}
74    case 1n+ +m:
75      draw_pos(w, x, 1n+m, m, Nat.is_eq(and_bits(w, 1n+m, m), 0n))
76
77def hit(m: Maybe<&2, Nat>, +k: Nat) -> Nat:
78  match m:
79    case None{}:
80      0n
81    case Some{v}:
82      b2n(Nat.is_eq(v, k))
83
84# the number of source outputs x < N that draw k
85def count(+w: Nat, +n: Nat, +k: Nat, N: Nat) -> Nat:
86  match N:
87    case 0n:
88      0n
89    case 1n+ +x:
90      Nat.add(count(w, n, k, x), hit(draw(w, x, n), k))
91
92# the bounded draw once a draw of x (with verdict m) has left state s: m if
93# accepted, else (while fuel lasts) the next draw; with no fuel left, the
94# last candidate floor(x n / 2^64). The verdict is matched before the
95# recursion, so a draw that is not yet decided keeps the rest unexpanded.
96def below_go(~S: Data, ~next: S -> W.U64 & S, fuel: Nat, +n: Nat, m: Maybe<&2, Nat>, +x: Nat, +s: S) -> Nat & S:
97  match fuel m:
98    case 0n Some{k}:
99      (k, s)
100    case 0n None{}:
101      (C.high(64n, Nat.mul(x, n)), s)
102    case 1n+f Some{k}:
103      (k, s)
104    case 1n+f None{}:
105      below_go(~S, ~next, f, n, draw(64n, SW.value(SRC.fst64(S, next(s))), n), SW.value(SRC.fst64(S, next(s))), SRC.snd64(S, next(s)))
106
107# the first accepted draw of 64-bit outputs among fuel + 1 draws
108def below(~S: Data, ~next: S -> W.U64 & S, fuel: Nat, +n: Nat, +s: S) -> Nat & S:
109  below_go(~S, ~next, fuel, n, draw(64n, SW.value(SRC.fst64(S, next(s))), n), SW.value(SRC.fst64(S, next(s))), SRC.snd64(S, next(s)))
110
111def occurrences(~A: Data, ~V: Data, ~rel: A -> V -> Bool, +v: V, xs: List<&2, A>) -> Nat:
112  match xs:
113    case Nil{}:
114      0n
115    case Con{+x, rest}:
116      Nat.add(b2n(rel(x, v)), occurrences(~A, ~V, ~rel, v, rest))