1 // Copyright (c) 2019-2020, 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 
10 use crate::context::*;
11 use crate::mc::MotionVector;
12 use crate::partition::*;
13 use crate::predict::PredictionMode;
14 use crate::transform::*;
15 
16 use std::cmp;
17 use std::marker::PhantomData;
18 use std::ops::{Index, IndexMut};
19 use std::slice;
20 
21 /// Tiled view of FrameBlocks
22 #[derive(Debug)]
23 pub 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)]
36 pub 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
49 macro_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         unsafe {
154           let ptr = self.data.add(index * self.frame_cols);
155           slice::from_raw_parts(ptr, self.cols)
156         }
157       }
158     }
159 
160     // for convenience, also index by TileBlockOffset
161     impl Index<TileBlockOffset> for $name<'_> {
162       type Output = Block;
163       #[inline(always)]
164       fn index(&self, bo: TileBlockOffset) -> &Self::Output {
165         &self[bo.0.y][bo.0.x]
166       }
167     }
168   }
169 }
170 
171 tile_blocks_common!(TileBlocks);
172 tile_blocks_common!(TileBlocksMut, mut);
173 
174 impl TileBlocksMut<'_> {
175   #[inline(always)]
as_const(&self) -> TileBlocks<'_>176   pub const fn as_const(&self) -> TileBlocks<'_> {
177     TileBlocks {
178       data: self.data,
179       x: self.x,
180       y: self.y,
181       cols: self.cols,
182       rows: self.rows,
183       frame_cols: self.frame_cols,
184       frame_rows: self.frame_rows,
185       phantom: PhantomData,
186     }
187   }
188 
subregion_mut( &mut self, x: usize, y: usize, cols: usize, rows: usize, ) -> TileBlocksMut<'_>189   pub fn subregion_mut(
190     &mut self, x: usize, y: usize, cols: usize, rows: usize,
191   ) -> TileBlocksMut<'_> {
192     TileBlocksMut {
193       data: &mut self[y][x],
194       x: self.x + x,
195       y: self.y + y,
196       cols: cmp::min(cols, self.cols - x),
197       rows: cmp::min(rows, self.rows - y),
198       frame_cols: self.frame_cols,
199       frame_rows: self.frame_rows,
200       phantom: PhantomData,
201     }
202   }
203 
204   #[inline(always)]
for_each<F>(&mut self, bo: TileBlockOffset, bsize: BlockSize, f: F) where F: Fn(&mut Block),205   pub fn for_each<F>(&mut self, bo: TileBlockOffset, bsize: BlockSize, f: F)
206   where
207     F: Fn(&mut Block),
208   {
209     let mut bw = bsize.width_mi();
210     let bh = bsize.height_mi();
211 
212     if bo.0.x + bw >= self.cols {
213       bw = self.cols - bo.0.x;
214     }
215     for y in 0..bh {
216       if bo.0.y + y >= self.rows {
217         continue;
218       }
219       for block in self[bo.0.y + y][bo.0.x..bo.0.x + bw].iter_mut() {
220         f(block);
221       }
222     }
223   }
224 
225   #[inline(always)]
set_mode( &mut self, bo: TileBlockOffset, bsize: BlockSize, mode: PredictionMode, )226   pub fn set_mode(
227     &mut self, bo: TileBlockOffset, bsize: BlockSize, mode: PredictionMode,
228   ) {
229     self.for_each(bo, bsize, |block| block.mode = mode);
230   }
231 
232   #[inline(always)]
set_block_size(&mut self, bo: TileBlockOffset, bsize: BlockSize)233   pub fn set_block_size(&mut self, bo: TileBlockOffset, bsize: BlockSize) {
234     let n4_w = bsize.width_mi() as u8;
235     let n4_h = bsize.height_mi() as u8;
236     self.for_each(bo, bsize, |block| {
237       block.bsize = bsize;
238       block.n4_w = n4_w;
239       block.n4_h = n4_h;
240     });
241   }
242 
243   #[inline(always)]
set_tx_size( &mut self, bo: TileBlockOffset, bsize: BlockSize, tx_size: TxSize, )244   pub fn set_tx_size(
245     &mut self, bo: TileBlockOffset, bsize: BlockSize, tx_size: TxSize,
246   ) {
247     self.for_each(bo, bsize, |block| block.txsize = tx_size);
248   }
249 
250   #[inline(always)]
set_skip( &mut self, bo: TileBlockOffset, bsize: BlockSize, skip: bool, )251   pub fn set_skip(
252     &mut self, bo: TileBlockOffset, bsize: BlockSize, skip: bool,
253   ) {
254     self.for_each(bo, bsize, |block| block.skip = skip);
255   }
256 
257   #[inline(always)]
set_segmentation_idx( &mut self, bo: TileBlockOffset, bsize: BlockSize, idx: u8, )258   pub fn set_segmentation_idx(
259     &mut self, bo: TileBlockOffset, bsize: BlockSize, idx: u8,
260   ) {
261     self.for_each(bo, bsize, |block| block.segmentation_idx = idx);
262   }
263 
264   #[inline(always)]
set_ref_frames( &mut self, bo: TileBlockOffset, bsize: BlockSize, r: [RefType; 2], )265   pub fn set_ref_frames(
266     &mut self, bo: TileBlockOffset, bsize: BlockSize, r: [RefType; 2],
267   ) {
268     self.for_each(bo, bsize, |block| block.ref_frames = r);
269   }
270 
271   #[inline(always)]
set_motion_vectors( &mut self, bo: TileBlockOffset, bsize: BlockSize, mvs: [MotionVector; 2], )272   pub fn set_motion_vectors(
273     &mut self, bo: TileBlockOffset, bsize: BlockSize, mvs: [MotionVector; 2],
274   ) {
275     self.for_each(bo, bsize, |block| block.mv = mvs);
276   }
277 
278   #[inline(always)]
set_cdef(&mut self, sbo: TileSuperBlockOffset, cdef_index: u8)279   pub fn set_cdef(&mut self, sbo: TileSuperBlockOffset, cdef_index: u8) {
280     let bo = sbo.block_offset(0, 0).0;
281     // Checkme: Is 16 still the right block unit for 128x128 superblocks?
282     let bw = cmp::min(bo.x + MIB_SIZE, self.cols);
283     let bh = cmp::min(bo.y + MIB_SIZE, self.rows);
284     for y in bo.y..bh {
285       for x in bo.x..bw {
286         self[y as usize][x as usize].cdef_index = cdef_index;
287       }
288     }
289   }
290 }
291 
292 impl IndexMut<usize> for TileBlocksMut<'_> {
293   #[inline(always)]
index_mut(&mut self, index: usize) -> &mut Self::Output294   fn index_mut(&mut self, index: usize) -> &mut Self::Output {
295     assert!(index < self.rows);
296     unsafe {
297       let ptr = self.data.add(index * self.frame_cols);
298       slice::from_raw_parts_mut(ptr, self.cols)
299     }
300   }
301 }
302 
303 impl IndexMut<TileBlockOffset> for TileBlocksMut<'_> {
304   #[inline(always)]
index_mut(&mut self, bo: TileBlockOffset) -> &mut Self::Output305   fn index_mut(&mut self, bo: TileBlockOffset) -> &mut Self::Output {
306     &mut self[bo.0.y][bo.0.x]
307   }
308 }
309