1	import Base
2	import ../../../src/math/u64.bend as W
3	import ../../../src/math/w64.bend as X
4	import ../../../src/math/random/rand.bend as R
5	import ../../../spec/math/random/rand.bend as SR
6	import ../../lib/logic.bend as L
7	import ../../lib/nat.bend as N
8	import ../../lib/u32alg.bend as A
9	
10	# Fisher-Yates (src/math/random/rand.bend shuffle and perm, Go's Shuffle and
11	# Perm) returns a permutation of its input, for every source and every draw
12	# (the indices j need not be in range: an out-of-range swap is the identity):
13	# every swap keeps the number of elements related to any v (the multiset,
14	# for an equality), as in the Coq/Isabelle proofs that Fisher-Yates permutes.
15	
16	def occ(~A: Data, ~V: Data, ~rel: A -> V -> Bool, +v: V, +xs: List<&2, A>) -> Nat:
17	  SR.occurrences(~A, ~V, ~rel, v, xs)
18	
19	def some_or(-A: Data, m: Maybe<&2, A>, +d: A) -> A:
20	  match m:
21	    case None{}:
22	      d
23	    case Some{x}:
24	      x
25	
26	def is_some(-A: Data, m: Maybe<&2, A>) -> Bool:
27	  match m:
28	    case None{}:
29	      False{}
30	    case Some{x}:
31	      True{}
32	
33	def none_some(-A: Data, +y: A, e: {None{} == Some{y} : Maybe<&2, A>}) -> Empty:
34	  %Equal.sym(Maybe<&2, A>, None{}, Some{y}, e) : L.truth(is_some(A, _))
35	  Unit{}
36	
37	def some_inj(-A: Data, +x: A, +y: A, e: {Some{x} == Some{y} : Maybe<&2, A>}) -> {x == y : A}:
38	  Equal.cong(Maybe<&2, A>, A, m => some_or(A, m, x), Some{x}, Some{y}, e)
39	
40	# (a + c) + b == (b + c) + a
41	def swap_ends(+a: Nat, +b: Nat, +c: Nat) -> {Nat.add(Nat.add(a, c), b) == Nat.add(Nat.add(b, c), a) : Nat}:
42	  Equal.trans(Nat, Nat.add(Nat.add(a, c), b), Nat.add(a, Nat.add(c, b)), Nat.add(Nat.add(b, c), a),
43	    N.add_assoc(a, c, b),
44	    Equal.trans(Nat, Nat.add(a, Nat.add(c, b)), Nat.add(a, Nat.add(b, c)), Nat.add(Nat.add(b, c), a),
45	      Equal.cong(Nat, Nat, t => Nat.add(a, t), Nat.add(c, b), Nat.add(b, c), N.add_comm(c, b)),
46	      N.add_comm(a, Nat.add(b, c))))
47	
48	# (h + x) + y == h + (x + y), stated for the step of put_occ
49	def assoc3(+h: Nat, +x: Nat, +y: Nat, +z: Nat, +e: {Nat.add(x, y) == z : Nat}) -> {Nat.add(Nat.add(h, x), y) == Nat.add(h, z) : Nat}:
50	  Equal.trans(Nat, Nat.add(Nat.add(h, x), y), Nat.add(h, Nat.add(x, y)), Nat.add(h, z), N.add_assoc(h, x, y),
51	    Equal.cong(Nat, Nat, t => Nat.add(h, t), Nat.add(x, y), z, e))
52	
53	# writing x over the element y at i: occ + [y] == old occ + [x]
54	def put_occ(~A: Data, ~V: Data, ~rel: A -> V -> Bool, +v: V, +xs: List<&2, A>, +i: Nat, +x: A, +y: A, +h: {R.nth(A, xs, i) == Some{y} : Maybe<&2, A>}) -> {Nat.add(occ(~A, ~V, ~rel, v, R.put(A, xs, i, x)), SR.b2n(rel(y, v))) == Nat.add(occ(~A, ~V, ~rel, v, xs), SR.b2n(rel(x, v))) : Nat}:
55	  match xs i:
56	    case Nil{} _:
57	      Empty.absurd({Nat.add(0n, SR.b2n(rel(y, v))) == Nat.add(0n, SR.b2n(rel(x, v))) : Nat}, none_some(A, y, h))
58	    case Con{+hd, +t} 0n:
59	      %some_inj(A, hd, y, h) : {Nat.add(Nat.add(SR.b2n(rel(x, v)), occ(~A, ~V, ~rel, v, t)), SR.b2n(rel(_, v))) == Nat.add(Nat.add(SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, t)), SR.b2n(rel(x, v))) : Nat}
60	      swap_ends(SR.b2n(rel(x, v)), SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, t))
61	    case Con{+hd, +t} 1n+p:
62	      Equal.trans(Nat, Nat.add(Nat.add(SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, R.put(A, t, p, x))), SR.b2n(rel(y, v))), Nat.add(SR.b2n(rel(hd, v)), Nat.add(occ(~A, ~V, ~rel, v, t), SR.b2n(rel(x, v)))), Nat.add(Nat.add(SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, t)), SR.b2n(rel(x, v))),
63	        assoc3(SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, R.put(A, t, p, x)), SR.b2n(rel(y, v)), Nat.add(occ(~A, ~V, ~rel, v, t), SR.b2n(rel(x, v))), put_occ(~A, ~V, ~rel, v, t, p, x, y, h)),
64	        Equal.sym(Nat, Nat.add(Nat.add(SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, t)), SR.b2n(rel(x, v))), Nat.add(SR.b2n(rel(hd, v)), Nat.add(occ(~A, ~V, ~rel, v, t), SR.b2n(rel(x, v)))), N.add_assoc(SR.b2n(rel(hd, v)), occ(~A, ~V, ~rel, v, t), SR.b2n(rel(x, v)))))
65	
66	# writing y at i keeps an element y at j
67	def nth_put(-A: Data, +xs: List<&2, A>, +i: Nat, +j: Nat, +y: A, +h: {R.nth(A, xs, j) == Some{y} : Maybe<&2, A>}) -> {R.nth(A, R.put(A, xs, i, y), j) == Some{y} : Maybe<&2, A>}:
68	  match xs i j:
69	    case Nil{} _ _:
70	      Empty.absurd({R.nth(A, R.put(A, Nil{}, i, y), j) == Some{y} : Maybe<&2, A>}, none_some(A, y, h))
71	    case Con{+hd, +t} 0n 0n:
72	      {==}
73	    case Con{+hd, +t} 0n 1n+q:
74	      h
75	    case Con{+hd, +t} 1n+r 0n:
76	      h
77	    case Con{+hd, +t} 1n+r 1n+q:
78	      nth_put(A, t, r, q, y, h)
79	
80	def swap_some(~A: Data, ~V: Data, ~rel: A -> V -> Bool, +v: V, +xs: List<&2, A>, +i: Nat, +j: Nat, +x: A, +y: A, +hx: {R.nth(A, xs, i) == Some{x} : Maybe<&2, A>}, +hy: {R.nth(A, xs, j) == Some{y} : Maybe<&2, A>}) -> {occ(~A, ~V, ~rel, v, R.put(A, R.put(A, xs, i, y), j, x)) == occ(~A, ~V, ~rel, v, xs) : Nat}:
81	  +ys = R.put(A, xs, i, y)
82	  A.add_cancel_r(occ(~A, ~V, ~rel, v, R.put(A, ys, j, x)), occ(~A, ~V, ~rel, v, xs), SR.b2n(rel(y, v)),
83	    Equal.trans(Nat, Nat.add(occ(~A, ~V, ~rel, v, R.put(A, ys, j, x)), SR.b2n(rel(y, v))), Nat.add(occ(~A, ~V, ~rel, v, ys), SR.b2n(rel(x, v))), Nat.add(occ(~A, ~V, ~rel, v, xs), SR.b2n(rel(y, v))),
84	      put_occ(~A, ~V, ~rel, v, ys, j, x, y, nth_put(A, xs, i, j, y, hy)),
85	      put_occ(~A, ~V, ~rel, v, xs, i, y, x, hx)))
86	
87	def swap_m_occ(~A: Data, ~V: Data, ~rel: A -> V -> Bool, +v: V, +xs: List<&2, A>, +i: Nat, +j: Nat, +a: Maybe<&2, A>, +b: Maybe<&2, A>, +ha: {R.nth(A, xs, i) == a : Maybe<&2, A>}, +hb: {R.nth(A, xs, j) == b : Maybe<&2, A>}) -> {occ(~A, ~V, ~rel, v, R.swap_m(A, xs, i, j, a, b)) == occ(~A, ~V, ~rel, v, xs) : Nat}:
88	  match a b:
89	    case Some{x} Some{y}:
90	      swap_some(~A, ~V, ~rel, v, xs, i, j, x, y, ha, hb)
91	    case None{} _:
92	      {==}
93	    case Some{x} None{}:
94	      {==}
95	
96	# THEOREM: a swap keeps every count
97	def swap_occ(~A: Data, ~V: Data, ~rel: A -> V -> Bool, +v: V, +xs: List<&2, A>, +i: Nat, +j: Nat) -> {occ(~A, ~V, ~rel, v, R.swap(A, xs, i, j)) == occ(~A, ~V, ~rel, v, xs) : Nat}:
98	  swap_m_occ(~A, ~V, ~rel, v, xs, i, j, R.nth(A, xs, i), R.nth(A, xs, j), {==}, {==})
99	
100	def swap_at_occ(~A: Data, ~V: Data, ~rel: A -> V -> Bool, -S: Data, +v: V, +xs: List<&2, A>, +i: Nat, r: W.U64 & S) -> {occ(~A, ~V, ~rel, v, R.fst(List<&2, A>, S, R.swap_at(A, S, xs, i, r))) == occ(~A, ~V, ~rel, v, xs) : Nat}:
101	  match r:
102	    case Tuple{+j, s}:
103	      swap_occ(~A, ~V, ~rel, v, xs, i, X.n48(j))
104	
105	def go_occ(~A: Data, ~V: Data, ~rel: A -> V -> Bool, ~S: Data, ~next: S -> W.U64 & S, +v: V, +k: Nat, +n: W.U64, st: List<&2, A> & S) -> {occ(~A, ~V, ~rel, v, R.fst(List<&2, A>, S, R.shuffle_go(~A, ~S, ~next, k, n, st))) == occ(~A, ~V, ~rel, v, R.fst(List<&2, A>, S, st)) : Nat}:
106	  match k st:
107	    case 0n Tuple{+xs, +s}:
108	      {==}
109	    case 1n+p Tuple{+xs, +s}:
110	      Equal.trans(Nat, occ(~A, ~V, ~rel, v, R.fst(List<&2, A>, S, R.shuffle_go(~A, ~S, ~next, p, X.sub(n, W.U64{1, 0}), R.shuffle_step(~A, ~S, ~next, 1n+p, n, (xs, s))))), occ(~A, ~V, ~rel, v, R.fst(List<&2, A>, S, R.shuffle_step(~A, ~S, ~next, 1n+p, n, (xs, s)))), occ(~A, ~V, ~rel, v, xs),
111	        go_occ(~A, ~V, ~rel, ~S, ~next, v, p, X.sub(n, W.U64{1, 0}), R.shuffle_step(~A, ~S, ~next, 1n+p, n, (xs, s))),
112	        swap_at_occ(~A, ~V, ~rel, S, v, xs, 1n+p, R.uint64n(~S, ~next, s, n)))
113	
114	# THEOREM (Shuffle.permutation)
115	def shuffle_perm(~A: Data, ~V: Data, ~rel: A -> V -> Bool, ~S: Data, ~next: S -> W.U64 & S, +s: S, +xs: List<&2, A>, +v: V) -> {SR.occurrences(~A, ~V, ~rel, v, R.fst(List<&2, A>, S, R.shuffle(~A, ~S, ~next, s, xs))) == SR.occurrences(~A, ~V, ~rel, v, xs) : Nat}:
116	  go_occ(~A, ~V, ~rel, ~S, ~next, v, Nat.sub(List.length(&2, A, xs), 1n), R.len64(A, xs), (xs, s))
117	
118	# [p == v] + [v < p] == [v < p + 1]
119	def lt_succ_split(+v: Nat, +p: Nat) -> {Nat.add(SR.b2n(Nat.is_eq(p, v)), SR.b2n(Nat.is_lt(v, p))) == SR.b2n(Nat.is_lt(v, 1n+p)) : Nat}:
120	  match v p:
121	    case 0n 0n:
122	      {==}
123	    case 0n 1n+q:
124	      {==}
125	    case 1n+u 0n:
126	      %Equal.sym(Bool, Nat.is_lt(u, 0n), False{}, N.not_lt_zero(u)) : {0n == SR.b2n(_) : Nat}
127	      {==}
128	    case 1n+u 1n+q:
129	      lt_succ_split(u, q)
130	
131	# range_go(n, acc) = [0, 1, ..., n - 1] ++ acc holds v as often as acc,
132	# plus once when v < n
133	def range_occ(+n: Nat, +acc: List<&2, Nat>, +v: Nat) -> {SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, R.range_go(n, acc)) == Nat.add(SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc), SR.b2n(Nat.is_lt(v, n))) : Nat}:
134	  match n:
135	    case 0n:
136	      %Equal.sym(Bool, Nat.is_lt(v, 0n), False{}, N.not_lt_zero(v)) : {SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc) == Nat.add(SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc), SR.b2n(_)) : Nat}
137	      Equal.sym(Nat, Nat.add(SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc), 0n), SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc), N.add_zero(SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc)))
138	    case 1n+p:
139	      +o = SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, acc)
140	      +e = SR.b2n(Nat.is_eq(p, v))
141	      +l = SR.b2n(Nat.is_lt(v, p))
142	      Equal.trans(Nat, SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, v, R.range_go(p, Con{p, acc})), Nat.add(Nat.add(e, o), l), Nat.add(o, SR.b2n(Nat.is_lt(v, 1n+p))),
143	        range_occ(p, Con{p, acc}, v),
144	        Equal.trans(Nat, Nat.add(Nat.add(e, o), l), Nat.add(o, Nat.add(e, l)), Nat.add(o, SR.b2n(Nat.is_lt(v, 1n+p))),
145	          Equal.trans(Nat, Nat.add(Nat.add(e, o), l), Nat.add(Nat.add(o, e), l), Nat.add(o, Nat.add(e, l)),
146	            Equal.cong(Nat, Nat, t => Nat.add(t, l), Nat.add(e, o), Nat.add(o, e), N.add_comm(e, o)),
147	            N.add_assoc(o, e, l)),
148	          Equal.cong(Nat, Nat, t => Nat.add(o, t), Nat.add(e, l), SR.b2n(Nat.is_lt(v, 1n+p)), lt_succ_split(v, p))))
149	
150	# THEOREM (Perm.permutation)
151	def perm_perm(~S: Data, ~next: S -> W.U64 & S, +s: S, +n: Nat, +i: Nat) -> {SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, i, R.fst(List<&2, Nat>, S, R.perm(~S, ~next, s, n))) == SR.b2n(Nat.is_lt(i, n)) : Nat}:
152	  Equal.trans(Nat, SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, i, R.fst(List<&2, Nat>, S, R.shuffle(~Nat, ~S, ~next, s, R.range(n)))), SR.occurrences(~Nat, ~Nat, ~Nat.is_eq, i, R.range_go(n, [])), SR.b2n(Nat.is_lt(i, n)),
153	    shuffle_perm(~Nat, ~Nat, ~Nat.is_eq, ~S, ~next, s, R.range(n), i),
154	    range_occ(n, [], i))