1// Copyright (c) 2019-2022, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9
10use crate::context::*;
11use crate::encoder::FrameInvariants;
12use crate::lrf::*;
13use crate::util::Pixel;
14
15use std::marker::PhantomData;
16use std::ops::{Index, IndexMut};
17use std::ptr;
18use std::slice;
19
20/// Tiled view of `RestorationUnits`
21#[derive(Debug)]
22pub struct TileRestorationUnits<'a> {
23 data: *const RestorationUnit,
24 // private to guarantee borrowing rules
25 x: usize,
26 y: usize,
27 cols: usize,
28 rows: usize,
29 /// number of cols in the underlying `FrameRestorationUnits`
30 stride: usize,
31 phantom: PhantomData<&'a RestorationUnit>,
32}
33
34/// Mutable tiled view of `RestorationUnits`
35#[derive(Debug)]
36pub struct TileRestorationUnitsMut<'a> {
37 data: *mut RestorationUnit,
38 // private to guarantee borrowing rules
39 x: usize,
40 y: usize,
41 cols: usize,
42 rows: usize,
43 /// number of cols in the underlying `FrameRestorationUnits`
44 stride: usize,
45 phantom: PhantomData<&'a mut RestorationUnit>,
46}
47
48// common impl for TileRestorationUnits and TileRestorationUnitsMut
49macro_rules! tile_restoration_units_common {
50 // $name: TileRestorationUnits or TileRestorationUnitsMut
51 // $null: null or null_mut
52 // $opt_mut: nothing or mut
53 ($name:ident, $null:ident $(,$opt_mut:tt)?) => {
54 impl<'a> $name<'a> {
55
56 #[inline(always)]
57 pub fn new(
58 frame_units: &'a $($opt_mut)? FrameRestorationUnits,
59 x: usize,
60 y: usize,
61 cols: usize,
62 rows: usize,
63 ) -> Self {
64 Self {
65 data: if x < frame_units.cols && y < frame_units.rows {
66 & $($opt_mut)? frame_units[y][x]
67 } else {
68 // on edges, a tile may contain no restoration units
69 ptr::$null()
70 },
71 x,
72 y,
73 cols,
74 rows,
75 stride: frame_units.cols,
76 phantom: PhantomData,
77 }
78 }
79
80 #[inline(always)]
81 pub const fn x(&self) -> usize {
82 self.x
83 }
84
85 #[inline(always)]
86 pub const fn y(&self) -> usize {
87 self.y
88 }
89
90 #[inline(always)]
91 pub const fn cols(&self) -> usize {
92 self.cols
93 }
94
95 #[inline(always)]
96 pub const fn rows(&self) -> usize {
97 self.rows
98 }
99 }
100
101 unsafe impl Send for $name<'_> {}
102 unsafe impl Sync for $name<'_> {}
103
104 impl Index<usize> for $name<'_> {
105 type Output = [RestorationUnit];
106
107 #[inline(always)]
108 fn index(&self, index: usize) -> &Self::Output {
109 assert!(index < self.rows);
110 // SAFETY: The above assert ensures we do not access OOB data.
111 unsafe {
112 let ptr = self.data.add(index * self.stride);
113 slice::from_raw_parts(ptr, self.cols)
114 }
115 }
116 }
117 }
118}
119
120tile_restoration_units_common!(TileRestorationUnits, null);
121tile_restoration_units_common!(TileRestorationUnitsMut, null_mut, mut);
122
123impl TileRestorationUnitsMut<'_> {
124 #[inline(always)]
125 pub const fn as_const(&self) -> TileRestorationUnits<'_> {
126 TileRestorationUnits {
127 data: self.data,
128 x: self.x,
129 y: self.y,
130 cols: self.cols,
131 rows: self.rows,
132 stride: self.stride,
133 phantom: PhantomData,
134 }
135 }
136}
137
138impl IndexMut<usize> for TileRestorationUnitsMut<'_> {
139 #[inline(always)]
140 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
141 assert!(index < self.rows);
142 // SAFETY: The above assert ensures we do not access OOB data.
143 unsafe {
144 let ptr: *mut RestorationUnit = self.data.add(count:index * self.stride);
145 slice::from_raw_parts_mut(data:ptr, self.cols)
146 }
147 }
148}
149
150/// Tiled view of `RestorationPlane`
151#[derive(Debug)]
152pub struct TileRestorationPlane<'a> {
153 pub rp_cfg: &'a RestorationPlaneConfig,
154 pub wiener_ref: [[i8; 3]; 2],
155 pub sgrproj_ref: [i8; 2],
156 pub units: TileRestorationUnits<'a>,
157}
158
159/// Mutable tiled view of `RestorationPlane`
160#[derive(Debug)]
161pub struct TileRestorationPlaneMut<'a> {
162 pub rp_cfg: &'a RestorationPlaneConfig,
163 pub wiener_ref: [[i8; 3]; 2],
164 pub sgrproj_ref: [i8; 2],
165 pub units: TileRestorationUnitsMut<'a>,
166}
167
168// common impl for TileRestorationPlane and TileRestorationPlaneMut
169macro_rules! tile_restoration_plane_common {
170 // $name: TileRestorationPlane or TileRestorationPlaneMut
171 // $tru_type: TileRestorationUnits or TileRestorationUnitsMut
172 // $opt_mut: nothing or mut
173 ($name:ident, $tru_type:ident $(,$opt_mut:tt)?) => {
174 impl<'a> $name<'a> {
175
176 #[inline(always)]
177 pub fn new(
178 rp: &'a $($opt_mut)? RestorationPlane,
179 units_x: usize,
180 units_y: usize,
181 units_cols: usize,
182 units_rows: usize,
183 ) -> Self {
184 Self {
185 rp_cfg: &rp.cfg,
186 wiener_ref: [WIENER_TAPS_MID; 2],
187 sgrproj_ref: SGRPROJ_XQD_MID,
188 units: $tru_type::new(& $($opt_mut)? rp.units, units_x, units_y, units_cols, units_rows),
189 }
190 }
191
192 // determines the loop restoration unit row and column a
193 // superblock belongs to. The stretch boolean indicates if a
194 // superblock that belongs to a stretched LRU should return an
195 // index (stretch == true) or None (stretch == false).
196 pub const fn restoration_unit_index(&self, sbo: TileSuperBlockOffset, stretch: bool)
197 -> Option<(usize, usize)> {
198 if self.units.rows > 0 && self.units.cols > 0 {
199 // is this a stretch block?
200 let x_stretch = sbo.0.x < self.rp_cfg.sb_cols &&
201 sbo.0.x >> self.rp_cfg.sb_h_shift >= self.units.cols;
202 let y_stretch = sbo.0.y < self.rp_cfg.sb_rows &&
203 sbo.0.y >> self.rp_cfg.sb_v_shift >= self.units.rows;
204 if (x_stretch || y_stretch) && !stretch {
205 None
206 } else {
207 let x = (sbo.0.x >> self.rp_cfg.sb_h_shift) - if x_stretch { 1 } else { 0 };
208 let y = (sbo.0.y >> self.rp_cfg.sb_v_shift) - if y_stretch { 1 } else { 0 };
209 if x < self.units.cols && y < self.units.rows {
210 Some((x, y))
211 } else {
212 None
213 }
214 }
215 } else {
216 None
217 }
218 }
219
220 pub fn restoration_unit_offset(&self, base: TileSuperBlockOffset,
221 offset: TileSuperBlockOffset, stretch: bool)
222 -> Option<(usize, usize)> {
223 let base_option = self.restoration_unit_index(base, stretch);
224 let delta_option = self.restoration_unit_index(base + offset, stretch);
225 if let (Some((base_x, base_y)), Some((delta_x, delta_y))) =
226 (base_option, delta_option)
227 {
228 Some ((delta_x - base_x, delta_y - base_y))
229 } else {
230 None
231 }
232 }
233
234 pub const fn restoration_unit_countable(&self, x: usize, y: usize) -> usize {
235 y * self.units.cols + x
236 }
237
238 // Is this the last sb (in scan order) in the restoration unit
239 // that we will be considering for RDO? This would be a
240 // straightforward calculation but for stretch; if the LRU
241 // stretches into a different tile, we don't consider those SBs
242 // in the other tile to be part of the LRU for RDO purposes.
243 pub fn restoration_unit_last_sb_for_rdo<T: Pixel>(
244 &self,
245 fi: &FrameInvariants<T>,
246 global_sbo: PlaneSuperBlockOffset,
247 tile_sbo: TileSuperBlockOffset,
248 ) -> bool {
249 // there is 1 restoration unit for (1 << sb_shift) super-blocks
250 let h_mask = (1 << self.rp_cfg.sb_h_shift) - 1;
251 let v_mask = (1 << self.rp_cfg.sb_v_shift) - 1;
252 // is this a stretch block?
253 let x_stretch = tile_sbo.0.x >> self.rp_cfg.sb_h_shift >= self.units.cols;
254 let y_stretch = tile_sbo.0.y >> self.rp_cfg.sb_v_shift >= self.units.rows;
255 // Need absolute superblock offsets for edge check, not local to the tile.
256 let sbx = global_sbo.0.x + tile_sbo.0.x;
257 let sby = global_sbo.0.y + tile_sbo.0.y;
258 // edge-of-tile check + edge-of-frame check
259 let last_x = (tile_sbo.0.x & h_mask == h_mask && !x_stretch) || sbx == fi.sb_width-1;
260 let last_y = (tile_sbo.0.y & v_mask == v_mask && !y_stretch) || sby == fi.sb_height-1;
261 last_x && last_y
262 }
263
264 #[inline(always)]
265 pub fn restoration_unit(&self, sbo: TileSuperBlockOffset, stretch: bool)
266 -> Option<&RestorationUnit> {
267 self.restoration_unit_index(sbo, stretch).map(|(x, y)| &self.units[y][x])
268 }
269 }
270 }
271}
272
273tile_restoration_plane_common!(TileRestorationPlane, TileRestorationUnits);
274tile_restoration_plane_common!(
275 TileRestorationPlaneMut,
276 TileRestorationUnitsMut,
277 mut
278);
279
280impl<'a> TileRestorationPlaneMut<'a> {
281 #[inline(always)]
282 pub fn restoration_unit_mut(
283 &mut self, sbo: TileSuperBlockOffset,
284 ) -> Option<&mut RestorationUnit> {
285 // cannot use map() due to lifetime constraints
286 if let Some((x: usize, y: usize)) = self.restoration_unit_index(sbo, stretch:true) {
287 Some(&mut self.units[y][x])
288 } else {
289 None
290 }
291 }
292
293 #[inline(always)]
294 pub const fn as_const(&self) -> TileRestorationPlane<'_> {
295 TileRestorationPlane {
296 rp_cfg: self.rp_cfg,
297 wiener_ref: self.wiener_ref,
298 sgrproj_ref: self.sgrproj_ref,
299 units: self.units.as_const(),
300 }
301 }
302}
303
304/// Tiled view of `RestorationState`
305#[derive(Debug)]
306pub struct TileRestorationState<'a> {
307 pub planes: [TileRestorationPlane<'a>; MAX_PLANES],
308}
309
310/// Mutable tiled view of `RestorationState`
311#[derive(Debug)]
312pub struct TileRestorationStateMut<'a> {
313 pub planes: [TileRestorationPlaneMut<'a>; MAX_PLANES],
314}
315
316// common impl for TileRestorationState and TileRestorationStateMut
317macro_rules! tile_restoration_state_common {
318 // $name: TileRestorationState or TileRestorationStateMut
319 // $trp_type: TileRestorationPlane or TileRestorationPlaneMut
320 // $iter: iter or iter_mut
321 // $opt_mut: nothing or mut
322 ($name:ident, $trp_type:ident, $iter:ident $(,$opt_mut:tt)?) => {
323 impl<'a> $name<'a> {
324
325 #[inline(always)]
326 pub fn new(
327 rs: &'a $($opt_mut)? RestorationState,
328 sbo: PlaneSuperBlockOffset,
329 sb_width: usize,
330 sb_height: usize,
331 ) -> Self {
332 let (units_x0, units_y0, units_cols0, units_rows0) =
333 Self::get_units_region(rs, sbo, sb_width, sb_height, 0);
334 let (units_x1, units_y1, units_cols1, units_rows1) =
335 Self::get_units_region(rs, sbo, sb_width, sb_height, 1);
336 let (units_x2, units_y2, units_cols2, units_rows2) =
337 Self::get_units_region(rs, sbo, sb_width, sb_height, 2);
338 // we cannot retrieve &mut of slice items directly and safely
339 let mut planes_iter = rs.planes.$iter();
340 Self {
341 planes: [
342 {
343 let plane = planes_iter.next().unwrap();
344 $trp_type::new(plane, units_x0, units_y0, units_cols0, units_rows0)
345 },
346 {
347 let plane = planes_iter.next().unwrap();
348 $trp_type::new(plane, units_x1, units_y1, units_cols1, units_rows1)
349 },
350 {
351 let plane = planes_iter.next().unwrap();
352 $trp_type::new(plane, units_x2, units_y2, units_cols2, units_rows2)
353 },
354 ],
355 }
356 }
357
358 #[inline(always)]
359 fn get_units_region(
360 rs: &RestorationState,
361 sbo: PlaneSuperBlockOffset,
362 sb_width: usize,
363 sb_height: usize,
364 pli: usize,
365 ) -> (usize, usize, usize, usize) {
366 let sb_h_shift = rs.planes[pli].cfg.sb_h_shift;
367 let sb_v_shift = rs.planes[pli].cfg.sb_v_shift;
368 // there may be several super-blocks per restoration unit
369 // the given super-block offset must match the start of a restoration unit
370 debug_assert!(sbo.0.x % (1 << sb_h_shift) == 0);
371 debug_assert!(sbo.0.y % (1 << sb_v_shift) == 0);
372
373 let units_x = sbo.0.x >> sb_h_shift;
374 let units_y = sbo.0.y >> sb_v_shift;
375 let units_cols = sb_width + (1 << sb_h_shift) - 1 >> sb_h_shift;
376 let units_rows = sb_height + (1 << sb_v_shift) - 1 >> sb_v_shift;
377
378 let FrameRestorationUnits { cols: rs_cols, rows: rs_rows, .. } = rs.planes[pli].units;
379 // +1 because the last super-block may use the "stretched" restoration unit
380 // from its neighbours
381 // <https://github.com/xiph/rav1e/issues/631#issuecomment-454419152>
382 debug_assert!(units_x < rs_cols + 1);
383 debug_assert!(units_y < rs_rows + 1);
384 debug_assert!(units_x + units_cols <= rs_cols + 1);
385 debug_assert!(units_y + units_rows <= rs_rows + 1);
386
387 let units_x = units_x.min(rs_cols);
388 let units_y = units_y.min(rs_rows);
389 let units_cols = units_cols.min(rs_cols - units_x);
390 let units_rows = units_rows.min(rs_rows - units_y);
391 (units_x, units_y, units_cols, units_rows)
392 }
393
394 #[inline(always)]
395 pub fn has_restoration_unit(&self, sbo: TileSuperBlockOffset, pli: usize, stretch: bool)
396 -> bool {
397 self.planes[pli].restoration_unit(sbo, stretch).is_some()
398 }
399 }
400 }
401}
402
403tile_restoration_state_common!(
404 TileRestorationState,
405 TileRestorationPlane,
406 iter
407);
408tile_restoration_state_common!(
409 TileRestorationStateMut,
410 TileRestorationPlaneMut,
411 iter_mut,
412 mut
413);
414
415impl<'a> TileRestorationStateMut<'a> {
416 #[inline(always)]
417 pub const fn as_const(&self) -> TileRestorationState {
418 TileRestorationState {
419 planes: [
420 self.planes[0].as_const(),
421 self.planes[1].as_const(),
422 self.planes[2].as_const(),
423 ],
424 }
425 }
426}
427