src/math/random/pcg.bend

bend-collections-laws-crypto@1.0.0.0 · 82 lines · 16 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
4
5# PCG: Go's math/rand/v2.PCG, bit for bit: a 128-bit linear congruential
6# generator (O'Neill's PCG family) with the DXSM ("double xorshift
7# multiply") output function.
8#
9#   new(seed1, seed2)  the state hi:lo = seed1:seed2 (Go's NewPCG)
10#   next(p)            the next 64-bit output and the advanced generator
11#                      (Go's (*PCG).Uint64); the Source interface of rand.bend
12#
13# One step is state = state * MUL + INC (mod 2^128) with Go's constants
14# MUL = mulHi:mulLo, INC = incHi:incLo; the output is DXSM of the new state:
15# hi ^= hi >> 32; hi *= 0xda942042e4dd58b5; hi ^= hi >> 48; hi *= lo | 1.
16# The 64-bit words are src/math/u64.bend's two-limb U64 with the proved
17# arithmetic of src/math/w64.bend.
18#
19# Proved equal to spec/math/random/pcg.bend (the step and the output on
20# naturals) for every state (proofs/math/random/pcg.bend); tested against
21# Go's vectors (tools/check_random.py). PCG is not a cryptographic generator:
22# its state follows from a few outputs.
23
24type PCG is Data:
25  P{hi: W.U64, lo: W.U64}
26
27def new(+seed1: W.U64, +seed2: W.U64) -> PCG:
28  P{seed1, seed2}
29
30# 2549297995355413924, 4865540595714422341
31def mul_hi() -> W.U64:
32  W.U64{533093796, 593554693}
33
34def mul_lo() -> W.U64:
35  W.U64{2681009733, 1132846948}
36
37# 6364136223846793005, 1442695040888963407
38def inc_hi() -> W.U64:
39  W.U64{1284865837, 1481765933}
40
41def inc_lo() -> W.U64:
42  W.U64{4150755663, 335903614}
43
44# 0xda942042e4dd58b5
45def cheap_mul() -> W.U64:
46  W.U64{3839711413, 3667140674}
47
48def xor64(+a: W.U64, +b: W.U64) -> W.U64:
49  W.U64{U32.xor(X.lo(a), X.lo(b)), U32.xor(X.hi(a), X.hi(b))}
50
51# the new state from the product's low word l and high word h
52def step_fin(+ih: W.U64, +il: W.U64, +l: W.U64, +h: W.U64) -> PCG:
53  P{X.add(X.add(h, ih), W.U64{X.b32(X.add_over(l, il)), 0}), X.add(l, il)}
54
55def step_mul(+mh: W.U64, +ml: W.U64, +ih: W.U64, +il: W.U64, +hi: W.U64, +lo: W.U64, p: W.U64 & W.U64) -> PCG:
56  (+l, +h) = p
57  step_fin(ih, il, l, X.add(h, X.add(X.mul(hi, ml), X.mul(lo, mh))))
58
59# state * (mh:ml) + (ih:il) mod 2^128 (Go's (*PCG).next with its constants
60# as parameters)
61def step_with(+mh: W.U64, +ml: W.U64, +ih: W.U64, +il: W.U64, p: PCG) -> PCG:
62  match p:
63    case P{+hi, +lo}:
64      step_mul(mh, ml, ih, il, hi, lo, X.mul128(lo, ml))
65
66def step(p: PCG) -> PCG:
67  step_with(mul_hi(), mul_lo(), inc_hi(), inc_lo(), p)
68
69def dxsm2(+h: W.U64, +lo: W.U64) -> W.U64:
70  X.mul(xor64(h, X.shr(h, 48n)), W.U64{U32.or(X.lo(lo), 1), X.hi(lo)})
71
72# the DXSM output of a state hi:lo with the multiplier cm
73def dxsm_with(+cm: W.U64, +hi: W.U64, +lo: W.U64) -> W.U64:
74  dxsm2(X.mul(xor64(hi, X.shr(hi, 32n)), cm), lo)
75
76def out(+p: PCG) -> W.U64 & PCG:
77  match p:
78    case P{+hi, +lo}:
79      (dxsm_with(cheap_mul(), hi, lo), P{hi, lo})
80
81def next(p: PCG) -> W.U64 & PCG:
82  out(step(p))