1	import Base
2	import ../shape/core.bend as Shape
3	import ./axis_spec.bend as Axis
4	
5	def axis_limit(rank: Nat) -> Nat:
6	  match rank:
7	    case 0n:
8	      1n
9	    case 1n+p:
10	      1n+p
11	
12	def axes_valid(+rank: Nat, +a: Nat, +b: Nat) -> Bool:
13	  Bool.and(Nat.is_lt(a, axis_limit(rank)), Nat.is_lt(b, axis_limit(rank)))
14	
15	def valid(+shape: Shape.Shape, +a: U32, +b: U32) -> Bool:
16	  axes_valid(Shape.rank_dimensions(Shape.dimensions(shape)), U32.to_nat(a), U32.to_nat(b))
17	
18	def replace(-A: Data, items: List<&2, A>, index: Nat, +value: A) -> List<&2, A>:
19	  match items index:
20	    case Nil{} index:
21	      Nil{}
22	    case head <> tail 0n:
23	      value <> tail
24	    case head <> tail 1n+p:
25	      head <> replace(A, tail, p, value)
26	
27	def swap_selected(-A: Data, +items: List<&2, A>, +a: Nat, +b: Nat, left: Maybe<&2, A>, right: Maybe<&2, A>) -> List<&2, A>:
28	  match left right:
29	    case Some{x} Some{y}:
30	      replace(A, replace(A, items, a, y), b, x)
31	    case _ _:
32	      items
33	
34	def swap(-A: Data, +items: List<&2, A>, +a: Nat, +b: Nat) -> List<&2, A>:
35	  left right = Axis.lookup(A, items, a) Axis.lookup(A, items, b)
36	  swap_selected(A, items, a, b, left, right)
37	
38	def natural_axes(+rank: Nat, +a: Nat, +b: Nat) -> List<&2, Nat>:
39	  swap(Nat, Axis.indices(rank, 0n), a, b)
40	
41	def requested_axes(+rank: Nat, +a: Nat, +b: Nat) -> List<&2, U32>:
42	  Shape.words(natural_axes(rank, a, b))
43	
44	# The permutation `transpose_copy(input, a, b)` runs: the identity of the
45	# input's rank with axes a and b swapped.
46	def swapped_axes(+input: Shape.Shape, +a: U32, +b: U32) -> List<&2, U32>:
47	  requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b))