1	# bend-ml-bpe-tokenizer: byte-level BPE with a proved roundtrip.
2	#
3	#   import bend-ml-bpe-tokenizer@0.1.3.0/main.bend as BPE
4	#
5	# A token is a raw byte (B{n}) or the result of a merge rule (M{id}).
6	# The table is a list of rules from the oldest to the newest, like GPT-2's
7	# merges.txt: the rule with id k joins the pair (a, b) into the token M{k}.
8	# `encode` applies the rules in order; `decode` expands each token back
9	# into bytes. The laws are at the end of the file, each with its proof.
10	#
11	# How to read the proofs: a proof in Bend is a function whose TYPE is the statement.
12	# `match` does case analysis, the recursive call is the induction hypothesis,
13	# `%e : P` rewrites the goal with the equality `e`, `{==}` closes when the
14	# two sides are already the same term. `+x` marks a variable that may be
15	# used more than once.
16	
17	import Base
18	import bend-ml-nat-lemmas@0.1.0.0/main.bend as NL
19	
20	# Token: a raw byte (B) or the result of a merge rule, by rule id (M).
21	type Tk is Data:
22	  B{n: Nat}
23	  M{k: Nat}
24	
25	# Merge rule: joins the pair (a, b) into the token M{id}.
26	type Rule is Data:
27	  Rule{id: Nat, a: Tk, b: Tk}
28	
29	def Tk.eq(+x: Tk, +y: Tk) -> Bool:
30	  match x y:
31	    case B{+m} B{+n}:
32	      Nat.is_eq(m, n)
33	    case M{+i} M{+j}:
34	      Nat.is_eq(i, j)
35	    case B{m} M{j}:
36	      False{}
37	    case M{i} B{n}:
38	      False{}
39	
40	# does xs start with exactly the pair (a, b)?
41	def peek(+a: Tk, +b: Tk, xs: List<&2, Tk>) -> Bool:
42	  match xs:
43	    case Nil{}:
44	      False{}
45	    case Con{+x, t}:
46	      match t:
47	        case Nil{}:
48	          False{}
49	        case Con{+y, u}:
50	          Bool.and(Tk.eq(x, a), Tk.eq(y, b))
51	
52	# Replaces each occurrence (non-overlapping, left to right) of the pair
53	# (a, b) with the token c. `hit` is always `peek(a, b, xs)`: the caller computes it,
54	# because a match can only open parameters.
55	def merge.go(xs: List<&2, Tk>, hit: Bool, +a: Tk, +b: Tk, +c: Tk) -> List<&2, Tk>:
56	  match xs hit:
57	    case Con{x, Con{y, +t}} True{}:
58	      c <> merge.go(t, peek(a, b, t), a, b, c)
59	    case Con{x, Nil{}} True{}:
60	      x <> Nil{}
61	    case Con{x, +t} False{}:
62	      x <> merge.go(t, peek(a, b, t), a, b, c)
63	    case Nil{} _:
64	      Nil{}
65	
66	def merge(+a: Tk, +b: Tk, +c: Tk, +xs: List<&2, Tk>) -> List<&2, Tk>:
67	  merge.go(xs, peek(a, b, xs), a, b, c)
68	
69	# applies the rules from the oldest to the newest; `done` keeps the ones already applied
70	# (newest first), which is exactly the table that decode needs. Only the proofs read it, so it
71	# is erased (`-done`): the checker keeps it and the compiled program never builds it. Building
72	# it cost half of the time to encode (found by Mathews Fernando's bend-lint, dead_accumulator).
73	def encode.go(todo: List<&2, Rule>, -done: List<&2, Rule>, xs: List<&2, Tk>) -> List<&2, Tk>:
74	  match todo:
75	    case Nil{}:
76	      xs
77	    case Con{+r, rest}:
78	      match r:
79	        case Rule{+i, +a, +b}:
80	          encode.go(rest, Rule{i, a, b} <> done, merge(a, b, M{i}, xs))
81	
82	# The table goes from the oldest rule to the newest (like merges.txt).
83	def encode(+table: List<&2, Rule>, xs: List<&2, Tk>) -> List<&2, Tk>:
84	  encode.go(table, Nil{}, xs)
85	
86	# is the newest rule in the table the one that creates the token tk?
87	def hits(table: List<&2, Rule>, +tk: Tk) -> Bool:
88	  match table:
89	    case Nil{}:
90	      False{}
91	    case Con{Rule{+i, a, b}, rest}:
92	      match tk:
93	        case B{n}:
94	          False{}
95	        case M{+j}:
96	          Nat.is_eq(i, j)
97	
98	# Expansion of a token into bytes. The table has the newest rule first and
99	# a rule only refers to older tokens, so expanding a and b looks only at
100	# the rest of the table. `hit` is always `hits(table, tk)`.
101	def exp.go(table: List<&2, Rule>, hit: Bool, tk: Tk) -> List<&2, Nat>:
102	  match table hit tk:
103	    case Nil{} _ B{n}:
104	      n <> Nil{}
105	    case Nil{} _ M{j}:
106	      Nil{}
107	    case Con{r, rest} _ B{n}:
108	      n <> Nil{}
109	    case Con{Rule{i, +a, +b}, +rest} True{} M{j}:
110	      List.append(&2, Nat, exp.go(rest, hits(rest, a), a), exp.go(rest, hits(rest, b), b))
111	    case Con{r, +rest} False{} M{+j}:
112	      exp.go(rest, hits(rest, M{j}), M{j})
113	
114	def exp(+table: List<&2, Rule>, +tk: Tk) -> List<&2, Nat>:
115	  exp.go(table, hits(table, tk), tk)
116	
117	# decodes with the already reversed table (newest first)
118	def dec(+table: List<&2, Rule>, ids: List<&2, Tk>) -> List<&2, Nat>:
119	  match ids:
120	    case Nil{}:
121	      Nil{}
122	    case Con{+x, t}:
123	      List.append(&2, Nat, exp(table, x), dec(table, t))
124	
125	def decode(+table: List<&2, Rule>, ids: List<&2, Tk>) -> List<&2, Nat>:
126	  dec(List.reverse(&2, Rule, table), ids)
127	
128	# bytes -> initial tokens
129	def lift(bs: List<&2, Nat>) -> List<&2, Tk>:
130	  match bs:
131	    case Nil{}:
132	      Nil{}
133	    case Con{n, t}:
134	      B{n} <> lift(t)
135	
136	# ---- predicates used in the laws (all computable) ----
137	
138	# does some rule in the table have the id i?
139	def has(+D: List<&2, Rule>, +i: Nat) -> Bool:
140	  match D:
141	    case Nil{}:
142	      False{}
143	    case Con{Rule{+j, a, b}, t}:
144	      Bool.or(Nat.is_eq(i, j), has(t, i))
145	
146	# is the token a byte, or was it created by a rule in the table?
147	def kn(+D: List<&2, Rule>, tk: Tk) -> Bool:
148	  match tk:
149	    case B{n}:
150	      True{}
151	    case M{+j}:
152	      has(D, j)
153	
154	def knl(+D: List<&2, Rule>, xs: List<&2, Tk>) -> Bool:
155	  match xs:
156	    case Nil{}:
157	      True{}
158	    case Con{x, t}:
159	      Bool.and(kn(D, x), knl(D, t))
160	
161	# well-formed table: no id repeats
162	def wf.go(todo: List<&2, Rule>, +done: List<&2, Rule>) -> Bool:
163	  match todo:
164	    case Nil{}:
165	      True{}
166	    case Con{Rule{+i, +a, +b}, rest}:
167	      Bool.and(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done))
168	
169	def wf(table: List<&2, Rule>) -> Bool:
170	  wf.go(table, Nil{})
171	
172	# =====================================================================
173	# Training (needs no proof: it generates tables; a table is valid if wf(table))
174	# =====================================================================
175	# ---------------------------------------------------------------
176	# training: count adjacent pairs, merge the most frequent one, repeat
177	# (tie-break: the first pair that appeared, as in minbpe)
178	# ---------------------------------------------------------------
179	type Cnt is Data:
180	  Cnt{a: Tk, b: Tk, n: Nat}
181	
182	# is the first Cnt of cs the pair (a, b)?
183	def same_head(cs: List<&2, Cnt>, +a: Tk, +b: Tk) -> Bool:
184	  match cs:
185	    case Nil{}:
186	      False{}
187	    case Con{Cnt{+x, +y, k}, rest}:
188	      Bool.and(Tk.eq(x, a), Tk.eq(y, b))
189	
190	# adds 1 to the counter of the pair (a, b), or creates it at the end of the list
191	def bump.go(cs: List<&2, Cnt>, hit: Bool, +a: Tk, +b: Tk) -> List<&2, Cnt>:
192	  match cs hit:
193	    case Nil{} _:
194	      Cnt{a, b, 1n} <> Nil{}
195	    case Con{Cnt{x, y, +k}, rest} True{}:
196	      Cnt{x, y, 1n+k} <> rest
197	    case Con{c, +rest} False{}:
198	      c <> bump.go(rest, same_head(rest, a, b), a, b)
199	
200	def bump(+cs: List<&2, Cnt>, +a: Tk, +b: Tk) -> List<&2, Cnt>:
201	  bump.go(cs, same_head(cs, a, b), a, b)
202	
203	def count.go(xs: List<&2, Tk>, cs: List<&2, Cnt>) -> List<&2, Cnt>:
204	  match xs:
205	    case Nil{}:
206	      cs
207	    case Con{+x, +t}:
208	      match t:
209	        case Nil{}:
210	          cs
211	        case Con{+y, u}:
212	          count.go(Con{y, u}, bump(cs, x, y))
213	
214	def pick2(better: Bool, c: Cnt, cur: Cnt) -> Cnt:
215	  match better:
216	    case True{}:
217	      c
218	    case False{}:
219	      cur
220	
221	def pick(c: Cnt, cur: Cnt) -> Cnt:
222	  match c cur:
223	    case Cnt{+a, +b, +k} Cnt{+x, +y, +m}:
224	      pick2(Nat.is_lt(m, k), Cnt{a, b, k}, Cnt{x, y, m})
225	
226	def best.go(cs: List<&2, Cnt>, cur: Cnt) -> Cnt:
227	  match cs:
228	    case Nil{}:
229	      cur
230	    case Con{c, rest}:
231	      best.go(rest, pick(c, cur))
232	
233	def best(cs: List<&2, Cnt>) -> Maybe<&1, Cnt>:
234	  match cs:
235	    case Nil{}:
236	      None{}
237	    case Con{c, rest}:
238	      Some{best.go(rest, c)}
239	
240	# up to n merges, ids from `next` on; m is the current best pair (best(count(xs)))
241	def train.go(n: Nat, +next: Nat, +xs: List<&2, Tk>, m: Maybe<&1, Cnt>) -> List<&2, Rule>:
242	  match n m:
243	    case 0n _:
244	      Nil{}
245	    case 1n+p None{}:
246	      Nil{}
247	    case 1n+p Some{Cnt{+a, +b, k}}:
248	      +ys = merge(a, b, M{next}, xs)
249	      Rule{next, a, b} <> train.go(p, 1n+next, ys, best(count.go(ys, Nil{})))
250	
251	# the table from the oldest rule to the newest
252	def train(n: Nat, +next: Nat, +xs: List<&2, Tk>) -> List<&2, Rule>:
253	  train.go(n, next, xs, best(count.go(xs, Nil{})))
254	
255	# =====================================================================
256	# Auxiliary lemmas for the proofs
257	# =====================================================================
258	
259	# a type chosen by a Bool: rewriting through it refutes True == False
260	def BD(b: Bool, t: Type, f: Type) -> Type:
261	  match b:
262	    case True{}:
263	      t
264	    case False{}:
265	      f
266	
267	def nat_eq_sound(a: Nat, b: Nat, e: {True{} == Nat.is_eq(a, b) : Bool}) -> {a == b : Nat}:
268	  match a b:
269	    case 0n 0n:
270	      {==}
271	    case 0n 1n+q:
272	      %e : BD(_, Unit, {0n == 1n+q : Nat})
273	      Unit{}
274	    case 1n+p 0n:
275	      %e : BD(_, Unit, {1n+p == 0n : Nat})
276	      Unit{}
277	    case 1n+p 1n+q:
278	      %nat_eq_sound(p, q, e) : {1n+p == 1n+_ : Nat}
279	      {==}
280	
281	def nat_eq_refl(a: Nat) -> {True{} == Nat.is_eq(a, a) : Bool}:
282	  match a:
283	    case 0n:
284	      {==}
285	    case 1n+p:
286	      nat_eq_refl(p)
287	
288	def and_true(p: Bool, q: Bool, e: {True{} == Bool.and(p, q) : Bool}) -> {True{} == p : Bool} & {True{} == q : Bool}:
289	  match p:
290	    case True{}:
291	      ({==}, e)
292	    case False{}:
293	      %e : BD(_, Unit, {True{} == False{} : Bool} & {True{} == q : Bool})
294	      Unit{}
295	
296	def and_intro(p: Bool, q: Bool, ep: {True{} == p : Bool}, eq: {True{} == q : Bool}) -> {True{} == Bool.and(p, q) : Bool}:
297	  match p:
298	    case True{}:
299	      eq
300	    case False{}:
301	      %ep : BD(_, Unit, {True{} == False{} : Bool})
302	      Unit{}
303	
304	def or_left(p: Bool, q: Bool, ep: {True{} == p : Bool}) -> {True{} == Bool.or(p, q) : Bool}:
305	  match p:
306	    case True{}:
307	      {==}
308	    case False{}:
309	      %ep : BD(_, Unit, {True{} == q : Bool})
310	      Unit{}
311	
312	def or_right(p: Bool, q: Bool, eq: {True{} == q : Bool}) -> {True{} == Bool.or(p, q) : Bool}:
313	  match p:
314	    case True{}:
315	      {==}
316	    case False{}:
317	      eq
318	
319	# not p is true  ->  p is false
320	def not_true(p: Bool, e: {True{} == Bool.not(p) : Bool}) -> {False{} == p : Bool}:
321	  match p:
322	    case True{}:
323	      %e : BD(_, Unit, {False{} == True{} : Bool})
324	      Unit{}
325	    case False{}:
326	      {==}
327	
328	def tk_eq_sound(+x: Tk, +y: Tk, e: {True{} == Tk.eq(x, y) : Bool}) -> {x == y : Tk}:
329	  match x y:
330	    case B{+m} B{+n}:
331	      %nat_eq_sound(m, n, e) : {B{m} == B{_} : Tk}
332	      {==}
333	    case M{+i} M{+j}:
334	      %nat_eq_sound(i, j, e) : {M{i} == M{_} : Tk}
335	      {==}
336	    case B{m} M{j}:
337	      %e : BD(_, Unit, {B{m} == M{j} : Tk})
338	      Unit{}
339	    case M{i} B{n}:
340	      %e : BD(_, Unit, {M{i} == B{n} : Tk})
341	      Unit{}
342	
343	# Adding the (new) rule i on top of the table does not change the expansion of a
344	# token M{j} that the old table already knows (j != i, because i is new).
345	def ext_m(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +j: Nat, h: Bool, eh: {h == Nat.is_eq(i, j) : Bool}, nf: {True{} == Bool.not(has(D, i)) : Bool}, kt: {True{} == has(D, j) : Bool}) -> {exp(Rule{i, a, b} <> D, M{j}) == exp(D, M{j}) : List<&2, Nat>}:
346	  match h:
347	    case True{}:
348	      # h = True means i == j; but j is known (has D j) and i is not (not has D i)
349	      +ij = nat_eq_sound(i, j, eh)
350	      +hij = Equal.cong(Nat, Bool, k => has(D, k), i, j, ij)
351	      +f1 = Equal.trans(Bool, False{}, has(D, i), has(D, j), not_true(has(D, i), nf), hij)
352	      +tf = Equal.trans(Bool, True{}, has(D, j), False{}, kt, Equal.sym(Bool, False{}, has(D, j), f1))
353	      %tf : BD(_, Unit, {exp.go(Rule{i, a, b} <> D, Cmp.is_eq(Nat.cmp(i, j)), M{j}) == exp.go(D, hits(D, M{j}), M{j}) : List<&2, Nat>})
354	      Unit{}
355	    case False{}:
356	      %eh : {exp.go(Rule{i, a, b} <> D, _, M{j}) == exp(D, M{j}) : List<&2, Nat>}
357	      {==}
358	
359	def and_true_l(p: Bool, q: Bool, e: {True{} == Bool.and(p, q) : Bool}) -> {True{} == p : Bool}:
360	  match p:
361	    case True{}:
362	      {==}
363	    case False{}:
364	      %e : BD(_, Unit, {True{} == False{} : Bool})
365	      Unit{}
366	
367	def and_true_r(p: Bool, q: Bool, e: {True{} == Bool.and(p, q) : Bool}) -> {True{} == q : Bool}:
368	  match p:
369	    case True{}:
370	      e
371	    case False{}:
372	      %e : BD(_, Unit, {True{} == q : Bool})
373	      Unit{}
374	
375	# the same for any known token (a byte never depends on the table)
376	def ext_tok(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, t: Tk, kt: {True{} == kn(D, t) : Bool}, nf: {True{} == Bool.not(has(D, i)) : Bool}) -> {exp(Rule{i, a, b} <> D, t) == exp(D, t) : List<&2, Nat>}:
377	  match D t:
378	    case Nil{} B{n}:
379	      {==}
380	    case Con{Rule{j, a2, b2}, rest} B{n}:
381	      {==}
382	    case _ M{+j}:
383	      ext_m(D, i, a, b, j, Nat.is_eq(i, j), {==}, nf, kt)
384	
385	def ext_dec.step(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +x: Tk, +t: List<&2, Tk>, k1: {True{} == kn(D, x) : Bool}, nf: {True{} == Bool.not(has(D, i)) : Bool}, ih: {dec(Rule{i, a, b} <> D, t) == dec(D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, x <> t) == dec(D, x <> t) : List<&2, Nat>}:
386	  %ext_tok(D, i, a, b, x, k1, nf) : {List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, _, dec(D, t)) : List<&2, Nat>}
387	  %ih : {List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), _) : List<&2, Nat>}
388	  {==}
389	
390	# decoding a list of known tokens does not change when the new rule is added
391	def ext_dec(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, xs: List<&2, Tk>, +kx: {True{} == knl(D, xs) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}) -> {dec(Rule{i, a, b} <> D, xs) == dec(D, xs) : List<&2, Nat>}:
392	  match xs:
393	    case Nil{}:
394	      {==}
395	    case Con{+x, +t}:
396	      ext_dec.step(D, i, a, b, x, t, and_true_l(kn(D, x), knl(D, t), kx), nf, ext_dec(D, i, a, b, t, and_true_r(kn(D, x), knl(D, t), kx), nf))
397	
398	# The new token M{i} expands to the expansion of a followed by that of b
399	# (looking only at the old table D).
400	def exp_new(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, h: Bool, eh: {h == Nat.is_eq(i, i) : Bool}) -> {exp(Rule{i, a, b} <> D, M{i}) == List.append(&2, Nat, exp(D, a), exp(D, b)) : List<&2, Nat>}:
401	  match h:
402	    case True{}:
403	      %eh : {exp.go(Rule{i, a, b} <> D, _, M{i}) == List.append(&2, Nat, exp(D, a), exp(D, b)) : List<&2, Nat>}
404	      {==}
405	    case False{}:
406	      +f = Equal.trans(Bool, True{}, Nat.is_eq(i, i), False{}, nat_eq_refl(i), Equal.sym(Bool, False{}, Nat.is_eq(i, i), eh))
407	      %f : BD(_, Unit, {exp(Rule{i, a, b} <> D, M{i}) == List.append(&2, Nat, exp(D, a), exp(D, b)) : List<&2, Nat>})
408	      Unit{}
409	
410	# Central step: in place of a b t we get M{i} followed by m, where m decodes
411	# like t. Both sides decode the same because M{i} expands to exp(a) ++ exp(b).
412	def mp_core(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +t: List<&2, Tk>, +m: List<&2, Tk>, ka: {True{} == kn(D, a) : Bool}, kb: {True{} == kn(D, b) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}, ihm: {dec(Rule{i, a, b} <> D, m) == dec(Rule{i, a, b} <> D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, a <> b <> t) : List<&2, Nat>}:
413	  %Equal.sym(List<&2, Nat>, exp(Rule{i, a, b} <> D, M{i}), List.append(&2, Nat, exp(D, a), exp(D, b)), exp_new(D, i, a, b, Nat.is_eq(i, i), {==})) : {List.append(&2, Nat, _, dec(Rule{i, a, b} <> D, m)) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, a), List.append(&2, Nat, exp(Rule{i, a, b} <> D, b), dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
414	  %Equal.sym(List<&2, Nat>, dec(Rule{i, a, b} <> D, m), dec(Rule{i, a, b} <> D, t), ihm) : {List.append(&2, Nat, List.append(&2, Nat, exp(D, a), exp(D, b)), _) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, a), List.append(&2, Nat, exp(Rule{i, a, b} <> D, b), dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
415	  %Equal.sym(List<&2, Nat>, exp(Rule{i, a, b} <> D, a), exp(D, a), ext_tok(D, i, a, b, a, ka, nf)) : {List.append(&2, Nat, List.append(&2, Nat, exp(D, a), exp(D, b)), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, _, List.append(&2, Nat, exp(Rule{i, a, b} <> D, b), dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
416	  %Equal.sym(List<&2, Nat>, exp(Rule{i, a, b} <> D, b), exp(D, b), ext_tok(D, i, a, b, b, kb, nf)) : {List.append(&2, Nat, List.append(&2, Nat, exp(D, a), exp(D, b)), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, exp(D, a), List.append(&2, Nat, _, dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
417	  NL.append_assoc(Nat, exp(D, a), exp(D, b), dec(Rule{i, a, b} <> D, t))
418	
419	# if x == a and x is known, a is known too
420	def kn_sub(+D: List<&2, Rule>, +x: Tk, +a: Tk, ex: {x == a : Tk}, kx: {True{} == kn(D, x) : Bool}) -> {True{} == kn(D, a) : Bool}:
421	  %ex : {True{} == kn(D, _) : Bool}
422	  kx
423	
424	# x stays out of the rewrite: if the rest decodes the same, x <> rest does too
425	def mp_false(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +x: Tk, +t: List<&2, Tk>, +mg: List<&2, Tk>, ih: {dec(Rule{i, a, b} <> D, mg) == dec(Rule{i, a, b} <> D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, x <> mg) == dec(Rule{i, a, b} <> D, x <> t) : List<&2, Nat>}:
426	  %ih : {List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), dec(Rule{i, a, b} <> D, mg)) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), _) : List<&2, Nat>}
427	  {==}
428	
429	# the pair is x y: the tokens were equal to a and b
430	def mp_hit(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +x: Tk, +y: Tk, +t: List<&2, Tk>, +m: List<&2, Tk>, +ex: {x == a : Tk}, +ey: {y == b : Tk}, +kxx: {True{} == kn(D, x) : Bool}, +kyy: {True{} == kn(D, y) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}, ihm: {dec(Rule{i, a, b} <> D, m) == dec(Rule{i, a, b} <> D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, x <> y <> t) : List<&2, Nat>}:
431	  %Equal.sym(Tk, x, a, ex) : {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, _ <> y <> t) : List<&2, Nat>}
432	  %Equal.sym(Tk, y, b, ey) : {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, a <> _ <> t) : List<&2, Nat>}
433	  mp_core(D, i, a, b, t, m, kn_sub(D, x, a, ex, kxx), kn_sub(D, y, b, ey, kyy), nf, ihm)
434	
435	# merge.go does not change what decode returns (with the new rule on top of the table)
436	def mp(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, xs: List<&2, Tk>, hit: Bool, +eh: {hit == peek(a, b, xs) : Bool}, +kx: {True{} == knl(D, xs) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}) -> {dec(Rule{i, a, b} <> D, merge.go(xs, hit, a, b, M{i})) == dec(Rule{i, a, b} <> D, xs) : List<&2, Nat>}:
437	  match xs hit:
438	    case Nil{} _:
439	      {==}
440	    case Con{x, Nil{}} True{}:
441	      %eh : BD(_, Unit, {dec(Rule{i, a, b} <> D, merge.go(x <> Nil{}, True{}, a, b, M{i})) == dec(Rule{i, a, b} <> D, x <> Nil{}) : List<&2, Nat>})
442	      Unit{}
443	    case Con{+x, Con{+y, +t}} True{}:
444	      mp_hit(D, i, a, b, x, y, t, merge.go(t, peek(a, b, t), a, b, M{i}), tk_eq_sound(x, a, and_true_l(Tk.eq(x, a), Tk.eq(y, b), eh)), tk_eq_sound(y, b, and_true_r(Tk.eq(x, a), Tk.eq(y, b), eh)), and_true_l(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx), and_true_l(kn(D, y), knl(D, t), and_true_r(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx)), nf, mp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, y), knl(D, t), and_true_r(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx)), nf))
445	    case Con{+x, +t} False{}:
446	      mp_false(D, i, a, b, x, t, merge.go(t, peek(a, b, t), a, b, M{i}), mp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, x), knl(D, t), kx), nf))
447	
448	# ---- known tokens stay known ----
449	
450	# adding a rule to the table only increases what is known
451	def kn_mono(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, t: Tk, kt: {True{} == kn(D, t) : Bool}) -> {True{} == kn(Rule{i, a, b} <> D, t) : Bool}:
452	  match t:
453	    case B{n}:
454	      {==}
455	    case M{+j}:
456	      or_right(Nat.is_eq(j, i), has(D, j), kt)
457	
458	# the token created by the new rule is known
459	def kn_new(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk) -> {True{} == kn(Rule{i, a, b} <> D, M{i}) : Bool}:
460	  or_left(Nat.is_eq(i, i), has(D, i), nat_eq_refl(i))
461	
462	def kp(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, xs: List<&2, Tk>, hit: Bool, +eh: {hit == peek(a, b, xs) : Bool}, +kx: {True{} == knl(D, xs) : Bool}) -> {True{} == knl(Rule{i, a, b} <> D, merge.go(xs, hit, a, b, M{i})) : Bool}:
463	  match xs hit:
464	    case Nil{} _:
465	      {==}
466	    case Con{x, Nil{}} True{}:
467	      %eh : BD(_, Unit, {True{} == knl(Rule{i, a, b} <> D, merge.go(x <> Nil{}, True{}, a, b, M{i})) : Bool})
468	      Unit{}
469	    case Con{+x, Con{+y, +t}} True{}:
470	      and_intro(kn(Rule{i, a, b} <> D, M{i}), knl(Rule{i, a, b} <> D, merge.go(t, peek(a, b, t), a, b, M{i})), kn_new(D, i, a, b), kp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, y), knl(D, t), and_true_r(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx))))
471	    case Con{+x, +t} False{}:
472	      and_intro(kn(Rule{i, a, b} <> D, x), knl(Rule{i, a, b} <> D, merge.go(t, peek(a, b, t), a, b, M{i})), kn_mono(D, i, a, b, x, and_true_l(kn(D, x), knl(D, t), kx)), kp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, x), knl(D, t), kx)))
473	
474	# ---- induction over the rules ----
475	
476	# every emitted token is known to the complete (reversed) table
477	def enc_known(todo: List<&2, Rule>, +done: List<&2, Rule>, +xs: List<&2, Tk>, +wfe: {True{} == wf.go(todo, done) : Bool}, +kx: {True{} == knl(done, xs) : Bool}) -> {True{} == knl(List.reverse.go(&2, Rule, todo, done), encode.go(todo, done, xs)) : Bool}:
478	  match todo:
479	    case Nil{}:
480	      kx
481	    case Con{Rule{+i, +a, +b}, +rest}:
482	      enc_known(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}), and_true_r(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe), kp(done, i, a, b, xs, peek(a, b, xs), {==}, kx))
483	
484	# applying all the rules does not change what decode returns
485	def enc_dec(todo: List<&2, Rule>, +done: List<&2, Rule>, +xs: List<&2, Tk>, +wfe: {True{} == wf.go(todo, done) : Bool}, +kx: {True{} == knl(done, xs) : Bool}) -> {dec(List.reverse.go(&2, Rule, todo, done), encode.go(todo, done, xs)) == dec(done, xs) : List<&2, Nat>}:
486	  match todo:
487	    case Nil{}:
488	      {==}
489	    case Con{Rule{+i, +a, +b}, +rest}:
490	      Equal.trans(List<&2, Nat>, dec(List.reverse.go(&2, Rule, rest, Rule{i, a, b} <> done), encode.go(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}))), dec(Rule{i, a, b} <> done, xs), dec(done, xs), Equal.trans(List<&2, Nat>, dec(List.reverse.go(&2, Rule, rest, Rule{i, a, b} <> done), encode.go(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}))), dec(Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i})), dec(Rule{i, a, b} <> done, xs), enc_dec(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}), and_true_r(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe), kp(done, i, a, b, xs, peek(a, b, xs), {==}, kx)), mp(done, i, a, b, xs, peek(a, b, xs), {==}, kx, and_true_l(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe))), ext_dec(done, i, a, b, xs, kx, and_true_l(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe)))
491	
492	# ---- base case and final laws ----
493	
494	# with no rules, the bytes decode to themselves
495	def lift_dec(bs: List<&2, Nat>) -> {dec(Nil{}, lift(bs)) == bs : List<&2, Nat>}:
496	  match bs:
497	    case Nil{}:
498	      {==}
499	    case Con{+n, +t}:
500	      %lift_dec(t) : {n <> dec(Nil{}, lift(t)) == n <> _ : List<&2, Nat>}
501	      {==}
502	
503	def lift_known(bs: List<&2, Nat>) -> {True{} == knl(Nil{}, lift(bs)) : Bool}:
504	  match bs:
505	    case Nil{}:
506	      {==}
507	    case Con{n, t}:
508	      lift_known(t)
509	
510	# =====================================================================
511	# LAWS (the package laws, explained in plain language in the README)
512	# =====================================================================
513	
514	# LAW 1 (roundtrip): with a well-formed table (all rule ids
515	# distinct), decoding what was encoded returns the original bytes.
516	law roundtrip:
517	  for +table: List<&2, Rule>
518	  for +bs: List<&2, Nat>
519	  for +w: {True{} == wf(table) : Bool}
520	  {decode(table, encode(table, lift(bs))) == bs : List<&2, Nat>}
521	
522	def roundtrip(table, bs, w):
523	  Equal.trans(List<&2, Nat>, decode(table, encode(table, lift(bs))), dec(Nil{}, lift(bs)), bs, enc_dec(table, Nil{}, lift(bs), w, lift_known(bs)), lift_dec(bs))
524	
525	# LAW 2 (vocabulary): every token emitted by encode is a raw byte or was
526	# created by one of the table's rules; an unknown id never appears.
527	law vocab_bound:
528	  for +table: List<&2, Rule>
529	  for +bs: List<&2, Nat>
530	  for +w: {True{} == wf(table) : Bool}
531	  {True{} == knl(List.reverse(&2, Rule, table), encode(table, lift(bs))) : Bool}
532	
533	def vocab_bound(table, bs, w):
534	  enc_known(table, Nil{}, lift(bs), w, lift_known(bs))
535	
536	# LAW 3 (concatenation): decoding two token lists together is decoding
537	# each one and joining the bytes.
538	law dec_append:
539	  for +D: List<&2, Rule>
540	  for xs: List<&2, Tk>
541	  for +ys: List<&2, Tk>
542	  {dec(D, List.append(&2, Tk, xs, ys)) == List.append(&2, Nat, dec(D, xs), dec(D, ys)) : List<&2, Nat>}
543	
544	def dec_append(D, xs, ys):
545	  match xs:
546	    case Nil{}:
547	      {==}
548	    case Con{+x, +t}:
549	      %Equal.sym(List<&2, Nat>, dec(D, List.append(&2, Tk, t, ys)), List.append(&2, Nat, dec(D, t), dec(D, ys)), dec_append(D, t, ys)) : {List.append(&2, Nat, exp(D, x), _) == List.append(&2, Nat, List.append(&2, Nat, exp(D, x), dec(D, t)), dec(D, ys)) : List<&2, Nat>}
550	      Equal.sym(List<&2, Nat>, List.append(&2, Nat, List.append(&2, Nat, exp(D, x), dec(D, t)), dec(D, ys)), List.append(&2, Nat, exp(D, x), List.append(&2, Nat, dec(D, t), dec(D, ys))), NL.append_assoc(Nat, exp(D, x), dec(D, t), dec(D, ys)))
551	
552	# =====================================================================
553	# LAW 4: train generates well-formed tables (closes the roundtrip gap)
554	# =====================================================================
555	#
556	# train creates rules with ids next, next+1, next+2, ... None repeats, so the
557	# table satisfies wf, which is the hypothesis of the roundtrip law. Proof by induction
558	# over train.go, keeping the invariant "all ids already used are smaller
559	# than the next id".
560	
561	# n < 1 + n
562	def lt_succ(n: Nat) -> {True{} == Nat.is_lt(n, 1n+n) : Bool}:
563	  match n:
564	    case 0n:
565	      {==}
566	    case 1n+p:
567	      lt_succ(p)
568	
569	# j < i  ->  i != j
570	def lt_ne(j: Nat, i: Nat, e: {True{} == Nat.is_lt(j, i) : Bool}) -> {False{} == Nat.is_eq(i, j) : Bool}:
571	  match j i:
572	    case 0n 0n:
573	      %e : BD(_, Unit, {False{} == Nat.is_eq(0n, 0n) : Bool})
574	      Unit{}
575	    case 0n 1n+q:
576	      {==}
577	    case 1n+p 0n:
578	      %e : BD(_, Unit, {False{} == Nat.is_eq(0n, 1n+p) : Bool})
579	      Unit{}
580	    case 1n+p 1n+q:
581	      lt_ne(p, q, e)
582	
583	# j < i  ->  j < 1 + i
584	def lt_up(j: Nat, i: Nat, e: {True{} == Nat.is_lt(j, i) : Bool}) -> {True{} == Nat.is_lt(j, 1n+i) : Bool}:
585	  match j i:
586	    case 0n 0n:
587	      %e : BD(_, Unit, {True{} == Nat.is_lt(0n, 1n) : Bool})
588	      Unit{}
589	    case 0n 1n+q:
590	      {==}
591	    case 1n+p 0n:
592	      %e : BD(_, Unit, {True{} == Nat.is_lt(1n+p, 1n) : Bool})
593	      Unit{}
594	    case 1n+p 1n+q:
595	      lt_up(p, q, e)
596	
597	# all ids in the rule list are smaller than i
598	def below(D: List<&2, Rule>, +i: Nat) -> Bool:
599	  match D:
600	    case Nil{}:
601	      True{}
602	    case Con{Rule{+j, a, b}, t}:
603	      Bool.and(Nat.is_lt(j, i), below(t, i))
604	
605	def below_has(D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {False{} == has(D, i) : Bool}:
606	  match D:
607	    case Nil{}:
608	      {==}
609	    case Con{Rule{+j, a, b}, +t}:
610	      %lt_ne(j, i, and_true_l(Nat.is_lt(j, i), below(t, i), e)) : {False{} == Bool.or(_, has(t, i)) : Bool}
611	      below_has(t, i, and_true_r(Nat.is_lt(j, i), below(t, i), e))
612	
613	def below_up(D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {True{} == below(D, 1n+i) : Bool}:
614	  match D:
615	    case Nil{}:
616	      {==}
617	    case Con{Rule{+j, a, b}, +t}:
618	      and_intro(Nat.is_lt(j, 1n+i), below(t, 1n+i), lt_up(j, i, and_true_l(Nat.is_lt(j, i), below(t, i), e)), below_up(t, i, and_true_r(Nat.is_lt(j, i), below(t, i), e)))
619	
620	# adding the rule with id `next` (greater than all those in `done`) keeps the table "below 1+next"
621	def below_cons(+a: Tk, +b: Tk, +D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {True{} == below(Rule{i, a, b} <> D, 1n+i) : Bool}:
622	  and_intro(Nat.is_lt(i, 1n+i), below(D, 1n+i), lt_succ(i), below_up(D, i, e))
623	
624	# not (has D i) when all the ids in D are smaller than i
625	def below_fresh(+D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {True{} == Bool.not(has(D, i)) : Bool}:
626	  %below_has(D, i, e) : {True{} == Bool.not(_) : Bool}
627	  {==}
628	
629	# train.go generates rules with ids next, next+1, ...: none repeats, whatever the corpus
630	def tw(n: Nat, +next: Nat, +xs: List<&2, Tk>, m: Maybe<&1, Cnt>, +done: List<&2, Rule>, +bl: {True{} == below(done, next) : Bool}) -> {True{} == wf.go(train.go(n, next, xs, m), done) : Bool}:
631	  match n m:
632	    case 0n _:
633	      {==}
634	    case 1n++p None{}:
635	      {==}
636	    case 1n++p Some{Cnt{+a, +b, k}}:
637	      and_intro(Bool.not(has(done, next)), wf.go(train.go(p, 1n+next, merge(a, b, M{next}, xs), best(count.go(merge(a, b, M{next}, xs), Nil{}))), Rule{next, a, b} <> done), below_fresh(done, next, bl), tw(p, 1n+next, merge(a, b, M{next}, xs), best(count.go(merge(a, b, M{next}, xs), Nil{})), Rule{next, a, b} <> done, below_cons(a, b, done, next, bl)))
638	
639	# LAW (train_wf): the table that train returns is always well formed, so the roundtrip holds for it
640	law train_wf:
641	  for n: Nat
642	  for +next: Nat
643	  for +xs: List<&2, Tk>
644	  {True{} == wf(train(n, next, xs)) : Bool}
645	
646	def train_wf(n, next, xs):
647	  tw(n, next, xs, best(count.go(xs, Nil{})), Nil{}, {==})
648	
649	# LAW 5 (roundtrip_trained): train a table on any corpus, with as many merges as you like:
650	# decoding what it encodes always returns the original bytes. No hypothesis at all.
651	law roundtrip_trained:
652	  for +n: Nat
653	  for +corpus: List<&2, Nat>
654	  for +bs: List<&2, Nat>
655	  {decode(train(n, 0n, lift(corpus)), encode(train(n, 0n, lift(corpus)), lift(bs))) == bs : List<&2, Nat>}
656	
657	def roundtrip_trained(n, corpus, bs):
658	  roundtrip(train(n, 0n, lift(corpus)), bs, train_wf(n, 0n, lift(corpus)))