1	# Least-significant axis first. Singleton axes are omitted. The two innermost
2	# axes are flat levels, so an element step changes registers only; outer axes
3	# stay in a list that is rebuilt only when every level wraps. Every step is a
4	# direct template or a tail call, so native code keeps the advance in one loop.
5	import Base
6	import ./tensor_view.bend as Views
7	import ./traversal_array.bend as Traversal
8	
9	type Axis is Data:
10	  Axis{extent: Nat,remaining: Nat,stride: U32,rewind: U32}
11	
12	def include_axis(active: Bool,+extent: U32,+stride: U32,axes: List<&2,Axis>) -> List<&2,Axis>:
13	  match active:
14	    case False{}: axes
15	    case True{}:
16	      +count = U32.to_nat(extent)
17	      Axis{count,count,stride,((extent - 1) * stride : U32)} <> axes
18	
19	def axes(source: List<&2,Views.Axis>,reversed: List<&2,Axis>) -> List<&2,Axis>:
20	  match source:
21	    case []: reversed
22	    case Views.Axis{+extent,stride} <> tail:
23	      axes(tail,include_axis(U32.is_lt(1,extent),extent,stride,reversed))
24	
25	# A level with no remaining count never moves; it stands for an absent axis.
26	type Level<-Rest: Data> is Data:
27	  Level{axis: Axis,rest: Rest}
28	
29	def absent() -> Axis:
30	  Axis{0n,0n,0,0}
31	
32	# Outer axes: reset axes are kept reversed and restored after the carry stops.
33	def restore(reset: List<&2,Axis>,axes: List<&2,Axis>) -> List<&2,Axis>:
34	  match reset:
35	    case []: axes
36	    case axis <> tail: restore(tail,axis <> axes)
37	
38	def carry(axes: List<&2,Axis>,address: U32,reset: List<&2,Axis>) -> List<&2,Axis> & U32:
39	  match axes:
40	    case []: (restore(reset,[]),address)
41	    case Axis{extent,0n,stride,rewind} <> tail: (restore(reset,Axis{extent,0n,stride,rewind} <> tail),address)
42	    case Axis{+extent,1n,stride,+rewind} <> tail: carry(tail,(address - rewind : U32),Axis{extent,extent,stride,rewind} <> reset)
43	    case Axis{extent,1n+1n+remaining,+stride,rewind} <> tail:
44	      (restore(reset,Axis{extent,1n+remaining,stride,rewind} <> tail),(address + stride : U32))
45	
46	def advance_outer(axes: List<&2,Axis>,address: U32) -> List<&2,Axis> & U32:
47	  carry(axes,address,[])
48	
49	def reset(~Rest: Data,+extent: Nat,stride: U32,rewind: U32,result: Rest & U32) -> Level<Rest> & U32:
50	  (rest,address) = result
51	  (Level{Axis{extent,extent,stride,rewind},rest},address)
52	
53	def step(~Rest: Data,~next: Rest -> U32 -> Rest & U32,level: Level<Rest>,address: U32) -> Level<Rest> & U32:
54	  Level{axis,rest} = level
55	  match axis:
56	    case Axis{extent,0n,stride,rewind}: (Level{Axis{extent,0n,stride,rewind},rest},address)
57	    case Axis{+extent,1n,stride,+rewind}: reset(~Rest,extent,stride,rewind,next(rest,(address - rewind : U32)))
58	    case Axis{extent,1n+1n+remaining,+stride,rewind}: (Level{Axis{extent,1n+remaining,stride,rewind},rest},(address + stride : U32))
59	
60	def Depth0() -> Data:
61	  List<&2,Axis>
62	
63	def Depth1() -> Data:
64	  Level<Depth0()>
65	
66	def Levels() -> Data:
67	  Level<Depth1()>
68	
69	def advance_levels(axes: Levels(),address: U32) -> Levels() & U32:
70	  step(~Depth1(),~step(~Depth0(),~advance_outer),axes,address)
71	
72	type Cursor is Data:
73	  Cursor{axes: Levels(),address: U32}
74	
75	def moved(result: Levels() & U32) -> Cursor:
76	  (axes,address) = result
77	  Cursor{axes,address}
78	
79	def advance(axes: Levels(),address: U32) -> Cursor:
80	  moved(advance_levels(axes,address))
81	
82	def split_level(~Rest: Data,~split_rest: List<&2,Axis> -> Rest,axes: List<&2,Axis>) -> Level<Rest>:
83	  match axes:
84	    case []: Level{absent(),split_rest([])}
85	    case axis <> tail: Level{axis,split_rest(tail)}
86	
87	def outer(axes: List<&2,Axis>) -> List<&2,Axis>:
88	  axes
89	
90	def split(axes: List<&2,Axis>) -> Levels():
91	  split_level(~Depth1(),~split_level(~Depth0(),~outer),axes)
92	
93	def prepare(view: Views.View) -> Cursor:
94	  Views.View{source,offset} = view
95	  Cursor{split(axes(source,[])),offset}
96	
97	def read_value(~Element: Data,axes: Levels(),address: U32,result: Array<Element> & Element) -> Array<Element> & Traversal.Visited<Cursor,Element>:
98	  (storage,value) = result
99	  (storage,Traversal.Visited{advance(axes,address),value})
100	
101	def read(~Element: Data,storage: Array<Element>,cursor: Cursor) -> Array<Element> & Traversal.Visited<Cursor,Element>:
102	  Cursor{axes,+address} = cursor
103	  read_value(~Element,axes,address,Array.get(Element,storage,address))
104	
105	# One run of the innermost axis reads address, address + stride, ... with no
106	# countdown. Its two words stay in registers in whatever program the cursor is
107	# compiled into; the full cursor is advanced once, at the end of the run.
108	type Stride is Data:
109	  Stride{address: U32,stride: U32}
110	
111	def stride_value(~Element: Data,address: U32,+stride: U32,result: Array<Element> & Element) -> Array<Element> & Traversal.Visited<Stride,Element>:
112	  (storage,value) = result
113	  (storage,Traversal.Visited{Stride{(address + stride : U32),stride},value})
114	
115	def read_stride(~Element: Data,storage: Array<Element>,cursor: Stride) -> Array<Element> & Traversal.Visited<Stride,Element>:
116	  Stride{+address,+stride} = cursor
117	  stride_value(~Element,address,stride,Array.get(Element,storage,address))
118	
119	# The address a scatter stores at, with the cursor of the next element.
120	def target(cursor: Cursor) -> Traversal.Target<Cursor>:
121	  Cursor{axes,+address} = cursor
122	  Traversal.Target{advance(axes,address),address}
123	
124	def target_stride(cursor: Stride) -> Traversal.Target<Stride>:
125	  Stride{+address,+stride} = cursor
126	  Traversal.Target{Stride{(address + stride : U32),stride},address}
127	
128	type CursorPair is Data:
129	  CursorPair{left: Cursor,right: Cursor}