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