src/autograd/reduce_fiber_proof.bend

bendtorch-core@0.1.0.0 · 870 lines · 68 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
4import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
5import bendtorch-core@0.1.0.0/src/shape/broadcast_proofs.bend as Broadcast
6import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
7import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
8import bendtorch-core@0.1.0.0/src/tensor/matrix_index_proof.bend as MatrixIndex
9import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
10
11# The fiber walk visits the fiber. For index rows of a broadcast whose left
12# (target) extents are each 1 or the output extent, and a target cell below
13# the target's element count, the contribution offsets that `Reduce.fiber`
14# reads are, in order, exactly the offsets `walk_right` names at the
15# positions p = 0, 1, ..., count - 1 whose `walk_left` is the cell: the
16# positions the masked fold of `ReduceLaws.fiber` adds, in the order it adds
17# them. Everything here is about natural numbers and lists of them.
18
19# Arithmetic.
20
21def add_comm(+a: Nat, +b: Nat) -> {Nat.add(a, b) == Nat.add(b, a) : Nat}:
22  %Broadcast.add_zero(b) : {Nat.add(a, _) == Nat.add(b, a) : Nat}
23  %Broadcast.add_zero(a) : {Nat.add(a, Nat.add(b, 0n)) == Nat.add(b, _) : Nat}
24  Nats.add_swap(a, b, 0n)
25
26def lt_succ_eq(+r: Nat, +size: Nat, below: {Nat.is_lt(r, size) == True{} : Bool},
27  not_below: {False{} == Nat.is_lt(1n+r, size) : Bool}) -> {size == 1n+r : Nat}:
28  match r size:
29    case 0n 0n:
30      Empty.absurd({0n == 1n : Nat}, Arith.clash(below))
31    case 0n 1n+s:
32      match s:
33        case 0n:
34          {==}
35        case 1n+t:
36          Empty.absurd({2n+t == 1n : Nat}, Arith.clash(not_below))
37    case 1n+p 0n:
38      Empty.absurd({0n == 2n+p : Nat}, Arith.clash(below))
39    case 1n+p 1n+s:
40      Equal.cong(Nat, Nat, n => 1n+n, s, 1n+p, lt_succ_eq(p, s, below, not_below))
41
42# t = quotient * size + remainder with remainder < size.
43type Split<-t: Nat, -size: Nat> is Data:
44  Split{quotient: Nat, remainder: Nat, below: {Nat.is_lt(remainder, size) == True{} : Bool},
45    sum: {Nat.add(Nat.mul(quotient, size), remainder) == t : Nat}}
46
47def split_step(+t: Nat, +size: Nat, +q: Nat, +r: Nat, below: {Nat.is_lt(r, size) == True{} : Bool},
48  sum: {Nat.add(Nat.mul(q, size), r) == t : Nat}, last: Shape.BoolCase<Nat.is_lt(1n+r, size)>) -> Split<1n+t, size>:
49  match last:
50    case Shape.BTrue{equal}:
51      Split{q, 1n+r, Equal.sym(Bool, True{}, Nat.is_lt(1n+r, size), equal),
52        Equal.trans(Nat, Nat.add(Nat.mul(q, size), 1n+r), 1n+Nat.add(Nat.mul(q, size), r), 1n+t,
53          Arith.add_succ(Nat.mul(q, size), r), Equal.cong(Nat, Nat, n => 1n+n, Nat.add(Nat.mul(q, size), r), t, sum))}
54    case Shape.BFalse{equal}:
55      +whole = lt_succ_eq(r, size, below, equal)
56      Split{1n+q, 0n,
57        %Equal.sym(Nat, size, 1n+r, whole) : {Nat.is_lt(0n, _) == True{} : Bool}
58        {==},
59        %Equal.sym(Nat, Nat.add(Nat.mul(1n+q, size), 0n), Nat.mul(1n+q, size), Broadcast.add_zero(Nat.mul(1n+q, size))) : {_ == 1n+t : Nat}
60        %Equal.sym(Nat, size, 1n+r, whole) : {Nat.add(_, Nat.mul(q, size)) == 1n+t : Nat}
61        %Equal.sym(Nat, Nat.add(r, Nat.mul(q, size)), Nat.add(Nat.mul(q, size), r), add_comm(r, Nat.mul(q, size))) : {1n+_ == 1n+t : Nat}
62        Equal.cong(Nat, Nat, n => 1n+n, Nat.add(Nat.mul(q, size), r), t, sum)}
63
64def split_next(+t: Nat, +size: Nat, previous: Split<t, size>) -> Split<1n+t, size>:
65  Split{+q, +r, below, sum} = previous
66  split_step(t, size, q, r, below, sum, Shape.bool_case(Nat.is_lt(1n+r, size)))
67
68def split(+t: Nat, +size: Nat, positive: {Nat.is_lt(0n, size) == True{} : Bool}) -> Split<t, size>:
69  match t:
70    case 0n:
71      Split{0n, 0n, positive, {==}}
72    case 1n+p:
73      split_next(p, size, split(p, size, positive))
74
75# The quotient and remainder of a split are Nat.div and Nat.mod.
76def split_divmod(-t: Nat, +size: Nat, +q: Nat, +r: Nat, below: {Nat.is_lt(r, size) == True{} : Bool},
77  sum: {Nat.add(Nat.mul(q, size), r) == t : Nat}) -> {(q, r) == Nat.divmod(t, size) : Nat & Nat}:
78  %sum : {(q, r) == Nat.divmod(_, size) : Nat & Nat}
79  Nats.divmod_exact(q, size, r, MatrixIndex.lt_from(r, size, below))
80
81def div_of(-t: Nat, +size: Nat, +q: Nat, +r: Nat, below: {Nat.is_lt(r, size) == True{} : Bool},
82  sum: {Nat.add(Nat.mul(q, size), r) == t : Nat}) -> {Nat.div(t, size) == q : Nat}:
83  Equal.sym(Nat, q, Nat.div(t, size),
84    Equal.cong(Nat & Nat, Nat, pair => Pair.fst(Nat, Nat, pair), (q, r), Nat.divmod(t, size), split_divmod(t, size, q, r, below, sum)))
85
86def mod_of(-t: Nat, +size: Nat, +q: Nat, +r: Nat, below: {Nat.is_lt(r, size) == True{} : Bool},
87  sum: {Nat.add(Nat.mul(q, size), r) == t : Nat}) -> {Nat.mod(t, size) == r : Nat}:
88  Equal.sym(Nat, r, Nat.mod(t, size),
89    Equal.cong(Nat & Nat, Nat, pair => Pair.snd(Nat, Nat, pair), (q, r), Nat.divmod(t, size), split_divmod(t, size, q, r, below, sum)))
90
91# t = (t / size) * size + t % size, and t % size < size.
92type Decomposed<-t: Nat, -size: Nat> is Data:
93  Decomposed{below: {Nat.is_lt(Nat.mod(t, size), size) == True{} : Bool},
94    sum: {Nat.add(Nat.mul(Nat.div(t, size), size), Nat.mod(t, size)) == t : Nat}}
95
96def decomposed(+t: Nat, +size: Nat, given: Split<t, size>) -> Decomposed<t, size>:
97  Split{+q, +r, +below, +sum} = given
98  +dq = div_of(t, size, q, r, below, sum)
99  +dr = mod_of(t, size, q, r, below, sum)
100  Decomposed{
101    %Equal.sym(Nat, Nat.mod(t, size), r, dr) : {Nat.is_lt(_, size) == True{} : Bool}
102    below,
103    %Equal.sym(Nat, Nat.div(t, size), q, dq) : {Nat.add(Nat.mul(_, size), Nat.mod(t, size)) == t : Nat}
104    %Equal.sym(Nat, Nat.mod(t, size), r, dr) : {Nat.add(Nat.mul(q, size), _) == t : Nat}
105    sum}
106
107def decompose(+t: Nat, +size: Nat, positive: {Nat.is_lt(0n, size) == True{} : Bool}) -> Decomposed<t, size>:
108  decomposed(t, size, split(t, size, positive))
109
110def lt_irrefl(+x: Nat) -> {Nat.is_lt(x, x) == False{} : Bool}:
111  match x:
112    case 0n:
113      {==}
114    case 1n+p:
115      lt_irrefl(p)
116
117def le_add_cong(+c: Nat, +x: Nat, +y: Nat, h: {Nat.is_le(x, y) == True{} : Bool}) ->
118  {Nat.is_le(Nat.add(c, x), Nat.add(c, y)) == True{} : Bool}:
119  match c:
120    case 0n:
121      h
122    case 1n+p:
123      le_add_cong(p, x, y, h)
124
125def le_trans(+a: Nat, +b: Nat, +c: Nat, ab: {Nat.is_le(a, b) == True{} : Bool}, bc: {Nat.is_le(b, c) == True{} : Bool}) ->
126  {Nat.is_le(a, c) == True{} : Bool}:
127  match a b c:
128    case 0n b c:
129      Arith.le_zero(c)
130    case 1n+p 0n c:
131      Empty.absurd({Nat.is_le(1n+p, c) == True{} : Bool}, Arith.clash(ab))
132    case 1n+p 1n+q 0n:
133      Empty.absurd({Nat.is_le(1n+p, 0n) == True{} : Bool}, Arith.clash(bc))
134    case 1n+p 1n+q 1n+r:
135      le_trans(p, q, r, ab, bc)
136
137def mul_le_mono(+a: Nat, +q: Nat, +size: Nat, h: {Nat.is_le(a, q) == True{} : Bool}) ->
138  {Nat.is_le(Nat.mul(a, size), Nat.mul(q, size)) == True{} : Bool}:
139  match a q:
140    case 0n q:
141      Arith.le_zero(Nat.mul(q, size))
142    case 1n+p 0n:
143      Empty.absurd({Nat.is_le(Nat.mul(1n+p, size), 0n) == True{} : Bool}, Arith.clash(h))
144    case 1n+p 1n+r:
145      le_add_cong(size, Nat.mul(p, size), Nat.mul(r, size), mul_le_mono(p, r, size, h))
146
147# Below a * size, the quotient by size is below a.
148def quotient_below(+t: Nat, +a: Nat, +size: Nat, below: {Nat.is_lt(t, Nat.mul(a, size)) == True{} : Bool},
149  decision: Shape.BoolCase<Nat.is_lt(Nat.div(t, size), a)>, given: Decomposed<t, size>) ->
150  {Nat.is_lt(Nat.div(t, size), a) == True{} : Bool}:
151  match decision:
152    case Shape.BTrue{equal}:
153      Equal.sym(Bool, True{}, Nat.is_lt(Nat.div(t, size), a), equal)
154    case Shape.BFalse{equal}:
155      Decomposed{rest, sum} = given
156      +q = Nat.div(t, size)
157      +r = Nat.mod(t, size)
158      +grown = le_trans(Nat.mul(a, size), Nat.mul(q, size), t,
159        mul_le_mono(a, q, size, Arith.not_lt_le(q, a, Equal.sym(Bool, False{}, Nat.is_lt(q, a), equal))),
160        %sum : {Nat.is_le(Nat.mul(q, size), _) == True{} : Bool}
161        Api.le_add_right(Nat.mul(q, size), r))
162      Empty.absurd({Nat.is_lt(q, a) == True{} : Bool},
163        Arith.clash(Equal.trans(Bool, False{}, Nat.is_lt(t, t), True{}, Equal.sym(Bool, Nat.is_lt(t, t), False{}, lt_irrefl(t)),
164          Arith.lt_le_trans(t, Nat.mul(a, size), t, below, grown))))
165
166def positive_of(+t: Nat, +a: Nat, +size: Nat, below: {Nat.is_lt(t, Nat.mul(a, size)) == True{} : Bool}) ->
167  {Nat.is_lt(0n, size) == True{} : Bool}:
168  match size:
169    case 0n:
170      Empty.absurd({Nat.is_lt(0n, 0n) == True{} : Bool},
171        Arith.lt_zero(t,
172          %Nats.mul_zero_right(a) : {Nat.is_lt(t, _) == True{} : Bool}
173          below))
174    case 1n+p:
175      {==}
176
177def div_below(+t: Nat, +a: Nat, +size: Nat, +below: {Nat.is_lt(t, Nat.mul(a, size)) == True{} : Bool}) ->
178  {Nat.is_lt(Nat.div(t, size), a) == True{} : Bool}:
179  quotient_below(t, a, size, below, Shape.bool_case(Nat.is_lt(Nat.div(t, size), a)), decompose(t, size, positive_of(t, a, size, below)))
180
181# a + x = a + y exactly when x = y.
182def eq_add_cancel(+a: Nat, +x: Nat, +y: Nat) -> {Nat.is_eq(Nat.add(a, x), Nat.add(a, y)) == Nat.is_eq(x, y) : Bool}:
183  match a:
184    case 0n:
185      {==}
186    case 1n+p:
187      eq_add_cancel(p, x, y)
188
189def ne_from(+d: Nat, +e: Nat, same: {e == d : Nat}, differ: {Nat.is_eq(d, e) == False{} : Bool}) -> Empty:
190  Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(d, e), True{}, Equal.sym(Bool, Nat.is_eq(d, e), False{}, differ),
191    %Equal.sym(Nat, e, d, same) : {Nat.is_eq(d, _) == True{} : Bool}
192    Arith.eq_refl(d)))
193
194# On a kept row of left stride `size`, the position with digit d and
195# walk w below size lands on t exactly when d is t's quotient and w its
196# remainder.
197def mask_at(+size: Nat, +d: Nat, +w: Nat, +t: Nat, given: Decomposed<t, size>, same: {d == Nat.div(t, size) : Nat}) ->
198  {Nat.is_eq(Nat.add(Nat.mul(size, d), w), t) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}:
199  Decomposed{below, sum} = given
200  %sum :
201    {Nat.is_eq(Nat.add(Nat.mul(size, d), w), _) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}
202  %same : {Nat.is_eq(Nat.add(Nat.mul(size, d), w), Nat.add(Nat.mul(_, size), Nat.mod(t, size))) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}
203  %Nats.mul_comm(size, d) : {Nat.is_eq(Nat.add(Nat.mul(size, d), w), Nat.add(_, Nat.mod(t, size))) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}
204  eq_add_cancel(Nat.mul(size, d), w, Nat.mod(t, size))
205
206def mask_off_decided(+size: Nat, +d: Nat, +w: Nat, +t: Nat, below: {Nat.is_lt(w, size) == True{} : Bool},
207  differ: {Nat.is_eq(d, Nat.div(t, size)) == False{} : Bool},
208  decision: Shape.BoolCase<Nat.is_eq(Nat.add(Nat.mul(size, d), w), t)>) ->
209  {Nat.is_eq(Nat.add(Nat.mul(size, d), w), t) == False{} : Bool}:
210  match decision:
211    case Shape.BFalse{equal}:
212      Equal.sym(Bool, False{}, Nat.is_eq(Nat.add(Nat.mul(size, d), w), t), equal)
213    case Shape.BTrue{equal}:
214      +hit = Arith.eq_true(Nat.add(Nat.mul(size, d), w), t, Equal.sym(Bool, True{}, Nat.is_eq(Nat.add(Nat.mul(size, d), w), t), equal))
215      Empty.absurd({Nat.is_eq(Nat.add(Nat.mul(size, d), w), t) == False{} : Bool},
216        ne_from(d, Nat.div(t, size),
217          div_of(t, size, d, w, below,
218            %Nats.mul_comm(size, d) : {Nat.add(_, w) == t : Nat}
219            hit),
220          differ))
221
222def mask_off(+size: Nat, +d: Nat, +w: Nat, +t: Nat, below: {Nat.is_lt(w, size) == True{} : Bool},
223  differ: {Nat.is_eq(d, Nat.div(t, size)) == False{} : Bool}) ->
224  {Nat.is_eq(Nat.add(Nat.mul(size, d), w), t) == False{} : Bool}:
225  mask_off_decided(size, d, w, t, below, differ, Shape.bool_case(Nat.is_eq(Nat.add(Nat.mul(size, d), w), t)))
226
227# The masked positions and the walk.
228
229def kept(+keep: Bool, +position: Nat, rest: List<&2, Nat>) -> List<&2, Nat>:
230  match keep:
231    case True{}:
232      position <> rest
233    case False{}:
234      rest
235
236# The contribution offsets base + walk_right(p) of the positions
237# p = cursor .. cursor + fuel - 1 whose walk_left is `target`, in order:
238# the words the masked fold of `ReduceLaws.fiber` adds.
239def masked(fuel: Nat, +rows: List<&2, Shape.Row>, +target: Nat, +base: Nat, +cursor: Nat) -> List<&2, Nat>:
240  match fuel:
241    case 0n:
242      Nil{}
243    case 1n+rest:
244      kept(Nat.is_eq(Shape.walk_left(rows, cursor), target), Nat.add(base, Shape.walk_right(rows, cursor)),
245        masked(rest, rows, target, base, 1n+cursor))
246
247# The contribution offsets `Reduce.fiber` reads, in order.
248def visited(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, fuel: Nat, +target: Nat, +base: Nat, +digit: Nat) -> List<&2, Nat>:
249  match rows extents fuel:
250    case Nil{} _ _:
251      [base]
252    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
253      Nil{}
254    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
255      +inner = Nat.mod(target, left)
256      List.append(&2, Nat,
257        visited(rest, others, Reduce.first_span(rest, others), inner, Nat.add(base, Nat.mul(right, digit)), Reduce.first_digit(rest, inner)),
258        visited(rows, extents, remaining, target, base, 1n+digit))
259    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
260      Nil{}
261
262def append_nil(+xs: List<&2, Nat>) -> {List.append(&2, Nat, xs, Nil{}) == xs : List<&2, Nat>}:
263  match xs:
264    case Nil{}:
265      {==}
266    case +x <> +rest:
267      Equal.cong(List<&2, Nat>, List<&2, Nat>, tail => x <> tail, List.append(&2, Nat, rest, Nil{}), rest, append_nil(rest))
268
269def kept_append(+keep: Bool, +position: Nat, -xs: List<&2, Nat>, -ys: List<&2, Nat>) ->
270  {kept(keep, position, List.append(&2, Nat, xs, ys)) == List.append(&2, Nat, kept(keep, position, xs), ys) : List<&2, Nat>}:
271  match keep:
272    case True{}:
273      {==}
274    case False{}:
275      {==}
276
277def masked_split(+a: Nat, +b: Nat, +rows: List<&2, Shape.Row>, +target: Nat, +base: Nat, +cursor: Nat) ->
278  {masked(Nat.add(a, b), rows, target, base, cursor) ==
279    List.append(&2, Nat, masked(a, rows, target, base, cursor), masked(b, rows, target, base, Nat.add(cursor, a))) : List<&2, Nat>}:
280  match a:
281    case 0n:
282      %Equal.sym(Nat, Nat.add(cursor, 0n), cursor, Broadcast.add_zero(cursor)) :
283        {masked(b, rows, target, base, cursor) == masked(b, rows, target, base, _) : List<&2, Nat>}
284      {==}
285    case 1n+p:
286      +keep = Nat.is_eq(Shape.walk_left(rows, cursor), target)
287      +position = Nat.add(base, Shape.walk_right(rows, cursor))
288      +front = masked(p, rows, target, base, 1n+cursor)
289      %Equal.sym(Nat, Nat.add(cursor, 1n+p), 1n+Nat.add(cursor, p), Arith.add_succ(cursor, p)) :
290        {masked(1n+Nat.add(p, b), rows, target, base, cursor) ==
291          List.append(&2, Nat, kept(keep, position, front), masked(b, rows, target, base, _)) : List<&2, Nat>}
292      Equal.trans(List<&2, Nat>, kept(keep, position, masked(Nat.add(p, b), rows, target, base, 1n+cursor)),
293        kept(keep, position, List.append(&2, Nat, front, masked(b, rows, target, base, Nat.add(1n+cursor, p)))),
294        List.append(&2, Nat, kept(keep, position, front), masked(b, rows, target, base, Nat.add(1n+cursor, p))),
295        Equal.cong(List<&2, Nat>, List<&2, Nat>, xs => kept(keep, position, xs),
296          masked(Nat.add(p, b), rows, target, base, 1n+cursor),
297          List.append(&2, Nat, front, masked(b, rows, target, base, Nat.add(1n+cursor, p))),
298          masked_split(p, b, rows, target, base, 1n+cursor)),
299        kept_append(keep, position, front, masked(b, rows, target, base, Nat.add(1n+cursor, p))))
300
301def le_step(+q: Nat, +m: Nat, +s: Nat, room: {Nat.is_le(Nat.add(q, 1n+m), s) == True{} : Bool}) ->
302  {Nat.is_le(Nat.add(1n+q, m), s) == True{} : Bool}:
303  %Arith.add_succ(q, m) : {Nat.is_le(_, s) == True{} : Bool}
304  room
305
306def lt_of_room(+q: Nat, +m: Nat, +s: Nat, room: {Nat.is_le(Nat.add(q, 1n+m), s) == True{} : Bool}) ->
307  {Nat.is_lt(q, s) == True{} : Bool}:
308  Arith.lt_le_trans(q, Nat.add(q, 1n+m), s,
309    %Equal.sym(Nat, Nat.add(q, 1n+m), 1n+Nat.add(q, m), Arith.add_succ(q, m)) : {Nat.is_lt(q, _) == True{} : Bool}
310    Arith.lt_succ(q, Nat.add(q, m), Api.le_add_right(q, m)),
311    room)
312
313def lt_succ_right(+b: Nat, +k: Nat, h: {Nat.is_lt(b, k) == True{} : Bool}) -> {Nat.is_lt(b, 1n+k) == True{} : Bool}:
314  match b k:
315    case 0n k:
316      {==}
317    case 1n+p 0n:
318      Empty.absurd({Nat.is_lt(1n+p, 1n) == True{} : Bool}, Arith.clash(h))
319    case 1n+p 1n+q:
320      lt_succ_right(p, q, h)
321
322def gt_not_eq(+a: Nat, +b: Nat, h: {Nat.is_lt(b, a) == True{} : Bool}) -> {Nat.is_eq(a, b) == False{} : Bool}:
323  match a b:
324    case 0n b:
325      Empty.absurd({Nat.is_eq(0n, b) == False{} : Bool}, Arith.lt_zero(b, h))
326    case 1n+p 0n:
327      {==}
328    case 1n+p 1n+q:
329      gt_not_eq(p, q, h)
330
331def le_ne_succ(+a: Nat, +b: Nat, le: {Nat.is_le(a, b) == True{} : Bool}, ne: {Nat.is_eq(a, b) == False{} : Bool}) ->
332  {Nat.is_le(1n+a, b) == True{} : Bool}:
333  match a b:
334    case 0n 0n:
335      Empty.absurd({Nat.is_le(1n, 0n) == True{} : Bool}, Arith.clash(Equal.sym(Bool, Nat.is_eq(0n, 0n), False{}, ne)))
336    case 0n 1n+q:
337      Arith.le_zero(q)
338    case 1n+p 0n:
339      Empty.absurd({Nat.is_le(2n+p, 0n) == True{} : Bool}, Arith.clash(le))
340    case 1n+p 1n+q:
341      le_ne_succ(p, q, le, ne)
342
343def le_lt_absurd(+a: Nat, +b: Nat, le: {Nat.is_le(a, b) == True{} : Bool}, lt: {Nat.is_lt(b, a) == True{} : Bool}) -> Empty:
344  match a b:
345    case 0n b:
346      Arith.lt_zero(b, lt)
347    case 1n+p 0n:
348      Arith.clash(le)
349    case 1n+p 1n+q:
350      le_lt_absurd(p, q, le, lt)
351
352def lt_self_succ(+k: Nat) -> {Nat.is_lt(k, 1n+k) == True{} : Bool}:
353  match k:
354    case 0n:
355      {==}
356    case 1n+p:
357      lt_self_succ(p)
358
359def axis_zero(+rh: Nat, +oh: Nat, axis: Shape.Axis<0n, rh, oh>) -> {oh == 0n : Nat}:
360  match axis:
361    case Shape.AxisEqual{left_eq, right_eq}:
362      left_eq
363    case Shape.AxisLeftOne{one, right_eq}:
364      Empty.absurd({oh == 0n : Nat}, Arith.clash(Equal.cong(Nat, Bool, n => Nat.is_eq(n, 0n), 1n, 0n, one)))
365    case Shape.AxisRightOne{left_eq, one}:
366      left_eq
367
368def axis_kept(+j: Nat, +rh: Nat, +oh: Nat, axis: Shape.Axis<2n+j, rh, oh>) -> {oh == 2n+j : Nat}:
369  match axis:
370    case Shape.AxisEqual{left_eq, right_eq}:
371      left_eq
372    case Shape.AxisLeftOne{one, right_eq}:
373      Empty.absurd({oh == 2n+j : Nat}, Arith.clash(Equal.cong(Nat, Bool, n => Nat.is_eq(n, 1n), 2n+j, 1n, Equal.sym(Nat, 1n, 2n+j, one))))
374    case Shape.AxisRightOne{left_eq, one}:
375      left_eq
376
377def walk_left_row(+lh: Nat, +rh: Nat, +oh: Nat, axis: Shape.Axis<lh, rh, oh>, +size: Nat, +d: Nat, +w: Nat,
378  d_below: {Nat.is_lt(d, oh) == True{} : Bool}, w_below: {Nat.is_lt(w, size) == True{} : Bool}) ->
379  {Nat.is_lt(Nat.add(Nat.mul(Shape.broadcast_index(lh, size), d), w), Nat.mul(lh, size)) == True{} : Bool}:
380  match lh:
381    case 0n:
382      Empty.absurd({Nat.is_lt(Nat.add(Nat.mul(size, d), w), 0n) == True{} : Bool},
383        Arith.lt_zero(d,
384          %axis_zero(rh, oh, axis) : {Nat.is_lt(d, _) == True{} : Bool}
385          d_below))
386    case 1n+k:
387      match k:
388        case 0n:
389          %Equal.sym(Nat, Nat.add(size, 0n), size, Broadcast.add_zero(size)) : {Nat.is_lt(w, _) == True{} : Bool}
390          w_below
391        case 1n+j:
392          %Nats.mul_comm(d, size) : {Nat.is_lt(Nat.add(_, w), Nat.mul(2n+j, size)) == True{} : Bool}
393          MatrixIndex.lt_mul_add(d, 2n+j, w, size,
394            %axis_kept(j, rh, oh, axis) : {Nat.is_lt(d, _) == True{} : Bool}
395            d_below,
396            w_below)
397
398# The left walk of a position below the output count lies below the
399# left count.
400def walk_left_below(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +aligned: Shape.Aligned(left, right, output),
401  +q: Nat, +below: {Nat.is_lt(q, Shape.numel(output)) == True{} : Bool}) ->
402  {Nat.is_lt(Shape.walk_left(Shape.index_rows(left, right, output), q), Shape.numel(left)) == True{} : Bool}:
403  match left right output:
404    case Nil{} Nil{} Nil{}:
405      {==}
406    case +lh <> +lt +rh <> +rt +oh <> +ot:
407      Shape.AStep{+axis, +tail} = aligned
408      +s = Shape.numel(ot)
409      walk_left_row(lh, rh, oh, axis, Shape.numel(lt), Nat.div(q, s), Shape.walk_left(Shape.index_rows(lt, rt, ot), Nat.mod(q, s)),
410        div_below(q, oh, s, below),
411        walk_left_below(lt, rt, ot, tail, Nat.mod(q, s),
412          {\{Decomposed: m_below => sum => m_below} : Decomposed<q, s> -> {Nat.is_lt(Nat.mod(q, s), s) == True{} : Bool}}(
413            decompose(q, s, positive_of(q, oh, s, below)))))
414    case Nil{} Nil{} +oh <> +ot:
415      Empty.absurd({Nat.is_lt(Shape.walk_left(Shape.index_rows(Nil{}, Nil{}, oh <> ot), q), 1n) == True{} : Bool}, aligned)
416    case Nil{} +rh <> +rt Nil{}:
417      Empty.absurd({Nat.is_lt(Shape.walk_left(Shape.index_rows(Nil{}, rh <> rt, Nil{}), q), 1n) == True{} : Bool}, aligned)
418    case +lh <> +lt Nil{} Nil{}:
419      Empty.absurd({Nat.is_lt(Shape.walk_left(Shape.index_rows(lh <> lt, Nil{}, Nil{}), q), Shape.numel(lh <> lt)) == True{} : Bool}, aligned)
420    case Nil{} +rh <> +rt +oh <> +ot:
421      Empty.absurd({Nat.is_lt(Shape.walk_left(Shape.index_rows(Nil{}, rh <> rt, oh <> ot), q), 1n) == True{} : Bool}, aligned)
422    case +lh <> +lt Nil{} +oh <> +ot:
423      Empty.absurd({Nat.is_lt(Shape.walk_left(Shape.index_rows(lh <> lt, Nil{}, oh <> ot), q), Shape.numel(lh <> lt)) == True{} : Bool}, aligned)
424    case +lh <> +lt +rh <> +rt Nil{}:
425      Empty.absurd({Nat.is_lt(Shape.walk_left(Shape.index_rows(lh <> lt, rh <> rt, Nil{}), q), Shape.numel(lh <> lt)) == True{} : Bool}, aligned)
426
427# Below a kept row of left stride `left`, every walk of the rows below
428# stays under `left`: the rows are the index rows of an aligned broadcast
429# whose left count is `left` and whose output count is `stride`.
430type BoundBy<-rest: List<&2, Shape.Row>, -left: Nat, -stride: Nat> is Data:
431  BoundBy{lt: List<&2, Nat>, rt: List<&2, Nat>, ot: List<&2, Nat>, aligned: Shape.Aligned(lt, rt, ot),
432    rows: {Shape.index_rows(lt, rt, ot) == rest : List<&2, Shape.Row>},
433    size: {Shape.numel(lt) == left : Nat}, count: {Shape.numel(ot) == stride : Nat}}
434
435def bound_at(-rest: List<&2, Shape.Row>, -left: Nat, -stride: Nat, +by: BoundBy<rest, left, stride>, +q: Nat,
436  +below: {Nat.is_lt(q, stride) == True{} : Bool}) -> {Nat.is_lt(Shape.walk_left(rest, q), left) == True{} : Bool}:
437  BoundBy{+lt, +rt, +ot, +aligned, +rows, +size, +count} = by
438  %rows : {Nat.is_lt(Shape.walk_left(_, q), left) == True{} : Bool}
439  %size : {Nat.is_lt(Shape.walk_left(Shape.index_rows(lt, rt, ot), q), _) == True{} : Bool}
440  walk_left_below(lt, rt, ot, aligned, q,
441    %Equal.sym(Nat, Shape.numel(ot), stride, count) : {Nat.is_lt(q, _) == True{} : Bool}
442    below)
443
444# How the rows below a block decide which positions land on the target: on
445# a reduced row (left stride 0) the rows below land on the target itself; on
446# the kept row's target digit they land on the target's remainder; on any
447# other digit of a kept row nothing lands.
448type MaskMode<-left: Nat, -digit: Nat, -target: Nat, -inner: Nat> is Data:
449  Free{zero: {0n == left : Nat}, same: {target == inner : Nat}}
450  Hit{given: Decomposed<target, left>, same: {digit == Nat.div(target, left) : Nat}, remainder: {Nat.mod(target, left) == inner : Nat}}
451
452def mode_mask(+left: Nat, +digit: Nat, +target: Nat, +inner: Nat, +mode: MaskMode<left, digit, target, inner>, +w: Nat) ->
453  {Nat.is_eq(w, inner) == Nat.is_eq(Nat.add(Nat.mul(left, digit), w), target) : Bool}:
454  match mode:
455    case Free{zero, same}:
456      %zero : {Nat.is_eq(w, inner) == Nat.is_eq(Nat.add(Nat.mul(_, digit), w), target) : Bool}
457      %Equal.sym(Nat, target, inner, same) : {Nat.is_eq(w, inner) == Nat.is_eq(w, _) : Bool}
458      {==}
459    case Hit{given, same, remainder}:
460      %remainder : {Nat.is_eq(w, _) == Nat.is_eq(Nat.add(Nat.mul(left, digit), w), target) : Bool}
461      Equal.sym(Bool, Nat.is_eq(Nat.add(Nat.mul(left, digit), w), target), Nat.is_eq(w, Nat.mod(target, left)),
462        mask_at(left, digit, w, target, given, same))
463
464type OffMode<-rest: List<&2, Shape.Row>, -left: Nat, -stride: Nat, -digit: Nat, -target: Nat> is Data:
465  Off{bound: BoundBy<rest, left, stride>, differ: {Nat.is_eq(digit, Nat.div(target, left)) == False{} : Bool}}
466
467def off_mask(+rest: List<&2, Shape.Row>, +left: Nat, +stride: Nat, +digit: Nat, +target: Nat, +mode: OffMode<rest, left, stride, digit, target>,
468  +q: Nat, +below: {Nat.is_lt(q, stride) == True{} : Bool}) ->
469  {False{} == Nat.is_eq(Nat.add(Nat.mul(left, digit), Shape.walk_left(rest, q)), target) : Bool}:
470  Off{+bound, +differ} = mode
471  Equal.sym(Bool, Nat.is_eq(Nat.add(Nat.mul(left, digit), Shape.walk_left(rest, q)), target), False{},
472    mask_off(left, digit, Shape.walk_left(rest, q), target, bound_at(rest, left, stride, bound, q, below), differ))
473
474# On a row of stride s, left stride L and right stride R, the positions
475# d * s + q .. d * s + q + n - 1 have digit d and walk the rows below at
476# q .. q + n - 1.
477def shifted(+n: Nat, +stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +inner: Nat, +base: Nat,
478  +digit: Nat, +q: Nat, +room: {Nat.is_le(Nat.add(q, n), stride) == True{} : Bool}, +mode: MaskMode<left, digit, target, inner>) ->
479  {masked(n, Shape.Row{stride, left, right} <> rest, target, base, Nat.add(Nat.mul(digit, stride), q)) ==
480    masked(n, rest, inner, Nat.add(base, Nat.mul(right, digit)), q) : List<&2, Nat>}:
481  match n:
482    case 0n:
483      {==}
484    case 1n+m:
485      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
486      +c = Nat.add(Nat.mul(digit, stride), q)
487      +below = lt_of_room(q, m, stride, room)
488      +dq = div_of(c, stride, digit, q, below, {==})
489      +mq = mod_of(c, stride, digit, q, below, {==})
490      +wl = Shape.walk_left(rest, q)
491      +wr = Shape.walk_right(rest, q)
492      +tail = masked(m, rows, target, base, 1n+c)
493      +goal_right = kept(Nat.is_eq(wl, inner), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr),
494        masked(m, rest, inner, Nat.add(base, Nat.mul(right, digit)), 1n+q))
495      %Equal.sym(Nat, Nat.div(c, stride), digit, dq) :
496        {kept(Nat.is_eq(Nat.add(Nat.mul(left, _), Shape.walk_left(rest, Nat.mod(c, stride))), target),
497          Nat.add(base, Nat.add(Nat.mul(right, _), Shape.walk_right(rest, Nat.mod(c, stride)))), tail) == goal_right : List<&2, Nat>}
498      %Equal.sym(Nat, Nat.mod(c, stride), q, mq) :
499        {kept(Nat.is_eq(Nat.add(Nat.mul(left, digit), Shape.walk_left(rest, _)), target),
500          Nat.add(base, Nat.add(Nat.mul(right, digit), Shape.walk_right(rest, _))), tail) == goal_right : List<&2, Nat>}
501      %mode_mask(left, digit, target, inner, mode, wl) :
502        {kept(_, Nat.add(base, Nat.add(Nat.mul(right, digit), wr)), tail) == goal_right : List<&2, Nat>}
503      %Equal.sym(Nat, Nat.add(base, Nat.add(Nat.mul(right, digit), wr)), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr),
504        Index.add_assoc(base, Nat.mul(right, digit), wr)) :
505        {kept(Nat.is_eq(wl, inner), _, tail) == goal_right : List<&2, Nat>}
506      %Arith.add_succ(Nat.mul(digit, stride), q) :
507        {kept(Nat.is_eq(wl, inner), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr), masked(m, rows, target, base, _)) == goal_right : List<&2, Nat>}
508      Equal.cong(List<&2, Nat>, List<&2, Nat>, xs => kept(Nat.is_eq(wl, inner), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr), xs),
509        masked(m, rows, target, base, Nat.add(Nat.mul(digit, stride), 1n+q)),
510        masked(m, rest, inner, Nat.add(base, Nat.mul(right, digit)), 1n+q),
511        shifted(m, stride, left, right, rest, target, inner, base, digit, 1n+q, le_step(q, m, stride, room), mode))
512
513def shifted_off(+n: Nat, +stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +base: Nat,
514  +digit: Nat, +q: Nat, +room: {Nat.is_le(Nat.add(q, n), stride) == True{} : Bool}, +mode: OffMode<rest, left, stride, digit, target>) ->
515  {masked(n, Shape.Row{stride, left, right} <> rest, target, base, Nat.add(Nat.mul(digit, stride), q)) == Nil{} : List<&2, Nat>}:
516  match n:
517    case 0n:
518      {==}
519    case 1n+m:
520      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
521      +c = Nat.add(Nat.mul(digit, stride), q)
522      +below = lt_of_room(q, m, stride, room)
523      +dq = div_of(c, stride, digit, q, below, {==})
524      +mq = mod_of(c, stride, digit, q, below, {==})
525      +tail = masked(m, rows, target, base, 1n+c)
526      +position = Nat.add(base, Shape.walk_right(rows, c))
527      %Equal.sym(Nat, Nat.div(c, stride), digit, dq) :
528        {kept(Nat.is_eq(Nat.add(Nat.mul(left, _), Shape.walk_left(rest, Nat.mod(c, stride))), target), position, tail) == Nil{} : List<&2, Nat>}
529      %Equal.sym(Nat, Nat.mod(c, stride), q, mq) :
530        {kept(Nat.is_eq(Nat.add(Nat.mul(left, digit), Shape.walk_left(rest, _)), target), position, tail) == Nil{} : List<&2, Nat>}
531      %off_mask(rest, left, stride, digit, target, mode, q, below) :
532        {kept(_, position, tail) == Nil{} : List<&2, Nat>}
533      %Arith.add_succ(Nat.mul(digit, stride), q) : {masked(m, rows, target, base, _) == Nil{} : List<&2, Nat>}
534      shifted_off(m, stride, left, right, rest, target, base, digit, 1n+q, le_step(q, m, stride, room), mode)
535
536# The first block of a row, at cursor digit * stride: the rows below at
537# base + right * digit.
538def block(+stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +inner: Nat, +base: Nat, +digit: Nat,
539  +mode: MaskMode<left, digit, target, inner>) ->
540  {masked(stride, Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(digit, stride)) ==
541    masked(stride, rest, inner, Nat.add(base, Nat.mul(right, digit)), 0n) : List<&2, Nat>}:
542  %Broadcast.add_zero(Nat.mul(digit, stride)) :
543    {masked(stride, Shape.Row{stride, left, right} <> rest, target, base, _) ==
544      masked(stride, rest, inner, Nat.add(base, Nat.mul(right, digit)), 0n) : List<&2, Nat>}
545  shifted(stride, stride, left, right, rest, target, inner, base, digit, 0n, Api.le_refl(stride), mode)
546
547def block_off(+stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +base: Nat, +digit: Nat,
548  +mode: OffMode<rest, left, stride, digit, target>) ->
549  {masked(stride, Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(digit, stride)) == Nil{} : List<&2, Nat>}:
550  %Broadcast.add_zero(Nat.mul(digit, stride)) :
551    {masked(stride, Shape.Row{stride, left, right} <> rest, target, base, _) == Nil{} : List<&2, Nat>}
552  shifted_off(stride, stride, left, right, rest, target, base, digit, 0n, Api.le_refl(stride), mode)
553
554def next_cursor(+digit: Nat, +stride: Nat) -> {Nat.mul(1n+digit, stride) == Nat.add(Nat.mul(digit, stride), stride) : Nat}:
555  add_comm(stride, Nat.mul(digit, stride))
556
557# Moving the base moves every offset: base + xs.
558def shift(+base: Nat, xs: List<&2, Nat>) -> List<&2, Nat>:
559  match xs:
560    case Nil{}:
561      Nil{}
562    case x <> rest:
563      Nat.add(base, x) <> shift(base, rest)
564
565def shift_append(+base: Nat, +xs: List<&2, Nat>, +ys: List<&2, Nat>) ->
566  {shift(base, List.append(&2, Nat, xs, ys)) == List.append(&2, Nat, shift(base, xs), shift(base, ys)) : List<&2, Nat>}:
567  match xs:
568    case Nil{}:
569      {==}
570    case +x <> +rest:
571      Equal.cong(List<&2, Nat>, List<&2, Nat>, tail => Nat.add(base, x) <> tail,
572        shift(base, List.append(&2, Nat, rest, ys)), List.append(&2, Nat, shift(base, rest), shift(base, ys)), shift_append(base, rest, ys))
573
574def shift_shift(+a: Nat, +b: Nat, +xs: List<&2, Nat>) -> {shift(a, shift(b, xs)) == shift(Nat.add(a, b), xs) : List<&2, Nat>}:
575  match xs:
576    case Nil{}:
577      {==}
578    case +x <> +rest:
579      %Equal.sym(Nat, Nat.add(a, Nat.add(b, x)), Nat.add(Nat.add(a, b), x), Index.add_assoc(a, b, x)) : {_ <> shift(a, shift(b, rest)) == Nat.add(Nat.add(a, b), x) <> shift(Nat.add(a, b), rest) : List<&2, Nat>}
580      Equal.cong(List<&2, Nat>, List<&2, Nat>, tail => Nat.add(Nat.add(a, b), x) <> tail,
581        shift(a, shift(b, rest)), shift(Nat.add(a, b), rest), shift_shift(a, b, rest))
582
583def kept_shift(+keep: Bool, +base: Nat, +x: Nat, +rest: List<&2, Nat>) ->
584  {kept(keep, Nat.add(base, x), shift(base, rest)) == shift(base, kept(keep, x, rest)) : List<&2, Nat>}:
585  match keep:
586    case True{}:
587      {==}
588    case False{}:
589      {==}
590
591def masked_base(+n: Nat, +rows: List<&2, Shape.Row>, +target: Nat, +base: Nat, +cursor: Nat) ->
592  {masked(n, rows, target, base, cursor) == shift(base, masked(n, rows, target, 0n, cursor)) : List<&2, Nat>}:
593  match n:
594    case 0n:
595      {==}
596    case 1n+m:
597      +keep = Nat.is_eq(Shape.walk_left(rows, cursor), target)
598      +x = Shape.walk_right(rows, cursor)
599      Equal.trans(List<&2, Nat>, kept(keep, Nat.add(base, x), masked(m, rows, target, base, 1n+cursor)),
600        kept(keep, Nat.add(base, x), shift(base, masked(m, rows, target, 0n, 1n+cursor))),
601        shift(base, kept(keep, x, masked(m, rows, target, 0n, 1n+cursor))),
602        Equal.cong(List<&2, Nat>, List<&2, Nat>, tail => kept(keep, Nat.add(base, x), tail),
603          masked(m, rows, target, base, 1n+cursor), shift(base, masked(m, rows, target, 0n, 1n+cursor)), masked_base(m, rows, target, base, 1n+cursor)),
604        kept_shift(keep, base, x, masked(m, rows, target, 0n, 1n+cursor)))
605
606def visited_base(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat) ->
607  {visited(rows, extents, fuel, target, base, digit) == shift(base, visited(rows, extents, fuel, target, 0n, digit)) : List<&2, Nat>}:
608  match rows extents fuel:
609    case Nil{} _ _:
610      %Equal.sym(Nat, Nat.add(base, 0n), base, Broadcast.add_zero(base)) : {[base] == [_] : List<&2, Nat>}
611      {==}
612    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
613      {==}
614    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
615      +inner = Nat.mod(target, left)
616      +offset = Nat.mul(right, digit)
617      +span = Reduce.first_span(rest, others)
618      +first = Reduce.first_digit(rest, inner)
619      +low = visited(rest, others, span, inner, 0n, first)
620      +high = visited(rows, extents, remaining, target, 0n, 1n+digit)
621      %Equal.sym(List<&2, Nat>, visited(rest, others, span, inner, Nat.add(base, offset), first), shift(Nat.add(base, offset), low),
622        visited_base(rest, others, span, inner, Nat.add(base, offset), first)) :
623        {List.append(&2, Nat, _, visited(rows, extents, remaining, target, base, 1n+digit)) ==
624          shift(base, List.append(&2, Nat, visited(rest, others, span, inner, offset, first), high)) : List<&2, Nat>}
625      %Equal.sym(List<&2, Nat>, visited(rows, extents, remaining, target, base, 1n+digit), shift(base, high),
626        visited_base(rows, extents, remaining, target, base, 1n+digit)) :
627        {List.append(&2, Nat, shift(Nat.add(base, offset), low), _) ==
628          shift(base, List.append(&2, Nat, visited(rest, others, span, inner, offset, first), high)) : List<&2, Nat>}
629      %Equal.sym(List<&2, Nat>, visited(rest, others, span, inner, offset, first), shift(offset, low),
630        visited_base(rest, others, span, inner, offset, first)) :
631        {List.append(&2, Nat, shift(Nat.add(base, offset), low), shift(base, high)) ==
632          shift(base, List.append(&2, Nat, _, high)) : List<&2, Nat>}
633      %Equal.sym(List<&2, Nat>, shift(base, List.append(&2, Nat, shift(offset, low), high)),
634        List.append(&2, Nat, shift(base, shift(offset, low)), shift(base, high)), shift_append(base, shift(offset, low), high)) :
635        {List.append(&2, Nat, shift(Nat.add(base, offset), low), shift(base, high)) == _ : List<&2, Nat>}
636      %Equal.sym(List<&2, Nat>, shift(base, shift(offset, low)), shift(Nat.add(base, offset), low), shift_shift(base, offset, low)) :
637        {List.append(&2, Nat, shift(Nat.add(base, offset), low), shift(base, high)) == List.append(&2, Nat, _, shift(base, high)) : List<&2, Nat>}
638      {==}
639    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
640      {==}
641
642# The walk below a row, at base 0, visits the masked positions below it.
643def RestWalk(+rest: List<&2, Shape.Row>, +others: List<&2, Nat>, +inner: Nat, +stride: Nat) -> Data:
644  {masked(stride, rest, inner, 0n, 0n) ==
645    visited(rest, others, Reduce.first_span(rest, others), inner, 0n, Reduce.first_digit(rest, inner)) : List<&2, Nat>}
646
647# ... and so at every base.
648def rest_at(+rest: List<&2, Shape.Row>, +others: List<&2, Nat>, +inner: Nat, +stride: Nat, +walk: RestWalk(rest, others, inner, stride),
649  +base: Nat) ->
650  {masked(stride, rest, inner, base, 0n) ==
651    visited(rest, others, Reduce.first_span(rest, others), inner, base, Reduce.first_digit(rest, inner)) : List<&2, Nat>}:
652  %Equal.sym(List<&2, Nat>, masked(stride, rest, inner, base, 0n), shift(base, masked(stride, rest, inner, 0n, 0n)),
653    masked_base(stride, rest, inner, base, 0n)) :
654    {_ == visited(rest, others, Reduce.first_span(rest, others), inner, base, Reduce.first_digit(rest, inner)) : List<&2, Nat>}
655  %Equal.sym(List<&2, Nat>, visited(rest, others, Reduce.first_span(rest, others), inner, base, Reduce.first_digit(rest, inner)),
656    shift(base, visited(rest, others, Reduce.first_span(rest, others), inner, 0n, Reduce.first_digit(rest, inner))),
657    visited_base(rest, others, Reduce.first_span(rest, others), inner, base, Reduce.first_digit(rest, inner))) :
658    {shift(base, masked(stride, rest, inner, 0n, 0n)) == _ : List<&2, Nat>}
659  Equal.cong(List<&2, Nat>, List<&2, Nat>, xs => shift(base, xs), masked(stride, rest, inner, 0n, 0n),
660    visited(rest, others, Reduce.first_span(rest, others), inner, 0n, Reduce.first_digit(rest, inner)), walk)
661
662# A reduced row (left stride 0) walks every digit: the blocks digit ..
663# digit + f - 1 of the masked positions are what the walk visits.
664def reduced_blocks(+f: Nat, +stride: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +extent: Nat, +others: List<&2, Nat>,
665  +target: Nat, +base: Nat, +digit: Nat, +walk: RestWalk(rest, others, target, stride)) ->
666  {masked(Nat.mul(f, stride), Shape.Row{stride, 0n, right} <> rest, target, base, Nat.mul(digit, stride)) ==
667    visited(Shape.Row{stride, 0n, right} <> rest, extent <> others, f, target, base, digit) : List<&2, Nat>}:
668  match f:
669    case 0n:
670      {==}
671    case 1n+g:
672      +rows : List<&2, Shape.Row> = Shape.Row{stride, 0n, right} <> rest
673      +c = Nat.mul(digit, stride)
674      +w = visited(rest, others, Reduce.first_span(rest, others), target, Nat.add(base, Nat.mul(right, digit)), Reduce.first_digit(rest, target))
675      +v = visited(rows, extent <> others, g, target, base, 1n+digit)
676      %Equal.sym(List<&2, Nat>, masked(Nat.add(stride, Nat.mul(g, stride)), rows, target, base, c),
677        List.append(&2, Nat, masked(stride, rows, target, base, c), masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))),
678        masked_split(stride, Nat.mul(g, stride), rows, target, base, c)) :
679        {_ == List.append(&2, Nat, w, v) : List<&2, Nat>}
680      %Equal.sym(List<&2, Nat>, masked(stride, rows, target, base, c), masked(stride, rest, target, Nat.add(base, Nat.mul(right, digit)), 0n),
681        block(stride, 0n, right, rest, target, target, base, digit, Free{{==}, {==}})) :
682        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == List.append(&2, Nat, w, v) : List<&2, Nat>}
683      %Equal.sym(List<&2, Nat>, masked(stride, rest, target, Nat.add(base, Nat.mul(right, digit)), 0n), w,
684        rest_at(rest, others, target, stride, walk, Nat.add(base, Nat.mul(right, digit)))) :
685        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == List.append(&2, Nat, w, v) : List<&2, Nat>}
686      %next_cursor(digit, stride) :
687        {List.append(&2, Nat, w, masked(Nat.mul(g, stride), rows, target, base, _)) == List.append(&2, Nat, w, v) : List<&2, Nat>}
688      %Equal.sym(List<&2, Nat>, masked(Nat.mul(g, stride), rows, target, base, Nat.mul(1n+digit, stride)), v,
689        reduced_blocks(g, stride, right, rest, extent, others, target, base, 1n+digit, walk)) :
690        {List.append(&2, Nat, w, _) == List.append(&2, Nat, w, v) : List<&2, Nat>}
691      {==}
692
693# On a kept row, the blocks after the target's digit hold no position.
694def kept_after(+f: Nat, +stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +base: Nat, +digit: Nat,
695  +after: {Nat.is_lt(Nat.div(target, left), digit) == True{} : Bool}, +bound: BoundBy<rest, left, stride>) ->
696  {masked(Nat.mul(f, stride), Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(digit, stride)) == Nil{} : List<&2, Nat>}:
697  match f:
698    case 0n:
699      {==}
700    case 1n+g:
701      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
702      +c = Nat.mul(digit, stride)
703      %Equal.sym(List<&2, Nat>, masked(Nat.add(stride, Nat.mul(g, stride)), rows, target, base, c),
704        List.append(&2, Nat, masked(stride, rows, target, base, c), masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))),
705        masked_split(stride, Nat.mul(g, stride), rows, target, base, c)) :
706        {_ == Nil{} : List<&2, Nat>}
707      %Equal.sym(List<&2, Nat>, masked(stride, rows, target, base, c), Nil{},
708        block_off(stride, left, right, rest, target, base, digit, Off{bound, gt_not_eq(digit, Nat.div(target, left), after)})) :
709        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == Nil{} : List<&2, Nat>}
710      %next_cursor(digit, stride) :
711        {masked(Nat.mul(g, stride), rows, target, base, _) == Nil{} : List<&2, Nat>}
712      kept_after(g, stride, left, right, rest, target, base, 1n+digit, lt_succ_right(Nat.div(target, left), digit, after), bound)
713
714def kept_hit_step(+g: Nat, +stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +base: Nat, +digit: Nat,
715  decision: Shape.BoolCase<Nat.is_eq(digit, Nat.div(target, left))>,
716  +from: {Nat.is_le(digit, Nat.div(target, left)) == True{} : Bool},
717  +until: {Nat.is_lt(Nat.div(target, left), Nat.add(digit, 1n+g)) == True{} : Bool},
718  +given: Decomposed<target, left>, +bound: BoundBy<rest, left, stride>,
719  rec: @+next: Nat -> @+from: {Nat.is_le(next, Nat.div(target, left)) == True{} : Bool} ->
720    @+until: {Nat.is_lt(Nat.div(target, left), Nat.add(next, g)) == True{} : Bool} ->
721    {masked(Nat.mul(g, stride), Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(next, stride)) ==
722      masked(stride, rest, Nat.mod(target, left), Nat.add(base, Nat.mul(right, Nat.div(target, left))), 0n) : List<&2, Nat>}) ->
723  {masked(Nat.mul(1n+g, stride), Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(digit, stride)) ==
724    masked(stride, rest, Nat.mod(target, left), Nat.add(base, Nat.mul(right, Nat.div(target, left))), 0n) : List<&2, Nat>}:
725  match decision:
726    case Shape.BTrue{equal}:
727      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
728      +c = Nat.mul(digit, stride)
729      +dt = Nat.div(target, left)
730      +tr = Nat.mod(target, left)
731      +hit = masked(stride, rest, tr, Nat.add(base, Nat.mul(right, dt)), 0n)
732      +same = Arith.eq_true(digit, dt, Equal.sym(Bool, True{}, Nat.is_eq(digit, dt), equal))
733      +mine = masked(stride, rest, tr, Nat.add(base, Nat.mul(right, digit)), 0n)
734      %Equal.sym(List<&2, Nat>, masked(Nat.add(stride, Nat.mul(g, stride)), rows, target, base, c),
735        List.append(&2, Nat, masked(stride, rows, target, base, c), masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))),
736        masked_split(stride, Nat.mul(g, stride), rows, target, base, c)) :
737        {_ == hit : List<&2, Nat>}
738      %Equal.sym(List<&2, Nat>, masked(stride, rows, target, base, c), mine,
739        block(stride, left, right, rest, target, tr, base, digit, Hit{given, same, {==}})) :
740        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == hit : List<&2, Nat>}
741      %next_cursor(digit, stride) :
742        {List.append(&2, Nat, mine, masked(Nat.mul(g, stride), rows, target, base, _)) == hit : List<&2, Nat>}
743      %Equal.sym(List<&2, Nat>, masked(Nat.mul(g, stride), rows, target, base, Nat.mul(1n+digit, stride)), Nil{},
744        kept_after(g, stride, left, right, rest, target, base, 1n+digit,
745          %same : {Nat.is_lt(_, 1n+digit) == True{} : Bool}
746          lt_self_succ(digit), bound)) :
747        {List.append(&2, Nat, mine, _) == hit : List<&2, Nat>}
748      %Equal.sym(List<&2, Nat>, List.append(&2, Nat, mine, Nil{}), mine, append_nil(mine)) : {_ == hit : List<&2, Nat>}
749      Equal.cong(Nat, List<&2, Nat>, d => masked(stride, rest, tr, Nat.add(base, Nat.mul(right, d)), 0n), digit, dt, same)
750    case Shape.BFalse{equal}:
751      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
752      +c = Nat.mul(digit, stride)
753      +dt = Nat.div(target, left)
754      +tr = Nat.mod(target, left)
755      +hit = masked(stride, rest, tr, Nat.add(base, Nat.mul(right, dt)), 0n)
756      +differ = Equal.sym(Bool, False{}, Nat.is_eq(digit, dt), equal)
757      %Equal.sym(List<&2, Nat>, masked(Nat.add(stride, Nat.mul(g, stride)), rows, target, base, c),
758        List.append(&2, Nat, masked(stride, rows, target, base, c), masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))),
759        masked_split(stride, Nat.mul(g, stride), rows, target, base, c)) :
760        {_ == hit : List<&2, Nat>}
761      %Equal.sym(List<&2, Nat>, masked(stride, rows, target, base, c), Nil{},
762        block_off(stride, left, right, rest, target, base, digit, Off{bound, differ})) :
763        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == hit : List<&2, Nat>}
764      %next_cursor(digit, stride) :
765        {masked(Nat.mul(g, stride), rows, target, base, _) == hit : List<&2, Nat>}
766      rec(1n+digit, le_ne_succ(digit, dt, from, differ),
767        %Arith.add_succ(digit, g) : {Nat.is_lt(dt, _) == True{} : Bool}
768        until)
769
770# On a kept row, the blocks from `digit` that cover the target's digit hold
771# exactly the rows below at the target's remainder.
772def kept_hit(+f: Nat, +stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +base: Nat, +digit: Nat,
773  +from: {Nat.is_le(digit, Nat.div(target, left)) == True{} : Bool},
774  +until: {Nat.is_lt(Nat.div(target, left), Nat.add(digit, f)) == True{} : Bool},
775  +given: Decomposed<target, left>, +bound: BoundBy<rest, left, stride>) ->
776  {masked(Nat.mul(f, stride), Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(digit, stride)) ==
777    masked(stride, rest, Nat.mod(target, left), Nat.add(base, Nat.mul(right, Nat.div(target, left))), 0n) : List<&2, Nat>}:
778  match f:
779    case 0n:
780      Empty.absurd({masked(0n, Shape.Row{stride, left, right} <> rest, target, base, Nat.mul(digit, stride)) ==
781        masked(stride, rest, Nat.mod(target, left), Nat.add(base, Nat.mul(right, Nat.div(target, left))), 0n) : List<&2, Nat>},
782        le_lt_absurd(digit, Nat.div(target, left), from,
783          %Broadcast.add_zero(digit) : {Nat.is_lt(Nat.div(target, left), _) == True{} : Bool}
784          until))
785    case 1n+g:
786      kept_hit_step(g, stride, left, right, rest, target, base, digit, Shape.bool_case(Nat.is_eq(digit, Nat.div(target, left))),
787        from, until, given, bound,
788        next => from => until => kept_hit(g, stride, left, right, rest, target, base, next, from, until, given, bound))
789
790# A kept row walks the one digit the target names there.
791def kept_row(+size: Nat, +stride: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +extent: Nat, +others: List<&2, Nat>,
792  +target: Nat, +base: Nat, +inside: {Nat.is_lt(target, Nat.mul(extent, size)) == True{} : Bool},
793  +walk: RestWalk(rest, others, Nat.mod(target, size), stride), +bound: BoundBy<rest, size, stride>) ->
794  {masked(Nat.mul(extent, stride), Shape.Row{stride, size, right} <> rest, target, base, 0n) ==
795    visited(Shape.Row{stride, size, right} <> rest, extent <> others, Reduce.span(size, extent), target, base, Nat.div(target, size)) :
796    List<&2, Nat>}:
797  match size:
798    case 0n:
799      Empty.absurd({masked(Nat.mul(extent, stride), Shape.Row{stride, 0n, right} <> rest, target, base, 0n) ==
800        visited(Shape.Row{stride, 0n, right} <> rest, extent <> others, Reduce.span(0n, extent), target, base, Nat.div(target, 0n)) :
801        List<&2, Nat>},
802        Arith.lt_zero(target,
803          %Nats.mul_zero_right(extent) : {Nat.is_lt(target, _) == True{} : Bool}
804          inside))
805    case 1n+l:
806      +dt = Nat.div(target, 1n+l)
807      +tr = Nat.mod(target, 1n+l)
808      +w = visited(rest, others, Reduce.first_span(rest, others), tr, Nat.add(base, Nat.mul(right, dt)), Reduce.first_digit(rest, tr))
809      %Equal.sym(List<&2, Nat>, masked(Nat.mul(extent, stride), Shape.Row{stride, 1n+l, right} <> rest, target, base, 0n),
810        masked(stride, rest, tr, Nat.add(base, Nat.mul(right, dt)), 0n),
811        kept_hit(extent, stride, 1n+l, right, rest, target, base, 0n, Arith.le_zero(dt),
812          div_below(target, extent, 1n+l, inside), decompose(target, 1n+l, {==}), bound)) :
813        {_ == List.append(&2, Nat, w, Nil{}) : List<&2, Nat>}
814      %Equal.sym(List<&2, Nat>, masked(stride, rest, tr, Nat.add(base, Nat.mul(right, dt)), 0n), w,
815        rest_at(rest, others, tr, stride, walk, Nat.add(base, Nat.mul(right, dt)))) :
816        {_ == List.append(&2, Nat, w, Nil{}) : List<&2, Nat>}
817      Equal.sym(List<&2, Nat>, List.append(&2, Nat, w, Nil{}), w, append_nil(w))
818
819def RowsWalk(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +target: Nat) -> Data:
820  RestWalk(Shape.index_rows(left, right, output), output, target, Shape.numel(output))
821
822def fiber_cons(+lh: Nat, +lt: List<&2, Nat>, +rh: Nat, +rt: List<&2, Nat>, +oh: Nat, +ot: List<&2, Nat>, +axis: Shape.Axis<lh, rh, oh>,
823  +tail: Shape.Aligned(lt, rt, ot), +target: Nat, +below: {Nat.is_lt(target, Nat.mul(lh, Shape.numel(lt))) == True{} : Bool},
824  rec: @+inner: Nat -> @+inside: {Nat.is_lt(inner, Shape.numel(lt)) == True{} : Bool} -> RowsWalk(lt, rt, ot, inner)) ->
825  RowsWalk(lh <> lt, rh <> rt, oh <> ot, target):
826  match lh:
827    case 0n:
828      Empty.absurd(RowsWalk(0n <> lt, rh <> rt, oh <> ot, target), Arith.lt_zero(target, below))
829    case 1n+k:
830      match k:
831        case 0n:
832          reduced_blocks(oh, Shape.numel(ot), Shape.broadcast_index(rh, Shape.numel(rt)), Shape.index_rows(lt, rt, ot), oh, ot, target, 0n, 0n,
833            rec(target,
834              %Broadcast.add_zero(Shape.numel(lt)) : {Nat.is_lt(target, _) == True{} : Bool}
835              below))
836        case 1n+j:
837          +size = Shape.numel(lt)
838          kept_row(size, Shape.numel(ot), Shape.broadcast_index(rh, Shape.numel(rt)), Shape.index_rows(lt, rt, ot), oh, ot, target, 0n,
839            %Equal.sym(Nat, oh, 2n+j, axis_kept(j, rh, oh, axis)) : {Nat.is_lt(target, Nat.mul(_, size)) == True{} : Bool}
840            below,
841            rec(Nat.mod(target, size),
842              {\{Decomposed: m_below => sum => m_below} : Decomposed<target, size> -> {Nat.is_lt(Nat.mod(target, size), size) == True{} : Bool}}(
843                decompose(target, size, positive_of(target, 2n+j, size, below)))),
844            BoundBy{lt, rt, ot, tail, {==}, {==}, {==}})
845
846# The fiber-walk lemma: for the index rows of an aligned broadcast and a
847# cell below the left count, the offsets the walk reads, in order, are the
848# masked positions.
849def fiber_walk(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +aligned: Shape.Aligned(left, right, output),
850  +target: Nat, +below: {Nat.is_lt(target, Shape.numel(left)) == True{} : Bool}) -> RowsWalk(left, right, output, target):
851  match left right output:
852    case Nil{} Nil{} Nil{}:
853      %Arith.lt_one(target, below) : {kept(Nat.is_eq(0n, _), 0n, Nil{}) == [0n] : List<&2, Nat>}
854      {==}
855    case +lh <> +lt +rh <> +rt +oh <> +ot:
856      Shape.AStep{+axis, +tail} = aligned
857      fiber_cons(lh, lt, rh, rt, oh, ot, axis, tail, target, below,
858        inner => inside => fiber_walk(lt, rt, ot, tail, inner, inside))
859    case Nil{} Nil{} +oh <> +ot:
860      Empty.absurd(RowsWalk(Nil{}, Nil{}, oh <> ot, target), aligned)
861    case Nil{} +rh <> +rt Nil{}:
862      Empty.absurd(RowsWalk(Nil{}, rh <> rt, Nil{}, target), aligned)
863    case +lh <> +lt Nil{} Nil{}:
864      Empty.absurd(RowsWalk(lh <> lt, Nil{}, Nil{}, target), aligned)
865    case Nil{} +rh <> +rt +oh <> +ot:
866      Empty.absurd(RowsWalk(Nil{}, rh <> rt, oh <> ot, target), aligned)
867    case +lh <> +lt Nil{} +oh <> +ot:
868      Empty.absurd(RowsWalk(lh <> lt, Nil{}, oh <> ot, target), aligned)
869    case +lh <> +lt +rh <> +rt Nil{}:
870      Empty.absurd(RowsWalk(lh <> lt, rh <> rt, Nil{}, target), aligned)