1	import Base
2	import ./core.bend as Shape
3	
4	def add_assoc(a: Nat, -b: Nat, -c: Nat) -> {Nat.add(a, Nat.add(b, c)) == Nat.add(Nat.add(a, b), c) : Nat}:
5	  match a:
6	    case 0n:
7	      {==}
8	    case 1n+p:
9	      %add_assoc(p, b, c) : {1n+Nat.add(p, Nat.add(b, c)) == 1n+_ : Nat}
10	      {==}
11	
12	def split_left(+a: Nat, +b: Nat, +index: Shape.Fin(a)) -> {Shape.Fin.split(a, b, Shape.Fin.left(a, b, index)) == Inl{index} : Either<Shape.Fin(a), Shape.Fin(b)>}:
13	  match a:
14	    case 0n:
15	      match index:
16	    case 1n+p:
17	      match index:
18	        case Shape.FZero{}:
19	          {==}
20	        case Shape.FNext{+previous}:
21	          %split_left(p, b, previous) : {Shape.Fin.split.step(p, b, Shape.Fin.split(p, b, Shape.Fin.left(p, b, previous))) == Shape.Fin.split.step(p, b, _) : Either<Shape.Fin(1n+p), Shape.Fin(b)>}
22	          {==}
23	
24	def split_shift(+a: Nat, +b: Nat, +index: Shape.Fin(b)) -> {Shape.Fin.split(a, b, Shape.Fin.shift(a, b, index)) == Inr{index} : Either<Shape.Fin(a), Shape.Fin(b)>}:
25	  match a:
26	    case 0n:
27	      {==}
28	    case 1n+p:
29	      %split_shift(p, b, index) : {Shape.Fin.split.step(p, b, Shape.Fin.split(p, b, Shape.Fin.shift(p, b, index))) == Shape.Fin.split.step(p, b, _) : Either<Shape.Fin(1n+p), Shape.Fin(b)>}
30	      {==}
31	
32	def merge_step(+a: Nat, +b: Nat, split: Either<Shape.Fin(a), Shape.Fin(b)>) -> {Shape.Fin.merge(1n+a, b, Shape.Fin.split.step(a, b, split)) == Shape.FNext{Shape.Fin.merge(a, b, split)} : Shape.Fin(1n+Nat.add(a, b))}:
33	  match split:
34	    case Inl{index}:
35	      {==}
36	    case Inr{index}:
37	      {==}
38	
39	def merge_split(+a: Nat, +b: Nat, +index: Shape.Fin(Nat.add(a, b))) -> {Shape.Fin.merge(a, b, Shape.Fin.split(a, b, index)) == index : Shape.Fin(Nat.add(a, b))}:
40	  match a:
41	    case 0n:
42	      {==}
43	    case 1n+p:
44	      match index:
45	        case Shape.FZero{}:
46	          {==}
47	        case Shape.FNext{+previous}:
48	          %merge_split(p, b, previous) : {Shape.Fin.merge(1n+p, b, Shape.Fin.split.step(p, b, Shape.Fin.split(p, b, previous))) == Shape.FNext{_} : Shape.Fin(1n+Nat.add(p, b))}
49	          merge_step(p, b, Shape.Fin.split(p, b, previous))
50	
51	def left_value(+a: Nat, +b: Nat, +index: Shape.Fin(a)) -> {Shape.Fin.value(Nat.add(a, b), Shape.Fin.left(a, b, index)) == Shape.Fin.value(a, index) : Nat}:
52	  match a:
53	    case 0n:
54	      match index:
55	    case 1n+p:
56	      match index:
57	        case Shape.FZero{}:
58	          {==}
59	        case Shape.FNext{+previous}:
60	          %left_value(p, b, previous) : {1n+Shape.Fin.value(Nat.add(p, b), Shape.Fin.left(p, b, previous)) == 1n+_ : Nat}
61	          {==}
62	
63	def shift_value(+a: Nat, +b: Nat, +index: Shape.Fin(b)) -> {Shape.Fin.value(Nat.add(a, b), Shape.Fin.shift(a, b, index)) == Nat.add(a, Shape.Fin.value(b, index)) : Nat}:
64	  match a:
65	    case 0n:
66	      {==}
67	    case 1n+p:
68	      %shift_value(p, b, index) : {1n+Shape.Fin.value(Nat.add(p, b), Shape.Fin.shift(p, b, index)) == 1n+_ : Nat}
69	      {==}
70	
71	def pack_value(+a: Nat, +b: Nat, +first: Shape.Fin(a), +last: Shape.Fin(b)) -> {Shape.Fin.value(Nat.mul(a, b), Shape.Fin.pack(a, b, (first, last))) == Nat.add(Nat.mul(Shape.Fin.value(a, first), b), Shape.Fin.value(b, last)) : Nat}:
72	  match a:
73	    case 0n:
74	      match first:
75	    case 1n+p:
76	      match first:
77	        case Shape.FZero{}:
78	          left_value(b, Nat.mul(p, b), last)
79	        case Shape.FNext{+previous}:
80	          %add_assoc(b, Nat.mul(Shape.Fin.value(p, previous), b), Shape.Fin.value(b, last)) : {Shape.Fin.value(Nat.add(b, Nat.mul(p, b)), Shape.Fin.shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last)))) == _ : Nat}
81	          %pack_value(p, b, previous, last) : {Shape.Fin.value(Nat.add(b, Nat.mul(p, b)), Shape.Fin.shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last)))) == Nat.add(b, _) : Nat}
82	          shift_value(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last)))
83	
84	def unpack_pack(+a: Nat, +b: Nat, +first: Shape.Fin(a), +last: Shape.Fin(b)) -> {Shape.Fin.unpack(a, b, Shape.Fin.pack(a, b, (first, last))) == (first, last) : Shape.Fin(a) & Shape.Fin(b)}:
85	  match a:
86	    case 0n:
87	      match first:
88	    case 1n+p:
89	      match first:
90	        case Shape.FZero{}:
91	          %split_left(b, Nat.mul(p, b), last) : {Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), Shape.Fin.split(b, Nat.mul(p, b), Shape.Fin.left(b, Nat.mul(p, b), last))) == Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), _) : Shape.Fin(1n+p) & Shape.Fin(b)}
92	          {==}
93	        case Shape.FNext{+previous}:
94	          %unpack_pack(p, b, previous, last) : {Shape.Fin.unpack(1n+p, b, Shape.Fin.pack(1n+p, b, (Shape.FNext{previous}, last))) == Shape.Fin.unpack.lift(p, b, _) : Shape.Fin(1n+p) & Shape.Fin(b)}
95	          %split_shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last))) : {Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), Shape.Fin.split(b, Nat.mul(p, b), Shape.Fin.shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last))))) == Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), _) : Shape.Fin(1n+p) & Shape.Fin(b)}
96	          {==}
97	
98	def pack_lift(+a: Nat, +b: Nat, coordinate: Shape.Fin(a) & Shape.Fin(b)) -> {Shape.Fin.pack(1n+a, b, Shape.Fin.unpack.lift(a, b, coordinate)) == Shape.Fin.shift(b, Nat.mul(a, b), Shape.Fin.pack(a, b, coordinate)) : Shape.Fin(Nat.add(b, Nat.mul(a, b)))}:
99	  (first, last) = coordinate
100	  {==}
101	
102	def pack_unpack_step(+a: Nat, +b: Nat, rec: @+index: Shape.Fin(Nat.mul(a, b)) -> {Shape.Fin.pack(a, b, Shape.Fin.unpack(a, b, index)) == index : Shape.Fin(Nat.mul(a, b))}, split: Either<Shape.Fin(b), Shape.Fin(Nat.mul(a, b))>) -> {Shape.Fin.pack(1n+a, b, Shape.Fin.unpack.step(a, b, index => Shape.Fin.unpack(a, b, index), split)) == Shape.Fin.merge(b, Nat.mul(a, b), split) : Shape.Fin(Nat.add(b, Nat.mul(a, b)))}:
103	  match split:
104	    case Inl{last}:
105	      {==}
106	    case Inr{+index}:
107	      %rec(index) : {Shape.Fin.pack(1n+a, b, Shape.Fin.unpack.lift(a, b, Shape.Fin.unpack(a, b, index))) == Shape.Fin.shift(b, Nat.mul(a, b), _) : Shape.Fin(Nat.add(b, Nat.mul(a, b)))}
108	      pack_lift(a, b, Shape.Fin.unpack(a, b, index))
109	
110	def pack_unpack(+a: Nat, +b: Nat, +index: Shape.Fin(Nat.mul(a, b))) -> {Shape.Fin.pack(a, b, Shape.Fin.unpack(a, b, index)) == index : Shape.Fin(Nat.mul(a, b))}:
111	  match a:
112	    case 0n:
113	      match index:
114	    case 1n+p:
115	      %merge_split(b, Nat.mul(p, b), index) : {Shape.Fin.pack(1n+p, b, Shape.Fin.unpack(1n+p, b, index)) == _ : Shape.Fin(Nat.add(b, Nat.mul(p, b)))}
116	      pack_unpack_step(p, b, +index => pack_unpack(p, b, index), Shape.Fin.split(b, Nat.mul(p, b), index))
117	
118	def flatten_unflatten_step(+head: Nat, +tail: List<&2, Nat>, rec: @+index: Shape.Fin(Shape.numel(tail)) -> {Shape.flatten(tail, Shape.unflatten(tail, index)) == index : Shape.Fin(Shape.numel(tail))}, coordinate: Shape.Fin(head) & Shape.Fin(Shape.numel(tail))) -> {Shape.flatten(head <> tail, Shape.unflatten.step(head, tail, rest => Shape.unflatten(tail, rest), coordinate)) == Shape.Fin.pack(head, Shape.numel(tail), coordinate) : Shape.Fin(Shape.numel(head <> tail))}:
119	  (first, +last) = coordinate
120	  %rec(last) : {Shape.Fin.pack(head, Shape.numel(tail), (first, Shape.flatten(tail, Shape.unflatten(tail, last)))) == Shape.Fin.pack(head, Shape.numel(tail), (first, _)) : Shape.Fin(Shape.numel(head <> tail))}
121	  {==}
122	
123	def flatten_unflatten(+dims: List<&2, Nat>, +index: Shape.Fin(Shape.numel(dims))) -> {Shape.flatten(dims, Shape.unflatten(dims, index)) == index : Shape.Fin(Shape.numel(dims))}:
124	  match dims:
125	    case Nil{}:
126	      match index:
127	        case Shape.FZero{}:
128	          {==}
129	        case Shape.FNext{previous}:
130	          match previous:
131	    case +head <> +tail:
132	      %pack_unpack(head, Shape.numel(tail), index) : {Shape.flatten(head <> tail, Shape.unflatten(head <> tail, index)) == _ : Shape.Fin(Shape.numel(head <> tail))}
133	      flatten_unflatten_step(head, tail, +last => flatten_unflatten(tail, last), Shape.Fin.unpack(head, Shape.numel(tail), index))
134	
135	def unflatten_flatten(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {Shape.unflatten(dims, Shape.flatten(dims, coordinate)) == coordinate : Shape.Coordinate(dims)}:
136	  match dims:
137	    case Nil{}:
138	      match coordinate:
139	        case Unit{}:
140	          {==}
141	    case +head <> +tail:
142	      Shape.CStep{+first, +rest} = coordinate
143	      %unflatten_flatten(tail, rest) : {Shape.unflatten(head <> tail, Shape.flatten(head <> tail, Shape.CStep{first, rest})) == Shape.CStep{first, _} : Shape.Coordinate(head <> tail)}
144	      %unpack_pack(head, Shape.numel(tail), first, Shape.flatten(tail, rest)) : {Shape.unflatten.step(head, tail, last => Shape.unflatten(tail, last), Shape.Fin.unpack(head, Shape.numel(tail), Shape.Fin.pack(head, Shape.numel(tail), (first, Shape.flatten(tail, rest))))) == Shape.unflatten.step(head, tail, last => Shape.unflatten(tail, last), _) : Shape.Coordinate(head <> tail)}
145	      {==}
146	
147	def flatten_row_major(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {Shape.Fin.value(Shape.numel(dims), Shape.flatten(dims, coordinate)) == Shape.row_major(dims, coordinate) : Nat}:
148	  match dims:
149	    case Nil{}:
150	      {==}
151	    case +head <> +tail:
152	      Shape.CStep{+first, +rest} = coordinate
153	      %flatten_row_major(tail, rest) : {Shape.Fin.value(Shape.numel(head <> tail), Shape.flatten(head <> tail, Shape.CStep{first, rest})) == Nat.add(Nat.mul(Shape.Fin.value(head, first), Shape.numel(tail)), _) : Nat}
154	      pack_value(head, Shape.numel(tail), first, Shape.flatten(tail, rest))