1	import Base
2	import ./core.bend as Shape
3	import ./nat_proofs.bend as Nats
4	import ./index_proofs.bend as Index
5	import ./broadcast_proofs.bend as Broadcast
6	
7	type Fits.Step<-dh: Nat, -dt: List<&2, Nat>, -vh: Nat, -vt: List<&2, Nat>> is Data:
8	  FStep{bound: {True{} == Nat.is_lt(vh, dh) : Bool}, tail: Fits(dt, vt)}
9	
10	def Fits(dims: List<&2, Nat>, values: List<&2, Nat>) -> Data:
11	  match dims values:
12	    case Nil{} Nil{}:
13	      Unit
14	    case dh <> dt vh <> vt:
15	      Fits.Step<dh, dt, vh, vt>
16	    case Nil{} vh <> vt:
17	      Empty
18	    case dh <> dt Nil{}:
19	      Empty
20	
21	def fin_from(+bound: Nat, +value: Nat, fits: {True{} == Nat.is_lt(value, bound) : Bool}) -> Shape.Fin(bound):
22	  match bound value:
23	    case 0n 0n:
24	      Empty.absurd(Shape.Fin(0n), Shape.clash(fits))
25	    case 0n 1n+q:
26	      Empty.absurd(Shape.Fin(0n), Shape.clash(fits))
27	    case 1n+p 0n:
28	      Shape.FZero{}
29	    case 1n+p 1n+q:
30	      Shape.FNext{fin_from(p, q, fits)}
31	
32	def fin_from_value(+bound: Nat, +value: Nat, fits: {True{} == Nat.is_lt(value, bound) : Bool}) -> {Shape.Fin.value(bound, fin_from(bound, value, fits)) == value : Nat}:
33	  match bound value:
34	    case 0n 0n:
35	      Empty.absurd({Shape.Fin.value(0n, fin_from(0n, 0n, fits)) == 0n : Nat}, Shape.clash(fits))
36	    case 0n 1n+q:
37	      Empty.absurd({Shape.Fin.value(0n, fin_from(0n, 1n+q, fits)) == 1n+q : Nat}, Shape.clash(fits))
38	    case 1n+p 0n:
39	      {==}
40	    case 1n+p 1n+q:
41	      Equal.cong(Nat, Nat, value => 1n+value, Shape.Fin.value(p, fin_from(p, q, fits)), q, fin_from_value(p, q, fits))
42	
43	def coordinate_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, fits: Fits(dims, values)) -> Shape.Coordinate(dims):
44	  match dims values:
45	    case Nil{} Nil{}:
46	      Unit{}
47	    case +dh <> +dt +vh <> +vt:
48	      FStep{bound, tail} = fits
49	      Shape.CStep{fin_from(dh, vh, bound), coordinate_of(dt, vt, tail)}
50	    case Nil{} vh <> vt:
51	      match fits:
52	    case dh <> dt Nil{}:
53	      match fits:
54	
55	def indices_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, +fits: Fits(dims, values)) -> {Shape.indices(dims, coordinate_of(dims, values, fits)) == values : List<&2, Nat>}:
56	  match dims values:
57	    case Nil{} Nil{}:
58	      {==}
59	    case +dh <> +dt +vh <> +vt:
60	      FStep{+bound, +tail} = fits
61	      %Equal.sym(Nat, Shape.Fin.value(dh, fin_from(dh, vh, bound)), vh, fin_from_value(dh, vh, bound)) : {_ <> Shape.indices(dt, coordinate_of(dt, vt, tail)) == vh <> vt : List<&2, Nat>}
62	      %Equal.sym(List<&2, Nat>, Shape.indices(dt, coordinate_of(dt, vt, tail)), vt, indices_of(dt, vt, tail)) : {vh <> _ == vh <> vt : List<&2, Nat>}
63	      {==}
64	    case Nil{} vh <> vt:
65	      match fits:
66	    case dh <> dt Nil{}:
67	      match fits:
68	
69	def offset_indices(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {Shape.offset(dims, Shape.indices(dims, coordinate)) == Shape.row_major(dims, coordinate) : Nat}:
70	  match dims:
71	    case Nil{}:
72	      {==}
73	    case +head <> +tail:
74	      Shape.CStep{+index, +rest} = coordinate
75	      Equal.cong(Nat, Nat, value => Nat.add(Nat.mul(Shape.Fin.value(head, index), Shape.numel(tail)), value), Shape.offset(tail, Shape.indices(tail, rest)), Shape.row_major(tail, rest), offset_indices(tail, rest))
76	
77	def row_major_lt(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> Shape.LT(Shape.row_major(dims, coordinate), Shape.numel(dims)):
78	  %Index.flatten_row_major(dims, coordinate) : Shape.LT(_, Shape.numel(dims))
79	  Shape.Fin.bound(Shape.numel(dims), Shape.flatten(dims, coordinate))
80	
81	def size_fits(+shape: Shape.Shape, -index: Nat, fits: {True{} == Nat.is_lt(index, Shape.elements(shape)) : Bool}) -> {True{} == Nat.is_lt(index, Shape.numel(Shape.extents(shape))) : Bool}:
82	  %Shape.elements_exact(shape) : {True{} == Nat.is_lt(index, _) : Bool}
83	  fits
84	
85	def size_misses(+shape: Shape.Shape, -index: Nat, misses: {False{} == Nat.is_lt(index, Shape.elements(shape)) : Bool}) -> {False{} == Nat.is_lt(index, Shape.numel(Shape.extents(shape))) : Bool}:
86	  %Shape.elements_exact(shape) : {False{} == Nat.is_lt(index, _) : Bool}
87	  misses
88	
89	def offset_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, +fits: Fits(dims, values)) -> {Shape.row_major(dims, coordinate_of(dims, values, fits)) == Shape.offset(dims, values) : Nat}:
90	  %indices_of(dims, values, fits) : {Shape.row_major(dims, coordinate_of(dims, values, fits)) == Shape.offset(dims, _) : Nat}
91	  Equal.sym(Nat, Shape.offset(dims, Shape.indices(dims, coordinate_of(dims, values, fits))), Shape.row_major(dims, coordinate_of(dims, values, fits)), offset_indices(dims, coordinate_of(dims, values, fits)))
92	
93	def digits_step(+head: Nat, +tail: List<&2, Nat>, rec: @+index: Shape.Fin(Shape.numel(tail)) -> {Shape.indices(tail, Shape.unflatten(tail, index)) == Shape.digits(tail, Shape.Fin.value(Shape.numel(tail), index)) : List<&2, Nat>}, coordinate: Shape.Fin(head) & Shape.Fin(Shape.numel(tail))) -> {Shape.indices(head <> tail, Shape.unflatten.step(head, tail, rest => Shape.unflatten(tail, rest), coordinate)) == Shape.digits(head <> tail, Shape.Fin.value(Nat.mul(head, Shape.numel(tail)), Shape.Fin.pack(head, Shape.numel(tail), coordinate))) : List<&2, Nat>}:
94	  (+first, +last) = coordinate
95	  +stride = Shape.numel(tail)
96	  +major = Shape.Fin.value(head, first)
97	  +minor = Shape.Fin.value(stride, last)
98	  %Equal.sym(Nat, Shape.Fin.value(Nat.mul(head, stride), Shape.Fin.pack(head, stride, (first, last))), Nat.add(Nat.mul(major, stride), minor), Index.pack_value(head, stride, first, last)) : {major <> Shape.indices(tail, Shape.unflatten(tail, last)) == Shape.digits(head <> tail, _) : List<&2, Nat>}
99	  %Nats.divmod_exact(major, stride, minor, Shape.Fin.bound(stride, last)) : {major <> Shape.indices(tail, Shape.unflatten(tail, last)) == Pair.fst(Nat, Nat, _) <> Shape.digits(tail, Pair.snd(Nat, Nat, _)) : List<&2, Nat>}
100	  %rec(last) : {major <> Shape.indices(tail, Shape.unflatten(tail, last)) == major <> _ : List<&2, Nat>}
101	  {==}
102	
103	def digits_unflatten(+dims: List<&2, Nat>, +index: Shape.Fin(Shape.numel(dims))) -> {Shape.indices(dims, Shape.unflatten(dims, index)) == Shape.digits(dims, Shape.Fin.value(Shape.numel(dims), index)) : List<&2, Nat>}:
104	  match dims:
105	    case Nil{}:
106	      match index:
107	        case Shape.FZero{}:
108	          {==}
109	        case Shape.FNext{previous}:
110	          match previous:
111	    case +head <> +tail:
112	      %Index.pack_unpack(head, Shape.numel(tail), index) : {Shape.indices(head <> tail, Shape.unflatten(head <> tail, index)) == Shape.digits(head <> tail, Shape.Fin.value(Shape.numel(head <> tail), _)) : List<&2, Nat>}
113	      digits_step(head, tail, +index => digits_unflatten(tail, index), Shape.Fin.unpack(head, Shape.numel(tail), index))
114	
115	def fits_of.step(-dh: Nat, -dt: List<&2, Nat>, -vh: Nat, -vt: List<&2, Nat>, verdict: Shape.BoolCase<Nat.is_lt(vh, dh)>, ok: {True{} == Bool.and(Nat.is_lt(vh, dh), Shape.fit(dt, vt)) : Bool}, rec: {True{} == Shape.fit(dt, vt) : Bool} -> Fits(dt, vt)) -> Fits(dh <> dt, vh <> vt):
116	  match verdict:
117	    case Shape.BFalse{equal}:
118	      Empty.absurd(Fits(dh <> dt, vh <> vt), Shape.clash(Equal.trans(Bool, True{}, Bool.and(Nat.is_lt(vh, dh), Shape.fit(dt, vt)), False{}, ok, Equal.cong(Bool, Bool, head => Bool.and(head, Shape.fit(dt, vt)), Nat.is_lt(vh, dh), False{}, Equal.sym(Bool, False{}, Nat.is_lt(vh, dh), equal)))))
119	    case Shape.BTrue{+equal}:
120	      FStep{equal, rec(Equal.trans(Bool, True{}, Bool.and(Nat.is_lt(vh, dh), Shape.fit(dt, vt)), Shape.fit(dt, vt), ok, Equal.cong(Bool, Bool, head => Bool.and(head, Shape.fit(dt, vt)), Nat.is_lt(vh, dh), True{}, Equal.sym(Bool, True{}, Nat.is_lt(vh, dh), equal))))}
121	
122	def fits_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, ok: {True{} == Shape.fit(dims, values) : Bool}) -> Fits(dims, values):
123	  match dims values:
124	    case Nil{} Nil{}:
125	      Unit{}
126	    case +dh <> +dt +vh <> +vt:
127	      fits_of.step(dh, dt, vh, vt, Shape.bool_case(Nat.is_lt(vh, dh)), ok, rest => fits_of(dt, vt, rest))
128	    case Nil{} vh <> vt:
129	      Empty.absurd(Fits(Nil{}, vh <> vt), Shape.clash(ok))
130	    case dh <> dt Nil{}:
131	      Empty.absurd(Fits(dh <> dt, Nil{}), Shape.clash(ok))
132	
133	def fit_indices(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {True{} == Shape.fit(dims, Shape.indices(dims, coordinate)) : Bool}:
134	  match dims:
135	    case Nil{}:
136	      match coordinate:
137	        case Unit{}:
138	          {==}
139	    case +head <> +tail:
140	      Shape.CStep{+index, +rest} = coordinate
141	      %Nats.lt_true(Shape.Fin.value(head, index), head, Shape.Fin.bound(head, index)) : {True{} == Bool.and(_, Shape.fit(tail, Shape.indices(tail, rest))) : Bool}
142	      fit_indices(tail, rest)
143	
144	def fit_words_spec(+dims: List<&2, U32>, +values: List<&2, U32>) -> {Shape.fit_words(dims, values) == Shape.fit(Shape.naturals(dims), Shape.naturals(values)) : Bool}:
145	  match dims values:
146	    case Nil{} Nil{}:
147	      {==}
148	    case +dh <> +dt +vh <> +vt:
149	      Equal.cong(Bool, Bool, rest => Bool.and(Nat.is_lt(U32.to_nat(vh), U32.to_nat(dh)), rest), Shape.fit_words(dt, vt), Shape.fit(Shape.naturals(dt), Shape.naturals(vt)), fit_words_spec(dt, vt))
150	    case Nil{} vh <> vt:
151	      {==}
152	    case dh <> dt Nil{}:
153	      {==}
154	
155	def digits_words_spec(+dims: List<&2, U32>, +index: Nat) -> {Shape.digits_words(dims, index) == Shape.digits(Shape.naturals(dims), index) : List<&2, Nat>}:
156	  match dims:
157	    case Nil{}:
158	      {==}
159	    case +head <> +tail:
160	      +stride = Shape.numel(Shape.naturals(tail))
161	      %Equal.sym(Nat, Shape.numel_words(tail), stride, Shape.numel_words_spec(tail)) : {Nat.div(index, _) <> Shape.digits_words(tail, Nat.mod(index, _)) == Shape.digits(Shape.naturals(head <> tail), index) : List<&2, Nat>}
162	      Equal.cong(List<&2, Nat>, List<&2, Nat>, rest => Nat.div(index, stride) <> rest, Shape.digits_words(tail, Nat.mod(index, stride)), Shape.digits(Shape.naturals(tail), Nat.mod(index, stride)), digits_words_spec(tail, Nat.mod(index, stride)))
163	
164	def projected_words_spec(+dims: List<&2, U32>, +values: List<&2, U32>) -> {Shape.projected_words(dims, values) == Shape.broadcast_indices(Shape.naturals(dims), Shape.naturals(values)) : List<&2, Nat>}:
165	  match dims values:
166	    case +head <> +tail +value <> +rest:
167	      Equal.cong(List<&2, Nat>, List<&2, Nat>, rest => Shape.broadcast_index(U32.to_nat(head), U32.to_nat(value)) <> rest, Shape.projected_words(tail, rest), Shape.broadcast_indices(Shape.naturals(tail), Shape.naturals(rest)), projected_words_spec(tail, rest))
168	    case Nil{} values:
169	      {==}
170	    case head <> tail Nil{}:
171	      {==}
172	
173	def naturals_drop(+values: List<&2, U32>, +count: Nat) -> {Shape.naturals(List.drop(&2, U32, values, count)) == List.drop(&2, Nat, Shape.naturals(values), count) : List<&2, Nat>}:
174	  match values count:
175	    case Nil{} count:
176	      {==}
177	    case head <> tail 0n:
178	      {==}
179	    case head <> tail 1n+rest:
180	      naturals_drop(tail, rest)
181	
182	def horner_spec(+dims: List<&2, U32>, +values: List<&2, U32>, +acc: Nat, +fits: Fits(Shape.naturals(dims), Shape.naturals(values))) -> {Shape.horner(dims, values, acc) == Nat.add(Nat.mul(acc, Shape.numel(Shape.naturals(dims))), Shape.offset(Shape.naturals(dims), Shape.naturals(values))) : Nat}:
183	  match dims values:
184	    case Nil{} Nil{}:
185	      Equal.trans(Nat, acc, Nat.mul(acc, 1n), Nat.add(Nat.mul(acc, 1n), 0n), Equal.sym(Nat, Nat.mul(acc, 1n), acc, Nats.mul_one_right(acc)), Equal.sym(Nat, Nat.add(Nat.mul(acc, 1n), 0n), Nat.mul(acc, 1n), Broadcast.add_zero(Nat.mul(acc, 1n))))
186	    case +dh <> +dt +vh <> +vt:
187	      FStep{bound, +tail} = fits
188	      +size = Shape.numel(Shape.naturals(dt))
189	      +rest = Shape.offset(Shape.naturals(dt), Shape.naturals(vt))
190	      +next = Nat.add(Nat.mul(acc, U32.to_nat(dh)), U32.to_nat(vh))
191	      %Equal.sym(Nat, Shape.horner(dt, vt, next), Nat.add(Nat.mul(next, size), rest), horner_spec(dt, vt, next, tail)) : {_ == Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)) : Nat}
192	      %Equal.sym(Nat, Nat.mul(next, size), Nat.add(Nat.mul(Nat.mul(acc, U32.to_nat(dh)), size), Nat.mul(U32.to_nat(vh), size)), Nats.mul_add_right(Nat.mul(acc, U32.to_nat(dh)), U32.to_nat(vh), size)) : {Nat.add(_, rest) == Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)) : Nat}
193	      %Equal.sym(Nat, Nat.mul(Nat.mul(acc, U32.to_nat(dh)), size), Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nats.mul_assoc(acc, U32.to_nat(dh), size)) : {Nat.add(Nat.add(_, Nat.mul(U32.to_nat(vh), size)), rest) == Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)) : Nat}
194	      Equal.sym(Nat, Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)), Nat.add(Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.mul(U32.to_nat(vh), size)), rest), Index.add_assoc(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.mul(U32.to_nat(vh), size), rest))
195	    case Nil{} vh <> vt:
196	      match fits:
197	    case dh <> dt Nil{}:
198	      match fits:
199	
200	def pair_fits(+shape: Shape.Shape, +first: U32, +second: U32, +x: U32, +y: U32,
201	  dimensions: {Shape.dimensions(shape) == [first, second] : List<&2, U32>},
202	  first_inside: {Nat.is_lt(U32.to_nat(x), U32.to_nat(first)) == True{} : Bool},
203	  second_inside: {Nat.is_lt(U32.to_nat(y), U32.to_nat(second)) == True{} : Bool}) -> Fits(Shape.extents(shape), [U32.to_nat(x), U32.to_nat(y)]):
204	  fits_of(Shape.extents(shape), [U32.to_nat(x), U32.to_nat(y)],
205	    %Equal.sym(List<&2, U32>, Shape.dimensions(shape), [first, second], dimensions) :
206	      {True{} == Shape.fit(Shape.naturals(_), [U32.to_nat(x), U32.to_nat(y)]) : Bool}
207	    %Equal.sym(Bool, Nat.is_lt(U32.to_nat(x), U32.to_nat(first)), True{}, first_inside) :
208	      {True{} == Bool.and(_, Bool.and(Nat.is_lt(U32.to_nat(y), U32.to_nat(second)), True{})) : Bool}
209	    %Equal.sym(Bool, Nat.is_lt(U32.to_nat(y), U32.to_nat(second)), True{}, second_inside) :
210	      {True{} == Bool.and(True{}, Bool.and(_, True{})) : Bool}
211	    {==})