main.bend

bend-ml-tensor-array@0.1.19.0 · 2,340 lines · 321 definitions

1# bend-ml-tensor-array: tensors over a flat Array<F32>, with the shape in the type.
2#
3#   import bend-ml-tensor-array@0.1.19.0/main.bend as TA
4#
5# Mat<r, c> stores r*c numbers in an Array<F32>, row by row (index i*c + j). The
6# dimensions are erased type parameters: the same guarantee as bend-ml-tensor
7# (a product with wrong dimensions does not compile), but ~50x faster than lists:
8# Array.get/set by index, and products in parallel blocks of rows.
9#
10# An Array is affine (a single owner): every operation that reads an operand returns it
11# together with the result (as Array.get does), in records MM, AR, ... Use Mat.clone
12# when you need two copies.
13#
14# What is NOT proved: the invariant "the Array has capacity >= r*c" holds because
15# the constructors (Mat.zeros, Mat.of_list) allocate the right size; and F32 is
16# validated by tests against PyTorch (reference/test_tensor_array.py), not by proof.
17
18import Base
19
20# ---------------------------------------------------------------------
21# helper types
22# ---------------------------------------------------------------------
23
24type R is Type:
25  R{a: Array<F32>, b: Array<F32>, x: F32}
26
27type G is Type:
28  G{a: Array<F32>, b: Array<F32>, c: Array<F32>}
29
30# A dot product with four partial sums. Each step adds its product to a0 and rotates the sums
31# (a1, a2, a3, a0), so each sum gets every fourth product and the processor can work on the next
32# addition before the last one has finished; at the end the sums are added in pairs. The order of
33# the additions is fixed, so the result is the same on every run and every thread count, but it is
34# not the order of a single running sum (NOTES.md, exp. 17).
35def dotr(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a0: F32, a1: F32, a2: F32, a3: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> R:
36  match k ra rb:
37    case 0n Tuple{a, x} Tuple{b, y}:
38      R{a, b, ((a0 + a1 : F32) + (a2 + a3 : F32) : F32)}
39    case 1n+p Tuple{a, +x} Tuple{b, +y}:
40      dotr(p, (ia + sa : U32), (ib + sb : U32), sa, sb, a1, a2, a3, (a0 + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32)))
41
42# dotr's state between two steps: the four partial sums and the two pending reads
43type D4 is Type:
44  D4{a0: F32, a1: F32, a2: F32, a3: F32, ra: Array<F32> & F32, rb: Array<F32> & F32}
45
46# one step of dotr: adds the pending product to a0, rotates the sums, reads the next pair at ia, ib
47def dstep(+ia: U32, +ib: U32, s: D4) -> D4:
48  match s:
49    case D4{a0, a1, a2, a3, ra, rb}:
50      match ra rb:
51        case Tuple{a, +x} Tuple{b, +y}:
52          D4{a1, a2, a3, (a0 + (x * y : F32) : F32), Array.get(F32, a, ia), Array.get(F32, b, ib)}
53
54def dfin(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, s: D4) -> R:
55  match s:
56    case D4{a0, a1, a2, a3, ra, rb}:
57      dotr(k, ia, ib, sa, sb, a0, a1, a2, a3, ra, rb)
58
59# dotr with its loop unrolled by four: q rounds of four dstep, then the last rem (< 4) steps by dotr itself.
60# Every product goes to the same sum in the same order as in dotr (four rotations bring the names back), so the
61# result has dotr's bits (tests/dot4.bend). The compiled loop runs ~16% fewer instructions; one thread on a
62# 50,257 x 768 matrix · vector goes from 26.7 to 20.0 ms, smaller ones and 16 threads stay within noise (NOTES.md,
63# 2026-11-01). Eight steps a round were no faster.
64def dot4(q: Nat, +rem: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, s: D4) -> R:
65  match q:
66    case 0n:
67      dfin(rem, ia, ib, sa, sb, s)
68    case 1n+p:
69      dot4(p, rem, (ia + (4 * sa : U32) : U32), (ib + (4 * sb : U32) : U32), sa, sb, dstep((ia + (4 * sa : U32) : U32), (ib + (4 * sb : U32) : U32), dstep((ia + (3 * sa : U32) : U32), (ib + (3 * sb : U32) : U32), dstep((ia + (2 * sa : U32) : U32), (ib + (2 * sb : U32) : U32), dstep((ia + sa : U32), (ib + sb : U32), s)))))
70
71def dotg(+k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> R:
72  dot4(Nat.div(k, 4n), Nat.sub(k, Nat.mul(4n, Nat.div(k, 4n))), ia, ib, sa, sb, D4{acc, 0.0, 0.0, 0.0, ra, rb})
73
74# writes one cell of C: c[idx] = x, or c[idx] = c[idx] + x when ac (an accumulating product, 0.1.10.0). x is computed
75# exactly as without ac and added once, so C + A·B has the bits of adding a separately computed A·B to C.
76
77def cwr_add(+idx: U32, +x: F32, r: Array<F32> & F32) -> Array<F32>:
78  match r:
79    case Tuple{c, +v}:
80      Array.set(F32, c, idx, (v + x : F32))
81
82def cwr(ac: Bool, +idx: U32, +x: F32, c: Array<F32>) -> Array<F32>:
83  match ac:
84    case True{}:
85      cwr_add(idx, x, Array.get(F32, c, idx))
86    case False{}:
87      Array.set(F32, c, idx, x)
88
89# row i of C, columns j.. ; r is the product of the current column (already computed)
90
91def gcols(+ac: Bool, mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array<F32>, r: R) -> G:
92  match mleft r:
93    case 0n R{a, b, +x}:
94      G{a, b, cwr(ac, cidx, x, c)}
95    case 1n+p R{a, b, +x}:
96      gcols(ac, p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, (cidx + 1 : U32), cwr(ac, cidx, x, c), dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
97
98def grows(+ac: Bool, nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G:
99  match nleft g:
100    case 0n G{a, b, c}:
101      G{a, b, c}
102    case 1n+p G{a, b, c}:
103      grows(ac, p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, mu, gcols(ac, mm1, kn, (i * arow : U32), 0, sa, sb, bcol, (i * mu : U32), c, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0))))
104
105# C (n x m) = A · B with k steps, or C + A · B when ac; c is the output array (with capacity n*m)
106
107def gemm_ac(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
108  grows(ac, n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c})
109
110def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
111  gemm_ac(False{}, n, m, k, arow, sa, sb, bcol, a, b, c)
112
113
114# ---------------------------------------------------------------------
115# parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B,
116# returning lists that are written into C at the end
117# ---------------------------------------------------------------------
118
119type S3 is Type:
120  S3{a: Array<F32>, b: Array<F32>, out: List<&2, F32>}
121
122def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3:
123  match mleft r:
124    case 0n R{a, b, +x}:
125      S3{a, b, x <> out}
126    case 1n+p R{a, b, +x}:
127      lcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, x <> out, dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
128
129def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3:
130  match nleft st:
131    case 0n S3{a, b, out}:
132      S3{a, b, out}
133    case 1n+p S3{a, b, out}:
134      lrows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, lcols(mm1, kn, (i * arow : U32), 0, sa, sb, bcol, out, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0))))
135
136def lblock_fin(st: S3) -> List<&2, F32>:
137  match st:
138    case S3{a, b, out}:
139      List.reverse(&2, F32, out)
140
141def lblock(cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
142  lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}}))
143
144def lpar(d: Nat, +cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array<F32> & Array<F32>, rb: Array<F32> & Array<F32>) -> List<&2, F32>:
145  match d ra rb:
146    case 0n Tuple{a, a2} Tuple{b, b2}:
147      lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b)
148    case 1n++q Tuple{a, a2} Tuple{b, b2}:
149      x y = lpar(q, Nat.div(cnt, 2n), mm1, kn, i, arow, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) lpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), mm1, kn, (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))
150      List.append(&2, F32, x, y)
151
152def write_l(xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
153  match xs:
154    case Nil{}:
155      c
156    case Con{h, t}:
157      write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h))
158
159# c[idx..] = list, or c[idx..] += list when ac
160
161def write_la(+ac: Bool, xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
162  match xs:
163    case Nil{}:
164      c
165    case Con{+h, t}:
166      write_la(ac, t, (idx + 1 : U32), cwr(ac, idx, h, c))
167
168def gp3(+ac: Bool, a: Array<F32>, b: Array<F32>, c: Array<F32>, xs: List<&2, F32>) -> G:
169  G{a, b, write_la(ac, xs, 0, c)}
170
171def gp2(d: Nat, +n: Nat, +mm1: Nat, +kn: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, c: Array<F32>, pa: Array<F32> & Array<F32>, pb: Array<F32> & Array<F32>) -> G:
172  match pa pb:
173    case Tuple{a1, a2} Tuple{b1, b2}:
174      gp3(False{}, a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
175
176# C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C
177
178def gemm_par(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
179  gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
180
181
182# ---------------------------------------------------------------------
183# parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a
184# range of columns of C, with its own copy of the operands
185# ---------------------------------------------------------------------
186
187def cblock_fin(st: S3) -> List<&2, F32>:
188  match st:
189    case S3{a, b, out}:
190      List.reverse(&2, F32, out)
191
192# cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0
193
194def cblock_go(cnt1: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
195  cblock_fin(lcols(cnt1, kn, ia0, (j0 * bcol : U32), sa, sb, bcol, Nil{}, dotg(kn, ia0, (j0 * bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (j0 * bcol : U32)))))
196
197def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
198  match cnt:
199    case 0n:
200      Nil{}
201    case 1n+p:
202      cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b)
203
204def cpar(d: Nat, +cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array<F32> & Array<F32>, rb: Array<F32> & Array<F32>) -> List<&2, F32>:
205  match d ra rb:
206    case 0n Tuple{a, a2} Tuple{b, b2}:
207      cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b)
208    case 1n++q Tuple{a, a2} Tuple{b, b2}:
209      x y = cpar(q, Nat.div(cnt, 2n), kn, ia0, j0, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) cpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), kn, ia0, (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))
210      List.append(&2, F32, x, y)
211
212def cp2(+ac: Bool, d: Nat, +m: Nat, +kn: Nat, +sa: U32, +sb: U32, +bcol: U32, c: Array<F32>, pa: Array<F32> & Array<F32>, pb: Array<F32> & Array<F32>) -> G:
213  match pa pb:
214    case Tuple{a1, a2} Tuple{b1, b2}:
215      gp3(ac, a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
216
217# C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel. Since 0.1.12.0 no product calls it (one row runs
218# the scalar loop, see gemm_any3): it clones both operands at every node. Kept for callers of earlier versions.
219
220def gemm_par_cols(+ac: Bool, d: Nat, +m: Nat, +k: Nat, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
221  cp2(ac, d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
222
223# ---------------------------------------------------------------------
224# element-wise operations (in place, by index)
225# ---------------------------------------------------------------------
226
227# ---------------------------------------------------------------------
228# gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol]
229# (nn: arow=K, sa=1, sb=M, bcol=1 | nt: B stored m x k: sb=1, bcol=K | tn: A stored k x n: arow=1, sa=N)
230# ---------------------------------------------------------------------
231
232# acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair
233
234type P2 is Type:
235  P2{a: Array<F32>, b: Array<F32>}
236
237type AL is Type:
238  AL{c: Array<F32>, l: List<&2, F32>}
239
240# list of mleft + 1 elements ending at idx (reads with Array.get, returning the array). It reads from
241# idx down and conses each number onto acc, so the list comes out in index order with no reverse
242# (NOTES.md, 2026-11-01): callers pass the last index and its read.
243
244def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array<F32> & F32) -> AL:
245  match mleft rc:
246    case 0n Tuple{c, +v}:
247      AL{c, v <> acc}
248    case 1n+p Tuple{c, +v}:
249      read_l(p, (idx - 1 : U32), v <> acc, Array.get(F32, c, (idx - 1 : U32)))
250
251# writes the list starting at idx
252
253def addl(xs: List<&2, F32>, +idx: U32, rc: Array<F32> & F32) -> Array<F32>:
254  match xs rc:
255    case Nil{} Tuple{c, v}:
256      c
257    case Con{+h, t} Tuple{c, +v}:
258      addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32)))
259
260# adds a bias (a list of m) to each of the n rows of c (n x m)
261
262def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array<F32>) -> Array<F32>:
263  match nleft:
264    case 0n:
265      c
266    case 1n+p:
267      bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32))))
268
269# in-place relu on the first n elements
270
271def relu_l(nleft: Nat, +t: U32, rc: Array<F32> & F32) -> Array<F32>:
272  match nleft rc:
273    case 0n Tuple{c, v}:
274      c
275    case 1n+p Tuple{c, +v}:
276      relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32)))
277
278def pos2(b: Bool) -> F32:
279  match b:
280    case True{}:
281      1.0
282    case False{}:
283      0.0
284
285def pos(+y: F32) -> F32:
286  pos2(F32.is_gt(y, 0.0))
287
288# d[t] *= (h[t] > 0): the relu gradient; returns both arrays
289
290def mask_l(nleft: Nat, +t: U32, rd: Array<F32> & F32, rh: Array<F32> & F32) -> P2:
291  match nleft rd rh:
292    case 0n Tuple{d, x} Tuple{h, y}:
293      P2{d, h}
294    case 1n+p Tuple{d, +x} Tuple{h, +y}:
295      mask_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, d, t, (x * pos(y) : F32)), (t + 1 : U32)), Array.get(F32, h, (t + 1 : U32)))
296
297# w[t] -= lr * dw[t]
298
299def maxl2(xs: List<&2, F32>, cur: F32) -> F32:
300  match xs:
301    case Nil{}:
302      cur
303    case Con{+h, t}:
304      maxl2(t, F32.max(cur, h))
305
306def maxl(xs: List<&2, F32>) -> F32:
307  match xs:
308    case Nil{}:
309      0.0
310    case Con{+h, t}:
311      maxl2(t, h)
312
313def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
314  match xs:
315    case Nil{}:
316      Nil{}
317    case Con{+h, t}:
318      F32.exp((h - m : F32)) <> expl(t, m)
319
320def suml(xs: List<&2, F32>, acc: F32) -> F32:
321  match xs:
322    case Nil{}:
323      acc
324    case Con{+h, t}:
325      suml(t, (acc + h : F32))
326
327def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
328  match xs:
329    case Nil{}:
330      Nil{}
331    case Con{+h, t}:
332      (h / s : F32) <> divl(t, s)
333
334def softmax2(+es: List<&2, F32>) -> List<&2, F32>:
335  divl(es, suml(es, 0.0))
336
337def softmax(+xs: List<&2, F32>) -> List<&2, F32>:
338  softmax2(expl(xs, maxl(xs)))
339
340# p[label]
341
342def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array<F32> & F32, rd: Array<F32> & F32) -> P2:
343  match nleft rw rd:
344    case 0n Tuple{w, x} Tuple{d, y}:
345      P2{w, d}
346    case 1n+p Tuple{w, +x} Tuple{d, +y}:
347      sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x - (lr * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, d, (t + 1 : U32)))
348
349# ---------------------------------------------------------------------
350# softmax + cross-entropy per row (lists of 10 elements: negligible cost)
351# ---------------------------------------------------------------------
352
353def pick(ps: List<&2, F32>, n: Nat) -> F32:
354  match ps n:
355    case Con{+h, t} 0n:
356      h
357    case Con{h, t} 1n+q:
358      pick(t, q)
359    case Nil{} _:
360      0.0
361
362# (p - onehot(label)) / n
363
364def hit(b: Bool) -> F32:
365  match b:
366    case True{}:
367      1.0
368    case False{}:
369      0.0
370
371def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>:
372  match ps:
373    case Nil{}:
374      Nil{}
375    case Con{+h, t}:
376      (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf)
377
378type L2 is Type:
379  L2{loss: F32, c: Array<F32>}
380
381def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array<F32>) -> L2:
382  L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)}
383
384def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2:
385  match al:
386    case AL{c, l}:
387      ce_row2(softmax(l), label, nf, idx, loss, c)
388
389# walks the rows of z2 (n x m): r carries the accumulated loss and the array; at the end the array has become the gradient
390
391type Mat<-r: Nat, -c: Nat> is Type:
392  Mat{d: Array<F32>}
393
394def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
395  match labels r:
396    case Nil{} L2{loss, c}:
397      L2{loss, c}
398    case Con{+y, t} L2{loss, c}:
399      ce_rows(t, mm1, mu, nf, (i + 1 : U32), ce_row(read_l(mm1, ((i * mu : U32) + U32.from_nat(mm1) : U32), Nil{}, Array.get(F32, c, ((i * mu : U32) + U32.from_nat(mm1) : U32))), y, nf, (i * mu : U32), loss))
400
401# ---- the same rows over the array itself: no lists ----
402# The list version above reads each row into a list and walks it five times. This one walks the row in place with tail
403# loops, with the same operations in the same order (the max from the first cell, exp(x - max) summed from +0.0 left to
404# right, p = e / s, the loss term read before the row is overwritten, (p - onehot) / n written back), so the same bits;
405# MetalBend measured 3.4x fewer cycles on 64 rows of 49,152 (2026-10-25). The list version stays as the reference
406# tests/ce_array.bend compares it with.
407
408type AF is Type:
409  AF{c: Array<F32>, x: F32}
410
411# the max of k + 1 cells from idx, from m (the first cell's value, already read in r)
412def cea_max(k: Nat, +idx: U32, +m: F32, r: Array<F32> & F32) -> AF:
413  match k r:
414    case 0n Tuple{c, +x}:
415      AF{c, F32.max(m, x)}
416    case 1n+p Tuple{c, +x}:
417      cea_max(p, (idx + 1 : U32), F32.max(m, x), Array.get(F32, c, (idx + 1 : U32)))
418
419# cells idx..idx + k become exp(cell - m); returns their sum from s, added left to right
420def cea_exp(k: Nat, +idx: U32, +m: F32, +s: F32, r: Array<F32> & F32) -> AF:
421  match k r:
422    case 0n Tuple{c, +x}:
423      +e = F32.exp((x - m : F32))
424      AF{Array.set(F32, c, idx, e), (s + e : F32)}
425    case 1n+p Tuple{c, +x}:
426      +e = F32.exp((x - m : F32))
427      cea_exp(p, (idx + 1 : U32), m, (s + e : F32), Array.get(F32, Array.set(F32, c, idx, e), (idx + 1 : U32)))
428
429# cell idx (column j) becomes ((e / s) - onehot) / nf, for k + 1 cells
430def cea_grad(k: Nat, +idx: U32, +j: Nat, +label: Nat, +s: F32, +nf: F32, r: Array<F32> & F32) -> Array<F32>:
431  match k r:
432    case 0n Tuple{c, +x}:
433      Array.set(F32, c, idx, ((((x / s : F32) - hit(Nat.is_eq(label, j)) : F32) / nf) : F32))
434    case 1n+p Tuple{c, +x}:
435      cea_grad(p, (idx + 1 : U32), 1n+j, label, s, nf, Array.get(F32, Array.set(F32, c, idx, ((((x / s : F32) - hit(Nat.is_eq(label, j)) : F32) / nf) : F32)), (idx + 1 : U32)))
436
437def cea_p2(+s: F32, r: Array<F32> & F32) -> AF:
438  match r:
439    case Tuple{c, +e}:
440      AF{c, (e / s : F32)}
441
442# p[label] = e[label] / s, or 0.0 for a label past the row (pick's Nil case)
443def cea_p(inside: Bool, +o: U32, +label: Nat, +s: F32, c: Array<F32>) -> AF:
444  match inside:
445    case True{}:
446      cea_p2(s, Array.get(F32, c, (o + U32.from_nat(label) : U32)))
447    case False{}:
448      AF{c, 0.0}
449
450def cea_row3(+mm1: Nat, +o: U32, +label: Nat, +s: F32, +nf: F32, +loss: F32, q: AF) -> L2:
451  match q:
452    case AF{c, +p}:
453      L2{(loss - F32.log(F32.max(p, 0.000000000001)) : F32), cea_grad(mm1, o, 0n, label, s, nf, Array.get(F32, c, o))}
454
455def cea_row2(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, q: AF) -> L2:
456  match q:
457    case AF{c, +s}:
458      cea_row3(mm1, o, label, s, nf, loss, cea_p(Nat.is_le(label, mm1), o, label, s, c))
459
460def cea_row1(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, q: AF) -> L2:
461  match q:
462    case AF{c, +m}:
463      cea_row2(mm1, o, label, nf, loss, cea_exp(mm1, o, m, 0.0, Array.get(F32, c, o)))
464
465# a row's max starts from its first cell (maxl), then folds over the k + 1 cells from it
466def cea_row0(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, r: Array<F32> & F32) -> L2:
467  match r:
468    case Tuple{c, +x0}:
469      cea_row1(mm1, o, label, nf, loss, cea_max(mm1, o, x0, Array.get(F32, c, o)))
470
471# one row per label, in order, as ce_rows
472def cea_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
473  match labels r:
474    case Nil{} L2{loss, c}:
475      L2{loss, c}
476    case Con{+y, t} L2{+loss, c}:
477      cea_rows(t, mm1, mu, nf, (i + 1 : U32), cea_row0(mm1, (i * mu : U32), y, nf, loss, Array.get(F32, c, (i * mu : U32))))
478
479# ---- helper: hits per row ----
480
481type CR is Type:
482  CR{c: Array<F32>, n: Nat}
483
484def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
485  match xs:
486    case Nil{}:
487      False{}
488    case Con{+h, t}:
489      F32.is_gt(h, v)
490
491def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat:
492  match xs g:
493    case Nil{} _:
494      bi
495    case Con{+h, +t} True{}:
496      argmax2(t, gt_head(t, h), 1n+i, i, h)
497    case Con{h, +t} False{}:
498      argmax2(t, gt_head(t, bv), 1n+i, bi, bv)
499
500def argmax(xs: List<&2, F32>) -> Nat:
501  match xs:
502    case Nil{}:
503      0n
504    case Con{+h, +t}:
505      argmax2(t, gt_head(t, h), 1n, 0n, h)
506
507def one(b: Bool) -> Nat:
508  match b:
509    case True{}:
510      1n
511    case False{}:
512      0n
513
514def count_rows3(+y: Nat, n: Nat, al: AL) -> CR:
515  match al:
516    case AL{c, l}:
517      CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))}
518
519def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR:
520  match labels r:
521    case Nil{} CR{c, n}:
522      CR{c, n}
523    case Con{+y, t} CR{c, n}:
524      count_rows(t, mm1, mu, (i + 1 : U32), count_rows3(y, n, read_l(mm1, ((i * mu : U32) + U32.from_nat(mm1) : U32), Nil{}, Array.get(F32, c, ((i * mu : U32) + U32.from_nat(mm1) : U32)))))
525
526
527# =====================================================================
528# Typed API: Mat<r, c> (the dimensions enter the type)
529# =====================================================================
530
531# =====================================================================
532# Array capacity: a PROVED depth for the number of slots
533# =====================================================================
534#
535# cap_depth(n) is the d such that an Array with 2^d slots holds n elements. It is computed by
536# structural recursion on n (not with U32.log2) so that the law cap_ok below can prove
537# `n <= 2^cap_depth(n)`. What stays trusted is only that Array.new(T, d, v) allocates 2^d slots.
538
539# a type chosen by a Bool: rewriting through it refutes True == False
540def BD(b: Bool, t: Type, f: Type) -> Type:
541  match b:
542    case True{}:
543      t
544    case False{}:
545      f
546
547# a <= b, by direct recursion
548def leb(a: Nat, b: Nat) -> Bool:
549  match a b:
550    case 0n _:
551      True{}
552    case 1n+p 0n:
553      False{}
554    case 1n+p 1n+q:
555      leb(p, q)
556
557def pow2(d: Nat) -> Nat:
558  match d:
559    case 0n:
560      1n
561    case 1n++q:
562      # bend-lint: disable-next bend-lint-rules/performance/repeated_call -- pow2 appears only in proofs' types not-expected-because: the package never runs it (NOTES 2026-10-17)
563      Nat.add(pow2(q), pow2(q))
564
565# smallest power of two >= n, by structural recursion on n: returns (depth, 2^depth)
566def cap_pick(ok: Bool, d: Nat, +p: Nat) -> Nat & Nat:
567  match ok:
568    case True{}:
569      (d, p)
570    case False{}:
571      (1n+d, Nat.add(p, p))
572
573def cap_step(r: Nat & Nat, +n: Nat) -> Nat & Nat:
574  match r:
575    case Tuple{+d, +p}:
576      cap_pick(leb(n, p), d, p)
577
578def cap_pair(n: Nat) -> Nat & Nat:
579  match n:
580    case 0n:
581      (0n, 1n)
582    case 1n++m:
583      cap_step(cap_pair(m), 1n+m)
584
585def fst_nat(r: Nat & Nat) -> Nat:
586  match r:
587    case Tuple{d, p}:
588      d
589
590def snd_nat(r: Nat & Nat) -> Nat:
591  match r:
592    case Tuple{d, p}:
593      p
594
595# leb is Base's Nat.is_le, case for case. leb walks its arguments one step at a time in the compiled program; Base's
596# comparison is one native comparison there (0.1.13.0).
597def le_leb(a: Nat, b: Nat) -> {Nat.is_le(a, b) == leb(a, b) : Bool}:
598  match a b:
599    case 0n 0n:
600      {==}
601    case 0n 1n+q:
602      {==}
603    case 1n+p 0n:
604      {==}
605    case 1n+p 1n+q:
606      le_leb(p, q)
607
608def leb_of(+a: Nat, +b: Nat, e: {True{} == Nat.is_le(a, b) : Bool}) -> {True{} == leb(a, b) : Bool}:
609  %le_leb(a, b) : {True{} == _ : Bool}
610  e
611
612# the same pair by doubling (0.1.10.0): p doubles until n <= p, log2(n) steps in a tail loop, where cap_pair takes n
613# steps of a recursion that is not a tail call, each with a comparison of up to n steps. Its fuel is n, which always
614# suffices (2^n >= n); if it ran out, the answer would be cap_pair's, so the law below needs no arithmetic about how
615# many doublings there are. ok is Nat.is_le(n, p), computed by the caller (match only inspects a parameter); with leb
616# it was n steps in all, 152 ms for n = 786,432 on one GPU lane (0.1.13.0).
617def cap_dbl(f: Nat, ok: Bool, +n: Nat, +d: Nat, +p: Nat) -> Nat & Nat:
618  match f ok:
619    case _ True{}:
620      (d, p)
621    case 0n False{}:
622      cap_pair(n)
623    case 1n+g False{}:
624      cap_dbl(g, Nat.is_le(n, Nat.add(p, p)), n, 1n+d, Nat.add(p, p))
625
626# the depth d such that an Array with 2^d slots holds n elements (the smallest such d)
627def cap_depth(+n: Nat) -> Nat:
628  fst_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))
629
630# ---- lemmas ----
631
632# a <= b  ->  a <= 1 + b
633def leb_succ(a: Nat, b: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, 1n+b) : Bool}:
634  match a b:
635    case 0n _:
636      {==}
637    case 1n+p 0n:
638      %e : BD(_, Unit, {True{} == leb(1n+p, 1n) : Bool})
639      Unit{}
640    case 1n+p 1n+q:
641      leb_succ(p, q, e)
642
643# a <= b  ->  a <= c + b
644def leb_up_l(+a: Nat, +b: Nat, c: Nat, +e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(c, b)) : Bool}:
645  match c:
646    case 0n:
647      e
648    case 1n++c2:
649      leb_succ(a, Nat.add(c2, b), leb_up_l(a, b, c2, e))
650
651# a <= b  ->  a <= b + c
652def leb_up(a: Nat, b: Nat, c: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(b, c)) : Bool}:
653  match a b:
654    case 0n _:
655      {==}
656    case 1n+p 0n:
657      %e : BD(_, Unit, {True{} == leb(1n+p, Nat.add(0n, c)) : Bool})
658      Unit{}
659    case 1n+p 1n+q:
660      leb_up(p, q, c, e)
661
662# a <= b  and  c <= d  ->  a + c <= b + d
663def leb_add(a: Nat, b: Nat, c: Nat, d: Nat, e1: {True{} == leb(a, b) : Bool}, e2: {True{} == leb(c, d) : Bool}) -> {True{} == leb(Nat.add(a, c), Nat.add(b, d)) : Bool}:
664  match a b:
665    case 0n _:
666      leb_up_l(c, d, b, e2)
667    case 1n+p 0n:
668      %e1 : BD(_, Unit, {True{} == leb(Nat.add(1n+p, c), Nat.add(0n, d)) : Bool})
669      Unit{}
670    case 1n+p 1n+q:
671      leb_add(p, q, c, d, e1, e2)
672
673# 1 <= 2^d
674def pow2_pos(d: Nat) -> {True{} == leb(1n, pow2(d)) : Bool}:
675  match d:
676    case 0n:
677      {==}
678    case 1n++q:
679      leb_up(1n, pow2(q), pow2(q), pow2_pos(q))
680
681# 1 <= p when p = 2^d
682def pos_p(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {True{} == leb(1n, p) : Bool}:
683  %Equal.sym(Nat, p, pow2(d), i1) : {True{} == leb(1n, _) : Bool}
684  pow2_pos(d)
685
686# p + p == 2^(1+d) when p == 2^d
687def step_eq(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {Nat.add(p, p) == pow2(1n+d) : Nat}:
688  %i1 : {Nat.add(p, p) == Nat.add(_, _) : Nat}
689  {==}
690
691# one step: given the invariant for m, the invariant for 1+m holds whichever way the comparison goes
692def step_ok(ok: Bool, +d: Nat, +p: Nat, +m: Nat, +e: {ok == leb(1n+m, p) : Bool}, +i1: {p == pow2(d) : Nat}, +i2: {True{} == leb(m, p) : Bool}) -> {snd_nat(cap_pick(ok, d, p)) == pow2(fst_nat(cap_pick(ok, d, p))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_pick(ok, d, p))) : Bool}:
693  match ok:
694    case True{}:
695      (i1, e)
696    case False{}:
697      (step_eq(d, p, i1), leb_add(1n, p, m, p, pos_p(d, p, i1), i2))
698
699
700# the invariant for cap_step on a pair, opened
701def step_pair_ok(+m: Nat, r: Nat & Nat, i1: {snd_nat(r) == pow2(fst_nat(r)) : Nat}, i2: {True{} == leb(m, snd_nat(r)) : Bool}) -> {snd_nat(cap_step(r, 1n+m)) == pow2(fst_nat(cap_step(r, 1n+m))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_step(r, 1n+m))) : Bool}:
702  match r:
703    case Tuple{+d, +p}:
704      step_ok(leb(1n+m, p), d, p, m, {==}, i1, i2)
705
706def cap_succ(+m: Nat, ih: {snd_nat(cap_pair(m)) == pow2(fst_nat(cap_pair(m))) : Nat} & {True{} == leb(m, snd_nat(cap_pair(m))) : Bool}) -> {snd_nat(cap_pair(1n+m)) == pow2(fst_nat(cap_pair(1n+m))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_pair(1n+m))) : Bool}:
707  match ih:
708    case Tuple{i1, i2}:
709      step_pair_ok(m, cap_pair(m), i1, i2)
710
711# the invariant of cap_pair: the second component is 2^(first) and is >= n
712def cap_inv(n: Nat) -> {snd_nat(cap_pair(n)) == pow2(fst_nat(cap_pair(n))) : Nat} & {True{} == leb(n, snd_nat(cap_pair(n))) : Bool}:
713  match n:
714    case 0n:
715      ({==}, {==})
716    case 1n++m:
717      cap_succ(m, cap_inv(m))
718
719# the invariant of cap_dbl: whichever way it ends (n <= p found, or the fuel out and cap_pair's answer), the second
720# component is 2^(first) and is >= n
721def dbl_inv(f: Nat, ok: Bool, +n: Nat, +d: Nat, +p: Nat, e: {ok == Nat.is_le(n, p) : Bool}, i1: {p == pow2(d) : Nat}) -> {snd_nat(cap_dbl(f, ok, n, d, p)) == pow2(fst_nat(cap_dbl(f, ok, n, d, p))) : Nat} & {True{} == leb(n, snd_nat(cap_dbl(f, ok, n, d, p))) : Bool}:
722  match f ok:
723    case _ True{}:
724      (i1, leb_of(n, p, e))
725    case 0n False{}:
726      cap_inv(n)
727    case 1n+g False{}:
728      dbl_inv(g, Nat.is_le(n, Nat.add(p, p)), n, 1n+d, Nat.add(p, p), {==}, step_eq(d, p, i1))
729
730def dbl_final(n: Nat, pr: {snd_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n)) == pow2(fst_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))) : Nat} & {True{} == leb(n, snd_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))) : Bool}) -> {True{} == leb(n, pow2(cap_depth(n))) : Bool}:
731  match pr:
732    case Tuple{i1, i2}:
733      %i1 : {True{} == leb(n, _) : Bool}
734      i2
735
736# LAW (cap_ok): an Array with 2^cap_depth(n) slots always has room for n elements.
737law cap_ok:
738  for +n: Nat
739  {True{} == leb(n, pow2(cap_depth(n))) : Bool}
740
741def cap_ok(n):
742  dbl_final(n, dbl_inv(n, Nat.is_le(n, 1n), n, 0n, 1n, {==}, {==}))
743
744# r x c of zeros, with the exact capacity (smallest power of 2)
745
746def Mat.zeros(+r: Nat, +c: Nat) -> Mat<r, c>:
747  Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)}
748
749def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat<r, c>:
750  Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)}
751
752def fill_list(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
753  match xs:
754    case Nil{}:
755      a
756    case Con{h, t}:
757      fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h))
758
759def len_is(xs: List<&2, F32>, n: Nat) -> Bool:
760  match xs n:
761    case Nil{} 0n:
762      True{}
763    case Con{h, t} 1n+p:
764      len_is(t, p)
765    case _ _:
766      False{}
767
768def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
769  match ok:
770    case True{}:
771      Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}}
772    case False{}:
773      None{}
774
775# Mat<r, c> from r*c numbers in row order; None if the size does not match
776
777def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
778  Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs)
779
780# reads the first n numbers, returning the matrix
781
782# Mat<r, c> without checking the list size (a short list leaves the rest at zero, a long one is cut
783# by the capacity): only for when the size is known by construction. Prefer Mat.of_list.
784
785def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat<r, c>:
786  Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))}
787
788# writes the numbers of xs into m starting at flat index i (row-major); numbers past the
789# capacity are dropped. For loading a large matrix block by block without building one big list.
790def Mat.fill_at(-r: Nat, -c: Nat, +i: U32, xs: List<&2, F32>, m: Mat<r, c>) -> Mat<r, c>:
791  match m:
792    case Mat{d}:
793      Mat{fill_list(xs, i, d)}
794
795type ML<-r: Nat, -c: Nat> is Type:
796  ML{m: Mat<r, c>, l: List<&2, F32>}
797
798def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML<r, c>:
799  match al:
800    case AL{d, l}:
801      ML{Mat{d}, l}
802
803def ML.len(-r: Nat, -c: Nat, ml: ML<r, c>) -> Nat:
804  match ml:
805    case ML{m, l}:
806      List.length(&2, F32, l)
807
808# n numbers from index 0, in row order; n = 0 gives the empty list (a Mat with no numbers)
809def Mat.to_list_count(-r: Nat, -c: Nat, d: Array<F32>, +n: Nat) -> ML<r, c>:
810  match n:
811    case 0n:
812      ML{Mat{d}, Nil{}}
813    case 1n+p:
814      Mat.to_list2(r, c, read_l(p, (U32.from_nat(n) - 1 : U32), Nil{}, Array.get(F32, d, (U32.from_nat(n) - 1 : U32))))
815
816# two independent copies
817
818def Mat.to_list(+r: Nat, +c: Nat, m: Mat<r, c>) -> ML<r, c>:
819  match m:
820    case Mat{d}:
821      Mat.to_list_count(r, c, d, Nat.mul(r, c))
822
823type MM2<-r: Nat, -c: Nat> is Type:
824  MM2{a: Mat<r, c>, b: Mat<r, c>}
825
826def Mat.clone2(-r: Nat, -c: Nat, p: Array<F32> & Array<F32>) -> MM2<r, c>:
827  match p:
828    case Tuple{x, y}:
829      MM2{Mat{x}, Mat{y}}
830
831# ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ----
832
833# C = A(n x k) · B(k x m)
834
835def Mat.clone(-r: Nat, -c: Nat, m: Mat<r, c>) -> MM2<r, c>:
836  match m:
837    case Mat{d}:
838      Mat.clone2(r, c, Array.clone(F32, d))
839
840type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
841  MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}
842
843# C = A(n x k) · Bᵀ, with B stored m x k
844
845type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type:
846  MMulNT{a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>}
847
848# C = Aᵀ · B, with A stored k x n
849
850type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type:
851  MMulTN{a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>}
852
853# ---------------------------------------------------------------------
854# gemm in tiles of 8 (0.1.8.0): for products with 2 or more rows, A's rows and B's columns are packed once into cells of
855# 8 F32 (a row of k numbers is ceil(k/8) cells, zeros after the last), so the inner loop reads 8 numbers per array
856# access instead of one. The order of the additions is that of `dotr`: product p of a dot goes to partial sum p mod 4,
857# in order, and the four sums are combined as dotr combines them after k steps (its rotation depends on k mod 4). The
858# padding adds +0 to a sum that starts at +0, which changes nothing, so the results are the same bits as the scalar
859# loop's. Products with one row (matrix · vector) keep the scalar loop: packing B would cost as much as the product.
860# ---------------------------------------------------------------------
861
862type T8 is Data:
863  T8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32}
864
865def t8_zero() -> T8:
866  T8{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}
867
868def pick0(b: Bool, +v: F32) -> F32:
869  match b:
870    case True{}:
871      v
872    case False{}:
873      0.0
874
875type PK is Type:
876  PK{a: Array<F32>, t: Array<T8>}
877
878# one cell: the 8 numbers of a at idx, idx + st, ..., those at positions q0 + q >= lim (past the end of the line) as 0
879def pk_g7(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +v5: F32, +v6: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
880  match r:
881    case Tuple{a, +v}:
882      PK{a, Array.set(T8, t, o, T8{v0, v1, v2, v3, v4, v5, v6, pick0(U32.is_lt((q0 + 7 : U32), lim), v)})}
883
884def pk_g6(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +v5: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
885  match r:
886    case Tuple{a, +v}:
887      pk_g7(v0, v1, v2, v3, v4, v5, pick0(U32.is_lt((q0 + 6 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 7 : U32) : U32)))
888
889def pk_g5(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
890  match r:
891    case Tuple{a, +v}:
892      pk_g6(v0, v1, v2, v3, v4, pick0(U32.is_lt((q0 + 5 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 6 : U32) : U32)))
893
894def pk_g4(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
895  match r:
896    case Tuple{a, +v}:
897      pk_g5(v0, v1, v2, v3, pick0(U32.is_lt((q0 + 4 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 5 : U32) : U32)))
898
899def pk_g3(+v0: F32, +v1: F32, +v2: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
900  match r:
901    case Tuple{a, +v}:
902      pk_g4(v0, v1, v2, pick0(U32.is_lt((q0 + 3 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 4 : U32) : U32)))
903
904def pk_g2(+v0: F32, +v1: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
905  match r:
906    case Tuple{a, +v}:
907      pk_g3(v0, v1, pick0(U32.is_lt((q0 + 2 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 3 : U32) : U32)))
908
909def pk_g1(+v0: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
910  match r:
911    case Tuple{a, +v}:
912      pk_g2(v0, pick0(U32.is_lt((q0 + 1 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 2 : U32) : U32)))
913
914def pk_g0(+q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
915  match r:
916    case Tuple{a, +v}:
917      pk_g1(pick0(U32.is_lt((q0 + 0 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 1 : U32) : U32)))
918
919# the cells of one line, cell o.. of t; q0 counts the line's numbers before the cell
920def pk_line(tl: Nat, +q0: U32, +lim: U32, +base: U32, +st: U32, +o: U32, r: PK) -> PK:
921  match tl r:
922    case 0n r:
923      r
924    case 1n+q PK{a, t}:
925      pk_line(q, (q0 + 8 : U32), lim, (base + (st * 8 : U32) : U32), st, (o + 1 : U32), pk_g0(q0, lim, base, st, o, t, Array.get(F32, a, base)))
926
927# lines count lines; source line i starts at i * lstep (in a), and goes to cell row o (cells o * tku .. in t)
928def pk_lines(cnt: Nat, +tk: Nat, +tku: U32, +ku: U32, +i: U32, +o: U32, +lstep: U32, +st: U32, r: PK) -> PK:
929  match cnt:
930    case 0n:
931      r
932    case 1n+p:
933      pk_lines(p, tk, tku, ku, (i + 1 : U32), (o + 1 : U32), lstep, st, pk_line(tk, 0, ku, (i * lstep : U32), st, (o * tku : U32), r))
934
935def pk_pack(+cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array<F32>) -> PK:
936  pk_lines(cnt, tk, U32.from_nat(tk), U32.from_nat(k), 0, 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())})
937
938# the dot of two lines of cells: tk cells from ia (in a) and ib (in b); ra, rb hold the arrays and the cells at ia, ib,
939# already read. Product q of a cell goes to sum q mod 4, the first four before the last four, as dotr adds them.
940def t8_dot(tk: Nat, +ia: U32, +ib: U32, s0: F32, s1: F32, s2: F32, s3: F32, ra: Array<T8> & T8, rb: Array<T8> & T8) -> TD:
941  match tk ra rb:
942    case 0n Tuple{a, x} Tuple{b, y}:
943      TD{a, b, s0, s1, s2, s3}
944    case 1n+p Tuple{a, T8{x0, x1, x2, x3, x4, x5, x6, x7}} Tuple{b, T8{y0, y1, y2, y3, y4, y5, y6, y7}}:
945      t8_dot(p, (ia + 1 : U32), (ib + 1 : U32), ((s0 + (x0 * y0 : F32) : F32) + (x4 * y4 : F32) : F32), ((s1 + (x1 * y1 : F32) : F32) + (x5 * y5 : F32) : F32), ((s2 + (x2 * y2 : F32) : F32) + (x6 * y6 : F32) : F32), ((s3 + (x3 * y3 : F32) : F32) + (x7 * y7 : F32) : F32), Array.get(T8, a, (ia + 1 : U32)), Array.get(T8, b, (ib + 1 : U32)))
946
947type TD is Type:
948  TD{a: Array<T8>, b: Array<T8>, s0: F32, s1: F32, s2: F32, s3: F32}
949
950# dotr's final sum after k products: the sums sit rotated by k mod 4 when its loop ends
951def t8_comb(+rot: Nat, +s0: F32, +s1: F32, +s2: F32, +s3: F32) -> F32:
952  match rot:
953    case 0n:
954      ((s0 + s1 : F32) + (s2 + s3 : F32) : F32)
955    case 1n:
956      ((s1 + s2 : F32) + (s3 + s0 : F32) : F32)
957    case 2n:
958      ((s2 + s3 : F32) + (s0 + s1 : F32) : F32)
959    case _:
960      ((s3 + s0 : F32) + (s1 + s2 : F32) : F32)
961
962type TR is Type:
963  TR{a: Array<T8>, b: Array<T8>, x: F32}
964
965def t8_fin(+rot: Nat, d: TD) -> TR:
966  match d:
967    case TD{a, b, +s0, +s1, +s2, +s3}:
968      TR{a, b, t8_comb(rot, s0, s1, s2, s3)}
969
970def t8_cell(+tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, a: Array<T8>, b: Array<T8>) -> TR:
971  t8_fin(rot, t8_dot(tk, (i * tku : U32), (j * tku : U32), 0.0, 0.0, 0.0, 0.0, Array.get(T8, a, (i * tku : U32)), Array.get(T8, b, (j * tku : U32))))
972
973type TG is Type:
974  TG{a: Array<T8>, b: Array<T8>, c: Array<F32>}
975
976def t8_set(+ac: Bool, +idx: U32, c: Array<F32>, r: TR) -> TG:
977  match r:
978    case TR{a, b, +x}:
979      TG{a, b, cwr(ac, idx, x, c)}
980
981def t8_cols(+ac: Bool, mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, +cidx: U32, g: TG) -> TG:
982  match mleft g:
983    case 0n g:
984      g
985    case 1n+p TG{a, b, c}:
986      t8_cols(ac, p, tk, tku, rot, i, (j + 1 : U32), (cidx + 1 : U32), t8_set(ac, cidx, c, t8_cell(tk, tku, rot, i, j, a, b)))
987
988def t8_rows(+ac: Bool, nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +mu: U32, g: TG) -> TG:
989  match nleft:
990    case 0n:
991      g
992    case 1n+p:
993      t8_rows(ac, p, m, tk, tku, rot, (i + 1 : U32), mu, t8_cols(ac, m, tk, tku, rot, i, 0, (i * mu : U32), g))
994
995# rows i0 .. i0+cnt-1 of C as a list, in order (the parallel version's leaves)
996type TL is Type:
997  TL{a: Array<T8>, b: Array<T8>, out: List<&2, F32>}
998
999# cells j.. of row i, m - j of them; r is cell j, already computed
1000def t8_lcols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, out: List<&2, F32>, r: TR) -> TL:
1001  match mleft r:
1002    case 0n TR{a, b, +x}:
1003      TL{a, b, x <> out}
1004    case 1n+p TR{a, b, +x}:
1005      t8_lcols(p, tk, tku, rot, i, (j + 1 : U32), x <> out, t8_cell(tk, tku, rot, i, (j + 1 : U32), a, b))
1006
1007def t8_lrows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, l: TL) -> TL:
1008  match nleft l:
1009    case 0n l:
1010      l
1011    case 1n+p TL{a, b, out}:
1012      t8_lrows(p, m, tk, tku, rot, (i + 1 : U32), t8_lcols(Nat.sub(m, 1n), tk, tku, rot, i, 0, out, t8_cell(tk, tku, rot, i, 0, a, b)))
1013
1014def t8_lfin(l: TL) -> List<&2, F32>:
1015  match l:
1016    case TL{a, b, out}:
1017      List.reverse(&2, F32, out)
1018
1019type PKK is Type:
1020  PKK{a: Array<F32>, b: Array<F32>, ta: Array<T8>, tb: Array<T8>}
1021
1022def t8_pk2(pa: PK, pb: PK) -> PKK:
1023  match pa pb:
1024    case PK{a, ta} PK{b, tb}:
1025      PKK{a, b, ta, tb}
1026
1027def t8_g(a: Array<F32>, b: Array<F32>, g: TG) -> G:
1028  match g:
1029    case TG{ta, tb, c}:
1030      G{a, b, c}
1031
1032def t8_out_seq(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, p: PKK) -> G:
1033  match p:
1034    case PKK{a, b, ta, tb}:
1035      t8_g(a, b, t8_rows(ac, n, m, tk, tku, rot, 0, U32.from_nat(m), TG{ta, tb, c}))
1036
1037# ---- the parallel path (0.1.9.0): the larger operand is packed into one piece per block, so a block owns its piece;
1038# the other operand is copied once per node of the fork tree (2^d - 1 copies), instead of both at every node ----
1039
1040type PT is Type:
1041  PLeaf{t: Array<T8>}
1042  PNode{l: PT, r: PT}
1043
1044type PKT is Type:
1045  PKT{a: Array<F32>, t: PT}
1046
1047def pkt_leaf(r: PK) -> PKT:
1048  match r:
1049    case PK{a, t}:
1050      PKT{a, PLeaf{t}}
1051
1052def pkt_then(r: PKT, k: Array<F32> -> PT -> PKT) -> PKT:
1053  match r:
1054    case PKT{a, t}:
1055      k(a, t)
1056
1057# lines i .. i+cnt-1 of a, packed into 2^d pieces (halves of the line count, as the fork tree splits them)
1058def pk_tree(+d: Nat, +cnt: Nat, +i: U32, +tk: Nat, +tku: U32, +ku: U32, +lstep: U32, +st: U32, a: Array<F32>) -> PKT:
1059  match d:
1060    case 0n:
1061      pkt_leaf(pk_lines(cnt, tk, tku, ku, i, 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())}))
1062    case 1n+q:
1063      pkt_then(pk_tree(q, Nat.div(cnt, 2n), i, tk, tku, ku, lstep, st, a), a2 => l => pkt_then(pk_tree(q, Nat.sub(cnt, Nat.div(cnt, 2n)), (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), tk, tku, ku, lstep, st, a2), a3 => r => PKT{a3, PNode{l, r}}))
1064
1065def t8_with2(-R: Type, p: Array<T8> & Array<T8>, k: Array<T8> -> Array<T8> -> R) -> R:
1066  match p:
1067    case Tuple{x, y}:
1068      k(x, y)
1069
1070# split by rows of C: the pieces are A's rows, B is shared; the result is C's rows in order
1071def t8_rpar(t: PT, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, b: Array<T8>) -> List<&2, F32>:
1072  match t:
1073    case PLeaf{a}:
1074      t8_lfin(t8_lrows(cnt, m, tk, tku, rot, 0, TL{a, b, Nil{}}))
1075    case PNode{l, r}:
1076      t8_with2(List<&2, F32>, Array.clone(T8, b), b1 => b2 => (
1077        x y = t8_rpar(l, Nat.div(cnt, 2n), m, tk, tku, rot, b1) t8_rpar(r, Nat.sub(cnt, Nat.div(cnt, 2n)), m, tk, tku, rot, b2)
1078        List.append(&2, F32, x, y)))
1079
1080# a block of C's columns j0 .. j0+w-1 for every row, row by row
1081type CB is Data:
1082  CB{j0: U32, w: Nat, xs: List<&2, F32>}
1083
1084# split by columns of C: the pieces are B's columns, A is shared
1085def t8_cpar(t: PT, +cnt: Nat, +j0: U32, +n: Nat, +tk: Nat, +tku: U32, +rot: Nat, a: Array<T8>) -> List<&2, CB>:
1086  match t:
1087    case PLeaf{b}:
1088      [CB{j0, cnt, t8_lfin(t8_lrows(n, cnt, tk, tku, rot, 0, TL{a, b, Nil{}}))}]
1089    case PNode{l, r}:
1090      t8_with2(List<&2, CB>, Array.clone(T8, a), a1 => a2 => (
1091        x y = t8_cpar(l, Nat.div(cnt, 2n), j0, n, tk, tku, rot, a1) t8_cpar(r, Nat.sub(cnt, Nat.div(cnt, 2n)), (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), n, tk, tku, rot, a2)
1092        List.append(&2, CB, x, y)))
1093
1094type RW is Type:
1095  RW{c: Array<F32>, xs: List<&2, F32>}
1096
1097# w numbers of xs into c from idx; what is left of xs comes back
1098def cb_row(+ac: Bool, w: Nat, +idx: U32, c: Array<F32>, xs: List<&2, F32>) -> RW:
1099  match w xs:
1100    case 1n+p Con{+h, t}:
1101      cb_row(ac, p, (idx + 1 : U32), cwr(ac, idx, h, c), t)
1102    case _ _:
1103      RW{c, xs}
1104
1105def cb_rows(+ac: Bool, n: Nat, +i: U32, +j0: U32, +w: Nat, +mu: U32, r: RW) -> Array<F32>:
1106  match n r:
1107    case 0n RW{c, xs}:
1108      c
1109    case 1n+p RW{c, xs}:
1110      cb_rows(ac, p, (i + 1 : U32), j0, w, mu, cb_row(ac, w, ((i * mu : U32) + j0 : U32), c, xs))
1111
1112def cb_write(+ac: Bool, bs: List<&2, CB>, +n: Nat, +mu: U32, c: Array<F32>) -> Array<F32>:
1113  match bs:
1114    case Nil{}:
1115      c
1116    case Con{CB{+j0, +w, xs}, t}:
1117      cb_write(ac, t, n, mu, cb_rows(ac, n, 0, j0, w, mu, RW{c, xs}))
1118
1119def t8_pr2(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, ra: PKT, rb: PK) -> G:
1120  match ra rb:
1121    case PKT{a, t} PK{b, tb}:
1122      G{a, b, write_la(ac, t8_rpar(t, n, m, tk, tku, rot, tb), 0, c)}
1123
1124def t8_pc2(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, ra: PK, rb: PKT) -> G:
1125  match ra rb:
1126    case PK{a, ta} PKT{b, t}:
1127      G{a, b, cb_write(ac, t8_cpar(t, m, 0, n, tk, tku, rot, ta), n, U32.from_nat(m), c)}
1128
1129# 2^d blocks along the larger of n (rows of C, pieces of A) and m (columns of C, pieces of B)
1130def t8_par2(+ac: Bool, byrows: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1131  match byrows:
1132    case True{}:
1133      t8_pr2(ac, n, m, tk, tku, rot, c, pk_tree(d, n, 0, tk, tku, U32.from_nat(k), arow, sa, a), pk_pack(m, tk, k, bcol, sb, b))
1134    case False{}:
1135      t8_pc2(ac, n, m, tk, tku, rot, c, pk_pack(n, tk, k, arow, sa, a), pk_tree(d, m, 0, tk, tku, U32.from_nat(k), bcol, sb, b))
1136
1137def gemm_t8b(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1138  match d:
1139    case 0n:
1140      t8_out_seq(ac, n, m, tk, tku, rot, c, t8_pk2(pk_pack(n, tk, k, arow, sa, a), pk_pack(m, tk, k, bcol, sb, b)))
1141    case 1n+q:
1142      t8_par2(ac, Nat.is_le(m, n), 1n+q, n, m, k, tk, tku, rot, arow, sa, sb, bcol, a, b, c)
1143
1144# C (n x m) = A · B (or C + A · B when ac) in cells of 8, sequential (d = 0) or with 2^d blocks of rows in parallel; A
1145# and B come back intact
1146def gemm_t8_ac(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1147  gemm_t8b(ac, d, n, m, k, Nat.div(Nat.add(k, 7n), 8n), U32.from_nat(Nat.div(Nat.add(k, 7n), 8n)), Nat.mod(k, 4n), arow, sa, sb, bcol, a, b, c)
1148
1149def gemm_t8(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1150  gemm_t8_ac(False{}, d, n, m, k, arow, sa, sb, bcol, a, b, c)
1151
1152# one row (matrix · vector): the scalar loop at every depth; two or more rows: cells of 8. Up to 0.1.11.0 one row in
1153# parallel ran gemm_par_cols, which clones both operands at every node of the fork tree: a 50,257 x 768 weight matrix at
1154# depth 4 ran out of the default heap (5.1 GB) and 8192 x 2048 took 124 ms against 10 ms sequentially. Packing the
1155# matrix into cells of 8 instead costs more than the product, which reads each weight once (64 ms against 24 at
1156# 50,257 x 768). Matrix · vector in parallel is Bands.matvec. Measured 2026-10-16 (NOTES.md).
1157def gemm_any3(+ac: Bool, one: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1158  match one:
1159    case True{}:
1160      gemm_ac(ac, n, m, k, arow, sa, sb, bcol, a, b, c)
1161    case False{}:
1162      gemm_t8_ac(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c)
1163
1164def gemm_any2(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1165  gemm_any3(ac, Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c)
1166
1167def gemm_any(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1168  gemm_any2(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c)
1169
1170def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul<n, k, m>:
1171  match g:
1172    case G{a, b, c}:
1173      MMul{Mat{a}, Mat{b}, Mat{c}}
1174
1175def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT<n, k, m>:
1176  match g:
1177    case G{a, b, c}:
1178      MMulNT{Mat{a}, Mat{b}, Mat{c}}
1179
1180def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN<n, k, m>:
1181  match g:
1182    case G{a, b, c}:
1183      MMulTN{Mat{a}, Mat{b}, Mat{c}}
1184
1185# ---- the split depth of a product (0.1.11.0): par is an upper bound, and a product too small to gain splits 0 times ----
1186#
1187# A product of n x k by k x m splits into 2^par blocks only when it does at least 2^23 multiply-adds and each cell of C
1188# sums at least 64 of them; otherwise it runs sequentially. Measured on 16 threads (2026-10-16, 18 shapes): below 2^23 the
1189# split costs more than it saves, and with a short inner dimension (a weight gradient over a few tokens) every split
1190# was 2-3.6x slower, because each block hands its rows of C back as a list that is then written into C, a cost per cell
1191# of C that k multiply-adds per cell do not pay for. The result does not depend on the depth: the same bits at every
1192# depth (tests/t8.bend, tests/acc.bend, tests/split.bend).
1193def split_pick(big: Bool, d: Nat) -> Nat:
1194  match big:
1195    case True{}:
1196      d
1197    case False{}:
1198      0n
1199
1200def split_depth(d: Nat, +n: Nat, +m: Nat, +k: Nat) -> Nat:
1201  split_pick(Bool.and(Nat.is_le(64n, k), Nat.is_le(8388608n, Nat.mul(Nat.mul(n, m), k))), d)
1202
1203def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMul<n, k, m>:
1204  mmul_fin(n, k, m, gemm_any(False{}, d, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
1205
1206# C = A · B   (n x k) · (k x m); par = log2 of the most parallel blocks (0 = sequential); split_depth picks 0 or par
1207
1208def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
1209  match a b:
1210    case Mat{x} Mat{y}:
1211      mm_a(n, k, m, split_depth(par, n, m, k), x, y)
1212
1213# C + A · B, into C (n x m); A and B come back with the result, as Mat.matmul returns them. Each cell of A · B is
1214# computed as Mat.matmul computes it and added once to C's cell: the same bits as Mat.add(c, Mat.matmul(a, b)'s c),
1215# without allocating the product or walking it a second time.
1216
1217def Mat.matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>) -> MMul<n, k, m>:
1218  match a b c:
1219    case Mat{x} Mat{y} Mat{z}:
1220      mmul_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z))
1221
1222def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulNT<n, k, m>:
1223  mmul_nt_fin(n, k, m, gemm_any(False{}, d, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
1224
1225# C = A · Bᵀ   (n x k) · (m x k)ᵀ: for weights stored with one row per output
1226
1227def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>) -> MMulNT<n, k, m>:
1228  match a b:
1229    case Mat{x} Mat{y}:
1230      mm_nt(n, k, m, split_depth(par, n, m, k), x, y)
1231
1232# C + A · Bᵀ, into C (n x m)
1233
1234def Mat.matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>) -> MMulNT<n, k, m>:
1235  match a b c:
1236    case Mat{x} Mat{y} Mat{z}:
1237      mmul_nt_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z))
1238
1239def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulTN<n, k, m>:
1240  mmul_tn_fin(n, k, m, gemm_any(False{}, d, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
1241
1242# C = Aᵀ · B   (k x n)ᵀ · (k x m): weight gradient, X^T · dY
1243
1244def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>) -> MMulTN<n, k, m>:
1245  match a b:
1246    case Mat{x} Mat{y}:
1247      mm_tn(n, k, m, split_depth(par, n, m, k), x, y)
1248
1249# C + Aᵀ · B, into C (n x m): a weight gradient added into the gradient it accumulates, dW += Xᵀ · dY
1250
1251def Mat.matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>) -> MMulTN<n, k, m>:
1252  match a b c:
1253    case Mat{x} Mat{y} Mat{z}:
1254      mmul_tn_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z))
1255
1256# ---- a precise product for long dots (0.1.14.0) ----
1257#
1258# Mat.matmul_precise is Mat.matmul with k summed in blocks of 1024 (the last one 1..1024): block b's columns of A and
1259# rows of B are copied out, the first block's product written into C, each next one's added to it as matmul_acc adds,
1260# so every cell is ((x0 + x1) + x2) .., x_b the block's dot with the usual four partial sums. A dot of 49,152 products
1261# as four running sums of 12,288 lost the small terms (SmolLM2's head: its gradient 1.8x further from float64 than
1262# PyTorch's float32, 0.03x in blocks; BendTorch NOTES, 2026-10-18). Opt-in: it copies the slices and runs one product
1263# a block, and the default products keep their speed and bits. Up to 1024 steps it is Mat.matmul itself.
1264
1265type SL is Type:
1266  SL{a: Array<F32>, s: Array<F32>}
1267
1268# d[di..] = src[si..], k numbers; r holds src[si], already read
1269def sl_row(k: Nat, +si: U32, +di: U32, d: Array<F32>, r: Array<F32> & F32) -> SL:
1270  match k r:
1271    case 0n Tuple{a, x}: SL{a, d}
1272    case 1n+p Tuple{a, +x}: sl_row(p, (si + 1 : U32), (di + 1 : U32), Array.set(F32, d, di, x), Array.get(F32, a, (si + 1 : U32)))
1273
1274# rows i.. of a slice: w numbers a row from src[i*st + o], into d[i*w]
1275def sl_rows(r: Nat, +i: U32, +w: Nat, +w32: U32, +st: U32, +o: U32, sl: SL) -> SL:
1276  match r sl:
1277    case 0n sl: sl
1278    case 1n+p SL{a, d}: sl_rows(p, (i + 1 : U32), w, w32, st, o, sl_row(w, ((i * st : U32) + o : U32), (i * w32 : U32), d, Array.get(F32, a, ((i * st : U32) + o : U32))))
1279
1280def blk_len(+bi: U32, +k: U32) -> Nat:
1281  U32.to_nat(Bool.pick(U32, U32.is_lt((k - (bi * 1024 : U32) : U32), 1024), (k - (bi * 1024 : U32) : U32), 1024))
1282
1283def mp_s4(a: Array<F32>, b: Array<F32>, g: G) -> G:
1284  match g:
1285    case G{x, y, c}: G{a, b, c}
1286
1287def mp_s3(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, a: Array<F32>, as: Array<F32>, c: Array<F32>, sb: SL) -> G:
1288  match sb:
1289    case SL{b, bs}: mp_s4(a, b, gemm_any(ac, split_depth(par, n, m, w), n, m, w, w32, 1, m32, 1, as, bs, c))
1290
1291def mp_s2(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, +o: U32, c: Array<F32>, b: Array<F32>, sa: SL) -> G:
1292  match sa:
1293    case SL{a, as}: mp_s3(ac, par, n, m, w, w32, m32, a, as, c, sl_row(Nat.mul(w, m), (o * m32 : U32), 0, Array.new(F32, cap_depth(Nat.mul(w, m)), 0.0), Array.get(F32, b, (o * m32 : U32))))
1294
1295# block bi: A's columns o..o+w (A is n x k) and B's rows o..o+w (B is k x m), multiplied into C (added when ac)
1296def mp_step(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +k32: U32, +bi: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1297  mp_s2(ac, par, n, m, blk_len(bi, k32), U32.from_nat(blk_len(bi, k32)), U32.from_nat(m), (bi * 1024 : U32), c, b, sl_rows(n, 0, blk_len(bi, k32), U32.from_nat(blk_len(bi, k32)), k32, (bi * 1024 : U32), SL{a, Array.new(F32, cap_depth(Nat.mul(n, blk_len(bi, k32))), 0.0)}))
1298
1299def mp_go(nb: Nat, +bi: U32, +par: Nat, +n: Nat, +m: Nat, +k32: U32, g: G) -> G:
1300  match nb g:
1301    case 0n g: g
1302    case 1n++q G{a, b, c}: mp_go(q, (bi + 1 : U32), par, n, m, k32, mp_step(U32.is_lt(0, bi), par, n, m, k32, bi, a, b, c))
1303
1304def mm_precise2(long: Bool, +n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
1305  match long a b:
1306    case True{} Mat{x} Mat{y}: mmul_fin(n, k, m, mp_go(U32.to_nat(U32.div((U32.from_nat(k) + 1023 : U32), 1024)), 0, par, n, m, U32.from_nat(k), G{x, y, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)}))
1307    case False{} a b: Mat.matmul(n, k, m, par, a, b)
1308
1309# C = A · B (n x k) · (k x m), k summed in blocks of 1024 (above); opt-in
1310def Mat.matmul_precise(+n: Nat, +k: Nat, +m: Nat, +par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
1311  mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)
1312
1313# ---- element-wise and training operations ----
1314
1315# adds the bias (one row, m numbers) to each row of C (n x m)
1316
1317type MBias<-n: Nat, -m: Nat> is Type:
1318  MBias{c: Mat<n, m>, b: Mat<1n, m>}
1319
1320def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array<F32>, al: AL) -> MBias<n, m>:
1321  match al:
1322    case AL{bd, bl}:
1323      MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}}
1324
1325def Mat.add_row(+n: Nat, +m: Nat, c: Mat<n, m>, b: Mat<1n, m>) -> MBias<n, m>:
1326  match c b:
1327    case Mat{cd} Mat{bd}:
1328      add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), U32.from_nat(Nat.sub(m, 1n)), Nil{}, Array.get(F32, bd, U32.from_nat(Nat.sub(m, 1n)))))
1329
1330# pair of matrices of the same shape (the result of an operation that reads two operands)
1331
1332type MPair<-r: Nat, -c: Nat> is Type:
1333  MPair{a: Mat<r, c>, b: Mat<r, c>}
1334
1335def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair<r, c>:
1336  match p:
1337    case P2{a, b}:
1338      MPair{Mat{a}, Mat{b}}
1339
1340def relu_fin(-r: Nat, -c: Nat, d: Array<F32>) -> Mat<r, c>:
1341  Mat{d}
1342
1343def Mat.relu(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<r, c>:
1344  match a:
1345    case Mat{d}:
1346      relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0)))
1347
1348# (dy * (h > 0), h): the gradient that passes through the relu
1349
1350def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat<r, c>, h: Mat<r, c>) -> MPair<r, c>:
1351  match dy h:
1352    case Mat{d} Mat{x}:
1353      mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0)))
1354
1355# (w - lr * dw, dw): one SGD step
1356
1357def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat<r, c>, dw: Mat<r, c>) -> MPair<r, c>:
1358  match w dw:
1359    case Mat{x} Mat{y}:
1360      mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0)))
1361
1362# (a + b, b): element-wise sum (a - (-1)*b is exactly a + b)
1363
1364def Mat.add(+r: Nat, +c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> MPair<r, c>:
1365  match a b:
1366    case Mat{x} Mat{y}:
1367      mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, (0.0 - 1.0 : F32), Array.get(F32, x, 0), Array.get(F32, y, 0)))
1368
1369# sum of each column (n x m -> 1 x m), returning the matrix
1370
1371type MColSum<-n: Nat, -m: Nat> is Type:
1372  MColSum{a: Mat<n, m>, s: Mat<1n, m>}
1373
1374def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum<n, m>:
1375  match g:
1376    case G{ones, a, s}:
1377      MColSum{Mat{a}, Mat{s}}
1378
1379def Mat.col_sums(+n: Nat, +m: Nat, a: Mat<n, m>) -> MColSum<n, m>:
1380  match a:
1381    case Mat{d}:
1382      colsum_fin(n, m, gemm(1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, Array.new(F32, cap_depth(m), 0.0)))
1383
1384# s + the sum of each column, into s (1 x m): a bias gradient added into its accumulator, with the bits of
1385# Mat.add(s, Mat.col_sums(a)'s s)
1386
1387def Mat.col_sums_acc(+n: Nat, +m: Nat, a: Mat<n, m>, s: Mat<1n, m>) -> MColSum<n, m>:
1388  match a s:
1389    case Mat{d} Mat{z}:
1390      colsum_fin(n, m, gemm_ac(True{}, 1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, z))
1391
1392# mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n.
1393# labels has one entry per row (n); the gradient replaces the logits.
1394
1395type MCE<-n: Nat, -c: Nat> is Type:
1396  MCE{loss: F32, grad: Mat<n, c>}
1397
1398def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE<n, c>:
1399  match r:
1400    case L2{loss, d}:
1401      MCE{(loss / nf : F32), Mat{d}}
1402
1403def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MCE<n, c>:
1404  match z:
1405    case Mat{d}:
1406      ce_fin(n, c, F32.from_nat(n), cea_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), F32.from_nat(n), 0, L2{0.0, d}))
1407
1408# argmax hits per row against labels
1409
1410type MHits<-n: Nat, -c: Nat> is Type:
1411  MHits{z: Mat<n, c>, hits: Nat}
1412
1413def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits<n, c>:
1414  match r:
1415    case CR{d, k}:
1416      MHits{Mat{d}, k}
1417
1418def Mat.count_correct(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MHits<n, c>:
1419  match z:
1420    case Mat{d}:
1421      hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n}))
1422
1423# a row as a list (returning the matrix) and the other way around
1424
1425type MRow<-r: Nat, -c: Nat> is Type:
1426  MRow{m: Mat<r, c>, row: List<&2, F32>}
1427
1428def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow<r, c>:
1429  match al:
1430    case AL{d, l}:
1431      MRow{Mat{d}, l}
1432
1433def Mat.read_row(+r: Nat, +c: Nat, m: Mat<r, c>, +i: U32) -> MRow<r, c>:
1434  match m:
1435    case Mat{d}:
1436      row_fin(r, c, read_l(Nat.sub(c, 1n), ((i * U32.from_nat(c) : U32) + U32.from_nat(Nat.sub(c, 1n)) : U32), Nil{}, Array.get(F32, d, ((i * U32.from_nat(c) : U32) + U32.from_nat(Nat.sub(c, 1n)) : U32))))
1437
1438def Mat.write_row(-r: Nat, +c: Nat, m: Mat<r, c>, +i: U32, xs: List<&2, F32>) -> Mat<r, c>:
1439  match m:
1440    case Mat{d}:
1441      Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)}
1442
1443# ---- variants that check the number of labels ----
1444
1445# does the list of labels have exactly n entries (one per row)?
1446def labels_len_is(xs: List<&2, Nat>, n: Nat) -> Bool:
1447  match xs n:
1448    case Nil{} 0n:
1449      True{}
1450    case Con{h, t} 1n+p:
1451      labels_len_is(t, p)
1452    case _ _:
1453      False{}
1454
1455def ce_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> Maybe<&1, MCE<n, c>>:
1456  match ok:
1457    case True{}:
1458      Some{Mat.softmax_ce(n, c, z, labels)}
1459    case False{}:
1460      None{}
1461
1462# like Mat.softmax_ce, but returns None unless labels has exactly n entries
1463def Mat.softmax_ce_checked(+n: Nat, +c: Nat, z: Mat<n, c>, +labels: List<&2, Nat>) -> Maybe<&1, MCE<n, c>>:
1464  ce_checked(labels_len_is(labels, n), n, c, z, labels)
1465
1466def hits_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> Maybe<&1, MHits<n, c>>:
1467  match ok:
1468    case True{}:
1469      Some{Mat.count_correct(n, c, z, labels)}
1470    case False{}:
1471      None{}
1472
1473# like Mat.count_correct, but returns None unless labels has exactly n entries
1474def Mat.count_correct_checked(+n: Nat, +c: Nat, z: Mat<n, c>, +labels: List<&2, Nat>) -> Maybe<&1, MHits<n, c>>:
1475  hits_checked(labels_len_is(labels, n), n, c, z, labels)
1476
1477# ---------------------------------------------------------------------
1478# Bands<r, c>: a matrix cut into bands of rows, for matrix · vector in parallel
1479# without copying the weights.
1480#
1481# Mat.matmul_nt with n = 1 can only run in parallel by cloning the whole weight
1482# matrix for every task, and in matrix · vector each weight is read once, so the
1483# copies cost more than the products (NOTES.md, exp. 9). Bands keeps each band in
1484# its own Array: a task takes its band and nothing is copied except the input vector.
1485#
1486# A node with r rows has children with half(r) and r - half(r) rows. The split is
1487# part of the type, so the row count of every band follows from r, and the law
1488# half_cover says that the two halves always cover the r rows exactly.
1489# ---------------------------------------------------------------------
1490
1491# n / 2 rounded down. A tail loop with an accumulator: a non-tail recursion here would put a
1492# continuation on every node of the band tree and keep the bands from running in parallel.
1493def half.go(n: Nat, acc: Nat) -> Nat:
1494  match n:
1495    case 0n:
1496      acc
1497    case 1n+p:
1498      match p:
1499        case 0n:
1500          acc
1501        case 1n+q:
1502          half.go(q, 1n+acc)
1503
1504def half(n: Nat) -> Nat:
1505  half.go(n, 0n)
1506
1507# a <= a
1508def leb_refl(+a: Nat) -> {True{} == leb(a, a) : Bool}:
1509  match a:
1510    case 0n:
1511      {==}
1512    case 1n+p:
1513      leb_refl(p)
1514
1515# 1 + (a + b) == a + (1 + b)
1516def add_succ_r(+a: Nat, -b: Nat) -> {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}:
1517  match a:
1518    case 0n:
1519      {==}
1520    case 1n+p:
1521      %add_succ_r(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat}
1522      {==}
1523
1524# a + 0 == a
1525def add_zero_r(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}:
1526  match a:
1527    case 0n:
1528      {==}
1529    case 1n+p:
1530      %add_zero_r(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}
1531      {==}
1532
1533# half.go(n, acc) <= n + acc
1534def half_go_le(+n: Nat, +acc: Nat) -> {True{} == leb(half.go(n, acc), Nat.add(n, acc)) : Bool}:
1535  match n:
1536    case 0n:
1537      leb_refl(acc)
1538    case 1n+p:
1539      match p:
1540        case 0n:
1541          leb_succ(acc, acc, leb_refl(acc))
1542        case 1n+q:
1543          %Equal.sym(Nat, 1n+Nat.add(q, acc), Nat.add(q, 1n+acc), add_succ_r(q, acc)) : {True{} == leb(half.go(q, 1n+acc), 1n+_) : Bool}
1544          leb_succ(half.go(q, 1n+acc), Nat.add(q, 1n+acc), half_go_le(q, 1n+acc))
1545
1546# half(n) <= n
1547def half_le(+n: Nat) -> {True{} == leb(half(n), n) : Bool}:
1548  %add_zero_r(n) : {True{} == leb(half(n), _) : Bool}
1549  half_go_le(n, 0n)
1550
1551# a <= r  ->  r == a + (r - a)
1552def add_sub(+a: Nat, +r: Nat, e: {True{} == leb(a, r) : Bool}) -> {r == Nat.add(a, Nat.sub(r, a)) : Nat}:
1553  match a r:
1554    case 0n 0n:
1555      {==}
1556    case 0n 1n+q:
1557      {==}
1558    case 1n+p 0n:
1559      %e : BD(_, Unit, {0n == Nat.add(1n+p, Nat.sub(0n, 1n+p)) : Nat})
1560      Unit{}
1561    case 1n+p 1n+q:
1562      %add_sub(p, q, e) : {1n+q == 1n+_ : Nat}
1563      {==}
1564
1565# LAW (half_cover): the two bands of a node, half(r) rows and r - half(r) rows,
1566# add up to exactly r rows: no row is lost and none is counted twice.
1567law half_cover:
1568  for +n: Nat
1569  {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat}
1570
1571def half_cover(n):
1572  add_sub(half(n), n, half_le(n))
1573
1574type Bands<-r: Nat, -c: Nat> is Type:
1575  BLeaf{m: Mat<r, c>}
1576  BNode{x: Bands<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>}
1577
1578# r x c of zeros in 2^d bands (d = 0: a single band, the same as a Mat)
1579def Bands.zeros(+d: Nat, +r: Nat, +c: Nat) -> Bands<r, c>:
1580  match d:
1581    case 0n:
1582      BLeaf{Mat.zeros(r, c)}
1583    case 1n+q:
1584      # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- bands are built once at load not-expected-because: forking them measured as no change (NOTES 2026-10-16)
1585      BNode{Bands.zeros(q, half(r), c), Bands.zeros(q, Nat.sub(r, half(r)), c)}
1586
1587# the row counts of the two bands follow the row count of the node: if rr == r at run
1588# time, the same holds for both halves (the bands carry r only in their type)
1589def eq_half(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {half(a) == half(b) : Nat}:
1590  %e : {half(a) == half(_) : Nat}
1591  {==}
1592
1593def eq_rest(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {Nat.sub(a, half(a)) == Nat.sub(b, half(b)) : Nat}:
1594  %e : {Nat.sub(a, half(a)) == Nat.sub(_, half(_)) : Nat}
1595  {==}
1596
1597# does a block of n numbers starting at flat index i end before the second band of a
1598# node with r rows? does it start in the second band?
1599def bfill_l(+r: Nat, +c: Nat, +i: Nat, +n: Nat) -> Bool:
1600  Nat.is_le(Nat.add(i, n), Nat.mul(half(r), c))
1601
1602def bfill_r(+r: Nat, +c: Nat, +i: Nat) -> Bool:
1603  Nat.is_le(Nat.mul(half(r), c), i)
1604
1605# writes the first n numbers of xs at index i (and nothing past them)
1606def fill_n(n: Nat, xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
1607  match n xs:
1608    case 1n+p Con{h, t}:
1609      fill_n(p, t, (i + 1 : U32), Array.set(F32, a, i, h))
1610    case _ _:
1611      a
1612
1613def fill_band(-r: Nat, -c: Nat, n: Nat, xs: List<&2, F32>, +i: U32, m: Mat<r, c>) -> Mat<r, c>:
1614  match m:
1615    case Mat{d}:
1616      Mat{fill_n(n, xs, i, d)}
1617
1618# writes the first n numbers of xs at flat index i; l and rt say which bands they touch
1619# (bfill_l, bfill_r); rr is the row count r at run time. A block that crosses into the second
1620# band is not copied: the first band writes only its part, the second gets the rest by List.drop.
1621def bfill(+c: Nat, -r: Nat, b: Bands<r, c>, l: Bool, rt: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands<r, c>:
1622  match b l rt:
1623    case BLeaf{m} _ _:
1624      BLeaf{fill_band(r, c, n, xs, U32.from_nat(i), m)}
1625    case BNode{x, y} True{} _:
1626      BNode{bfill(c, half(r), x, bfill_l(half(rr), c, i, n), bfill_r(half(rr), c, i), half(rr), eq_half(rr, r, e), i, n, xs), y}
1627    case BNode{x, y} False{} True{}:
1628      BNode{x, bfill(c, Nat.sub(r, half(r)), y, bfill_l(Nat.sub(rr, half(rr)), c, Nat.sub(i, Nat.mul(half(rr), c)), n), bfill_r(Nat.sub(rr, half(rr)), c, Nat.sub(i, Nat.mul(half(rr), c))), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Nat.sub(i, Nat.mul(half(rr), c)), n, xs)}
1629    case BNode{x, y} False{} False{}:
1630      # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- bands are built once at load not-expected-because: forking them measured as no change (NOTES 2026-10-16)
1631      BNode{bfill(c, half(r), x, bfill_l(half(rr), c, i, Nat.sub(Nat.mul(half(rr), c), i)), bfill_r(half(rr), c, i), half(rr), eq_half(rr, r, e), i, Nat.sub(Nat.mul(half(rr), c), i), xs), bfill(c, Nat.sub(r, half(r)), y, bfill_l(Nat.sub(rr, half(rr)), c, 0n, Nat.sub(n, Nat.sub(Nat.mul(half(rr), c), i))), bfill_r(Nat.sub(rr, half(rr)), c, 0n), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), 0n, Nat.sub(n, Nat.sub(Nat.mul(half(rr), c), i)), List.drop(&2, F32, xs, Nat.sub(Nat.mul(half(rr), c), i)))}
1632
1633def bfill_top(ok: Bool, +r: Nat, +c: Nat, b: Bands<r, c>, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands<r, c>:
1634  match ok:
1635    case True{}:
1636      bfill(c, r, b, bfill_l(r, c, i, n), bfill_r(r, c, i), r, {==}, i, n, xs)
1637    case False{}:
1638      bfill(c, r, b, bfill_l(r, c, i, Nat.sub(Nat.mul(r, c), i)), bfill_r(r, c, i), r, {==}, i, Nat.sub(Nat.mul(r, c), i), xs)
1639
1640# writes the n numbers of xs into b starting at flat index i (row-major, as Mat.fill_at);
1641# numbers past r*c are dropped. For loading a large matrix block by block.
1642def Bands.fill_at(+r: Nat, +c: Nat, +i: Nat, +n: Nat, +xs: List<&2, F32>, b: Bands<r, c>) -> Bands<r, c>:
1643  bfill_top(Nat.is_le(Nat.add(i, n), Nat.mul(r, c)), r, c, b, i, n, xs)
1644
1645# ---- y = B · x for B with r rows of c numbers (a weight matrix stored out x in) ----
1646
1647# B (intact) and y = B · x, typed like Mat.matmul_nt: x is 1 x c, y is 1 x r
1648type BV<-r: Nat, -c: Nat> is Type:
1649  BV{b: Bands<r, c>, y: Mat<1n, r>}
1650
1651# the bands and y as a list in row order (the result of Bands.matvec_l)
1652type BL<-r: Nat, -c: Nat> is Type:
1653  BL{b: Bands<r, c>, y: List<&2, F32>}
1654
1655def bv_g2(-r: Nat, -c: Nat, w: Array<F32>, al: AL) -> BL<r, c>:
1656  match al:
1657    case AL{d, l}:
1658      BL{BLeaf{Mat{w}}, l}
1659
1660# the band's 1 + m rows came out in an Array (one number per row); read them back as a list
1661def bv_g(-r: Nat, -c: Nat, +m: Nat, g: G) -> BL<r, c>:
1662  match g:
1663    case G{a, b, y}:
1664      bv_g2(r, c, b, read_l(m, U32.from_nat(m), Nil{}, Array.get(F32, y, U32.from_nat(m))))
1665
1666# the rows of one band, one dot product each, with the gemm of Mat.matmul_nt (n = 1), so the
1667# results are identical. Writing the results into an Array and reading them back is faster than
1668# building the list row by row (NOTES.md, exp. 14).
1669def bv_gemm(-r: Nat, +c: Nat, rows: Nat, +rr: Nat, x: Array<F32>, w: Array<F32>) -> BL<r, c>:
1670  match rows:
1671    case 0n:
1672      BL{BLeaf{Mat{w}}, Nil{}}
1673    case 1n+p:
1674      bv_g(r, c, p, gemm(1n, rr, c, U32.from_nat(c), 1, 1, U32.from_nat(c), x, w, Array.new(F32, cap_depth(rr), 0.0)))
1675
1676def bv_mat(-r: Nat, +c: Nat, +rr: Nat, x: Array<F32>, m: Mat<r, c>) -> BL<r, c>:
1677  match m:
1678    case Mat{w}:
1679      bv_gemm(r, c, rr, rr, x, w)
1680
1681def bv_join(-r: Nat, -c: Nat, p: BL<half(r), c>, q: BL<Nat.sub(r, half(r)), c>) -> BL<r, c>:
1682  match p q:
1683    case BL{bx, yx} BL{by, yy}:
1684      BL{BNode{bx, by}, List.append(&2, F32, yx, yy)}
1685
1686# the two bands of a node run in parallel, each with its own copy of x (c numbers);
1687# xp is a pair of copies of x (a leaf uses one). xp comes after the erased e: up to Bend 2.0.36 an erased
1688# parameter at the end of the list turned the parallel let into two sequential calls (NOTES.md, exp. 11; fixed in
1689# 2.0.37 by bendlang/bend#1377). The order is kept, and harmless.
1690def bmv(+c: Nat, -r: Nat, b: Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}, xp: Array<F32> & Array<F32>) -> BL<r, c>:
1691  match b xp:
1692    case BLeaf{m} Tuple{x, x2}:
1693      bv_mat(r, c, rr, x, m)
1694    case BNode{bx, by} Tuple{x1, x2}:
1695      p q = bmv(c, half(r), bx, half(rr), eq_half(rr, r, e), Array.clone(F32, x1)) bmv(c, Nat.sub(r, half(r)), by, Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Array.clone(F32, x2))
1696      bv_join(r, c, p, q)
1697
1698def bv_fin(+r: Nat, -c: Nat, bl: BL<r, c>) -> BV<r, c>:
1699  match bl:
1700    case BL{b, y}:
1701      BV{b, Mat.from_list(1n, r, y)}
1702
1703# y = B · x for B (r x c, one row per output) and x (1 x c); returns B intact and y (1 x r).
1704# The same product as Mat.matmul_nt(1n, c, r, ..), and the same numbers, bit for bit.
1705def Bands.matvec(+r: Nat, +c: Nat, b: Bands<r, c>, x: Mat<1n, c>) -> BV<r, c>:
1706  match x:
1707    case Mat{xa}:
1708      bv_fin(r, c, bmv(c, r, b, r, {==}, Array.clone(F32, xa)))
1709
1710# the same product with x and y as lists (x: c numbers, y: r numbers in row order), without the
1711# Mat conversions of Bands.matvec: for callers that already hold the vector as a list
1712def Bands.matvec_l(+r: Nat, +c: Nat, b: Bands<r, c>, +x: List<&2, F32>) -> BL<r, c>:
1713  bmv(c, r, b, r, {==}, Array.clone(F32, fill_list(x, 0, Array.new(F32, cap_depth(c), 0.0))))
1714
1715# ---- reading rows and the whole matrix ----
1716
1717type BRow<-r: Nat, -c: Nat> is Type:
1718  BRow{b: Bands<r, c>, row: List<&2, F32>}
1719
1720def brow_fin(-r: Nat, -c: Nat, al: AL) -> BRow<r, c>:
1721  match al:
1722    case AL{d, l}:
1723      BRow{BLeaf{Mat{d}}, l}
1724
1725# n numbers starting at flat index i of one band (a band may have no rows)
1726def brow_leaf2(-r: Nat, -c: Nat, +n: Nat, +i: U32, d: Array<F32>) -> BRow<r, c>:
1727  match n:
1728    case 0n:
1729      BRow{BLeaf{Mat{d}}, Nil{}}
1730    case 1n+p:
1731      brow_fin(r, c, read_l(p, ((i + U32.from_nat(n) : U32) - 1 : U32), Nil{}, Array.get(F32, d, ((i + U32.from_nat(n) : U32) - 1 : U32))))
1732
1733def brow_leaf(-r: Nat, -c: Nat, m: Mat<r, c>, +i: U32, n: Nat) -> BRow<r, c>:
1734  match m:
1735    case Mat{d}:
1736      brow_leaf2(r, c, n, i, d)
1737
1738def brow_l(-r: Nat, -c: Nat, p: BRow<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1739  match p:
1740    case BRow{bx, row}:
1741      BRow{BNode{bx, y}, row}
1742
1743def brow_r(-r: Nat, -c: Nat, x: Bands<half(r), c>, q: BRow<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1744  match q:
1745    case BRow{by, row}:
1746      BRow{BNode{x, by}, row}
1747
1748# row i; left says whether i is in the first band of the node (i < half(rr))
1749def brow(+c: Nat, -r: Nat, b: Bands<r, c>, left: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat) -> BRow<r, c>:
1750  match b left:
1751    case BLeaf{m} _:
1752      brow_leaf(r, c, m, U32.from_nat(Nat.mul(i, c)), c)
1753    case BNode{x, y} True{}:
1754      brow_l(r, c, brow(c, half(r), x, Nat.is_lt(i, half(half(rr))), half(rr), eq_half(rr, r, e), i), y)
1755    case BNode{x, y} False{}:
1756      brow_r(r, c, x, brow(c, Nat.sub(r, half(r)), y, Nat.is_lt(Nat.sub(i, half(rr)), half(Nat.sub(rr, half(rr)))), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Nat.sub(i, half(rr))))
1757
1758# row i (c numbers), for example an embedding lookup
1759def Bands.read_row(+r: Nat, +c: Nat, b: Bands<r, c>, +i: Nat) -> BRow<r, c>:
1760  brow(c, r, b, Nat.is_lt(i, half(r)), r, {==}, i)
1761
1762def bl_join(-r: Nat, -c: Nat, p: BRow<half(r), c>, q: BRow<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1763  match p q:
1764    case BRow{bx, xs} BRow{by, ys}:
1765      BRow{BNode{bx, by}, List.append(&2, F32, xs, ys)}
1766
1767def blist(+c: Nat, -r: Nat, b: Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}) -> BRow<r, c>:
1768  match b:
1769    case BLeaf{m}:
1770      brow_leaf(r, c, m, 0, Nat.mul(rr, c))
1771    case BNode{x, y}:
1772      # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- bands are built once at load not-expected-because: forking them measured as no change (NOTES 2026-10-16)
1773      bl_join(r, c, blist(c, half(r), x, half(rr), eq_half(rr, r, e)), blist(c, Nat.sub(r, half(r)), y, Nat.sub(rr, half(rr)), eq_rest(rr, r, e)))
1774
1775# all r*c numbers in row order (in BRow.row)
1776def Bands.to_list(+r: Nat, +c: Nat, b: Bands<r, c>) -> BRow<r, c>:
1777  blist(c, r, b, r, {==})
1778
1779# ---- the list path returns exactly r numbers ----
1780#
1781# The numbers do not matter here, only how many there are: each lemma below counts the list a
1782# function builds, following the function's own recursion.
1783
1784def bl_len(-r: Nat, -c: Nat, bl: BL<r, c>) -> Nat:
1785  match bl:
1786    case BL{b, y}:
1787      List.length(&2, F32, y)
1788
1789# s == a + (1 + b)  ->  s == 1 + (a + b)
1790def succ_out(-s: Nat, +a: Nat, +b: Nat, e: {s == Nat.add(a, 1n+b) : Nat}) -> {s == 1n+Nat.add(a, b) : Nat}:
1791  %Equal.sym(Nat, 1n+Nat.add(a, b), Nat.add(a, 1n+b), add_succ_r(a, b)) : {s == _ : Nat}
1792  e
1793
1794# s == 1 + (p + (1 + l))  ->  s == 2 + (p + l)
1795def len_step(-s: Nat, +p: Nat, +l: Nat, e: {s == 1n+Nat.add(p, 1n+l) : Nat}) -> {s == 2n+Nat.add(p, l) : Nat}:
1796  %Equal.sym(Nat, 1n+Nat.add(p, l), Nat.add(p, 1n+l), add_succ_r(p, l)) : {s == 1n+_ : Nat}
1797  e
1798
1799# s == n + 0  ->  s == n
1800def zero_out(-s: Nat, +n: Nat, e: {s == Nat.add(n, 0n) : Nat}) -> {s == n : Nat}:
1801  %add_zero_r(n) : {s == _ : Nat}
1802  e
1803
1804# appending adds the lengths
1805def app_len(+xs: List<&2, F32>, +ys: List<&2, F32>) -> {List.length(&2, F32, List.append(&2, F32, xs, ys)) == Nat.add(List.length(&2, F32, xs), List.length(&2, F32, ys)) : Nat}:
1806  match xs:
1807    case Nil{}:
1808      {==}
1809    case Con{h, t}:
1810      %app_len(t, ys) : {1n+List.length(&2, F32, List.append(&2, F32, t, ys)) == 1n+_ : Nat}
1811      {==}
1812
1813# the same count for the list of a whole Mat: to_list2 wraps read_l's answer, and ML.len counts it
1814def read_ml_len(-r: Nat, -c: Nat, +mleft: Nat, +idx: U32, +acc: List<&2, F32>, rc: Array<F32> & F32) -> {ML.len(r, c, Mat.to_list2(r, c, read_l(mleft, idx, acc, rc))) == Nat.add(1n+mleft, List.length(&2, F32, acc)) : Nat}:
1815  match mleft rc:
1816    case 0n Tuple{d, +v}:
1817      {==}
1818    case 1n+p Tuple{d, +v}:
1819      len_step(ML.len(r, c, Mat.to_list2(r, c, read_l(p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32))))), p, List.length(&2, F32, acc), read_ml_len(r, c, p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32))))
1820
1821# a leaf's list: read_l reads mleft + 1 numbers onto acc. By induction on mleft: with 0 left it
1822# reads one number onto acc (1 + the length of acc, by the definitions); with 1 + p left it reads
1823# one and recurses, and len_step moves the 1 out of the sum.
1824law leaf_len:
1825  for -r: Nat
1826  for -c: Nat
1827  for w: Array<F32>
1828  for +mleft: Nat
1829  for +idx: U32
1830  for +acc: List<&2, F32>
1831  for rc: Array<F32> & F32
1832  {bl_len(r, c, bv_g2(r, c, w, read_l(mleft, idx, acc, rc))) == Nat.add(1n+mleft, List.length(&2, F32, acc)) : Nat}
1833
1834def leaf_len(r, c, w, mleft, idx, acc, rc):
1835  match mleft rc:
1836    case 0n Tuple{d, +v}:
1837      {==}
1838    case 1n+p Tuple{d, +v}:
1839      len_step(bl_len(r, c, bv_g2(r, c, w, read_l(p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32))))), p, List.length(&2, F32, acc), leaf_len(r, c, w, p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32))))
1840
1841# whatever the gemm wrote, reading 1 + p rows back gives 1 + p numbers (G has one constructor, so
1842# matching g opens it, and leaf_len counts the list)
1843def g_len(-r: Nat, -c: Nat, +p: Nat, g: G) -> {bl_len(r, c, bv_g(r, c, p, g)) == 1n+p : Nat}:
1844  match g:
1845    case G{a, b, y}:
1846      zero_out(bl_len(r, c, bv_g2(r, c, b, read_l(p, U32.from_nat(p), Nil{}, Array.get(F32, y, U32.from_nat(p))))), 1n+p, leaf_len(r, c, b, p, U32.from_nat(p), Nil{}, Array.get(F32, y, U32.from_nat(p))))
1847
1848# LAW (band_len): the product of one band with `rows` rows gives exactly `rows` numbers.
1849# 0 rows give the empty list; 1 + p rows go through the gemm, and g_len counts what is read back.
1850law band_len:
1851  for -r: Nat
1852  for +c: Nat
1853  for +rows: Nat
1854  for x: Array<F32>
1855  for w: Array<F32>
1856  {bl_len(r, c, bv_gemm(r, c, rows, rows, x, w)) == rows : Nat}
1857
1858def band_len(r, c, rows, x, w):
1859  match rows:
1860    case 0n:
1861      {==}
1862    case 1n+p:
1863      g_len(r, c, p, gemm(1n, 1n+p, c, U32.from_nat(c), 1, 1, U32.from_nat(c), x, w, Array.new(F32, cap_depth(1n+p), 0.0)))
1864
1865# joining two halves of h1 and h2 numbers gives h1 + h2 numbers. (The statement for the whole tree,
1866# "Bands.matvec_l gives r numbers", would apply this to the results of the two recursive calls and
1867# to the induction hypotheses about the same calls; that uses each band's Array twice, which Bend's
1868# affine rules refuse. It stays open, covered by tests: NOTES.md, v3.1.)
1869def join_len(-r: Nat, -c: Nat, -h1: Nat, -h2: Nat, p: BL<half(r), c>, q: BL<Nat.sub(r, half(r)), c>, hp: {bl_len(half(r), c, p) == h1 : Nat}, hq: {bl_len(Nat.sub(r, half(r)), c, q) == h2 : Nat}) -> {bl_len(r, c, bv_join(r, c, p, q)) == Nat.add(h1, h2) : Nat}:
1870  match p q:
1871    case BL{bx, +yx} BL{by, +yy}:
1872      %hp : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(_, h2) : Nat}
1873      %hq : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(List.length(&2, F32, yx), _) : Nat}
1874      app_len(yx, yy)
1875
1876# Mat.to_list_count reads exactly n numbers, and n = 0 reads none.
1877def to_list_count_len(-r: Nat, -c: Nat, d: Array<F32>, +n: Nat) -> {ML.len(r, c, Mat.to_list_count(r, c, d, n)) == n : Nat}:
1878  match n:
1879    case 0n:
1880      {==}
1881    case 1n+p:
1882      zero_out(ML.len(r, c, Mat.to_list_count(r, c, d, 1n+p)), 1n+p, read_ml_len(r, c, p, (U32.from_nat(1n+p) - 1 : U32), Nil{}, Array.get(F32, d, (U32.from_nat(1n+p) - 1 : U32))))
1883
1884# LAW (to_list_len): Mat.to_list reads exactly r*c numbers, no more and no fewer.
1885law to_list_len:
1886  for +r: Nat
1887  for +c: Nat
1888  for m: Mat<r, c>
1889  {ML.len(r, c, Mat.to_list(r, c, m)) == Nat.mul(r, c) : Nat}
1890
1891def to_list_len(r, c, m):
1892  match m:
1893    case Mat{d}:
1894      to_list_count_len(r, c, d, Nat.mul(r, c))
1895
1896# ---------------------------------------------------------------------
1897# Base's Array: three laws about Array.set and Array.get
1898# ---------------------------------------------------------------------
1899#
1900# - size_set: Array.set keeps Array.size;
1901# - get_after_set: what Array.set writes at i, Array.get reads back at i;
1902# - get_other_set: a write at i leaves alone the cell that a read at k reaches, when the two descents part.
1903# They are about Base's Array, for any element type, and live here until Base states them
1904# (bendlang/bend#1934; nuxyel/bend-ml#26). The helpers are named ASet.*.
1905#
1906# The first two hold for every i (the two descents wrap i the same way) and for every tree Base can build,
1907# balanced or not, since a write replaces one leaf and leaves the tree's shape alone. "The size" is the
1908# second component of the pair Array.size returns, read with Base's Pair.snd (ASet.sz); the value is
1909# Pair.snd of the pair Array.get returns.
1910#
1911# The third law's hypothesis says that the descent Array.set makes for i and the descent Array.get makes
1912# for k take different branches at some node before either reaches a leaf (ASet.Apart). It is not "i and
1913# k differ after the mask", because that is false on an unbalanced tree: Array.size follows the leftmost
1914# path, so in ANode{ANode{ALeaf, ALeaf}, ALeaf} (size 4) the indices 2 and 3 differ but both reach the
1915# last leaf, and a write at 2 changes what a read at 3 returns (tests/array_laws.bend prints it: 9.0). On a
1916# perfect tree of 2^d leaves, which Array.new and literals build, the descents part exactly when the masked
1917# indices differ; that equivalence is not proved here, since it needs lemmas about U32 bits. For a concrete
1918# tree ASet.Apart computes to Unit, so Unit{} proves it (tests/array_laws.bend).
1919#
1920# Why the proofs are shaped this way: an Array is affine, and so is it inside a proof, so a proof may use a
1921# given array once. The obvious lemmas each need the same subtree again (one to say Array.size of it is
1922# (it, its size), one for the induction step, one to open the pair the descent returns), so they cannot
1923# be combined. Instead, one pass over the tree (ASet.facts) proves all six facts the step needs at once;
1924# each fact is an equation, and equations can be reused (+). A second pass (ASet.copy) hands back the same
1925# array and its size as a plain value, which the top-level proof needs to pick the first branch. The third
1926# law uses one pass too (ASet.both): it proves the six facts and four more about the read along k, and on
1927# the half the write does not enter it runs a smaller pass (ASet.side).
1928#
1929# The trailing Unit in ASet.Facts, ASet.Copy, ASet.Side and ASet.Reads works around bendlang/bend#1077: a +
1930# binder on the last field of a nested pair fails to check, so no fact sits in the last field.
1931
1932
1933# ASet.sz(T, a): the size of a as a plain U32, the second component of the pair Array.size returns
1934def ASet.sz(-T: Data, a: Array<T>) -> U32:
1935  Pair.snd(Array<T>, U32, Array.size(T, a))
1936
1937# the three results the descents build: the swapped pair, the written array, and the read pair
1938def ASet.G(-T: Data, a: Array<T>, +n: U32, +j: U32, +v: T, +z: Bool) -> Array<T> & T:
1939  Array.swap.go(T, a, n, j, v, z)
1940
1941def ASet.NA(-T: Data, a: Array<T>, +n: U32, +j: U32, +v: T, +z: Bool) -> Array<T>:
1942  Array.set.fin(T, ASet.G(T, a, n, j, v, z))
1943
1944def ASet.K(-T: Data, a: Array<T>, +n: U32, +j: U32, +v: T, +z: Bool) -> Array<T> & T:
1945  Array.get.go(T, ASet.NA(T, a, n, j, v, z), n, j, z)
1946
1947# ASet.Facts: what one pass over the tree a proves, for any n, j, v, z (n need not be a's real size):
1948# 1. Array.size(a) is the pair (a, ASet.sz(a));
1949# 2. the pair the write's descent returns is (the written array, the old value);
1950# 3. Array.size of the written array is the pair (written array, its sz);
1951# 4. the written array has a's sz;
1952# 5. reading the written array along the same path gives v;
1953# 6. the pair the read returns is (its first component, its second).
1954def ASet.Facts(-T: Data, -a: Array<T>, -n: U32, -j: U32, -v: T, -z: Bool) -> Type:
1955  {(a, ASet.sz(T, a)) == Array.size(T, a) : Array<T> & U32} & {(ASet.NA(T, a, n, j, v, z), Pair.snd(Array<T>, T, ASet.G(T, a, n, j, v, z))) == ASet.G(T, a, n, j, v, z) : Array<T> & T} & {(ASet.NA(T, a, n, j, v, z), ASet.sz(T, ASet.NA(T, a, n, j, v, z))) == Array.size(T, ASet.NA(T, a, n, j, v, z)) : Array<T> & U32} & {ASet.sz(T, ASet.NA(T, a, n, j, v, z)) == ASet.sz(T, a) : U32} & {Pair.snd(Array<T>, T, ASet.K(T, a, n, j, v, z)) == v : T} & {(Pair.fst(Array<T>, T, ASet.K(T, a, n, j, v, z)), Pair.snd(Array<T>, T, ASet.K(T, a, n, j, v, z))) == ASet.K(T, a, n, j, v, z) : Array<T> & T} & Unit
1956
1957# Array.size hands back the array it measured, next to ASet.sz(T, a)
1958def ASet.size_pair(-T: Data, a: Array<T>) -> {(a, ASet.sz(T, a)) == Array.size(T, a) : Array<T> & U32}:
1959  match a:
1960    case ALeaf{x}:
1961      {==}
1962    case ANode{xs, ys}:
1963      %ASet.size_pair(T, xs) : {(ANode{xs, ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}
1964      {==}
1965
1966def ASet.node_t(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +j: U32, +v: T, f: ASet.Facts(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))) -> ASet.Facts(T, ANode{xs, ys}, n, j, v, True{}):
1967  match f:
1968    case Tuple{+c1, r1}:
1969      match r1:
1970        case Tuple{+c2, r2}:
1971          match r2:
1972            case Tuple{+c3, r3}:
1973              match r3:
1974                case Tuple{+c4, r4}:
1975                  match r4:
1976                    case Tuple{+c5, r5}:
1977                      match r5:
1978                        case Tuple{+c6, u}:
1979                          ({(%c1 : {(ANode{xs, ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}; {==}) : {(ANode{xs, ys}, ASet.sz(T, ANode{xs, ys})) == Array.size(T, ANode{xs, ys}) : Array<T> & U32}}, {(%c2 : {(Array.set.fin(T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), Pair.snd(Array<T>, T, ASet.G(T, ANode{xs, ys}, n, j, v, True{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, True{}) : Array<T> & T}}, {(%c2 : {(Array.set.fin(T, Array.swap.lo(T, ys, _)), ASet.sz(T, Array.set.fin(T, Array.swap.lo(T, ys, _)))) == Array.size(T, Array.set.fin(T, Array.swap.lo(T, ys, _))) : Array<T> & U32}; %c3 : {(ANode{Array.set.fin(T, ASet.G(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))), ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}))) == Array.size(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{})) : Array<T> & U32}}, {(%c2 : {ASet.sz(T, Array.set.fin(T, Array.swap.lo(T, ys, _))) == ASet.sz(T, ANode{xs, ys}) : U32}; %c3 : {Pair.snd(Array<T>, U32, Array.size.node(T, ys, _)) == ASet.sz(T, ANode{xs, ys}) : U32}; %c1 : {U32.shl(ASet.sz(T, Array.set.fin(T, ASet.G(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))))) == Pair.snd(Array<T>, U32, Array.size.node(T, ys, _)) : U32}; %c4 : {U32.shl(ASet.sz(T, Array.set.fin(T, ASet.G(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))))) == U32.shl(_) : U32}; {==}) : {ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{})) == ASet.sz(T, ANode{xs, ys}) : U32}}, {(%c2 : {Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, j, True{})) == v : T}; %c6 : {Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _)) == v : T}; c5) : {Pair.snd(Array<T>, T, ASet.K(T, ANode{xs, ys}, n, j, v, True{})) == v : T}}, {(%c2 : {(Pair.fst(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, j, True{})), Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, j, True{}))) == Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, j, True{}) : Array<T> & T}; %c6 : {(Pair.fst(Array<T>, T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, ASet.K(T, ANode{xs, ys}, n, j, v, True{})), Pair.snd(Array<T>, T, ASet.K(T, ANode{xs, ys}, n, j, v, True{}))) == ASet.K(T, ANode{xs, ys}, n, j, v, True{}) : Array<T> & T}}, Unit{})
1980
1981def ASet.node_f(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +j: U32, +v: T, +s: {(xs, ASet.sz(T, xs)) == Array.size(T, xs) : Array<T> & U32}, f: ASet.Facts(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))) -> ASet.Facts(T, ANode{xs, ys}, n, j, v, False{}):
1982  match f:
1983    case Tuple{+c1, r1}:
1984      match r1:
1985        case Tuple{+c2, r2}:
1986          match r2:
1987            case Tuple{+c3, r3}:
1988              match r3:
1989                case Tuple{+c4, r4}:
1990                  match r4:
1991                    case Tuple{+c5, r5}:
1992                      match r5:
1993                        case Tuple{+c6, u}:
1994                          ({(%s : {(ANode{xs, ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}; {==}) : {(ANode{xs, ys}, ASet.sz(T, ANode{xs, ys})) == Array.size(T, ANode{xs, ys}) : Array<T> & U32}}, {(%c2 : {(Array.set.fin(T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), Pair.snd(Array<T>, T, ASet.G(T, ANode{xs, ys}, n, j, v, False{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, False{}) : Array<T> & T}}, {(%c2 : {(Array.set.fin(T, Array.swap.hi(T, xs, _)), ASet.sz(T, Array.set.fin(T, Array.swap.hi(T, xs, _)))) == Array.size(T, Array.set.fin(T, Array.swap.hi(T, xs, _))) : Array<T> & U32}; %s : {(ANode{xs, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n)))))}, Pair.snd(Array<T>, U32, Array.size.node(T, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))), _))) == Array.size.node(T, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))), _) : Array<T> & U32}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}))) == Array.size(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{})) : Array<T> & U32}}, {(%c2 : {ASet.sz(T, Array.set.fin(T, Array.swap.hi(T, xs, _))) == ASet.sz(T, ANode{xs, ys}) : U32}; %s : {Pair.snd(Array<T>, U32, Array.size.node(T, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))), _)) == Pair.snd(Array<T>, U32, Array.size.node(T, ys, _)) : U32}; {==}) : {ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{})) == ASet.sz(T, ANode{xs, ys}) : U32}}, {(%c2 : {Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, j, False{})) == v : T}; %c6 : {Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _)) == v : T}; c5) : {Pair.snd(Array<T>, T, ASet.K(T, ANode{xs, ys}, n, j, v, False{})) == v : T}}, {(%c2 : {(Pair.fst(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, j, False{})), Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, j, False{}))) == Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, j, False{}) : Array<T> & T}; %c6 : {(Pair.fst(Array<T>, T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, ASet.K(T, ANode{xs, ys}, n, j, v, False{})), Pair.snd(Array<T>, T, ASet.K(T, ANode{xs, ys}, n, j, v, False{}))) == ASet.K(T, ANode{xs, ys}, n, j, v, False{}) : Array<T> & T}}, Unit{})
1995
1996def ASet.facts(-T: Data, a: Array<T>, +n: U32, +j: U32, +v: T, +z: Bool) -> ASet.Facts(T, a, n, j, v, z):
1997  match a z:
1998    case ALeaf{x} _:
1999      ({==}, {==}, {==}, {==}, {==}, {==}, Unit{})
2000    case ANode{xs, ys} True{}:
2001      ASet.node_t(T, xs, ys, n, j, v, ASet.facts(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n)))))
2002    case ANode{xs, ys} False{}:
2003      ASet.node_f(T, xs, ys, n, j, v, ASet.size_pair(T, xs), ASet.facts(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n)))))
2004
2005# ASet.copy(T, a): one pass over a hands back the same array b, its size n as a value, and the facts
2006# that tie them to a: b is a, n is ASet.sz(T, b), and Array.size(b) is the pair (b, ASet.sz(T, b))
2007def ASet.Copy(-T: Data, -a: Array<T>) -> Type:
2008  &b: Array<T> -> &n: U32 -> {b == a : Array<T>} & {n == ASet.sz(T, b) : U32} & {(b, ASet.sz(T, b)) == Array.size(T, b) : Array<T> & U32} & Unit
2009
2010def ASet.copy_node(-T: Data, -xs: Array<T>, ys: Array<T>, r: ASet.Copy(T, xs)) -> ASet.Copy(T, ANode{xs, ys}):
2011  match r:
2012    case Tuple{b, r1}:
2013      match r1:
2014        case Tuple{+m, r2}:
2015          match r2:
2016            case Tuple{+p, r3}:
2017              match r3:
2018                case Tuple{+q, r4}:
2019                  match r4:
2020                    case Tuple{+s, u}:
2021                      (ANode{b, ys}, U32.shl(m),
2022                        {(%p : {ANode{b, ys} == ANode{_, ys} : Array<T>}; {==}) : {ANode{b, ys} == ANode{xs, ys} : Array<T>}},
2023                        {(%s : {U32.shl(m) == Pair.snd(Array<T>, U32, Array.size.node(T, ys, _)) : U32}; %q : {U32.shl(m) == U32.shl(_) : U32}; {==}) : {U32.shl(m) == ASet.sz(T, ANode{b, ys}) : U32}},
2024                        {(%s : {(ANode{b, ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}; {==}) : {(ANode{b, ys}, ASet.sz(T, ANode{b, ys})) == Array.size(T, ANode{b, ys}) : Array<T> & U32}},
2025                        Unit{})
2026
2027def ASet.copy(-T: Data, a: Array<T>) -> ASet.Copy(T, a):
2028  match a:
2029    case ALeaf{x}:
2030      (ALeaf{x}, 1, {==}, {==}, {==}, Unit{})
2031    case ANode{xs, ys}:
2032      ASet.copy_node(T, xs, ys, ASet.copy(T, xs))
2033
2034# the masked index and the first branch, as Array.swap.at and Array.get.at compute them from a size n
2035def ASet.J(+i: U32, +n: U32) -> U32:
2036  U32.and(i, U32.sub(n, 1))
2037
2038def ASet.Z(+i: U32, +n: U32) -> Bool:
2039  U32.is_lt(ASet.J(i, n), U32.shr(n))
2040
2041def ASet.close_size(-T: Data, -b: Array<T>, +n: U32, +i: U32, +v: T, +q: {n == ASet.sz(T, b) : U32}, f: ASet.Facts(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))) -> {ASet.sz(T, Array.set(T, b, i, v)) == ASet.sz(T, b) : U32}:
2042  match f:
2043    case Tuple{+c1, r1}:
2044      match r1:
2045        case Tuple{+c2, r2}:
2046          match r2:
2047            case Tuple{+c3, r3}:
2048              match r3:
2049                case Tuple{+c4, r4}:
2050                  %c1 : {ASet.sz(T, Array.set.fin(T, Array.swap.at(T, i, v, _))) == ASet.sz(T, b) : U32}
2051                  %q : {ASet.sz(T, Array.set.fin(T, Array.swap.go(T, b, _, ASet.J(i, _), v, ASet.Z(i, _)))) == _ : U32}
2052                  Equal.trans(U32, ASet.sz(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))), ASet.sz(T, b), n, c4, Equal.sym(U32, n, ASet.sz(T, b), q))
2053
2054def ASet.close_get(-T: Data, -b: Array<T>, +n: U32, +i: U32, +v: T, +q: {n == ASet.sz(T, b) : U32}, f: ASet.Facts(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))) -> {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, b, i, v), i)) == v : T}:
2055  match f:
2056    case Tuple{+c1, r1}:
2057      match r1:
2058        case Tuple{+c2, r2}:
2059          match r2:
2060            case Tuple{+c3, r3}:
2061              match r3:
2062                case Tuple{+c4, r4}:
2063                  match r4:
2064                    case Tuple{+c5, r5}:
2065                      %c1 : {Pair.snd(Array<T>, T, Array.get.at(T, i, Array.size(T, Array.set.fin(T, Array.swap.at(T, i, v, _))))) == v : T}
2066                      %q : {Pair.snd(Array<T>, T, Array.get.at(T, i, Array.size(T, Array.set.fin(T, Array.swap.go(T, b, _, ASet.J(i, _), v, ASet.Z(i, _)))))) == v : T}
2067                      %c3 : {Pair.snd(Array<T>, T, Array.get.at(T, i, _)) == v : T}
2068                      %Equal.trans(U32, n, ASet.sz(T, b), ASet.sz(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))), q, Equal.sym(U32, ASet.sz(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))), ASet.sz(T, b), c4)) : {Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n)), _, ASet.J(i, _), ASet.Z(i, _))) == v : T}
2069                      c5
2070
2071def ASet.top_size(-T: Data, -a: Array<T>, +i: U32, +v: T, r: ASet.Copy(T, a)) -> {ASet.sz(T, Array.set(T, a, i, v)) == ASet.sz(T, a) : U32}:
2072  match r:
2073    case Tuple{b, r1}:
2074      match r1:
2075        case Tuple{+n, r2}:
2076          match r2:
2077            case Tuple{+p, r3}:
2078              match r3:
2079                case Tuple{+q, r4}:
2080                  %p : {ASet.sz(T, Array.set(T, _, i, v)) == ASet.sz(T, _) : U32}
2081                  ASet.close_size(T, b, n, i, v, q, ASet.facts(T, b, n, ASet.J(i, n), v, ASet.Z(i, n)))
2082
2083def ASet.top_get(-T: Data, -a: Array<T>, +i: U32, +v: T, r: ASet.Copy(T, a)) -> {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, a, i, v), i)) == v : T}:
2084  match r:
2085    case Tuple{b, r1}:
2086      match r1:
2087        case Tuple{+n, r2}:
2088          match r2:
2089            case Tuple{+p, r3}:
2090              match r3:
2091                case Tuple{+q, r4}:
2092                  %p : {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, _, i, v), i)) == v : T}
2093                  ASet.close_get(T, b, n, i, v, q, ASet.facts(T, b, n, ASet.J(i, n), v, ASet.Z(i, n)))
2094
2095# LAW: writing v at i keeps the array's size (any i: the index wraps around)
2096law size_set:
2097  for -T: Data
2098  for a: Array<T>
2099  for +i: U32
2100  for +v: T
2101  {ASet.sz(T, Array.set(T, a, i, v)) == ASet.sz(T, a) : U32}
2102
2103def size_set(T, a, i, v):
2104  ASet.top_size(T, a, i, v, ASet.copy(T, a))
2105
2106# LAW: what Array.set writes at i, Array.get reads at i (any i: both wrap it the same way)
2107law get_after_set:
2108  for -T: Data
2109  for a: Array<T>
2110  for +i: U32
2111  for +v: T
2112  {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, a, i, v), i)) == v : T}
2113
2114def get_after_set(T, a, i, v):
2115  ASet.top_get(T, a, i, v, ASet.copy(T, a))
2116
2117# ASet.Apart(T, a, n, j, z, k, w): the write's descent (index j, first branch z) and the read's descent (k, w),
2118# both started from the size n, take different branches at some node of a before either reaches a leaf.
2119# It is Empty when both reach the same leaf, and Unit at the node where they split.
2120def ASet.Apart(-T: Data, a: Array<T>, +n: U32, +j: U32, +z: Bool, +k: U32, +w: Bool) -> Type:
2121  match a z w:
2122    case ALeaf{x} _ _:
2123      Empty
2124    case ANode{xs, ys} True{} True{}:
2125      ASet.Apart(T, xs, U32.shr(n), j, U32.is_lt(j, U32.shr(U32.shr(n))), k, U32.is_lt(k, U32.shr(U32.shr(n))))
2126    case ANode{xs, ys} False{} False{}:
2127      ASet.Apart(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n))))
2128    case ANode{xs, ys} True{} False{}:
2129      Unit
2130    case ANode{xs, ys} False{} True{}:
2131      Unit
2132
2133# ASet.Side: what a pass over a subtree the write does not enter proves: Array.size(a) is (a, ASet.sz(a)), and the
2134# pair a read along (n, k, w) returns is (its first component, its second)
2135def ASet.Side(-T: Data, -a: Array<T>, -n: U32, -k: U32, -w: Bool) -> Type:
2136  {(a, ASet.sz(T, a)) == Array.size(T, a) : Array<T> & U32} & {(Pair.fst(Array<T>, T, Array.get.go(T, a, n, k, w)), Pair.snd(Array<T>, T, Array.get.go(T, a, n, k, w))) == Array.get.go(T, a, n, k, w) : Array<T> & T} & Unit
2137
2138def ASet.side_node(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +k: U32, +w: Bool, sx: ASet.Side(T, xs, U32.shr(n), k, U32.is_lt(k, U32.shr(U32.shr(n)))), sy: ASet.Side(T, ys, U32.shr(n), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n))))) -> ASet.Side(T, ANode{xs, ys}, n, k, w):
2139  match w:
2140    case True{}:
2141      match sx:
2142        case Tuple{+s1, r1}:
2143          match r1:
2144            case Tuple{+g1, u1}:
2145              match sy:
2146                case Tuple{+s2, r2}:
2147                  match r2:
2148                    case Tuple{+g2, u2}:
2149                      ({(%s1 : {(ANode{xs, ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}; {==}) : {(ANode{xs, ys}, ASet.sz(T, ANode{xs, ys})) == Array.size(T, ANode{xs, ys}) : Array<T> & U32}}, {(%g1 : {(Pair.fst(Array<T>, T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{}))) == Array.get.go(T, ANode{xs, ys}, n, k, True{}) : Array<T> & T}}, Unit{})
2150    case False{}:
2151      match sx:
2152        case Tuple{+s1, r1}:
2153          match r1:
2154            case Tuple{+g1, u1}:
2155              match sy:
2156                case Tuple{+s2, r2}:
2157                  match r2:
2158                    case Tuple{+g2, u2}:
2159                      ({(%s1 : {(ANode{xs, ys}, Pair.snd(Array<T>, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array<T> & U32}; {==}) : {(ANode{xs, ys}, ASet.sz(T, ANode{xs, ys})) == Array.size(T, ANode{xs, ys}) : Array<T> & U32}}, {(%g2 : {(Pair.fst(Array<T>, T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{}))) == Array.get.go(T, ANode{xs, ys}, n, k, False{}) : Array<T> & T}}, Unit{})
2160
2161def ASet.side(-T: Data, a: Array<T>, +n: U32, +k: U32, +w: Bool) -> ASet.Side(T, a, n, k, w):
2162  match a:
2163    case ALeaf{x}:
2164      ({==}, {==}, Unit{})
2165    case ANode{xs, ys}:
2166      # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- a proof by one pass over the tree not-expected-because: proofs are checked, never run
2167      ASet.side_node(T, xs, ys, n, k, w, ASet.side(T, xs, U32.shr(n), k, U32.is_lt(k, U32.shr(U32.shr(n)))), ASet.side(T, ys, U32.shr(n), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n)))))
2168
2169# ASet.Reads: what one pass over the tree a proves about a write along (n, j, v, z) and a read along (n, k, w):
2170# 1. the pair the write's descent returns is (the written array, the old value);
2171# 2. the read of a returns a pair (first component, second);
2172# 3. so does the read of the written array;
2173# 4. if the descents split, the read finds the same value in the written array as in a.
2174def ASet.Reads(-T: Data, -a: Array<T>, -n: U32, -j: U32, -v: T, -z: Bool, -k: U32, -w: Bool) -> Type:
2175  {(ASet.NA(T, a, n, j, v, z), Pair.snd(Array<T>, T, ASet.G(T, a, n, j, v, z))) == ASet.G(T, a, n, j, v, z) : Array<T> & T} & {(Pair.fst(Array<T>, T, Array.get.go(T, a, n, k, w)), Pair.snd(Array<T>, T, Array.get.go(T, a, n, k, w))) == Array.get.go(T, a, n, k, w) : Array<T> & T} & {(Pair.fst(Array<T>, T, Array.get.go(T, ASet.NA(T, a, n, j, v, z), n, k, w)), Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, a, n, j, v, z), n, k, w))) == Array.get.go(T, ASet.NA(T, a, n, j, v, z), n, k, w) : Array<T> & T} & (ASet.Apart(T, a, n, j, z, k, w) -> {Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, a, n, j, v, z), n, k, w)) == Pair.snd(Array<T>, T, Array.get.go(T, a, n, k, w)) : T}) & Unit
2176
2177# the written low half ASet.WL and the written high half ASet.WH of a node of size n
2178def ASet.WL(-T: Data, xs: Array<T>, +n: U32, +j: U32, +v: T) -> Array<T>:
2179  ASet.NA(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))
2180
2181def ASet.WH(-T: Data, ys: Array<T>, +n: U32, +j: U32, +v: T) -> Array<T>:
2182  ASet.NA(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))
2183
2184# the read's descent into the low half and into the high half of a node of size n
2185def ASet.RL(-T: Data, xs: Array<T>, +n: U32, +k: U32) -> Array<T> & T:
2186  Array.get.go(T, xs, U32.shr(n), k, U32.is_lt(k, U32.shr(U32.shr(n))))
2187
2188def ASet.RH(-T: Data, ys: Array<T>, +n: U32, +k: U32) -> Array<T> & T:
2189  Array.get.go(T, ys, U32.shr(n), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n))))
2190
2191def ASet.no_leaf(-T: Data, -x: T, -n: U32, -j: U32, -v: T, -z: Bool, -k: U32, -w: Bool, h: ASet.Apart(T, ALeaf{x}, n, j, z, k, w)) -> {Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, ALeaf{x}, n, j, v, z), n, k, w)) == Pair.snd(Array<T>, T, Array.get.go(T, ALeaf{x}, n, k, w)) : T}:
2192  match h:
2193
2194# the write goes into the low half xs
2195def ASet.reads_t(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +j: U32, +v: T, +k: U32, +w: Bool, +g: {(Pair.fst(Array<T>, T, ASet.RH(T, ys, n, k)), Pair.snd(Array<T>, T, ASet.RH(T, ys, n, k))) == ASet.RH(T, ys, n, k) : Array<T> & T}, r: ASet.Reads(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))), k, U32.is_lt(k, U32.shr(U32.shr(n))))) -> ASet.Reads(T, ANode{xs, ys}, n, j, v, True{}, k, w):
2196  match w:
2197    case True{}:
2198      match r:
2199        case Tuple{+d, r1}:
2200          match r1:
2201            case Tuple{+e1, r2}:
2202              match r2:
2203                case Tuple{+e2, r3}:
2204                  match r3:
2205                    case Tuple{e3, u}:
2206                      ({(%d : {(Array.set.fin(T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), Pair.snd(Array<T>, T, ASet.G(T, ANode{xs, ys}, n, j, v, True{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, True{}) : Array<T> & T}},
2207                        {(%e1 : {(Pair.fst(Array<T>, T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{}))) == Array.get.go(T, ANode{xs, ys}, n, k, True{}) : Array<T> & T}},
2208                        {(%d : {(Pair.fst(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, True{}))) == Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, True{}) : Array<T> & T}; %e2 : {(Pair.fst(Array<T>, T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), n, k, True{}))) == Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), n, k, True{}) : Array<T> & T}},
2209                        (h => (%d : {Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, True{})) == Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{})) : T}; %e2 : {Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _)) == Pair.snd(Array<T>, T, Array.swap.lo(T, ys, ASet.RL(T, xs, n, k))) : T}; %e1 : {Pair.snd(Array<T>, T, ASet.RL(T, ASet.WL(T, xs, n, j, v), n, k)) == Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _)) : T}; e3(h))),
2210                        Unit{})
2211    case False{}:
2212      match r:
2213        case Tuple{+d, r1}:
2214          match r1:
2215            case Tuple{+e1, r2}:
2216              match r2:
2217                case Tuple{+e2, r3}:
2218                  match r3:
2219                    case Tuple{e3, u}:
2220                      ({(%d : {(Array.set.fin(T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), Pair.snd(Array<T>, T, ASet.G(T, ANode{xs, ys}, n, j, v, True{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, True{}) : Array<T> & T}},
2221                        {(%g : {(Pair.fst(Array<T>, T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{}))) == Array.get.go(T, ANode{xs, ys}, n, k, False{}) : Array<T> & T}},
2222                        {(%d : {(Pair.fst(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, False{}))) == Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, False{}) : Array<T> & T}; %g : {(Pair.fst(Array<T>, T, Array.swap.hi(T, ASet.WL(T, xs, n, j, v), _)), Pair.snd(Array<T>, T, Array.swap.hi(T, ASet.WL(T, xs, n, j, v), _))) == Array.swap.hi(T, ASet.WL(T, xs, n, j, v), _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), n, k, False{}))) == Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), n, k, False{}) : Array<T> & T}},
2223                        (h => (%d : {Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, k, False{})) == Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{})) : T}; %g : {Pair.snd(Array<T>, T, Array.swap.hi(T, ASet.WL(T, xs, n, j, v), _)) == Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _)) : T}; {==})),
2224                        Unit{})
2225
2226# the write goes into the high half ys
2227def ASet.reads_f(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +j: U32, +v: T, +k: U32, +w: Bool, +g: {(Pair.fst(Array<T>, T, ASet.RL(T, xs, n, k)), Pair.snd(Array<T>, T, ASet.RL(T, xs, n, k))) == ASet.RL(T, xs, n, k) : Array<T> & T}, r: ASet.Reads(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n))))) -> ASet.Reads(T, ANode{xs, ys}, n, j, v, False{}, k, w):
2228  match w:
2229    case True{}:
2230      match r:
2231        case Tuple{+d, r1}:
2232          match r1:
2233            case Tuple{+e1, r2}:
2234              match r2:
2235                case Tuple{+e2, r3}:
2236                  match r3:
2237                    case Tuple{e3, u}:
2238                      ({(%d : {(Array.set.fin(T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), Pair.snd(Array<T>, T, ASet.G(T, ANode{xs, ys}, n, j, v, False{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, False{}) : Array<T> & T}},
2239                        {(%g : {(Pair.fst(Array<T>, T, Array.swap.lo(T, ys, _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{}))) == Array.get.go(T, ANode{xs, ys}, n, k, True{}) : Array<T> & T}},
2240                        {(%d : {(Pair.fst(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, True{}))) == Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, True{}) : Array<T> & T}; %g : {(Pair.fst(Array<T>, T, Array.swap.lo(T, ASet.WH(T, ys, n, j, v), _)), Pair.snd(Array<T>, T, Array.swap.lo(T, ASet.WH(T, ys, n, j, v), _))) == Array.swap.lo(T, ASet.WH(T, ys, n, j, v), _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), n, k, True{})), Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), n, k, True{}))) == Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), n, k, True{}) : Array<T> & T}},
2241                        (h => (%d : {Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, True{})) == Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, True{})) : T}; %g : {Pair.snd(Array<T>, T, Array.swap.lo(T, ASet.WH(T, ys, n, j, v), _)) == Pair.snd(Array<T>, T, Array.swap.lo(T, ys, _)) : T}; {==})),
2242                        Unit{})
2243    case False{}:
2244      match r:
2245        case Tuple{+d, r1}:
2246          match r1:
2247            case Tuple{+e1, r2}:
2248              match r2:
2249                case Tuple{+e2, r3}:
2250                  match r3:
2251                    case Tuple{e3, u}:
2252                      ({(%d : {(Array.set.fin(T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), Pair.snd(Array<T>, T, ASet.G(T, ANode{xs, ys}, n, j, v, False{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, False{}) : Array<T> & T}},
2253                        {(%e1 : {(Pair.fst(Array<T>, T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{}))) == Array.get.go(T, ANode{xs, ys}, n, k, False{}) : Array<T> & T}},
2254                        {(%d : {(Pair.fst(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, False{}))) == Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, False{}) : Array<T> & T}; %e2 : {(Pair.fst(Array<T>, T, Array.swap.hi(T, xs, _)), Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array<T> & T}; {==}) : {(Pair.fst(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), n, k, False{})), Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), n, k, False{}))) == Array.get.go(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), n, k, False{}) : Array<T> & T}},
2255                        (h => (%d : {Pair.snd(Array<T>, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, k, False{})) == Pair.snd(Array<T>, T, Array.get.go(T, ANode{xs, ys}, n, k, False{})) : T}; %e2 : {Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _)) == Pair.snd(Array<T>, T, Array.swap.hi(T, xs, ASet.RH(T, ys, n, k))) : T}; %e1 : {Pair.snd(Array<T>, T, ASet.RH(T, ASet.WH(T, ys, n, j, v), n, k)) == Pair.snd(Array<T>, T, Array.swap.hi(T, xs, _)) : T}; e3(h))),
2256                        Unit{})
2257
2258# ASet.Both: the six facts of the first two laws and the four of the third, from one pass
2259def ASet.Both(-T: Data, -a: Array<T>, -n: U32, -j: U32, -v: T, -z: Bool, -k: U32, -w: Bool) -> Type:
2260  ASet.Facts(T, a, n, j, v, z) & ASet.Reads(T, a, n, j, v, z, k, w)
2261
2262def ASet.both_t(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +j: U32, +v: T, +k: U32, +w: Bool, sy: ASet.Side(T, ys, U32.shr(n), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n)))), f: ASet.Both(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))), k, U32.is_lt(k, U32.shr(U32.shr(n))))) -> ASet.Both(T, ANode{xs, ys}, n, j, v, True{}, k, w):
2263  match sy:
2264    case Tuple{+s, r1}:
2265      match r1:
2266        case Tuple{+g, u}:
2267          match f:
2268            case Tuple{o, r}:
2269              (ASet.node_t(T, xs, ys, n, j, v, o), ASet.reads_t(T, xs, ys, n, j, v, k, w, g, r))
2270
2271def ASet.both_f(-T: Data, -xs: Array<T>, -ys: Array<T>, +n: U32, +j: U32, +v: T, +k: U32, +w: Bool, sx: ASet.Side(T, xs, U32.shr(n), k, U32.is_lt(k, U32.shr(U32.shr(n)))), f: ASet.Both(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n))))) -> ASet.Both(T, ANode{xs, ys}, n, j, v, False{}, k, w):
2272  match sx:
2273    case Tuple{+s, r1}:
2274      match r1:
2275        case Tuple{+g, u}:
2276          match f:
2277            case Tuple{o, r}:
2278              (ASet.node_f(T, xs, ys, n, j, v, s, o), ASet.reads_f(T, xs, ys, n, j, v, k, w, g, r))
2279
2280def ASet.both(-T: Data, a: Array<T>, +n: U32, +j: U32, +v: T, +z: Bool, +k: U32, +w: Bool) -> ASet.Both(T, a, n, j, v, z, k, w):
2281  match a z:
2282    case ALeaf{x} _:
2283      (({==}, {==}, {==}, {==}, {==}, {==}, Unit{}), ({==}, {==}, {==}, (h => ASet.no_leaf(T, x, n, j, v, z, k, w, h)), Unit{}))
2284    case ANode{xs, ys} True{}:
2285      ASet.both_t(T, xs, ys, n, j, v, k, w, ASet.side(T, ys, U32.shr(n), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n)))), ASet.both(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))), k, U32.is_lt(k, U32.shr(U32.shr(n)))))
2286    case ANode{xs, ys} False{}:
2287      ASet.both_f(T, xs, ys, n, j, v, k, w, ASet.side(T, xs, U32.shr(n), k, U32.is_lt(k, U32.shr(U32.shr(n)))), ASet.both(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))), U32.sub(k, U32.shr(n)), U32.is_lt(U32.sub(k, U32.shr(n)), U32.shr(U32.shr(n)))))
2288
2289
2290def ASet.close_other(-T: Data, -b: Array<T>, +n: U32, +i: U32, +k: U32, +v: T, +q: {n == ASet.sz(T, b) : U32}, f: ASet.Both(T, b, n, ASet.J(i, n), v, ASet.Z(i, n), ASet.J(k, n), ASet.Z(k, n))) -> ASet.Apart(T, b, ASet.sz(T, b), ASet.J(i, ASet.sz(T, b)), ASet.Z(i, ASet.sz(T, b)), ASet.J(k, ASet.sz(T, b)), ASet.Z(k, ASet.sz(T, b))) -> {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, b, i, v), k)) == Pair.snd(Array<T>, T, Array.get(T, b, k)) : T}:
2291  match f:
2292    case Tuple{o, r}:
2293      match o:
2294        case Tuple{+c1, o1}:
2295          match o1:
2296            case Tuple{+c2, o2}:
2297              match o2:
2298                case Tuple{+c3, o3}:
2299                  match o3:
2300                    case Tuple{+c4, o4}:
2301                      match r:
2302                        case Tuple{+d, r1}:
2303                          match r1:
2304                            case Tuple{+e1, r2}:
2305                              match r2:
2306                                case Tuple{+e2, r3}:
2307                                  match r3:
2308                                    case Tuple{e3, u}:
2309                                      %c1 : (ASet.Apart(T, b, ASet.sz(T, b), ASet.J(i, ASet.sz(T, b)), ASet.Z(i, ASet.sz(T, b)), ASet.J(k, ASet.sz(T, b)), ASet.Z(k, ASet.sz(T, b))) -> {Pair.snd(Array<T>, T, Array.get.at(T, k, Array.size(T, Array.set.fin(T, Array.swap.at(T, i, v, _))))) == Pair.snd(Array<T>, T, Array.get.at(T, k, _)) : T})
2310                                      %q : (ASet.Apart(T, b, _, ASet.J(i, _), ASet.Z(i, _), ASet.J(k, _), ASet.Z(k, _)) -> {Pair.snd(Array<T>, T, Array.get.at(T, k, Array.size(T, Array.set.fin(T, Array.swap.go(T, b, _, ASet.J(i, _), v, ASet.Z(i, _)))))) == Pair.snd(Array<T>, T, Array.get.go(T, b, _, ASet.J(k, _), ASet.Z(k, _))) : T})
2311                                      %c3 : (ASet.Apart(T, b, n, ASet.J(i, n), ASet.Z(i, n), ASet.J(k, n), ASet.Z(k, n)) -> {Pair.snd(Array<T>, T, Array.get.at(T, k, _)) == Pair.snd(Array<T>, T, Array.get.go(T, b, n, ASet.J(k, n), ASet.Z(k, n))) : T})
2312                                      %Equal.trans(U32, n, ASet.sz(T, b), ASet.sz(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))), q, Equal.sym(U32, ASet.sz(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n))), ASet.sz(T, b), c4)) : (ASet.Apart(T, b, n, ASet.J(i, n), ASet.Z(i, n), ASet.J(k, n), ASet.Z(k, n)) -> {Pair.snd(Array<T>, T, Array.get.go(T, ASet.NA(T, b, n, ASet.J(i, n), v, ASet.Z(i, n)), _, ASet.J(k, _), ASet.Z(k, _))) == Pair.snd(Array<T>, T, Array.get.go(T, b, n, ASet.J(k, n), ASet.Z(k, n))) : T})
2313                                      e3
2314
2315
2316def ASet.top_other(-T: Data, -a: Array<T>, +i: U32, +k: U32, +v: T, r: ASet.Copy(T, a)) -> ASet.Apart(T, a, ASet.sz(T, a), ASet.J(i, ASet.sz(T, a)), ASet.Z(i, ASet.sz(T, a)), ASet.J(k, ASet.sz(T, a)), ASet.Z(k, ASet.sz(T, a))) -> {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, a, i, v), k)) == Pair.snd(Array<T>, T, Array.get(T, a, k)) : T}:
2317  match r:
2318    case Tuple{b, r1}:
2319      match r1:
2320        case Tuple{+n, r2}:
2321          match r2:
2322            case Tuple{+p, r3}:
2323              match r3:
2324                case Tuple{+q, r4}:
2325                  %p : (ASet.Apart(T, _, ASet.sz(T, _), ASet.J(i, ASet.sz(T, _)), ASet.Z(i, ASet.sz(T, _)), ASet.J(k, ASet.sz(T, _)), ASet.Z(k, ASet.sz(T, _))) -> {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, _, i, v), k)) == Pair.snd(Array<T>, T, Array.get(T, _, k)) : T})
2326                  ASet.close_other(T, b, n, i, k, v, q, ASet.both(T, b, n, ASet.J(i, n), v, ASet.Z(i, n), ASet.J(k, n), ASet.Z(k, n)))
2327
2328# LAW: a write at i leaves alone the cell a read at k reaches, when the two descents part before a leaf
2329# (any i and k: both wrap them; h is what makes them different cells, see the top of the file)
2330law get_other_set:
2331  for -T: Data
2332  for a: Array<T>
2333  for +i: U32
2334  for +k: U32
2335  for +v: T
2336  for h: ASet.Apart(T, a, ASet.sz(T, a), ASet.J(i, ASet.sz(T, a)), ASet.Z(i, ASet.sz(T, a)), ASet.J(k, ASet.sz(T, a)), ASet.Z(k, ASet.sz(T, a)))
2337  {Pair.snd(Array<T>, T, Array.get(T, Array.set(T, a, i, v), k)) == Pair.snd(Array<T>, T, Array.get(T, a, k)) : T}
2338
2339def get_other_set(T, a, i, k, v, h):
2340  ASet.top_other(T, a, i, k, v, ASet.copy(T, a))(h)