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::mc::MotionVector;
12use crate::partition::*;
13use crate::predict::PredictionMode;
14use crate::transform::*;
15
16use std::cmp;
17use std::marker::PhantomData;
18use std::ops::{Index, IndexMut};
19use std::slice;
20
21/// Tiled view of `FrameBlocks`
22#[derive(Debug)]
23pub struct TileBlocks<'a> {
24 data: *const Block,
25 x: usize,
26 y: usize,
27 cols: usize,
28 rows: usize,
29 frame_cols: usize,
30 frame_rows: usize,
31 phantom: PhantomData<&'a Block>,
32}
33
34/// Mutable tiled view of `FrameBlocks`
35#[derive(Debug)]
36pub struct TileBlocksMut<'a> {
37 data: *mut Block,
38 // private to guarantee borrowing rules
39 x: usize,
40 y: usize,
41 cols: usize,
42 rows: usize,
43 frame_cols: usize,
44 frame_rows: usize,
45 phantom: PhantomData<&'a mut Block>,
46}
47
48// common impl for TileBlocks and TileBlocksMut
49macro_rules! tile_blocks_common {
50 // $name: TileBlocks or TileBlocksMut
51 // $opt_mut: nothing or mut
52 ($name:ident $(,$opt_mut:tt)?) => {
53 impl<'a> $name<'a> {
54
55 #[inline(always)]
56 pub fn new(
57 frame_blocks: &'a $($opt_mut)? FrameBlocks,
58 x: usize,
59 y: usize,
60 cols: usize,
61 rows: usize,
62 ) -> Self {
63 Self {
64 data: & $($opt_mut)? frame_blocks[y][x],
65 x,
66 y,
67 cols,
68 rows,
69 frame_cols: frame_blocks.cols,
70 frame_rows: frame_blocks.rows,
71 phantom: PhantomData,
72 }
73 }
74
75 pub fn subregion(
76 &mut self,
77 x: usize,
78 y: usize,
79 cols: usize,
80 rows: usize,
81 ) -> TileBlocks<'_> {
82 TileBlocks {
83 data: &self[y][x],
84 x: self.x+x,
85 y: self.y+y,
86 cols: cmp::min(cols, self.cols - x),
87 rows: cmp::min(rows, self.rows - y),
88 frame_cols: self.frame_cols,
89 frame_rows: self.frame_rows,
90 phantom: PhantomData,
91 }
92 }
93
94 #[inline(always)]
95 pub const fn x(&self) -> usize {
96 self.x
97 }
98
99 #[inline(always)]
100 pub const fn y(&self) -> usize {
101 self.y
102 }
103
104 #[inline(always)]
105 pub const fn cols(&self) -> usize {
106 self.cols
107 }
108
109 #[inline(always)]
110 pub const fn rows(&self) -> usize {
111 self.rows
112 }
113
114 #[inline(always)]
115 pub const fn frame_cols(&self) -> usize {
116 self.frame_cols
117 }
118
119 #[inline(always)]
120 pub const fn frame_rows(&self) -> usize {
121 self.frame_rows
122 }
123
124 #[inline(always)]
125 pub fn above_of(&self, bo: TileBlockOffset) -> &Block {
126 &self[bo.0.y - 1][bo.0.x]
127 }
128
129 #[inline(always)]
130 pub fn left_of(&self, bo: TileBlockOffset) -> &Block {
131 &self[bo.0.y][bo.0.x - 1]
132 }
133
134 #[inline(always)]
135 pub fn above_left_of(&self, bo: TileBlockOffset) -> &Block {
136 &self[bo.0.y - 1][bo.0.x - 1]
137 }
138
139 pub fn get_cdef(&self, sbo: TileSuperBlockOffset) -> u8 {
140 let bo = sbo.block_offset(0, 0).0;
141 self[bo.y][bo.x].cdef_index
142 }
143 }
144
145 unsafe impl Send for $name<'_> {}
146 unsafe impl Sync for $name<'_> {}
147
148 impl Index<usize> for $name<'_> {
149 type Output = [Block];
150 #[inline(always)]
151 fn index(&self, index: usize) -> &Self::Output {
152 assert!(index < self.rows);
153 // SAFETY: The above assert ensures we do not access OOB data.
154 unsafe {
155 let ptr = self.data.add(index * self.frame_cols);
156 slice::from_raw_parts(ptr, self.cols)
157 }
158 }
159 }
160
161 // for convenience, also index by TileBlockOffset
162 impl Index<TileBlockOffset> for $name<'_> {
163 type Output = Block;
164 #[inline(always)]
165 fn index(&self, bo: TileBlockOffset) -> &Self::Output {
166 &self[bo.0.y][bo.0.x]
167 }
168 }
169 }
170}
171
172tile_blocks_common!(TileBlocks);
173tile_blocks_common!(TileBlocksMut, mut);
174
175impl TileBlocksMut<'_> {
176 #[inline(always)]
177 pub const fn as_const(&self) -> TileBlocks<'_> {
178 TileBlocks {
179 data: self.data,
180 x: self.x,
181 y: self.y,
182 cols: self.cols,
183 rows: self.rows,
184 frame_cols: self.frame_cols,
185 frame_rows: self.frame_rows,
186 phantom: PhantomData,
187 }
188 }
189
190 pub fn subregion_mut(
191 &mut self, x: usize, y: usize, cols: usize, rows: usize,
192 ) -> TileBlocksMut<'_> {
193 TileBlocksMut {
194 data: &mut self[y][x],
195 x: self.x + x,
196 y: self.y + y,
197 cols: cmp::min(cols, self.cols - x),
198 rows: cmp::min(rows, self.rows - y),
199 frame_cols: self.frame_cols,
200 frame_rows: self.frame_rows,
201 phantom: PhantomData,
202 }
203 }
204
205 #[inline(always)]
206 pub fn for_each<F>(&mut self, bo: TileBlockOffset, bsize: BlockSize, f: F)
207 where
208 F: Fn(&mut Block),
209 {
210 let mut bw = bsize.width_mi();
211 let bh = bsize.height_mi();
212
213 if bo.0.x + bw >= self.cols {
214 bw = self.cols - bo.0.x;
215 }
216 for y in 0..bh {
217 if bo.0.y + y >= self.rows {
218 continue;
219 }
220 for block in self[bo.0.y + y][bo.0.x..bo.0.x + bw].iter_mut() {
221 f(block);
222 }
223 }
224 }
225
226 #[inline(always)]
227 pub fn set_mode(
228 &mut self, bo: TileBlockOffset, bsize: BlockSize, mode: PredictionMode,
229 ) {
230 self.for_each(bo, bsize, |block| block.mode = mode);
231 }
232
233 #[inline(always)]
234 pub fn set_block_size(&mut self, bo: TileBlockOffset, bsize: BlockSize) {
235 let n4_w = bsize.width_mi() as u8;
236 let n4_h = bsize.height_mi() as u8;
237 self.for_each(bo, bsize, |block| {
238 block.bsize = bsize;
239 block.n4_w = n4_w;
240 block.n4_h = n4_h;
241 });
242 }
243
244 #[inline(always)]
245 pub fn set_tx_size(
246 &mut self, bo: TileBlockOffset, bsize: BlockSize, tx_size: TxSize,
247 ) {
248 self.for_each(bo, bsize, |block| block.txsize = tx_size);
249 }
250
251 #[inline(always)]
252 pub fn set_skip(
253 &mut self, bo: TileBlockOffset, bsize: BlockSize, skip: bool,
254 ) {
255 self.for_each(bo, bsize, |block| block.skip = skip);
256 }
257
258 #[inline(always)]
259 pub fn set_segmentation_idx(
260 &mut self, bo: TileBlockOffset, bsize: BlockSize, idx: u8,
261 ) {
262 self.for_each(bo, bsize, |block| block.segmentation_idx = idx);
263 }
264
265 #[inline(always)]
266 pub fn set_ref_frames(
267 &mut self, bo: TileBlockOffset, bsize: BlockSize, r: [RefType; 2],
268 ) {
269 self.for_each(bo, bsize, |block| block.ref_frames = r);
270 }
271
272 #[inline(always)]
273 pub fn set_motion_vectors(
274 &mut self, bo: TileBlockOffset, bsize: BlockSize, mvs: [MotionVector; 2],
275 ) {
276 self.for_each(bo, bsize, |block| block.mv = mvs);
277 }
278
279 #[inline(always)]
280 pub fn set_cdef(&mut self, sbo: TileSuperBlockOffset, cdef_index: u8) {
281 let bo = sbo.block_offset(0, 0).0;
282 // Checkme: Is 16 still the right block unit for 128x128 superblocks?
283 let bw = cmp::min(bo.x + MIB_SIZE, self.cols);
284 let bh = cmp::min(bo.y + MIB_SIZE, self.rows);
285 for y in bo.y..bh {
286 for x in bo.x..bw {
287 self[y][x].cdef_index = cdef_index;
288 }
289 }
290 }
291}
292
293impl IndexMut<usize> for TileBlocksMut<'_> {
294 #[inline(always)]
295 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
296 assert!(index < self.rows);
297 // SAFETY: The above assert ensures we do not access OOB data.
298 unsafe {
299 let ptr: *mut Block = self.data.add(count:index * self.frame_cols);
300 slice::from_raw_parts_mut(data:ptr, self.cols)
301 }
302 }
303}
304
305impl IndexMut<TileBlockOffset> for TileBlocksMut<'_> {
306 #[inline(always)]
307 fn index_mut(&mut self, bo: TileBlockOffset) -> &mut Self::Output {
308 &mut self[bo.0.y][bo.0.x]
309 }
310}
311