rav1e/tiling/
tile_blocks.rs

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 = self.data.add(index * self.frame_cols);
300      slice::from_raw_parts_mut(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}
rav1e/tiling/tile_blocks.rs

rav1e/tiling/
tile_blocks.rs
