1	import Base
2	import ../shape/core.bend as Shape
3	import ../shape/nat_proofs.bend as Nats
4	import ../storage/arith.bend as Arith
5	import ./elementwise.bend as Elementwise
6	
7	# Proofs that a build whose inputs read output cell p at their own slot p or
8	# at slot 0 (`Elementwise.left_slot`, `Elementwise.right_slot`) computes every
9	# cell `Elementwise.cell` computes: there the broadcast walk of every position
10	# lands on that slot.
11	
12	def add_zero_left(+number: Nat) -> {number == Nat.add(number, 0n) : Nat}:
13	  match number:
14	    case 0n:
15	      {==}
16	    case 1n+p:
17	      Equal.cong(Nat, Nat, value => 1n+value, p, Nat.add(p, 0n), add_zero_left(p))
18	
19	def add_comm(+a: Nat, +b: Nat) -> {Nat.add(a, b) == Nat.add(b, a) : Nat}:
20	  match a:
21	    case 0n:
22	      add_zero_left(b)
23	    case 1n+p:
24	      Equal.trans(Nat, 1n+Nat.add(p, b), 1n+Nat.add(b, p), Nat.add(b, 1n+p),
25	        Equal.cong(Nat, Nat, n => 1n+n, Nat.add(p, b), Nat.add(b, p), add_comm(p, b)),
26	        Equal.sym(Nat, Nat.add(b, 1n+p), 1n+Nat.add(b, p), Arith.add_succ(b, p)))
27	
28	# The division loop keeps quotient * (1 + bp) + remainder + what is left to
29	# count fixed, as long as the remainder and the countdown add up to bp.
30	def go_total(+n: Nat, +m: Nat, +d: Nat, +r: Nat, +bp: Nat, +fits: {Nat.add(m, r) == bp : Nat}) ->
31	  {Nat.add(Nat.mul(Pair.fst(Nat, Nat, Nat.divmod.go(n, m, d, r)), 1n+bp), Pair.snd(Nat, Nat, Nat.divmod.go(n, m, d, r))) ==
32	    Nat.add(Nat.add(Nat.mul(d, 1n+bp), r), n) : Nat}:
33	  match n m:
34	    case 0n m:
35	      add_zero_left(Nat.add(Nat.mul(d, 1n+bp), r))
36	    case 1n+np 0n:
37	      +total = Nat.mul(d, 1n+bp)
38	      Equal.trans(Nat,
39	        Nat.add(Nat.mul(Pair.fst(Nat, Nat, Nat.divmod.go(np, r, 1n+d, 0n)), 1n+bp), Pair.snd(Nat, Nat, Nat.divmod.go(np, r, 1n+d, 0n))),
40	        1n+Nat.add(Nat.add(Nat.add(bp, total), 0n), np),
41	        Nat.add(Nat.add(total, r), 1n+np),
42	        go_total(np, r, 1n+d, 0n, bp, Equal.trans(Nat, Nat.add(r, 0n), r, bp, Equal.sym(Nat, r, Nat.add(r, 0n), add_zero_left(r)), fits)),
43	        Equal.trans(Nat, 1n+Nat.add(Nat.add(Nat.add(bp, total), 0n), np), 1n+Nat.add(Nat.add(total, r), np), Nat.add(Nat.add(total, r), 1n+np),
44	          Equal.cong(Nat, Nat, x => 1n+Nat.add(x, np), Nat.add(Nat.add(bp, total), 0n), Nat.add(total, r),
45	            Equal.trans(Nat, Nat.add(Nat.add(bp, total), 0n), Nat.add(bp, total), Nat.add(total, r),
46	              Equal.sym(Nat, Nat.add(bp, total), Nat.add(Nat.add(bp, total), 0n), add_zero_left(Nat.add(bp, total))),
47	              Equal.trans(Nat, Nat.add(bp, total), Nat.add(total, bp), Nat.add(total, r), add_comm(bp, total),
48	                Equal.cong(Nat, Nat, x => Nat.add(total, x), bp, r, Equal.sym(Nat, r, bp, fits))))),
49	          Equal.sym(Nat, Nat.add(Nat.add(total, r), 1n+np), 1n+Nat.add(Nat.add(total, r), np), Arith.add_succ(Nat.add(total, r), np))))
50	    case 1n+np 1n+mp:
51	      +total = Nat.mul(d, 1n+bp)
52	      Equal.trans(Nat,
53	        Nat.add(Nat.mul(Pair.fst(Nat, Nat, Nat.divmod.go(np, mp, d, 1n+r)), 1n+bp), Pair.snd(Nat, Nat, Nat.divmod.go(np, mp, d, 1n+r))),
54	        Nat.add(Nat.add(total, 1n+r), np),
55	        Nat.add(Nat.add(total, r), 1n+np),
56	        go_total(np, mp, d, 1n+r, bp, Equal.trans(Nat, Nat.add(mp, 1n+r), 1n+Nat.add(mp, r), bp, Arith.add_succ(mp, r), fits)),
57	        Equal.trans(Nat, Nat.add(Nat.add(total, 1n+r), np), 1n+Nat.add(Nat.add(total, r), np), Nat.add(Nat.add(total, r), 1n+np),
58	          Equal.cong(Nat, Nat, x => Nat.add(x, np), Nat.add(total, 1n+r), 1n+Nat.add(total, r), Arith.add_succ(total, r)),
59	          Equal.sym(Nat, Nat.add(Nat.add(total, r), 1n+np), 1n+Nat.add(Nat.add(total, r), np), Arith.add_succ(Nat.add(total, r), np))))
60	
61	# s * (i / s) + i % s = i, for every s; at s = 0, Nat.div gives 0 and
62	# Nat.mod gives i.
63	def div_mod(+i: Nat, +s: Nat) -> {Nat.add(Nat.mul(s, Nat.div(i, s)), Nat.mod(i, s)) == i : Nat}:
64	  match s:
65	    case 0n:
66	      {==}
67	    case 1n+bp:
68	      %Nats.mul_comm(Nat.div(i, 1n+bp), 1n+bp) : {Nat.add(_, Nat.mod(i, 1n+bp)) == i : Nat}
69	      go_total(i, bp, 0n, 0n, bp, Equal.sym(Nat, bp, Nat.add(bp, 0n), add_zero_left(bp)))
70	
71	def mod_one_go(+n: Nat, +d: Nat) -> {Pair.snd(Nat, Nat, Nat.divmod.go(n, 0n, d, 0n)) == 0n : Nat}:
72	  match n:
73	    case 0n:
74	      {==}
75	    case 1n+np:
76	      mod_one_go(np, 1n+d)
77	
78	def mod_one(+i: Nat) -> {Nat.mod(i, 1n) == 0n : Nat}:
79	  mod_one_go(i, 0n)
80	
81	# Rows whose left stride is the output's on every axis, the last one 1:
82	# the left walk of every position is the position.
83	def walk_left_same(+rows: List<&2, Shape.Row>, +same: {Elementwise.same_left(rows) == True{} : Bool}, +i: Nat) ->
84	  {Shape.walk_left(rows, i) == i : Nat}:
85	  match rows:
86	    case Shape.Row{+stride, +left, +right} <> Nil{}:
87	      +unit = Arith.eq_true(stride, 1n, Arith.and_left(Nat.is_eq(stride, 1n), Nat.is_eq(left, 1n), same))
88	      +step = Arith.eq_true(left, 1n, Arith.and_right(Nat.is_eq(stride, 1n), Nat.is_eq(left, 1n), same))
89	      %Equal.sym(Nat, left, stride, Equal.trans(Nat, left, 1n, stride, step, Equal.sym(Nat, stride, 1n, unit))) :
90	        {Nat.add(Nat.mul(_, Nat.div(i, stride)), 0n) == i : Nat}
91	      %Equal.sym(Nat, stride, 1n, unit) : {Nat.add(Nat.mul(_, Nat.div(i, _)), 0n) == i : Nat}
92	      %mod_one(i) : {Nat.add(Nat.mul(1n, Nat.div(i, 1n)), _) == i : Nat}
93	      div_mod(i, 1n)
94	    case Shape.Row{+stride, +left, +right} <> +next <> +rest:
95	      +equal = Arith.eq_true(left, stride, Arith.and_left(Nat.is_eq(left, stride), Elementwise.same_left(next <> rest), same))
96	      %Equal.sym(Nat, left, stride, equal) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_left(next <> rest, Nat.mod(i, stride))) == i : Nat}
97	      %Equal.sym(Nat, Shape.walk_left(next <> rest, Nat.mod(i, stride)), Nat.mod(i, stride),
98	        walk_left_same(next <> rest, Arith.and_right(Nat.is_eq(left, stride), Elementwise.same_left(next <> rest), same), Nat.mod(i, stride))) :
99	        {Nat.add(Nat.mul(stride, Nat.div(i, stride)), _) == i : Nat}
100	      div_mod(i, stride)
101	    case Nil{}:
102	      Empty.absurd({Shape.walk_left(Nil{}, i) == i : Nat}, Arith.clash(same))
103	
104	def walk_right_same(+rows: List<&2, Shape.Row>, +same: {Elementwise.same_right(rows) == True{} : Bool}, +i: Nat) ->
105	  {Shape.walk_right(rows, i) == i : Nat}:
106	  match rows:
107	    case Shape.Row{+stride, +left, +right} <> Nil{}:
108	      +unit = Arith.eq_true(stride, 1n, Arith.and_left(Nat.is_eq(stride, 1n), Nat.is_eq(right, 1n), same))
109	      +step = Arith.eq_true(right, 1n, Arith.and_right(Nat.is_eq(stride, 1n), Nat.is_eq(right, 1n), same))
110	      %Equal.sym(Nat, right, stride, Equal.trans(Nat, right, 1n, stride, step, Equal.sym(Nat, stride, 1n, unit))) :
111	        {Nat.add(Nat.mul(_, Nat.div(i, stride)), 0n) == i : Nat}
112	      %Equal.sym(Nat, stride, 1n, unit) : {Nat.add(Nat.mul(_, Nat.div(i, _)), 0n) == i : Nat}
113	      %mod_one(i) : {Nat.add(Nat.mul(1n, Nat.div(i, 1n)), _) == i : Nat}
114	      div_mod(i, 1n)
115	    case Shape.Row{+stride, +left, +right} <> +next <> +rest:
116	      +equal = Arith.eq_true(right, stride, Arith.and_left(Nat.is_eq(right, stride), Elementwise.same_right(next <> rest), same))
117	      %Equal.sym(Nat, right, stride, equal) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_right(next <> rest, Nat.mod(i, stride))) == i : Nat}
118	      %Equal.sym(Nat, Shape.walk_right(next <> rest, Nat.mod(i, stride)), Nat.mod(i, stride),
119	        walk_right_same(next <> rest, Arith.and_right(Nat.is_eq(right, stride), Elementwise.same_right(next <> rest), same), Nat.mod(i, stride))) :
120	        {Nat.add(Nat.mul(stride, Nat.div(i, stride)), _) == i : Nat}
121	      div_mod(i, stride)
122	    case Nil{}:
123	      Empty.absurd({Shape.walk_right(Nil{}, i) == i : Nat}, Arith.clash(same))
124	
125	# Rows whose left stride is 0 on every axis: every position walks to 0.
126	def walk_left_zero(+rows: List<&2, Shape.Row>, +zero: {Elementwise.zero_left(rows) == True{} : Bool}, +i: Nat) ->
127	  {Shape.walk_left(rows, i) == 0n : Nat}:
128	  match rows:
129	    case Shape.Row{+stride, +left, +right} <> +rest:
130	      +none = Arith.eq_true(left, 0n, Arith.and_left(Nat.is_eq(left, 0n), Elementwise.zero_left(rest), zero))
131	      %Equal.sym(Nat, left, 0n, none) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_left(rest, Nat.mod(i, stride))) == 0n : Nat}
132	      walk_left_zero(rest, Arith.and_right(Nat.is_eq(left, 0n), Elementwise.zero_left(rest), zero), Nat.mod(i, stride))
133	    case Nil{}:
134	      {==}
135	
136	def walk_right_zero(+rows: List<&2, Shape.Row>, +zero: {Elementwise.zero_right(rows) == True{} : Bool}, +i: Nat) ->
137	  {Shape.walk_right(rows, i) == 0n : Nat}:
138	  match rows:
139	    case Shape.Row{+stride, +left, +right} <> +rest:
140	      +none = Arith.eq_true(right, 0n, Arith.and_left(Nat.is_eq(right, 0n), Elementwise.zero_right(rest), zero))
141	      %Equal.sym(Nat, right, 0n, none) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_right(rest, Nat.mod(i, stride))) == 0n : Nat}
142	      walk_right_zero(rest, Arith.and_right(Nat.is_eq(right, 0n), Elementwise.zero_right(rest), zero), Nat.mod(i, stride))
143	    case Nil{}:
144	      {==}
145	
146	# The slot an input reads for output cell `position`.
147	def left_at(+slot: Elementwise.Slot, -left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
148	  +map: Shape.Index<left, right, output>, +position: Nat) -> Nat:
149	  match slot:
150	    case Elementwise.OwnSlot{}:
151	      position
152	    case Elementwise.FirstSlot{}:
153	      0n
154	    case Elementwise.WalkedSlot{}:
155	      Elementwise.left_walk(left, right, output, map, position)
156	
157	def right_at(+slot: Elementwise.Slot, -left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
158	  +map: Shape.Index<left, right, output>, +position: Nat) -> Nat:
159	  match slot:
160	    case Elementwise.OwnSlot{}:
161	      position
162	    case Elementwise.FirstSlot{}:
163	      0n
164	    case Elementwise.WalkedSlot{}:
165	      Elementwise.right_walk(left, right, output, map, position)
166	
167	# The slot `left_slot` picks is where the walk of every position lands.
168	def left_pick(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +rows: List<&2, Shape.Row>,
169	  +exact: {rows == Shape.index_table(le, re, oe) : List<&2, Shape.Row>},
170	  same: Bool, +same_exact: {same == Elementwise.same_left(rows) : Bool},
171	  zero: Bool, +zero_exact: {zero == Elementwise.zero_left(rows) : Bool}, +i: Nat) ->
172	  {left_at(Elementwise.slot_of(same, zero), le, re, oe, Shape.Indexed{rows, exact}, i) == Shape.walk_left(rows, i) : Nat}:
173	  match same zero:
174	    case True{} zero:
175	      Equal.sym(Nat, Shape.walk_left(rows, i), i, walk_left_same(rows, Equal.sym(Bool, True{}, Elementwise.same_left(rows), same_exact), i))
176	    case False{} True{}:
177	      Equal.sym(Nat, Shape.walk_left(rows, i), 0n, walk_left_zero(rows, Equal.sym(Bool, True{}, Elementwise.zero_left(rows), zero_exact), i))
178	    case False{} False{}:
179	      {==}
180	
181	def right_pick(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +rows: List<&2, Shape.Row>,
182	  +exact: {rows == Shape.index_table(le, re, oe) : List<&2, Shape.Row>},
183	  same: Bool, +same_exact: {same == Elementwise.same_right(rows) : Bool},
184	  zero: Bool, +zero_exact: {zero == Elementwise.zero_right(rows) : Bool}, +i: Nat) ->
185	  {right_at(Elementwise.slot_of(same, zero), le, re, oe, Shape.Indexed{rows, exact}, i) == Shape.walk_right(rows, i) : Nat}:
186	  match same zero:
187	    case True{} zero:
188	      Equal.sym(Nat, Shape.walk_right(rows, i), i, walk_right_same(rows, Equal.sym(Bool, True{}, Elementwise.same_right(rows), same_exact), i))
189	    case False{} True{}:
190	      Equal.sym(Nat, Shape.walk_right(rows, i), 0n, walk_right_zero(rows, Equal.sym(Bool, True{}, Elementwise.zero_right(rows), zero_exact), i))
191	    case False{} False{}:
192	      {==}
193	
194	def left_index(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +map: Shape.Index<le, re, oe>, +i: Nat) ->
195	  {left_at(Elementwise.left_slot(le, re, oe, map), le, re, oe, map, i) == Elementwise.left_walk(le, re, oe, map, i) : Nat}:
196	  match map:
197	    case Shape.Indexed{+rows, +exact}:
198	      left_pick(le, re, oe, rows, exact, Elementwise.same_left(rows), {==}, Elementwise.zero_left(rows), {==}, i)
199	    case Shape.Unindexed{+zero}:
200	      {==}
201	
202	def right_index(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +map: Shape.Index<le, re, oe>, +i: Nat) ->
203	  {right_at(Elementwise.right_slot(le, re, oe, map), le, re, oe, map, i) == Elementwise.right_walk(le, re, oe, map, i) : Nat}:
204	  match map:
205	    case Shape.Indexed{+rows, +exact}:
206	      right_pick(le, re, oe, rows, exact, Elementwise.same_right(rows), {==}, Elementwise.zero_right(rows), {==}, i)
207	    case Shape.Unindexed{+zero}:
208	      {==}
209	
210	# An inside cell whose slots are known is the cell with the walk.
211	def direct_same(+operation: Elementwise.Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>,
212	  +map: Shape.Index<le, re, oe>, +left: Elementwise.Slot, +left_exact: {left == Elementwise.left_slot(le, re, oe, map) : Elementwise.Slot},
213	  +right: Elementwise.Slot, +right_exact: {right == Elementwise.right_slot(le, re, oe, map) : Elementwise.Slot},
214	  +position: Nat, -la: Array<F32>, -ra: Array<F32>) ->
215	  {Elementwise.direct_cell(operation, left_at(left, le, re, oe, map, position), right_at(right, le, re, oe, map, position), la, ra) ==
216	    Elementwise.cell(operation, le, re, oe, map, position, la, ra, True{}) : F32 & (Array<F32> & Array<F32>)}:
217	  +lw = Elementwise.left_walk(le, re, oe, map, position)
218	  +rw = Elementwise.right_walk(le, re, oe, map, position)
219	  %Equal.sym(Nat, left_at(left, le, re, oe, map, position), lw,
220	    %Equal.sym(Elementwise.Slot, left, Elementwise.left_slot(le, re, oe, map), left_exact) :
221	      {left_at(_, le, re, oe, map, position) == lw : Nat}
222	    left_index(le, re, oe, map, position)) :
223	    {Elementwise.cell_left(operation, ra, U32.from_nat(right_at(right, le, re, oe, map, position)), Array.get(F32, la, U32.from_nat(_))) ==
224	      Elementwise.cell_left(operation, ra, U32.from_nat(rw), Array.get(F32, la, U32.from_nat(lw))) : F32 & (Array<F32> & Array<F32>)}
225	  %Equal.sym(Nat, right_at(right, le, re, oe, map, position), rw,
226	    %Equal.sym(Elementwise.Slot, right, Elementwise.right_slot(le, re, oe, map), right_exact) :
227	      {right_at(_, le, re, oe, map, position) == rw : Nat}
228	    right_index(le, re, oe, map, position)) :
229	    {Elementwise.cell_left(operation, ra, U32.from_nat(_), Array.get(F32, la, U32.from_nat(lw))) ==
230	      Elementwise.cell_left(operation, ra, U32.from_nat(rw), Array.get(F32, la, U32.from_nat(lw))) : F32 & (Array<F32> & Array<F32>)}
231	  {==}
232	
233	# With the slots `left_slot` and `right_slot` pick, the cell of a build is
234	# the cell with the walk.
235	def cell_by_same(+operation: Elementwise.Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>,
236	  +map: Shape.Index<le, re, oe>, inside: Bool, left: Elementwise.Slot, +left_exact: {left == Elementwise.left_slot(le, re, oe, map) : Elementwise.Slot},
237	  right: Elementwise.Slot, +right_exact: {right == Elementwise.right_slot(le, re, oe, map) : Elementwise.Slot},
238	  +position: Nat, -la: Array<F32>, -ra: Array<F32>) ->
239	  {Elementwise.cell_by(inside, left, right, operation, le, re, oe, map, position, la, ra) ==
240	    Elementwise.cell(operation, le, re, oe, map, position, la, ra, inside) : F32 & (Array<F32> & Array<F32>)}:
241	  match inside left right:
242	    case False{} left right:
243	      {==}
244	    case True{} Elementwise.OwnSlot{} Elementwise.OwnSlot{}:
245	      direct_same(operation, le, re, oe, map, Elementwise.OwnSlot{}, left_exact, Elementwise.OwnSlot{}, right_exact, position, la, ra)
246	    case True{} Elementwise.OwnSlot{} Elementwise.FirstSlot{}:
247	      direct_same(operation, le, re, oe, map, Elementwise.OwnSlot{}, left_exact, Elementwise.FirstSlot{}, right_exact, position, la, ra)
248	    case True{} Elementwise.OwnSlot{} Elementwise.WalkedSlot{}:
249	      direct_same(operation, le, re, oe, map, Elementwise.OwnSlot{}, left_exact, Elementwise.WalkedSlot{}, right_exact, position, la, ra)
250	    case True{} Elementwise.FirstSlot{} Elementwise.OwnSlot{}:
251	      direct_same(operation, le, re, oe, map, Elementwise.FirstSlot{}, left_exact, Elementwise.OwnSlot{}, right_exact, position, la, ra)
252	    case True{} Elementwise.FirstSlot{} Elementwise.FirstSlot{}:
253	      direct_same(operation, le, re, oe, map, Elementwise.FirstSlot{}, left_exact, Elementwise.FirstSlot{}, right_exact, position, la, ra)
254	    case True{} Elementwise.FirstSlot{} Elementwise.WalkedSlot{}:
255	      direct_same(operation, le, re, oe, map, Elementwise.FirstSlot{}, left_exact, Elementwise.WalkedSlot{}, right_exact, position, la, ra)
256	    case True{} Elementwise.WalkedSlot{} Elementwise.OwnSlot{}:
257	      direct_same(operation, le, re, oe, map, Elementwise.WalkedSlot{}, left_exact, Elementwise.OwnSlot{}, right_exact, position, la, ra)
258	    case True{} Elementwise.WalkedSlot{} Elementwise.FirstSlot{}:
259	      direct_same(operation, le, re, oe, map, Elementwise.WalkedSlot{}, left_exact, Elementwise.FirstSlot{}, right_exact, position, la, ra)
260	    case True{} Elementwise.WalkedSlot{} Elementwise.WalkedSlot{}:
261	      {==}
262	
263	# The cell function of a build is `Elementwise.cell` of its plan.
264	def cell_at_walk(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +position: Nat, -la: Array<F32>, -ra: Array<F32>,
265	  inside: Bool) ->
266	  {Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), position, (la, ra), inside) ==
267	    Elementwise.cell(operation, Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
268	      Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position, la, ra, inside) : F32 & (Array<F32> & Array<F32>)}:
269	  +le = Shape.extents(Shape.left_shape(plan))
270	  +re = Shape.extents(Shape.right_shape(plan))
271	  +oe = Shape.extents(Shape.output_shape(plan))
272	  cell_by_same(operation, le, re, oe, Shape.index_map(plan), inside, Elementwise.left_slot(le, re, oe, Shape.index_map(plan)), {==},
273	    Elementwise.right_slot(le, re, oe, Shape.index_map(plan)), {==}, position, la, ra)