1	import Base
2	import ../../lib/common.bend as C
3	import ../../../src/math/u64.bend as W
4	import ../w64.bend as SW
5	import ./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	
31	def b2n(b: Bool) -> Nat:
32	  match b:
33	    case True{}:
34	      1n
35	    case False{}:
36	      0n
37	
38	def 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	
45	def 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	
52	def 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
60	def 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	
63	def 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	
70	def 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	
77	def 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
85	def 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.
96	def 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
108	def 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	
111	def 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))