1	import Base
2	import ../shape/core.bend as Shape
3	import ./core.bend as Core
4	import ./layout.bend as Layout
5	import ./axis_spec.bend as Axis
6	
7	def dimensions_pair(pair: Core.Dynamic & Shape.Shape) -> Maybe<&2, List<&2, U32>>:
8	  (dynamic, shape) = pair
9	  Some{Shape.dimensions(shape)}
10	
11	def count_pair(pair: Core.Dynamic & U32) -> Maybe<&2, U32>:
12	  (dynamic, count) = pair
13	  Some{count}
14	
15	def dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
16	  match result:
17	    case Fail{error}:
18	      None{}
19	    case Done{tensor}:
20	      dimensions_pair(Core.shape(Core.pack(shape, tensor)))
21	
22	def count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
23	  match result:
24	    case Fail{error}:
25	      None{}
26	    case Done{tensor}:
27	      count_pair(Core.count(Core.pack(shape, tensor)))
28	
29	def point(-shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
30	  match result:
31	    case Fail{error}:
32	      None{}
33	    case Done{tensor}:
34	      Core.read_value(shape, Core.read_flat(shape, tensor, index))
35	
36	def after_value_pair(-shape: Shape.Shape, +other: U32,
37	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Maybe<&2, F32>:
38	  (tensor, result) = pair
39	  Core.read_value(shape, Core.read_flat(shape, tensor, other))
40	
41	def after_value(-shape: Shape.Shape, +index: U32, +other: U32,
42	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
43	  match result:
44	    case Fail{error}:
45	      None{}
46	    case Done{tensor}:
47	      after_value_pair(shape, other, Core.read_flat(shape, tensor, index))
48	
49	def after_length_pair(-shape: Shape.Shape,
50	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Nat:
51	  (tensor, result) = pair
52	  Core.logical_length(shape, tensor)
53	
54	def after_length(-shape: Shape.Shape, +index: U32,
55	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
56	  match result:
57	    case Fail{error}:
58	      None{}
59	    case Done{tensor}:
60	      Some{after_length_pair(shape, Core.read_flat(shape, tensor, index))}
61	
62	def after_dimensions_pair(+shape: Shape.Shape,
63	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Maybe<&2, List<&2, U32>>:
64	  (tensor, result) = pair
65	  dimensions_pair(Core.shape(Core.pack(shape, tensor)))
66	
67	def after_dimensions(+shape: Shape.Shape, +index: U32,
68	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
69	  match result:
70	    case Fail{error}:
71	      None{}
72	    case Done{tensor}:
73	      after_dimensions_pair(shape, Core.read_flat(shape, tensor, index))
74	
75	def after_count_pair(+shape: Shape.Shape,
76	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Maybe<&2, U32>:
77	  (tensor, result) = pair
78	  count_pair(Core.count(Core.pack(shape, tensor)))
79	
80	def after_count(+shape: Shape.Shape, +index: U32,
81	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
82	  match result:
83	    case Fail{error}:
84	      None{}
85	    case Done{tensor}:
86	      after_count_pair(shape, Core.read_flat(shape, tensor, index))
87	
88	def write_done_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Bool:
89	  (tensor, result) = pair
90	  Result.is_done(&2, &2, Core.Error, Unit, result)
91	
92	def write_done(-shape: Shape.Shape, +index: U32, +value: F32,
93	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Bool>:
94	  match result:
95	    case Fail{error}:
96	      None{}
97	    case Done{tensor}:
98	      Some{write_done_pair(shape, Core.write_flat(shape, tensor, index, value))}
99	
100	def written_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) ->
101	  Maybe<&2, F32>:
102	  (tensor, result) = pair
103	  Core.read_value(shape, Core.read_flat(shape, tensor, other))
104	
105	def written_value(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
106	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
107	  match result:
108	    case Fail{error}:
109	      None{}
110	    case Done{tensor}:
111	      written_value_pair(shape, other, Core.write_flat(shape, tensor, index, value))
112	
113	def written_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Nat:
114	  (tensor, result) = pair
115	  Core.logical_length(shape, tensor)
116	
117	def written_length(-shape: Shape.Shape, +index: U32, +value: F32,
118	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
119	  match result:
120	    case Fail{error}:
121	      None{}
122	    case Done{tensor}:
123	      Some{written_length_pair(shape, Core.write_flat(shape, tensor, index, value))}
124	
125	def written_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) ->
126	  Maybe<&2, List<&2, U32>>:
127	  (tensor, result) = pair
128	  dimensions_pair(Core.shape(Core.pack(shape, tensor)))
129	
130	def written_dimensions(+shape: Shape.Shape, +index: U32, +value: F32,
131	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
132	  match result:
133	    case Fail{error}:
134	      None{}
135	    case Done{tensor}:
136	      written_dimensions_pair(shape, Core.write_flat(shape, tensor, index, value))
137	
138	def written_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Maybe<&2, U32>:
139	  (tensor, result) = pair
140	  count_pair(Core.count(Core.pack(shape, tensor)))
141	
142	def written_count(+shape: Shape.Shape, +index: U32, +value: F32,
143	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
144	  match result:
145	    case Fail{error}:
146	      None{}
147	    case Done{tensor}:
148	      written_count_pair(shape, Core.write_flat(shape, tensor, index, value))
149	
150	def read_done_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Bool:
151	  (tensor, result) = pair
152	  Result.is_done(&2, &2, Core.Error, F32, result)
153	
154	def read_done(-shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Bool>:
155	  match result:
156	    case Fail{error}:
157	      None{}
158	    case Done{tensor}:
159	      Some{read_done_pair(shape, Core.read_flat(shape, tensor, index))}
160	
161	def read_status(result: Result<&2, &2, Core.Error, F32>) -> U32:
162	  match result:
163	    case Fail{error}:
164	      Core.error_code(error)
165	    case Done{value}:
166	      0
167	
168	def read_code_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> U32:
169	  (tensor, result) = pair
170	  read_status(result)
171	
172	def read_code(-shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
173	  match result:
174	    case Fail{error}:
175	      None{}
176	    case Done{tensor}:
177	      Some{read_code_pair(shape, Core.read_flat(shape, tensor, index))}
178	
179	def write_status(result: Result<&2, &2, Core.Error, Unit>) -> U32:
180	  match result:
181	    case Fail{error}:
182	      Core.error_code(error)
183	    case Done{unit}:
184	      0
185	
186	def write_code_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> U32:
187	  (tensor, result) = pair
188	  write_status(result)
189	
190	def write_code(-shape: Shape.Shape, +index: U32, +value: F32,
191	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
192	  match result:
193	    case Fail{error}:
194	      None{}
195	    case Done{tensor}:
196	      Some{write_code_pair(shape, Core.write_flat(shape, tensor, index, value))}
197	
198	def coordinate_index(-shape: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>) -> Maybe<&2, U32>:
199	  match result:
200	    case Fail{error}:
201	      None{}
202	    case Done{index}:
203	      Some{Shape.index_value(shape, index)}
204	
205	def coordinate_value_pair(+shape: Shape.Shape, +coordinates: List<&2, U32>,
206	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Maybe<&2, F32>:
207	  (tensor, result) = pair
208	  Core.read_value(shape, Core.read_coordinate(shape, tensor, coordinates))
209	
210	def coordinate_value(+shape: Shape.Shape, +index: U32, +value: F32, +coordinates: List<&2, U32>,
211	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
212	  match result:
213	    case Fail{error}:
214	      None{}
215	    case Done{tensor}:
216	      coordinate_value_pair(shape, coordinates, Core.write_flat(shape, tensor, index, value))
217	
218	def coordinate_written(+shape: Shape.Shape, +coordinates: List<&2, U32>, +value: F32, +other: U32,
219	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
220	  match result:
221	    case Fail{error}:
222	      None{}
223	    case Done{tensor}:
224	      written_value_pair(shape, other, Core.write_coordinate(shape, tensor, coordinates, value))
225	
226	def dynamic_point(+shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
227	  match result:
228	    case Fail{error}:
229	      None{}
230	    case Done{tensor}:
231	      Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, tensor}, index))
232	
233	def clone_first_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
234	  (first, second) = pair
235	  Core.read_value(shape, Core.read_flat(shape, first, other))
236	
237	def clone_second_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
238	  (first, second) = pair
239	  Core.read_value(shape, Core.read_flat(shape, second, other))
240	
241	def clone_first_value(-shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
242	  match result:
243	    case Fail{error}:
244	      None{}
245	    case Done{tensor}:
246	      clone_first_value_pair(shape, other, Core.clone(shape, tensor))
247	
248	def clone_second_value(-shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
249	  match result:
250	    case Fail{error}:
251	      None{}
252	    case Done{tensor}:
253	      clone_second_value_pair(shape, other, Core.clone(shape, tensor))
254	
255	def clone_first_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Nat:
256	  (first, second) = pair
257	  Core.logical_length(shape, first)
258	
259	def clone_second_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Nat:
260	  (first, second) = pair
261	  Core.logical_length(shape, second)
262	
263	def clone_first_length(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
264	  match result:
265	    case Fail{error}:
266	      None{}
267	    case Done{tensor}:
268	      Some{clone_first_length_pair(shape, Core.clone(shape, tensor))}
269	
270	def clone_second_length(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
271	  match result:
272	    case Fail{error}:
273	      None{}
274	    case Done{tensor}:
275	      Some{clone_second_length_pair(shape, Core.clone(shape, tensor))}
276	
277	def clone_first_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, List<&2, U32>>:
278	  (first, second) = pair
279	  dimensions_pair(Core.shape(Core.pack(shape, first)))
280	
281	def clone_second_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, List<&2, U32>>:
282	  (first, second) = pair
283	  dimensions_pair(Core.shape(Core.pack(shape, second)))
284	
285	def clone_first_dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
286	  match result:
287	    case Fail{error}:
288	      None{}
289	    case Done{tensor}:
290	      clone_first_dimensions_pair(shape, Core.clone(shape, tensor))
291	
292	def clone_second_dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
293	  match result:
294	    case Fail{error}:
295	      None{}
296	    case Done{tensor}:
297	      clone_second_dimensions_pair(shape, Core.clone(shape, tensor))
298	
299	def clone_first_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, U32>:
300	  (first, second) = pair
301	  count_pair(Core.count(Core.pack(shape, first)))
302	
303	def clone_second_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, U32>:
304	  (first, second) = pair
305	  count_pair(Core.count(Core.pack(shape, second)))
306	
307	def clone_first_count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
308	  match result:
309	    case Fail{error}:
310	      None{}
311	    case Done{tensor}:
312	      clone_first_count_pair(shape, Core.clone(shape, tensor))
313	
314	def clone_second_count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
315	  match result:
316	    case Fail{error}:
317	      None{}
318	    case Done{tensor}:
319	      clone_second_count_pair(shape, Core.clone(shape, tensor))
320	
321	def clone_written_first_pair(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
322	  pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
323	  (first, second) = pair
324	  written_value_pair(shape, other, Core.write_flat(shape, first, index, value))
325	
326	def clone_written_first(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
327	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
328	  match result:
329	    case Fail{error}:
330	      None{}
331	    case Done{tensor}:
332	      clone_written_first_pair(shape, index, value, other, Core.clone(shape, tensor))
333	
334	def untouched(-shape: Shape.Shape, +other: U32, written: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>,
335	  second: Core.Tensor<shape>) -> Maybe<&2, F32>:
336	  (tensor, result) = written
337	  Core.read_value(shape, Core.read_flat(shape, second, other))
338	
339	def clone_written_second_pair(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
340	  pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
341	  (first, second) = pair
342	  untouched(shape, other, Core.write_flat(shape, first, index, value), second)
343	
344	def clone_written_second(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
345	  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
346	  match result:
347	    case Fail{error}:
348	      None{}
349	    case Done{tensor}:
350	      clone_written_second_pair(shape, index, value, other, Core.clone(shape, tensor))
351	
352	def length_number(-shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Nat:
353	  (tensor, n) = pair
354	  n
355	
356	def length_value(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
357	  match result:
358	    case Fail{error}:
359	      None{}
360	    case Done{tensor}:
361	      Some{length_number(shape, Core.length(shape, tensor))}
362	
363	def length_owner_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Nat) -> Maybe<&2, F32>:
364	  (tensor, n) = pair
365	  Core.read_value(shape, Core.read_flat(shape, tensor, other))
366	
367	def length_owner_value(-shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
368	  match result:
369	    case Fail{error}:
370	      None{}
371	    case Done{tensor}:
372	      length_owner_value_pair(shape, other, Core.length(shape, tensor))
373	
374	def length_owner_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Nat:
375	  (tensor, n) = pair
376	  Core.logical_length(shape, tensor)
377	
378	def length_owner_length(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
379	  match result:
380	    case Fail{error}:
381	      None{}
382	    case Done{tensor}:
383	      Some{length_owner_length_pair(shape, Core.length(shape, tensor))}
384	
385	def length_owner_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Maybe<&2, List<&2, U32>>:
386	  (tensor, n) = pair
387	  dimensions_pair(Core.shape(Core.pack(shape, tensor)))
388	
389	def length_owner_dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
390	  match result:
391	    case Fail{error}:
392	      None{}
393	    case Done{tensor}:
394	      length_owner_dimensions_pair(shape, Core.length(shape, tensor))
395	
396	def length_owner_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Maybe<&2, U32>:
397	  (tensor, n) = pair
398	  count_pair(Core.count(Core.pack(shape, tensor)))
399	
400	def length_owner_count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
401	  match result:
402	    case Fail{error}:
403	      None{}
404	    case Done{tensor}:
405	      length_owner_count_pair(shape, Core.length(shape, tensor))
406	
407	def dynamic_owner_value(+other: U32, dynamic: Core.Dynamic) -> Maybe<&2, F32>:
408	  Core.dynamic_read_value(Core.read_dynamic(dynamic, other))
409	
410	def shape_owner_pair(+other: U32, pair: Core.Dynamic & Shape.Shape) -> Maybe<&2, F32>:
411	  (dynamic, shape) = pair
412	  dynamic_owner_value(other, dynamic)
413	
414	def shape_owner_value(+shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
415	  match result:
416	    case Fail{error}:
417	      None{}
418	    case Done{tensor}:
419	      shape_owner_pair(other, Core.shape(Core.pack(shape, tensor)))
420	
421	def count_owner_pair(+other: U32, pair: Core.Dynamic & U32) -> Maybe<&2, F32>:
422	  (dynamic, count) = pair
423	  dynamic_owner_value(other, dynamic)
424	
425	def count_owner_value(+shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
426	  match result:
427	    case Fail{error}:
428	      None{}
429	    case Done{tensor}:
430	      count_owner_pair(other, Core.count(Core.pack(shape, tensor)))
431	
432	def listed(values: List<&2, F32>, found: Maybe<&2, U32>) -> Maybe<&2, F32>:
433	  match found:
434	    case None{}:
435	      None{}
436	    case Some{index}:
437	      List.get(&2, F32, values, U32.to_nat(index))
438	
439	def permute_source(+input: Shape.Shape, +axes: List<&2, U32>, +index: U32) -> Maybe<&2, U32>:
440	  Layout.reference_source_coordinates(Shape.dimensions(input),
441	    Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(input)), Shape.naturals(axes),
442	      Shape.words(Shape.digits(Shape.naturals(Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))), U32.to_nat(index))), 0n))
443	
444	# The row-major source index that `permute_source` names for output index
445	# `index`, as a natural (0 where it names none).
446	def permute_source_index(+input: Shape.Shape, +axes: List<&2, U32>, +index: Nat) -> Nat:
447	  U32.to_nat(Maybe.default(&2, U32, permute_source(input, axes, U32.from_nat(index)), 0))