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