1	import Base
2	import ../shape/core.bend as Shape
3	import ./axis_spec.bend as Axis
4	
5	type Error is Data:
6	  ShapeError{error: Shape.Error}
7	  AxesLengthMismatch{}
8	  AxisOutOfBounds{}
9	  DuplicateAxis{}
10	  ElementCountMismatch{}
11	  MappingError{}
12	
13	def contains(+axis: U32, axes: List<&2, U32>) -> Bool:
14	  match axes:
15	    case Nil{}:
16	      False{}
17	    case +head <> tail:
18	      Bool.or(U32.is_eq(axis, head), contains(axis, tail))
19	
20	def unique(axes: List<&2, U32>) -> Bool:
21	  match axes:
22	    case Nil{}:
23	      True{}
24	    case +head <> +tail:
25	      Bool.and(Bool.not(contains(head, tail)), unique(tail))
26	
27	def bounded(axes: List<&2, U32>, +rank: U32) -> Bool:
28	  match axes:
29	    case Nil{}:
30	      True{}
31	    case head <> tail:
32	      Bool.and(U32.is_lt(head, rank), bounded(tail, rank))
33	
34	def bijective(+shape: Shape.Shape, +axes: List<&2, U32>) -> Bool:
35	  Bool.and(Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(shape))),
36	    Bool.and(bounded(axes, Shape.rank(shape)), unique(axes)))
37	
38	def lookup(values: List<&2, U32>, index: Nat) -> Maybe<&2, U32>:
39	  match values index:
40	    case Nil{} _:
41	      None{}
42	    case head <> tail 0n:
43	      Some{head}
44	    case head <> tail 1n+p:
45	      lookup(tail, p)
46	
47	def reorder_tail(head: U32, result: Maybe<&2, List<&2, U32>>) -> Maybe<&2, List<&2, U32>>:
48	  match result:
49	    case None{}:
50	      None{}
51	    case Some{tail}:
52	      Some{head <> tail}
53	
54	def reorder_head(rec: U32 -> Maybe<&2, List<&2, U32>>, head: Maybe<&2, U32>) -> Maybe<&2, List<&2, U32>>:
55	  match head:
56	    case None{}:
57	      None{}
58	    case Some{value}:
59	      rec(value)
60	
61	def reorder(+values: List<&2, U32>, axes: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
62	  match axes:
63	    case Nil{}:
64	      Some{Nil{}}
65	    case axis <> tail:
66	      reorder_head(value => reorder_tail(value, reorder(values, tail)), lookup(values, U32.to_nat(axis)))
67	
68	def coordinate_for(axes: List<&2, U32>, coordinates: List<&2, U32>, +axis: U32) -> Maybe<&2, U32>:
69	  match axes coordinates:
70	    case Nil{} _:
71	      None{}
72	    case _ Nil{}:
73	      None{}
74	    case +head <> tail value <> rest:
75	      Bool.pick(Maybe<&2, U32>, U32.is_eq(head, axis), Some{value}, coordinate_for(tail, rest, axis))
76	
77	def inverse_coordinates(dimensions: List<&2, U32>, +axes: List<&2, U32>,
78	  +coordinates: List<&2, U32>, +axis: U32) -> Maybe<&2, List<&2, U32>>:
79	  match dimensions:
80	    case Nil{}:
81	      Some{Nil{}}
82	    case dimension <> tail:
83	      reorder_head(value => reorder_tail(value, inverse_coordinates(tail, axes, coordinates, U32.inc(axis))),
84	        coordinate_for(axes, coordinates, axis))
85	
86	type Plan<-input: Shape.Shape, -axes: List<&2, U32>> is Data:
87	  Permutation{output: Shape.Shape,
88	    valid: {True{} == bijective(input, axes) : Bool},
89	    ordered: {Some{Shape.dimensions(output)} == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>},
90	    count: {True{} == Nat.is_eq(Shape.elements(input), Shape.elements(output)) : Bool}}
91	
92	def output(-input: Shape.Shape, -axes: List<&2, U32>, plan: Plan<input, axes>) -> Shape.Shape:
93	  Permutation{shape, valid, ordered, count} = plan
94	  shape
95	
96	def plan_counted(+input: Shape.Shape, +axes: List<&2, U32>, +out: Shape.Shape,
97	  valid: {True{} == bijective(input, axes) : Bool},
98	  ordered: {Some{Shape.dimensions(out)} == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>},
99	  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(input), Shape.elements(out))>) ->
100	  Result<&2, &2, Error, Plan<input, axes>>:
101	  match decision:
102	    case Shape.BFalse{equal}:
103	      Fail{ElementCountMismatch{}}
104	    case Shape.BTrue{equal}:
105	      Done{Permutation{out, valid, ordered, equal}}
106	
107	def plan_built(+input: Shape.Shape, +axes: List<&2, U32>, -dimensions: List<&2, U32>,
108	  valid: {True{} == bijective(input, axes) : Bool},
109	  selected: {Some{dimensions} == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>},
110	  built: Shape.ConstructResult<dimensions>) -> Result<&2, &2, Error, Plan<input, axes>>:
111	  match built:
112	    case Shape.SError{error}:
113	      Fail{ShapeError{error}}
114	    case Shape.SBuilt{+out, preserved}:
115	      plan_counted(input, axes, out, valid,
116	        Equal.trans(Maybe<&2, List<&2, U32>>, Some{Shape.dimensions(out)}, Some{dimensions},
117	          reorder(Shape.dimensions(input), axes),
118	          Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, xs => Some{xs}, Shape.dimensions(out), dimensions, preserved), selected),
119	        Shape.bool_case(Nat.is_eq(Shape.elements(input), Shape.elements(out))))
120	
121	def plan_selected(+input: Shape.Shape, +axes: List<&2, U32>, valid: {True{} == bijective(input, axes) : Bool},
122	  selected: Maybe<&2, List<&2, U32>>,
123	  same: {selected == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>}) -> Result<&2, &2, Error, Plan<input, axes>>:
124	  match selected:
125	    case None{}:
126	      Fail{MappingError{}}
127	    case Some{+dimensions}:
128	      plan_built(input, axes, dimensions, valid, same, Shape.construct(dimensions))
129	
130	def plan_decided(+input: Shape.Shape, +axes: List<&2, U32>, decision: Shape.BoolCase<bijective(input, axes)>) ->
131	  Result<&2, &2, Error, Plan<input, axes>>:
132	  match decision:
133	    case Shape.BFalse{equal}:
134	      Fail{MappingError{}}
135	    case Shape.BTrue{equal}:
136	      plan_selected(input, axes, equal, reorder(Shape.dimensions(input), axes), {==})
137	
138	def plan_unique(+input: Shape.Shape, +axes: List<&2, U32>, decision: Shape.BoolCase<unique(axes)>) ->
139	  Result<&2, &2, Error, Plan<input, axes>>:
140	  match decision:
141	    case Shape.BFalse{equal}:
142	      Fail{DuplicateAxis{}}
143	    case Shape.BTrue{equal}:
144	      plan_decided(input, axes, Shape.bool_case(bijective(input, axes)))
145	
146	def plan_bounded(+input: Shape.Shape, +axes: List<&2, U32>, decision: Shape.BoolCase<bounded(axes, Shape.rank(input))>) ->
147	  Result<&2, &2, Error, Plan<input, axes>>:
148	  match decision:
149	    case Shape.BFalse{equal}:
150	      Fail{AxisOutOfBounds{}}
151	    case Shape.BTrue{equal}:
152	      plan_unique(input, axes, Shape.bool_case(unique(axes)))
153	
154	def plan_ranked(+input: Shape.Shape, +axes: List<&2, U32>,
155	  decision: Shape.BoolCase<Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input)))>) ->
156	  Result<&2, &2, Error, Plan<input, axes>>:
157	  match decision:
158	    case Shape.BFalse{equal}:
159	      Fail{AxesLengthMismatch{}}
160	    case Shape.BTrue{equal}:
161	      plan_bounded(input, axes, Shape.bool_case(bounded(axes, Shape.rank(input))))
162	
163	def permutation(+input: Shape.Shape, +axes: List<&2, U32>) -> Result<&2, &2, Error, Plan<input, axes>>:
164	  plan_ranked(input, axes,
165	    Shape.bool_case(Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input)))))
166	
167	def transpose_axes(dimensions: List<&2, U32>, +a: U32, +b: U32, +i: U32) -> List<&2, U32>:
168	  match dimensions:
169	    case Nil{}:
170	      Nil{}
171	    case dimension <> tail:
172	      Bool.pick(U32, U32.is_eq(i, a), b, Bool.pick(U32, U32.is_eq(i, b), a, i)) <>
173	        transpose_axes(tail, a, b, U32.inc(i))
174	
175	def word_same(n: Nat, a: Word(n), b: Word(n)) -> Bool:
176	  match n:
177	    case 0n:
178	      True{}
179	    case 1n+p:
180	      WCon{ab, at} = a
181	      WCon{bb, bt} = b
182	      Bool.and(Bool.not(Bool.xor(ab, bb)), word_same(p, at, bt))
183	
184	def word_value_same(a: U32, b: U32) -> Bool:
185	  match a b:
186	    case U32{x} U32{y}:
187	      word_same(32n, x, y)
188	
189	def words_same(a: List<&2, U32>, b: List<&2, U32>) -> Bool:
190	  match a b:
191	    case Nil{} Nil{}:
192	      True{}
193	    case ah <> at bh <> bt:
194	      Bool.and(word_value_same(ah, bh), words_same(at, bt))
195	    case _ _:
196	      False{}
197	
198	def maybe_same(a: Maybe<&2, List<&2, U32>>, b: Maybe<&2, List<&2, U32>>) -> Bool:
199	  match a b:
200	    case None{} None{}:
201	      True{}
202	    case Some{xs} Some{ys}:
203	      words_same(xs, ys)
204	    case _ _:
205	      False{}
206	
207	type WordsValidStep<-dimension: U32, -dt: List<&2, U32>, -coordinate: U32, -ct: List<&2, U32>> is Data:
208	  WStep{bound: {True{} == U32.is_lt(coordinate, dimension) : Bool}, tail: WordsValid(dt, ct)}
209	
210	def WordsValid(dims: List<&2, U32>, coordinates: List<&2, U32>) -> Data:
211	  match dims coordinates:
212	    case Nil{} Nil{}:
213	      Unit
214	    case dh <> dt ch <> ct:
215	      WordsValidStep<dh, dt, ch, ct>
216	    case _ _:
217	      Empty
218	
219	def words_finish(-dimension: U32, -dt: List<&2, U32>, -coordinate: U32, -ct: List<&2, U32>,
220	  bound: {True{} == U32.is_lt(coordinate, dimension) : Bool}, result: Result<&2, &2, Shape.Error, WordsValid(dt, ct)>) ->
221	  Result<&2, &2, Shape.Error, WordsValid(dimension <> dt, coordinate <> ct)>:
222	  match result:
223	    case Fail{error}:
224	      Fail{error}
225	    case Done{tail}:
226	      Done{WStep{bound, tail}}
227	
228	def words_decided(+dimension: U32, -dt: List<&2, U32>, +coordinate: U32, -ct: List<&2, U32>,
229	  rec: Unit -> Result<&2, &2, Shape.Error, WordsValid(dt, ct)>, decision: Shape.BoolCase<U32.is_lt(coordinate, dimension)>) ->
230	  Result<&2, &2, Shape.Error, WordsValid(dimension <> dt, coordinate <> ct)>:
231	  match decision:
232	    case Shape.BFalse{equal}:
233	      Fail{Shape.CoordinateOutOfBounds{}}
234	    case Shape.BTrue{equal}:
235	      words_finish(dimension, dt, coordinate, ct, equal, rec(Unit{}))
236	
237	def words_check(+dims: List<&2, U32>, +coordinates: List<&2, U32>) -> Result<&2, &2, Shape.Error, WordsValid(dims, coordinates)>:
238	  match dims coordinates:
239	    case Nil{} Nil{}:
240	      Done{Unit{}}
241	    case +dh <> +dt +ch <> +ct:
242	      words_decided(dh, dt, ch, ct, unit => words_check(dt, ct), Shape.bool_case(U32.is_lt(ch, dh)))
243	    case _ _:
244	      Fail{Shape.CoordinateRankMismatch{}}
245	
246	type CheckedCoordinate<-shape: Shape.Shape> is Data:
247	  CWord{values: List<&2, U32>, valid: WordsValid(Shape.dimensions(shape), values)}
248	
249	def coordinate_values(-shape: Shape.Shape, coordinate: CheckedCoordinate<shape>) -> List<&2, U32>:
250	  CWord{values, valid} = coordinate
251	  values
252	
253	type FlatIndex<-shape: Shape.Shape> is Data:
254	  FIndex{value: U32, valid: {True{} == U32.is_lt(value, Shape.count(shape)) : Bool}}
255	
256	def index_value(-shape: Shape.Shape, index: FlatIndex<shape>) -> U32:
257	  FIndex{value, valid} = index
258	  value
259	
260	def flatten_words(dims: List<&2, U32>, coordinates: List<&2, U32>, accumulator: U32) -> U32:
261	  U32.from_nat(Shape.horner(dims, coordinates, U32.to_nat(accumulator)))
262	
263	def reverse_words_go(values: List<&2, U32>, accumulator: List<&2, U32>) -> List<&2, U32>:
264	  match values:
265	    case Nil{}:
266	      accumulator
267	    case head <> tail:
268	      reverse_words_go(tail, head <> accumulator)
269	
270	def reverse_words(values: List<&2, U32>) -> List<&2, U32>:
271	  reverse_words_go(values, Nil{})
272	
273	def unflatten_words(reversed_dimensions: List<&2, U32>, +index: U32) -> List<&2, U32>:
274	  reverse_words(Shape.words(Shape.digits_words(reverse_words(reversed_dimensions), U32.to_nat(index))))
275	
276	def reference_source_coordinates(+dimensions: List<&2, U32>, coordinates: Maybe<&2, List<&2, U32>>) -> Maybe<&2, U32>:
277	  match coordinates:
278	    case None{}:
279	      None{}
280	    case Some{values}:
281	      Some{U32.from_nat(Shape.offset(Shape.naturals(dimensions), Shape.naturals(values)))}
282	
283	def reference_source(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, index: U32) -> Maybe<&2, U32>:
284	  reference_source_coordinates(Shape.dimensions(input),
285	    Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(input)), Shape.naturals(axes),
286	      Shape.words(Shape.digits(Shape.extents(output(input, axes, plan)), U32.to_nat(index))), 0n))
287	
288	type Mapping<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32> is Data:
289	  Mapped{output_index: FlatIndex<output(input, axes, plan)>,
290	    output_coordinate: CheckedCoordinate<output(input, axes, plan)>,
291	    input_coordinate: CheckedCoordinate<input>, index: FlatIndex<input>,
292	    requested: {index_value(output(input, axes, plan), output_index) == i : U32},
293	    decoded: {coordinate_values(output(input, axes, plan), output_coordinate) ==
294	      reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
295	    inverse: {Some{coordinate_values(input, input_coordinate)} ==
296	      inverse_coordinates(Shape.dimensions(input), axes,
297	        coordinate_values(output(input, axes, plan), output_coordinate), 0) : Maybe<&2, List<&2, U32>>},
298	    flattened: {index_value(input, index) ==
299	      flatten_words(Shape.dimensions(input), coordinate_values(input, input_coordinate), 0) : U32},
300	    roundtrip: {True{} == maybe_same(reorder(coordinate_values(input, input_coordinate), axes),
301	      Some{coordinate_values(output(input, axes, plan), output_coordinate)}) : Bool}}
302	
303	def source_index(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32,
304	  mapping: Mapping<input, axes, plan, i>) -> U32:
305	  Mapped{oi, oc, ic, index, requested, decoded, inverse, flattened, roundtrip} = mapping
306	  index_value(input, index)
307	
308	def output_words(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32,
309	  mapping: Mapping<input, axes, plan, i>) -> List<&2, U32>:
310	  Mapped{oi, oc, ic, index, requested, decoded, inverse, flattened, roundtrip} = mapping
311	  coordinate_values(output(input, axes, plan), oc)
312	
313	def input_words(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32,
314	  mapping: Mapping<input, axes, plan, i>) -> List<&2, U32>:
315	  Mapped{oi, oc, ic, index, requested, decoded, inverse, flattened, roundtrip} = mapping
316	  coordinate_values(input, ic)
317	
318	def map_flat(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
319	  oi: FlatIndex<output(input, axes, plan)>, oc: CheckedCoordinate<output(input, axes, plan)>,
320	  +ic: CheckedCoordinate<input>,
321	  requested: {index_value(output(input, axes, plan), oi) == i : U32},
322	  decoded: {coordinate_values(output(input, axes, plan), oc) ==
323	    reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
324	  inverse: {Some{coordinate_values(input, ic)} ==
325	    inverse_coordinates(Shape.dimensions(input), axes, coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
326	  roundtrip: {True{} == maybe_same(reorder(coordinate_values(input, ic), axes),
327	    Some{coordinate_values(output(input, axes, plan), oc)}) : Bool},
328	  +j: U32, flat: {j == flatten_words(Shape.dimensions(input), coordinate_values(input, ic), 0) : U32},
329	  decision: Shape.BoolCase<U32.is_lt(j, Shape.count(input))>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
330	  match decision:
331	    case Shape.BFalse{equal}:
332	      Fail{MappingError{}}
333	    case Shape.BTrue{equal}:
334	      Done{Mapped{oi, oc, ic, FIndex{j, equal}, requested, decoded, inverse, flat, roundtrip}}
335	
336	def map_roundtrip(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
337	  oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
338	  +ic: CheckedCoordinate<input>,
339	  requested: {index_value(output(input, axes, plan), oi) == i : U32},
340	  decoded: {coordinate_values(output(input, axes, plan), oc) ==
341	    reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
342	  inverse: {Some{coordinate_values(input, ic)} ==
343	    inverse_coordinates(Shape.dimensions(input), axes, coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
344	  decision: Shape.BoolCase<maybe_same(reorder(coordinate_values(input, ic), axes),
345	    Some{coordinate_values(output(input, axes, plan), oc)})>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
346	  match decision:
347	    case Shape.BFalse{equal}:
348	      Fail{MappingError{}}
349	    case Shape.BTrue{equal}:
350	      +j = flatten_words(Shape.dimensions(input), coordinate_values(input, ic), 0)
351	      map_flat(input, axes, plan, i, oi, oc, ic, requested, decoded, inverse, equal, j, {==},
352	        Shape.bool_case(U32.is_lt(j, Shape.count(input))))
353	
354	def map_input(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
355	  oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
356	  requested: {index_value(output(input, axes, plan), oi) == i : U32},
357	  decoded: {coordinate_values(output(input, axes, plan), oc) ==
358	    reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
359	  +words: List<&2, U32>,
360	  inverse: {Some{words} == inverse_coordinates(Shape.dimensions(input), axes,
361	    coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
362	  checked: Result<&2, &2, Shape.Error, WordsValid(Shape.dimensions(input), words)>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
363	  match checked:
364	    case Fail{error}:
365	      Fail{ShapeError{error}}
366	    case Done{valid}:
367	      +ic = {CWord{words, valid} : CheckedCoordinate<input>}
368	      map_roundtrip(input, axes, plan, i, oi, oc, ic, requested, decoded, inverse,
369	        Shape.bool_case(maybe_same(reorder(words, axes), Some{coordinate_values(output(input, axes, plan), oc)})))
370	
371	def map_inverse(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
372	  oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
373	  requested: {index_value(output(input, axes, plan), oi) == i : U32},
374	  decoded: {coordinate_values(output(input, axes, plan), oc) ==
375	    reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
376	  result: Maybe<&2, List<&2, U32>>,
377	  same: {result == inverse_coordinates(Shape.dimensions(input), axes,
378	    coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>}) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
379	  match result:
380	    case None{}:
381	      Fail{MappingError{}}
382	    case Some{+words}:
383	      map_input(input, axes, plan, i, oi, oc, requested, decoded, words, same,
384	        words_check(Shape.dimensions(input), words))
385	
386	def map_output(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
387	  oi: FlatIndex<output(input, axes, plan)>,
388	  requested: {index_value(output(input, axes, plan), oi) == i : U32},
389	  +words: List<&2, U32>,
390	  decoded: {words == reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
391	  checked: Result<&2, &2, Shape.Error, WordsValid(Shape.dimensions(output(input, axes, plan)), words)>) ->
392	  Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
393	  match checked:
394	    case Fail{error}:
395	      Fail{ShapeError{error}}
396	    case Done{valid}:
397	      +oc = {CWord{words, valid} : CheckedCoordinate<output(input, axes, plan)>}
398	      map_inverse(input, axes, plan, i, oi, oc, requested, decoded,
399	        inverse_coordinates(Shape.dimensions(input), axes, words, 0), {==})
400	
401	def map_decided(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
402	  decision: Shape.BoolCase<U32.is_lt(i, Shape.count(output(input, axes, plan)))>) ->
403	  Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
404	  match decision:
405	    case Shape.BFalse{equal}:
406	      Fail{ShapeError{Shape.FlatIndexOutOfBounds{}}}
407	    case Shape.BTrue{equal}:
408	      +words = reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i))
409	      map_output(input, axes, plan, i, FIndex{i, equal}, {==}, words, {==},
410	        words_check(Shape.dimensions(output(input, axes, plan)), words))
411	
412	def map_index(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32) ->
413	  Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
414	  map_decided(input, axes, plan, i, Shape.bool_case(U32.is_lt(i, Shape.count(output(input, axes, plan)))))