1	# A view of the storage of a geometry, seen from its bands. A view that keeps
2	# the images and the rows of the geometry in place has axes before the rows,
3	# repeats of each row and axes after the rows, each moving by whole channels
4	# and columns; every band reads its cells through the same axes over its own
5	# rows.
6	import Base
7	import ./tensor_view.bend as Views
8	import ./tensor_shape.bend as Shapes
9	import ./tensor_band.bend as Bands
10	
11	# The axes before the rows, the repeats of each row, the axes after the rows,
12	# and the offset of the view, as the view gives them.
13	type Banded is Data:
14	  Banded{before: List<&2,Views.Axis>,repeats: List<&2,U32>,after: List<&2,Views.Axis>,offset: U32}
15	
16	# ---- The view of a band ----
17	
18	# The channel of plane words a word of the storage of a geometry lies in, or
19	# its column within that channel.
20	def component(+plane: U32,column: Bool,+word: U32) -> U32:
21	  match column:
22	    case False{}: (word / plane : U32)
23	    case True{}: (word % plane : U32)
24	
25	# A word of the storage of the geometry, a channel of plane words and a column
26	# within it, as the same channel and column in a band of band words per
27	# channel.
28	def localized(+plane: U32,+band: U32,+word: U32) -> U32:
29	  (component(plane,False{},word) * band + component(plane,True{},word) : U32)
30	
31	def localized_axes(+plane: U32,+band: U32,axes: List<&2,Views.Axis>) -> List<&2,Views.Axis>:
32	  match axes:
33	    case []: []
34	    case Views.Axis{extent,stride} <> rest: Views.Axis{extent,localized(plane,band,stride)} <> localized_axes(plane,band,rest)
35	
36	def repeated(repeats: List<&2,U32>) -> List<&2,Views.Axis>:
37	  match repeats:
38	    case []: []
39	    case extent <> rest: Views.Axis{extent,0} <> repeated(rest)
40	
41	# The rows axis, which a view of a geometry of one row per image leaves out,
42	# and the repeats of each row.
43	def rows_axes(rowed: Bool,+rows: U32,+width: U32,repeats: List<&2,U32>) -> List<&2,Views.Axis>:
44	  match rowed:
45	    case True{}: Views.Axis{rows,width} <> repeated(repeats)
46	    case False{}: repeated(repeats)
47	
48	# The axes of the view over rows rows of each channel. A band of a geometry
49	# of one row per image has at most one row.
50	def axes(+banded: Banded,+geometry: Bands.Geometry,+rows: U32) -> List<&2,Views.Axis>:
51	  Banded{before,repeats,after,offset} = banded
52	  Bands.Geometry{batch,channels,+height,+width} = geometry
53	  +plane = (height * width : U32)
54	  +band = (rows * width : U32)
55	  List.append(&2,Views.Axis,localized_axes(plane,band,before),List.append(&2,Views.Axis,rows_axes(True{},rows,width,repeats),localized_axes(plane,band,after)))
56	
57	# The view of the local storage of a band of rows rows.
58	def local(+banded: Banded,+geometry: Bands.Geometry,+rows: U32) -> Views.View:
59	  Banded{before,repeats,after,+offset} = banded
60	  Bands.Geometry{batch,channels,+height,+width} = geometry
61	  Views.View{axes(banded,geometry,rows),localized((height * width : U32),(rows * width : U32),offset)}
62	
63	# The extents of the view over rows rows.
64	def extents(+banded: Banded,+rows: U32) -> List<&2,U32>:
65	  Banded{before,repeats,after,offset} = banded
66	  List.append(&2,U32,Views.axis_extents(before),List.append(&2,U32,rows <> repeats,Views.axis_extents(after)))
67	
68	def known(count: Maybe<&2,U32>) -> U32:
69	  match count:
70	    case Some{count}: count
71	    case None{}: 0
72	
73	def product_or_zero(extents: List<&2,U32>) -> U32:
74	  known(Shapes.product(extents,Some{1}))
75	
76	# The rows of the output for each row of the input.
77	def scale(+banded: Banded) -> U32:
78	  Banded{before,repeats,after,offset} = banded
79	  product_or_zero(repeats)
80	
81	# The geometry of the output: the same images, the axes before the rows as
82	# channels, each row repeated, the axes after the rows as columns.
83	def output(+banded: Banded,+geometry: Bands.Geometry) -> Bands.Geometry:
84	  Banded{before,repeats,after,offset} = banded
85	  Bands.Geometry{batch,channels,height,width} = geometry
86	  Bands.Geometry{batch,product_or_zero(Views.axis_extents(before)),(height * product_or_zero(repeats) : U32),product_or_zero(Views.axis_extents(after))}
87	
88	# ---- Directions ----
89	
90	def count_when(present: Bool) -> Nat:
91	  match present:
92	    case True{}: 1n
93	    case False{}: 0n
94	
95	def local_when(valid: Bool,flags: List<&2,Bool>) -> Maybe<&2,List<&2,Bool>>:
96	  match valid:
97	    case True{}: Some{flags}
98	    case False{}: None{}
99	
100	# The flags of the axes of a local view: those of the axes before the rows,
101	# the rows axis and its repeats forward, those of the axes after the rows.
102	def local_directions(+before: List<&2,Views.Axis>,+repeats: List<&2,U32>,+after: List<&2,Views.Axis>,+before_flags: List<&2,Bool>,+after_flags: List<&2,Bool>) -> List<&2,Bool>:
103	  List.append(&2,Bool,Views.fitted_flags(List.length(&2,Views.Axis,before),before_flags),
104	    List.append(&2,Bool,Views.fitted_flags(1n+List.length(&2,U32,repeats),[]),Views.fitted_flags(List.length(&2,Views.Axis,after),after_flags)))
105	
106	# The flags of the local view of a band, when the view reads its images axis,
107	# its rows axis and the repeats of each row forward. An axis without a flag is
108	# read forward.
109	def local_flags(+banded: Banded,+geometry: Bands.Geometry,+flags: List<&2,Bool>) -> Maybe<&2,List<&2,Bool>>:
110	  Banded{+before,+repeats,+after,offset} = banded
111	  Bands.Geometry{+batch,channels,+height,width} = geometry
112	  +images = count_when(U32.is_lt(1,batch))
113	  +rows = Nat.add(count_when(U32.is_lt(1,height)),List.length(&2,U32,repeats))
114	  +later = Views.drop_flags(Nat.add(images,List.length(&2,Views.Axis,before)),flags)
115	  local_when(Views.forward(images,flags) && Views.forward(rows,later),local_directions(before,repeats,after,Views.drop_flags(images,flags),Views.drop_flags(rows,later)))
116	
117	# ---- Classification ----
118	
119	# The axes after the images axis, which a view of more than one image starts
120	# with.
121	def images_kept(kept: Bool,rest: List<&2,Views.Axis>) -> Maybe<&2,List<&2,Views.Axis>>:
122	  match kept:
123	    case True{}: Some{rest}
124	    case False{}: None{}
125	
126	def after_images(many: Bool,+batch: U32,+image: U32,axes: List<&2,Views.Axis>) -> Maybe<&2,List<&2,Views.Axis>>:
127	  match many axes:
128	    case False{} axes: Some{axes}
129	    case True{} Nil{}: None{}
130	    case True{} Views.Axis{+extent,+stride} <> rest: images_kept(U32.is_eq(extent,batch) && U32.is_eq(stride,image),rest)
131	
132	# The axes before the rows axis and those after it.
133	type Split is Data:
134	  Split{before: List<&2,Views.Axis>,rest: List<&2,Views.Axis>}
135	
136	def prepend_before(axis: Views.Axis,split: Maybe<&2,Split>) -> Maybe<&2,Split>:
137	  match split:
138	    case Some{Split{before,rest}}: Some{Split{axis <> before,rest}}
139	    case None{}: None{}
140	
141	def split_from(rest: List<&2,Views.Axis>,found: Bool,+height: U32,+width: U32,+extent: U32,+stride: U32) -> Maybe<&2,Split>:
142	  match rest found:
143	    case rest True{}: Some{Split{[],rest}}
144	    case Nil{} False{}: None{}
145	    case Views.Axis{+next_extent,+next_stride} <> tail False{}:
146	      prepend_before(Views.Axis{extent,stride},split_from(tail,U32.is_eq(next_extent,height) && U32.is_eq(next_stride,width),height,width,next_extent,next_stride))
147	
148	def split_rows(+height: U32,+width: U32,axes: List<&2,Views.Axis>) -> Maybe<&2,Split>:
149	  match axes:
150	    case []: None{}
151	    case Views.Axis{+extent,+stride} <> rest: split_from(rest,U32.is_eq(extent,height) && U32.is_eq(stride,width),height,width,extent,stride)
152	
153	# The leading axes of stride zero repeat each row.
154	def prepend_repeat(extent: U32,split: List<&2,U32> & List<&2,Views.Axis>) -> List<&2,U32> & List<&2,Views.Axis>:
155	  (repeats,rest) = split
156	  (extent <> repeats,rest)
157	
158	def repeats_from(rest: List<&2,Views.Axis>,zero: Bool,extent: U32,stride: U32) -> List<&2,U32> & List<&2,Views.Axis>:
159	  match rest zero:
160	    case rest False{}: ([],Views.Axis{extent,stride} <> rest)
161	    case Nil{} True{}: ([extent],[])
162	    case Views.Axis{next_extent,+next_stride} <> tail True{}: prepend_repeat(extent,repeats_from(tail,U32.is_zero(next_stride),next_extent,next_stride))
163	
164	def split_repeats(axes: List<&2,Views.Axis>) -> List<&2,U32> & List<&2,Views.Axis>:
165	  match axes:
166	    case []: ([],[])
167	    case Views.Axis{extent,+stride} <> rest: repeats_from(rest,U32.is_zero(stride),extent,stride)
168	
169	# The largest channel, or column, the axes reach from a start.
170	def reach(+plane: U32,+column: Bool,axes: List<&2,Views.Axis>,start: Maybe<&2,U32>) -> Maybe<&2,U32>:
171	  match axes start:
172	    case [] start: start
173	    case Views.Axis{+extent,stride} <> rest Some{+start}:
174	      reach(plane,column,rest,Shapes.checked_add_to(Shapes.checked_multiply((extent - 1 : U32),component(plane,column,stride)),start))
175	    case axes None{}: None{}
176	
177	def below(reach: Maybe<&2,U32>,+limit: U32) -> Bool:
178	  match reach:
179	    case Some{reach}: U32.is_lt(reach,limit)
180	    case None{}: False{}
181	
182	def counted(extents: List<&2,U32>) -> Bool:
183	  Views.has_count(Shapes.product(extents,Some{1}))
184	
185	# A classification is kept when the geometry has images and rows, every axis
186	# stays inside the channels and the columns of a band and every group of
187	# extents has a count.
188	def kept_when(valid: Bool,banded: Banded) -> Maybe<&2,Banded>:
189	  match valid:
190	    case True{}: Some{banded}
191	    case False{}: None{}
192	
193	def fits(+banded: Banded,+geometry: Bands.Geometry) -> Bool:
194	  Banded{+before,+repeats,+after,+offset} = banded
195	  Bands.Geometry{+batch,+channels,+height,+width} = geometry
196	  +plane = (height * width : U32)
197	  U32.is_lt(0,batch) && U32.is_lt(0,height) && below(reach(plane,False{},after,reach(plane,False{},before,Some{component(plane,False{},offset)})),channels) &&
198	    below(reach(plane,True{},after,reach(plane,True{},before,Some{component(plane,True{},offset)})),width) &&
199	    counted(Views.axis_extents(before)) && counted(Views.axis_extents(after)) && counted(repeats) && Views.has_count(Shapes.checked_multiply(height,product_or_zero(repeats)))
200	
201	def kept(+banded: Banded,+geometry: Bands.Geometry) -> Maybe<&2,Banded>:
202	  kept_when(fits(banded,geometry),banded)
203	
204	def split_kept(+offset: U32,+geometry: Bands.Geometry,before: List<&2,Views.Axis>,split: List<&2,U32> & List<&2,Views.Axis>) -> Maybe<&2,Banded>:
205	  (repeats,after) = split
206	  kept(Banded{before,repeats,after,offset},geometry)
207	
208	def rowed_split(+offset: U32,+geometry: Bands.Geometry,split: Maybe<&2,Split>) -> Maybe<&2,Banded>:
209	  match split:
210	    case Some{Split{before,rest}}: split_kept(offset,geometry,before,split_repeats(rest))
211	    case None{}: None{}
212	
213	def split_geometry(+geometry: Bands.Geometry,axes: List<&2,Views.Axis>) -> Maybe<&2,Split>:
214	  Bands.Geometry{batch,channels,+height,+width} = geometry
215	  split_rows(height,width,axes)
216	
217	def classify_rows(rowed: Bool,axes: Maybe<&2,List<&2,Views.Axis>>,+offset: U32,+geometry: Bands.Geometry) -> Maybe<&2,Banded>:
218	  match rowed axes:
219	    case True{} Some{axes}: rowed_split(offset,geometry,split_geometry(geometry,axes))
220	    case False{} Some{axes}: kept(Banded{[],[],axes,offset},geometry)
221	    case rowed None{}: None{}
222	
223	# The banded form of a view, given by its axes and offset, of the storage of a
224	# geometry, when the view starts with the images axis of a batch of more than
225	# one image, has the rows axis of a geometry of more than one row, and stays
226	# inside the channels and columns. A geometry of one image of one row, as a
227	# whole tensor has, takes every view whose addresses fit it.
228	def classify(axes: List<&2,Views.Axis>,+offset: U32,+geometry: Bands.Geometry) -> Maybe<&2,Banded>:
229	  Bands.Geometry{+batch,+channels,+height,+width} = geometry
230	  classify_rows(U32.is_lt(1,height),after_images(U32.is_lt(1,batch),batch,(channels * (height * width : U32) : U32),axes),offset,geometry)