rav1e/
partition.rs

1// Copyright (c) 2017-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
10#![allow(non_camel_case_types)]
11#![allow(dead_code)]
12
13use self::BlockSize::*;
14use self::TxSize::*;
15use crate::context::*;
16use crate::frame::*;
17use crate::predict::*;
18use crate::recon_intra::*;
19use crate::serialize::{Deserialize, Serialize};
20use crate::tiling::*;
21use crate::transform::TxSize;
22use crate::util::*;
23use thiserror::Error;
24
25use std::mem::transmute;
26use std::mem::MaybeUninit;
27
28// LAST_FRAME through ALTREF_FRAME correspond to slots 0-6.
29#[derive(PartialEq, Eq, PartialOrd, Copy, Clone, Debug)]
30pub enum RefType {
31  INTRA_FRAME = 0,
32  LAST_FRAME = 1,
33  LAST2_FRAME = 2,
34  LAST3_FRAME = 3,
35  GOLDEN_FRAME = 4,
36  BWDREF_FRAME = 5,
37  ALTREF2_FRAME = 6,
38  ALTREF_FRAME = 7,
39  NONE_FRAME = 8,
40}
41
42impl RefType {
43  /// convert to a ref list index, 0-6 (`INTER_REFS_PER_FRAME`)
44  ///
45  /// # Panics
46  ///
47  /// - If the ref type is a None or Intra frame
48  #[inline]
49  pub fn to_index(self) -> usize {
50    match self {
51      NONE_FRAME => {
52        panic!("Tried to get slot of NONE_FRAME");
53      }
54      INTRA_FRAME => {
55        panic!("Tried to get slot of INTRA_FRAME");
56      }
57      _ => (self as usize) - 1,
58    }
59  }
60  #[inline]
61  pub const fn is_fwd_ref(self) -> bool {
62    (self as usize) < 5
63  }
64  #[inline]
65  pub const fn is_bwd_ref(self) -> bool {
66    (self as usize) >= 5
67  }
68}
69
70use self::RefType::*;
71use std::fmt;
72use std::fmt::Display;
73
74pub const ALL_INTER_REFS: [RefType; 7] = [
75  LAST_FRAME,
76  LAST2_FRAME,
77  LAST3_FRAME,
78  GOLDEN_FRAME,
79  BWDREF_FRAME,
80  ALTREF2_FRAME,
81  ALTREF_FRAME,
82];
83
84pub const LAST_LAST2_FRAMES: usize = 0; // { LAST_FRAME, LAST2_FRAME }
85pub const LAST_LAST3_FRAMES: usize = 1; // { LAST_FRAME, LAST3_FRAME }
86pub const LAST_GOLDEN_FRAMES: usize = 2; // { LAST_FRAME, GOLDEN_FRAME }
87pub const BWDREF_ALTREF_FRAMES: usize = 3; // { BWDREF_FRAME, ALTREF_FRAME }
88pub const LAST2_LAST3_FRAMES: usize = 4; // { LAST2_FRAME, LAST3_FRAME }
89pub const LAST2_GOLDEN_FRAMES: usize = 5; // { LAST2_FRAME, GOLDEN_FRAME }
90pub const LAST3_GOLDEN_FRAMES: usize = 6; // { LAST3_FRAME, GOLDEN_FRAME }
91pub const BWDREF_ALTREF2_FRAMES: usize = 7; // { BWDREF_FRAME, ALTREF2_FRAME }
92pub const ALTREF2_ALTREF_FRAMES: usize = 8; // { ALTREF2_FRAME, ALTREF_FRAME }
93pub const TOTAL_UNIDIR_COMP_REFS: usize = 9;
94
95// NOTE: UNIDIR_COMP_REFS is the number of uni-directional reference pairs
96//       that are explicitly signaled.
97pub const UNIDIR_COMP_REFS: usize = BWDREF_ALTREF_FRAMES + 1;
98
99pub const FWD_REFS: usize = 4;
100pub const BWD_REFS: usize = 3;
101pub const SINGLE_REFS: usize = 7;
102pub const TOTAL_REFS_PER_FRAME: usize = 8;
103pub const INTER_REFS_PER_FRAME: usize = 7;
104pub const TOTAL_COMP_REFS: usize =
105  FWD_REFS * BWD_REFS + TOTAL_UNIDIR_COMP_REFS;
106
107pub const REF_FRAMES_LOG2: usize = 3;
108pub const REF_FRAMES: usize = 1 << REF_FRAMES_LOG2;
109
110pub const REF_CONTEXTS: usize = 3;
111pub const MVREF_ROW_COLS: usize = 3;
112
113#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Debug)]
114pub enum PartitionType {
115  PARTITION_NONE,
116  PARTITION_HORZ,
117  PARTITION_VERT,
118  PARTITION_SPLIT,
119  PARTITION_HORZ_A, // HORZ split and the top partition is split again
120  PARTITION_HORZ_B, // HORZ split and the bottom partition is split again
121  PARTITION_VERT_A, // VERT split and the left partition is split again
122  PARTITION_VERT_B, // VERT split and the right partition is split again
123  PARTITION_HORZ_4, // 4:1 horizontal partition
124  PARTITION_VERT_4, // 4:1 vertical partition
125  PARTITION_INVALID,
126}
127
128#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
129#[cfg_attr(test, derive(Default))]
130pub enum BlockSize {
131  BLOCK_4X4,
132  BLOCK_4X8,
133  BLOCK_8X4,
134  BLOCK_8X8,
135  BLOCK_8X16,
136  BLOCK_16X8,
137  BLOCK_16X16,
138  BLOCK_16X32,
139  BLOCK_32X16,
140  BLOCK_32X32,
141  BLOCK_32X64,
142  BLOCK_64X32,
143  #[cfg_attr(test, default)]
144  BLOCK_64X64,
145  BLOCK_64X128,
146  BLOCK_128X64,
147  BLOCK_128X128,
148  BLOCK_4X16,
149  BLOCK_16X4,
150  BLOCK_8X32,
151  BLOCK_32X8,
152  BLOCK_16X64,
153  BLOCK_64X16,
154}
155
156#[derive(Debug, Error, Copy, Clone, Eq, PartialEq)]
157pub struct InvalidBlockSize;
158
159impl Display for InvalidBlockSize {
160  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161    f.write_str("invalid block size")
162  }
163}
164
165impl PartialOrd for BlockSize {
166  #[inline(always)]
167  fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
168    use std::cmp::Ordering::{Equal, Greater, Less};
169    match (
170      self.width().cmp(&other.width()),
171      self.height().cmp(&other.height()),
172    ) {
173      (Greater, Less) | (Less, Greater) => None,
174      (Equal, Equal) => Some(Equal),
175      (Greater, _) | (_, Greater) => Some(Greater),
176      (Less, _) | (_, Less) => Some(Less),
177    }
178  }
179}
180
181impl BlockSize {
182  pub const BLOCK_SIZES_ALL: usize = 22;
183  pub const BLOCK_SIZES: usize = BlockSize::BLOCK_SIZES_ALL - 6; // BLOCK_SIZES_ALL minus 4:1 non-squares, six of them
184
185  #[inline]
186  /// # Errors
187  ///
188  /// - Returns `InvalidBlockSize` if the given `w` and `h` do not produce
189  ///   a valid block size.
190  pub fn from_width_and_height_opt(
191    w: usize, h: usize,
192  ) -> Result<BlockSize, InvalidBlockSize> {
193    match (w, h) {
194      (4, 4) => Ok(BLOCK_4X4),
195      (4, 8) => Ok(BLOCK_4X8),
196      (4, 16) => Ok(BLOCK_4X16),
197      (8, 4) => Ok(BLOCK_8X4),
198      (8, 8) => Ok(BLOCK_8X8),
199      (8, 16) => Ok(BLOCK_8X16),
200      (8, 32) => Ok(BLOCK_8X32),
201      (16, 4) => Ok(BLOCK_16X4),
202      (16, 8) => Ok(BLOCK_16X8),
203      (16, 16) => Ok(BLOCK_16X16),
204      (16, 32) => Ok(BLOCK_16X32),
205      (16, 64) => Ok(BLOCK_16X64),
206      (32, 8) => Ok(BLOCK_32X8),
207      (32, 16) => Ok(BLOCK_32X16),
208      (32, 32) => Ok(BLOCK_32X32),
209      (32, 64) => Ok(BLOCK_32X64),
210      (64, 16) => Ok(BLOCK_64X16),
211      (64, 32) => Ok(BLOCK_64X32),
212      (64, 64) => Ok(BLOCK_64X64),
213      (64, 128) => Ok(BLOCK_64X128),
214      (128, 64) => Ok(BLOCK_128X64),
215      (128, 128) => Ok(BLOCK_128X128),
216      _ => Err(InvalidBlockSize),
217    }
218  }
219
220  /// # Panics
221  ///
222  /// - If the given `w` and `h` do not produce a valid block size.
223  pub fn from_width_and_height(w: usize, h: usize) -> BlockSize {
224    Self::from_width_and_height_opt(w, h).unwrap()
225  }
226
227  #[inline]
228  pub fn cfl_allowed(self) -> bool {
229    // TODO: fix me when enabling EXT_PARTITION_TYPES
230    self <= BlockSize::BLOCK_32X32
231  }
232
233  #[inline]
234  pub const fn width(self) -> usize {
235    1 << self.width_log2()
236  }
237
238  /// width * height
239  #[inline]
240  pub const fn area(self) -> usize {
241    self.width() * self.height()
242  }
243
244  #[inline]
245  pub const fn width_log2(self) -> usize {
246    match self {
247      BLOCK_4X4 | BLOCK_4X8 | BLOCK_4X16 => 2,
248      BLOCK_8X4 | BLOCK_8X8 | BLOCK_8X16 | BLOCK_8X32 => 3,
249      BLOCK_16X4 | BLOCK_16X8 | BLOCK_16X16 | BLOCK_16X32 | BLOCK_16X64 => 4,
250      BLOCK_32X8 | BLOCK_32X16 | BLOCK_32X32 | BLOCK_32X64 => 5,
251      BLOCK_64X16 | BLOCK_64X32 | BLOCK_64X64 | BLOCK_64X128 => 6,
252      BLOCK_128X64 | BLOCK_128X128 => 7,
253    }
254  }
255
256  #[inline]
257  pub const fn width_mi_log2(self) -> usize {
258    self.width_log2() - 2
259  }
260
261  #[inline]
262  pub const fn width_mi(self) -> usize {
263    self.width() >> MI_SIZE_LOG2
264  }
265
266  #[inline]
267  pub fn width_imp_b(self) -> usize {
268    (self.width() >> (IMPORTANCE_BLOCK_TO_BLOCK_SHIFT + BLOCK_TO_PLANE_SHIFT))
269      .max(1)
270  }
271
272  #[inline]
273  pub const fn height(self) -> usize {
274    1 << self.height_log2()
275  }
276
277  #[inline]
278  pub const fn height_log2(self) -> usize {
279    match self {
280      BLOCK_4X4 | BLOCK_8X4 | BLOCK_16X4 => 2,
281      BLOCK_4X8 | BLOCK_8X8 | BLOCK_16X8 | BLOCK_32X8 => 3,
282      BLOCK_4X16 | BLOCK_8X16 | BLOCK_16X16 | BLOCK_32X16 | BLOCK_64X16 => 4,
283      BLOCK_8X32 | BLOCK_16X32 | BLOCK_32X32 | BLOCK_64X32 => 5,
284      BLOCK_16X64 | BLOCK_32X64 | BLOCK_64X64 | BLOCK_128X64 => 6,
285      BLOCK_64X128 | BLOCK_128X128 => 7,
286    }
287  }
288
289  #[inline]
290  pub const fn height_mi_log2(self) -> usize {
291    self.height_log2() - 2
292  }
293
294  #[inline]
295  pub const fn height_mi(self) -> usize {
296    self.height() >> MI_SIZE_LOG2
297  }
298
299  #[inline]
300  pub fn height_imp_b(self) -> usize {
301    (self.height() >> (IMPORTANCE_BLOCK_TO_BLOCK_SHIFT + BLOCK_TO_PLANE_SHIFT))
302      .max(1)
303  }
304
305  #[inline]
306  pub const fn tx_size(self) -> TxSize {
307    match self {
308      BLOCK_4X4 => TX_4X4,
309      BLOCK_4X8 => TX_4X8,
310      BLOCK_8X4 => TX_8X4,
311      BLOCK_8X8 => TX_8X8,
312      BLOCK_8X16 => TX_8X16,
313      BLOCK_16X8 => TX_16X8,
314      BLOCK_16X16 => TX_16X16,
315      BLOCK_16X32 => TX_16X32,
316      BLOCK_32X16 => TX_32X16,
317      BLOCK_32X32 => TX_32X32,
318      BLOCK_32X64 => TX_32X64,
319      BLOCK_64X32 => TX_64X32,
320      BLOCK_4X16 => TX_4X16,
321      BLOCK_16X4 => TX_16X4,
322      BLOCK_8X32 => TX_8X32,
323      BLOCK_32X8 => TX_32X8,
324      BLOCK_16X64 => TX_16X64,
325      BLOCK_64X16 => TX_64X16,
326      _ => TX_64X64,
327    }
328  }
329
330  /// Source: `Subsampled_Size` (AV1 specification section 5.11.38)
331  ///
332  /// # Errors
333  ///
334  /// - Returns `InvalidBlockSize` if the given block size cannot
335  ///   be subsampled in the requested way.
336  #[inline]
337  pub const fn subsampled_size(
338    self, xdec: usize, ydec: usize,
339  ) -> Result<BlockSize, InvalidBlockSize> {
340    Ok(match (xdec, ydec) {
341      (0, 0) /* 4:4:4 */ => self,
342      (1, 0) /* 4:2:2 */ => match self {
343        BLOCK_4X4 | BLOCK_8X4 => BLOCK_4X4,
344        BLOCK_8X8 => BLOCK_4X8,
345        BLOCK_16X4 => BLOCK_8X4,
346        BLOCK_16X8 => BLOCK_8X8,
347        BLOCK_16X16 => BLOCK_8X16,
348        BLOCK_32X8 => BLOCK_16X8,
349        BLOCK_32X16 => BLOCK_16X16,
350        BLOCK_32X32 => BLOCK_16X32,
351        BLOCK_64X16 => BLOCK_32X16,
352        BLOCK_64X32 => BLOCK_32X32,
353        BLOCK_64X64 => BLOCK_32X64,
354        BLOCK_128X64 => BLOCK_64X64,
355        BLOCK_128X128 => BLOCK_64X128,
356        _ => return Err(InvalidBlockSize),
357      },
358      (1, 1) /* 4:2:0 */ => match self {
359        BLOCK_4X4 | BLOCK_4X8 | BLOCK_8X4 | BLOCK_8X8 => BLOCK_4X4,
360        BLOCK_4X16 | BLOCK_8X16 => BLOCK_4X8,
361        BLOCK_8X32 => BLOCK_4X16,
362        BLOCK_16X4 | BLOCK_16X8 => BLOCK_8X4,
363        BLOCK_16X16 => BLOCK_8X8,
364        BLOCK_16X32 => BLOCK_8X16,
365        BLOCK_16X64 => BLOCK_8X32,
366        BLOCK_32X8 => BLOCK_16X4,
367        BLOCK_32X16 => BLOCK_16X8,
368        BLOCK_32X32 => BLOCK_16X16,
369        BLOCK_32X64 => BLOCK_16X32,
370        BLOCK_64X16 => BLOCK_32X8,
371        BLOCK_64X32 => BLOCK_32X16,
372        BLOCK_64X64 => BLOCK_32X32,
373        BLOCK_64X128 => BLOCK_32X64,
374        BLOCK_128X64 => BLOCK_64X32,
375        BLOCK_128X128 => BLOCK_64X64,
376      },
377      _ => return Err(InvalidBlockSize),
378    })
379  }
380
381  /// # Panics
382  ///
383  /// Will panic if the subsampling is not possible
384  #[inline]
385  pub fn largest_chroma_tx_size(self, xdec: usize, ydec: usize) -> TxSize {
386    let plane_bsize = self
387      .subsampled_size(xdec, ydec)
388      .expect("invalid block size for this subsampling mode");
389
390    let chroma_tx_size = max_txsize_rect_lookup[plane_bsize as usize];
391
392    av1_get_coded_tx_size(chroma_tx_size)
393  }
394
395  #[inline]
396  pub const fn is_sqr(self) -> bool {
397    self.width_log2() == self.height_log2()
398  }
399
400  #[inline]
401  pub const fn is_sub8x8(self, xdec: usize, ydec: usize) -> bool {
402    xdec != 0 && self.width_log2() == 2 || ydec != 0 && self.height_log2() == 2
403  }
404
405  #[inline]
406  pub const fn sub8x8_offset(
407    self, xdec: usize, ydec: usize,
408  ) -> (isize, isize) {
409    let offset_x = if xdec != 0 && self.width_log2() == 2 { -1 } else { 0 };
410    let offset_y = if ydec != 0 && self.height_log2() == 2 { -1 } else { 0 };
411
412    (offset_x, offset_y)
413  }
414
415  /// # Errors
416  ///
417  /// - Returns `InvalidBlockSize` if the block size cannot be split
418  ///   in the requested way.
419  pub const fn subsize(
420    self, partition: PartitionType,
421  ) -> Result<BlockSize, InvalidBlockSize> {
422    use PartitionType::*;
423
424    Ok(match partition {
425      PARTITION_NONE => self,
426      PARTITION_SPLIT => match self {
427        BLOCK_8X8 => BLOCK_4X4,
428        BLOCK_16X16 => BLOCK_8X8,
429        BLOCK_32X32 => BLOCK_16X16,
430        BLOCK_64X64 => BLOCK_32X32,
431        BLOCK_128X128 => BLOCK_64X64,
432        _ => return Err(InvalidBlockSize),
433      },
434      PARTITION_HORZ | PARTITION_HORZ_A | PARTITION_HORZ_B => match self {
435        BLOCK_8X8 => BLOCK_8X4,
436        BLOCK_16X16 => BLOCK_16X8,
437        BLOCK_32X32 => BLOCK_32X16,
438        BLOCK_64X64 => BLOCK_64X32,
439        BLOCK_128X128 => BLOCK_128X64,
440        _ => return Err(InvalidBlockSize),
441      },
442      PARTITION_VERT | PARTITION_VERT_A | PARTITION_VERT_B => match self {
443        BLOCK_8X8 => BLOCK_4X8,
444        BLOCK_16X16 => BLOCK_8X16,
445        BLOCK_32X32 => BLOCK_16X32,
446        BLOCK_64X64 => BLOCK_32X64,
447        BLOCK_128X128 => BLOCK_64X128,
448        _ => return Err(InvalidBlockSize),
449      },
450      PARTITION_HORZ_4 => match self {
451        BLOCK_16X16 => BLOCK_16X4,
452        BLOCK_32X32 => BLOCK_32X8,
453        BLOCK_64X64 => BLOCK_64X16,
454        _ => return Err(InvalidBlockSize),
455      },
456      PARTITION_VERT_4 => match self {
457        BLOCK_16X16 => BLOCK_4X16,
458        BLOCK_32X32 => BLOCK_8X32,
459        BLOCK_64X64 => BLOCK_16X64,
460        _ => return Err(InvalidBlockSize),
461      },
462      _ => return Err(InvalidBlockSize),
463    })
464  }
465
466  pub const fn is_rect_tx_allowed(self) -> bool {
467    !matches!(
468      self,
469      BLOCK_4X4
470        | BLOCK_8X8
471        | BLOCK_16X16
472        | BLOCK_32X32
473        | BLOCK_64X64
474        | BLOCK_64X128
475        | BLOCK_128X64
476        | BLOCK_128X128
477    )
478  }
479}
480
481impl fmt::Display for BlockSize {
482  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
483    write!(
484      f,
485      "{}",
486      match self {
487        BlockSize::BLOCK_4X4 => "4x4",
488        BlockSize::BLOCK_4X8 => "4x8",
489        BlockSize::BLOCK_8X4 => "8x4",
490        BlockSize::BLOCK_8X8 => "8x8",
491        BlockSize::BLOCK_8X16 => "8x16",
492        BlockSize::BLOCK_16X8 => "16x8",
493        BlockSize::BLOCK_16X16 => "16x16",
494        BlockSize::BLOCK_16X32 => "16x32",
495        BlockSize::BLOCK_32X16 => "32x16",
496        BlockSize::BLOCK_32X32 => "32x32",
497        BlockSize::BLOCK_32X64 => "32x64",
498        BlockSize::BLOCK_64X32 => "64x32",
499        BlockSize::BLOCK_64X64 => "64x64",
500        BlockSize::BLOCK_64X128 => "64x128",
501        BlockSize::BLOCK_128X64 => "128x64",
502        BlockSize::BLOCK_128X128 => "128x128",
503        BlockSize::BLOCK_4X16 => "4x16",
504        BlockSize::BLOCK_16X4 => "16x4",
505        BlockSize::BLOCK_8X32 => "8x32",
506        BlockSize::BLOCK_32X8 => "32x8",
507        BlockSize::BLOCK_16X64 => "16x64",
508        BlockSize::BLOCK_64X16 => "64x16",
509      }
510    )
511  }
512}
513
514pub const NEWMV_MODE_CONTEXTS: usize = 7;
515pub const GLOBALMV_MODE_CONTEXTS: usize = 2;
516pub const REFMV_MODE_CONTEXTS: usize = 6;
517pub const INTER_COMPOUND_MODES: usize = 8;
518
519pub const REFMV_OFFSET: usize = 4;
520pub const GLOBALMV_OFFSET: usize = 3;
521pub const NEWMV_CTX_MASK: usize = (1 << GLOBALMV_OFFSET) - 1;
522pub const GLOBALMV_CTX_MASK: usize =
523  (1 << (REFMV_OFFSET - GLOBALMV_OFFSET)) - 1;
524pub const REFMV_CTX_MASK: usize = (1 << (8 - REFMV_OFFSET)) - 1;
525
526pub static RAV1E_PARTITION_TYPES: &[PartitionType] = &[
527  PartitionType::PARTITION_NONE,
528  PartitionType::PARTITION_HORZ,
529  PartitionType::PARTITION_VERT,
530  PartitionType::PARTITION_SPLIT,
531];
532
533#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd)]
534pub enum GlobalMVMode {
535  IDENTITY = 0,    // identity transformation, 0-parameter
536  TRANSLATION = 1, // translational motion 2-parameter
537  ROTZOOM = 2,     // simplified affine with rotation + zoom only, 4-parameter
538  AFFINE = 3,      // affine, 6-parameter
539}
540
541#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd)]
542pub enum MvSubpelPrecision {
543  MV_SUBPEL_NONE = -1,
544  MV_SUBPEL_LOW_PRECISION = 0,
545  MV_SUBPEL_HIGH_PRECISION,
546}
547
548/* Symbols for coding which components are zero jointly */
549pub const MV_JOINTS: usize = 4;
550
551#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd)]
552pub enum MvJointType {
553  MV_JOINT_ZERO = 0,   /* Zero vector */
554  MV_JOINT_HNZVZ = 1,  /* Vert zero, hor nonzero */
555  MV_JOINT_HZVNZ = 2,  /* Hor zero, vert nonzero */
556  MV_JOINT_HNZVNZ = 3, /* Both components nonzero */
557}
558
559fn supersample_chroma_bsize(
560  bsize: BlockSize, ss_x: usize, ss_y: usize,
561) -> BlockSize {
562  debug_assert!(ss_x < 2);
563  debug_assert!(ss_y < 2);
564
565  match bsize {
566    BLOCK_4X4 => match (ss_x, ss_y) {
567      (1, 1) => BLOCK_8X8,
568      (1, 0) => BLOCK_8X4,
569      (0, 1) => BLOCK_4X8,
570      _ => bsize,
571    },
572    BLOCK_4X8 => match (ss_x, ss_y) {
573      (1, 1) => BLOCK_8X8,
574      (1, 0) => BLOCK_8X8,
575      (0, 1) => BLOCK_4X8,
576      _ => bsize,
577    },
578    BLOCK_8X4 => match (ss_x, ss_y) {
579      (1, 1) => BLOCK_8X8,
580      (1, 0) => BLOCK_8X4,
581      (0, 1) => BLOCK_8X8,
582      _ => bsize,
583    },
584    BLOCK_4X16 => match (ss_x, ss_y) {
585      (1, 1) => BLOCK_8X16,
586      (1, 0) => BLOCK_8X16,
587      (0, 1) => BLOCK_4X16,
588      _ => bsize,
589    },
590    BLOCK_16X4 => match (ss_x, ss_y) {
591      (1, 1) => BLOCK_16X8,
592      (1, 0) => BLOCK_16X4,
593      (0, 1) => BLOCK_16X8,
594      _ => bsize,
595    },
596    _ => bsize,
597  }
598}
599
600type IntraEdgeBuffer<T> = Aligned<[MaybeUninit<T>; 4 * MAX_TX_SIZE + 1]>;
601
602#[cfg(any(test, feature = "bench"))]
603type IntraEdgeMock<T> = Aligned<[T; 4 * MAX_TX_SIZE + 1]>;
604
605pub struct IntraEdge<'a, T: Pixel>(&'a [T], &'a [T], &'a [T]);
606
607impl<'a, T: Pixel> IntraEdge<'a, T> {
608  fn new(
609    edge_buf: &'a mut IntraEdgeBuffer<T>, init_left: usize, init_above: usize,
610  ) -> Self {
611    // SAFETY: Initialized in `get_intra_edges`.
612    let left = unsafe {
613      let begin_left = 2 * MAX_TX_SIZE - init_left;
614      let end_above = 2 * MAX_TX_SIZE + 1 + init_above;
615      slice_assume_init_mut(&mut edge_buf.data[begin_left..end_above])
616    };
617    let (left, top_left) = left.split_at(init_left);
618    let (top_left, above) = top_left.split_at(1);
619    Self(left, top_left, above)
620  }
621
622  pub const fn as_slices(&self) -> (&'a [T], &'a [T], &'a [T]) {
623    (self.0, self.1, self.2)
624  }
625
626  pub const fn top_left_ptr(&self) -> *const T {
627    self.1.as_ptr()
628  }
629
630  #[cfg(any(test, feature = "bench"))]
631  pub fn mock(edge_buf: &'a IntraEdgeMock<T>) -> Self {
632    let left = &edge_buf.data[..];
633    let (left, top_left) = left.split_at(2 * MAX_TX_SIZE);
634    let (top_left, above) = top_left.split_at(1);
635    Self(left, top_left, above)
636  }
637}
638
639pub fn get_intra_edges<'a, T: Pixel>(
640  edge_buf: &'a mut IntraEdgeBuffer<T>,
641  dst: &PlaneRegion<'_, T>,
642  partition_bo: TileBlockOffset, // partition bo, BlockOffset
643  bx: usize,
644  by: usize,
645  partition_size: BlockSize, // partition size, BlockSize
646  po: PlaneOffset,
647  tx_size: TxSize,
648  bit_depth: usize,
649  opt_mode: Option<PredictionMode>,
650  enable_intra_edge_filter: bool,
651  intra_param: IntraParam,
652) -> IntraEdge<'a, T> {
653  let mut init_left: usize = 0;
654  let mut init_above: usize = 0;
655
656  let plane_cfg = &dst.plane_cfg;
657
658  let base = 128u16 << (bit_depth - 8);
659
660  {
661    // left pixels are ordered from bottom to top and right-aligned
662    let (left, not_left) = edge_buf.data.split_at_mut(2 * MAX_TX_SIZE);
663    let (top_left, above) = not_left.split_at_mut(1);
664
665    let x = po.x as usize;
666    let y = po.y as usize;
667
668    let mut needs_left = true;
669    let mut needs_topleft = true;
670    let mut needs_top = true;
671    let mut needs_topright = true;
672    let mut needs_bottomleft = true;
673    let mut needs_topleft_filter = false;
674
675    if let Some(mut mode) = opt_mode {
676      mode = match mode {
677        PredictionMode::PAETH_PRED => match (x, y) {
678          (0, 0) => PredictionMode::DC_PRED,
679          (0, _) => PredictionMode::V_PRED,
680          (_, 0) => PredictionMode::H_PRED,
681          _ => PredictionMode::PAETH_PRED,
682        },
683        _ => mode,
684      };
685
686      let p_angle = intra_mode_to_angle(mode)
687        + match intra_param {
688          IntraParam::AngleDelta(val) => (val * ANGLE_STEP) as isize,
689          _ => 0,
690        };
691
692      let dc_or_cfl =
693        mode == PredictionMode::DC_PRED || mode == PredictionMode::UV_CFL_PRED;
694
695      needs_left = (!dc_or_cfl || x != 0) || (p_angle > 90 && p_angle != 180);
696      needs_topleft = mode == PredictionMode::PAETH_PRED
697        || (mode.is_directional() && p_angle != 90 && p_angle != 180);
698      needs_top = (!dc_or_cfl || y != 0) || (p_angle != 90 && p_angle < 180);
699      needs_topright = mode.is_directional() && p_angle < 90;
700      needs_bottomleft = mode.is_directional() && p_angle > 180;
701      needs_topleft_filter =
702        enable_intra_edge_filter && p_angle > 90 && p_angle < 180;
703    }
704
705    let rect_w =
706      dst.rect().width.min(dst.plane_cfg.width - dst.rect().x as usize);
707    let rect_h =
708      dst.rect().height.min(dst.plane_cfg.height - dst.rect().y as usize);
709
710    // Needs left
711    if needs_left {
712      let txh = if y + tx_size.height() > rect_h {
713        rect_h - y
714      } else {
715        tx_size.height()
716      };
717      if x != 0 {
718        for i in 0..txh {
719          debug_assert!(y + i < rect_h);
720          left[2 * MAX_TX_SIZE - 1 - i].write(dst[y + i][x - 1]);
721        }
722        if txh < tx_size.height() {
723          let val = dst[y + txh - 1][x - 1];
724          for i in txh..tx_size.height() {
725            left[2 * MAX_TX_SIZE - 1 - i].write(val);
726          }
727        }
728      } else {
729        let val = if y != 0 { dst[y - 1][0] } else { T::cast_from(base + 1) };
730        for v in left[2 * MAX_TX_SIZE - tx_size.height()..].iter_mut() {
731          v.write(val);
732        }
733      }
734      init_left += tx_size.height();
735    }
736
737    // Needs top
738    if needs_top {
739      let txw = if x + tx_size.width() > rect_w {
740        rect_w - x
741      } else {
742        tx_size.width()
743      };
744      if y != 0 {
745        above[..txw].copy_from_slice(
746          // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
747          unsafe {
748            transmute::<&[T], &[MaybeUninit<T>]>(&dst[y - 1][x..x + txw])
749          },
750        );
751        if txw < tx_size.width() {
752          let val = dst[y - 1][x + txw - 1];
753          for i in txw..tx_size.width() {
754            above[i].write(val);
755          }
756        }
757      } else {
758        let val = if x != 0 { dst[0][x - 1] } else { T::cast_from(base - 1) };
759        for v in above[..tx_size.width()].iter_mut() {
760          v.write(val);
761        }
762      }
763      init_above += tx_size.width();
764    }
765
766    let bx4 = bx * (tx_size.width() >> MI_SIZE_LOG2); // bx,by are in tx block indices
767    let by4 = by * (tx_size.height() >> MI_SIZE_LOG2);
768
769    let have_top = by4 != 0
770      || if plane_cfg.ydec != 0 {
771        partition_bo.0.y > 1
772      } else {
773        partition_bo.0.y > 0
774      };
775    let have_left = bx4 != 0
776      || if plane_cfg.xdec != 0 {
777        partition_bo.0.x > 1
778      } else {
779        partition_bo.0.x > 0
780      };
781
782    let right_available = x + tx_size.width() < rect_w;
783    let bottom_available = y + tx_size.height() < rect_h;
784
785    let scaled_partition_size =
786      supersample_chroma_bsize(partition_size, plane_cfg.xdec, plane_cfg.ydec);
787
788    // Needs top right
789    if needs_topright {
790      debug_assert!(plane_cfg.xdec <= 1 && plane_cfg.ydec <= 1);
791
792      let num_avail = if y != 0
793        && has_top_right(
794          scaled_partition_size,
795          partition_bo,
796          have_top,
797          right_available,
798          tx_size,
799          by4,
800          bx4,
801          plane_cfg.xdec,
802          plane_cfg.ydec,
803        ) {
804        tx_size.width().min(rect_w - x - tx_size.width())
805      } else {
806        0
807      };
808      if num_avail > 0 {
809        above[tx_size.width()..][..num_avail].copy_from_slice(
810          // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
811          unsafe {
812            transmute::<&[T], &[MaybeUninit<T>]>(
813              &dst[y - 1][x + tx_size.width()..][..num_avail],
814            )
815          },
816        );
817      }
818      if num_avail < tx_size.height() {
819        let val = above[tx_size.width() + num_avail - 1];
820        for v in above
821          [tx_size.width() + num_avail..tx_size.width() + tx_size.height()]
822          .iter_mut()
823        {
824          *v = val;
825        }
826      }
827      init_above += tx_size.height();
828    }
829
830    // SAFETY: The blocks above have initialized the first `init_above` items.
831    let above = unsafe { slice_assume_init_mut(&mut above[..init_above]) };
832
833    // Needs bottom left
834    if needs_bottomleft {
835      debug_assert!(plane_cfg.xdec <= 1 && plane_cfg.ydec <= 1);
836
837      let num_avail = if x != 0
838        && has_bottom_left(
839          scaled_partition_size,
840          partition_bo,
841          bottom_available,
842          have_left,
843          tx_size,
844          by4,
845          bx4,
846          plane_cfg.xdec,
847          plane_cfg.ydec,
848        ) {
849        tx_size.height().min(rect_h - y - tx_size.height())
850      } else {
851        0
852      };
853      if num_avail > 0 {
854        for i in 0..num_avail {
855          left[2 * MAX_TX_SIZE - tx_size.height() - 1 - i]
856            .write(dst[y + tx_size.height() + i][x - 1]);
857        }
858      }
859      if num_avail < tx_size.width() {
860        let val = left[2 * MAX_TX_SIZE - tx_size.height() - num_avail];
861        for v in left[(2 * MAX_TX_SIZE - tx_size.height() - tx_size.width())
862          ..(2 * MAX_TX_SIZE - tx_size.height() - num_avail)]
863          .iter_mut()
864        {
865          *v = val;
866        }
867      }
868      init_left += tx_size.width();
869    }
870
871    // SAFETY: The blocks above have initialized last `init_left` items.
872    let left = unsafe {
873      slice_assume_init_mut(&mut left[2 * MAX_TX_SIZE - init_left..])
874    };
875
876    // Needs top-left
877    if needs_topleft {
878      let top_left = top_left[0].write(match (x, y) {
879        (0, 0) => T::cast_from(base),
880        (_, 0) => dst[0][x - 1],
881        (0, _) => dst[y - 1][0],
882        _ => dst[y - 1][x - 1],
883      });
884
885      let (w, h) = (tx_size.width(), tx_size.height());
886      if needs_topleft_filter && w + h >= 24 {
887        let (l, a, tl): (u32, u32, u32) =
888          (left[left.len() - 1].into(), above[0].into(), (*top_left).into());
889        let s = l * 5 + tl * 6 + a * 5;
890
891        *top_left = T::cast_from((s + (1 << 3)) >> 4);
892      }
893    } else {
894      top_left[0].write(T::cast_from(base));
895    }
896  }
897  IntraEdge::new(edge_buf, init_left, init_above)
898}
899
900pub fn has_tr(bo: TileBlockOffset, bsize: BlockSize) -> bool {
901  let sb_mi_size = BLOCK_64X64.width_mi(); /* Assume 64x64 for now */
902  let mask_row = bo.0.y & LOCAL_BLOCK_MASK;
903  let mask_col = bo.0.x & LOCAL_BLOCK_MASK;
904  let target_n4_w = bsize.width_mi();
905  let target_n4_h = bsize.height_mi();
906
907  let mut bs = target_n4_w.max(target_n4_h);
908
909  if bs > BLOCK_64X64.width_mi() {
910    return false;
911  }
912
913  let mut has_tr = !((mask_row & bs) != 0 && (mask_col & bs) != 0);
914
915  /* TODO: assert its a power of two */
916
917  while bs < sb_mi_size {
918    if (mask_col & bs) != 0 {
919      if (mask_col & (2 * bs) != 0) && (mask_row & (2 * bs) != 0) {
920        has_tr = false;
921        break;
922      }
923    } else {
924      break;
925    }
926    bs <<= 1;
927  }
928
929  /* The left hand of two vertical rectangles always has a top right (as the
930   * block above will have been decoded) */
931  if (target_n4_w < target_n4_h) && (bo.0.x & target_n4_w) == 0 {
932    has_tr = true;
933  }
934
935  /* The bottom of two horizontal rectangles never has a top right (as the block
936   * to the right won't have been decoded) */
937  if (target_n4_w > target_n4_h) && (bo.0.y & target_n4_h) != 0 {
938    has_tr = false;
939  }
940
941  /* The bottom left square of a Vertical A (in the old format) does
942   * not have a top right as it is decoded before the right hand
943   * rectangle of the partition */
944  /*
945    if blk.partition == PartitionType::PARTITION_VERT_A {
946      if blk.n4_w == blk.n4_h {
947        if (mask_row & bs) != 0 {
948          has_tr = false;
949        }
950      }
951    }
952  */
953
954  has_tr
955}
956
957pub fn has_bl(bo: TileBlockOffset, bsize: BlockSize) -> bool {
958  let sb_mi_size = BLOCK_64X64.width_mi(); /* Assume 64x64 for now */
959  let mask_row = bo.0.y & LOCAL_BLOCK_MASK;
960  let mask_col = bo.0.x & LOCAL_BLOCK_MASK;
961  let target_n4_w = bsize.width_mi();
962  let target_n4_h = bsize.height_mi();
963
964  let mut bs = target_n4_w.max(target_n4_h);
965
966  if bs > BLOCK_64X64.width_mi() {
967    return false;
968  }
969
970  let mut has_bl =
971    (mask_row & bs) == 0 && (mask_col & bs) == 0 && bs < sb_mi_size;
972
973  /* TODO: assert its a power of two */
974
975  while 2 * bs < sb_mi_size {
976    if (mask_col & bs) == 0 {
977      if (mask_col & (2 * bs) == 0) && (mask_row & (2 * bs) == 0) {
978        has_bl = true;
979        break;
980      }
981    } else {
982      break;
983    }
984    bs <<= 1;
985  }
986
987  /* The right hand of two vertical rectangles never has a bottom left (as the
988   * block below won't have been decoded) */
989  if (target_n4_w < target_n4_h) && (bo.0.x & target_n4_w) != 0 {
990    has_bl = false;
991  }
992
993  /* The top of two horizontal rectangles always has a bottom left (as the block
994   * to the left will have been decoded) */
995  if (target_n4_w > target_n4_h) && (bo.0.y & target_n4_h) == 0 {
996    has_bl = true;
997  }
998
999  /* The bottom left square of a Vertical A (in the old format) does
1000   * not have a top right as it is decoded before the right hand
1001   * rectangle of the partition */
1002  /*
1003    if blk.partition == PartitionType::PARTITION_VERT_A {
1004      if blk.n4_w == blk.n4_h {
1005        if (mask_row & bs) != 0 {
1006          has_tr = false;
1007        }
1008      }
1009    }
1010  */
1011
1012  has_bl
1013}
1014
1015#[cfg(test)]
1016mod tests {
1017  use crate::partition::BlockSize::*;
1018  use crate::partition::{BlockSize, InvalidBlockSize};
1019
1020  #[test]
1021  fn from_wh_matches_naive() {
1022    fn from_wh_opt_naive(
1023      w: usize, h: usize,
1024    ) -> Result<BlockSize, InvalidBlockSize> {
1025      match (w, h) {
1026        (4, 4) => Ok(BLOCK_4X4),
1027        (4, 8) => Ok(BLOCK_4X8),
1028        (8, 4) => Ok(BLOCK_8X4),
1029        (8, 8) => Ok(BLOCK_8X8),
1030        (8, 16) => Ok(BLOCK_8X16),
1031        (16, 8) => Ok(BLOCK_16X8),
1032        (16, 16) => Ok(BLOCK_16X16),
1033        (16, 32) => Ok(BLOCK_16X32),
1034        (32, 16) => Ok(BLOCK_32X16),
1035        (32, 32) => Ok(BLOCK_32X32),
1036        (32, 64) => Ok(BLOCK_32X64),
1037        (64, 32) => Ok(BLOCK_64X32),
1038        (64, 64) => Ok(BLOCK_64X64),
1039        (64, 128) => Ok(BLOCK_64X128),
1040        (128, 64) => Ok(BLOCK_128X64),
1041        (128, 128) => Ok(BLOCK_128X128),
1042        (4, 16) => Ok(BLOCK_4X16),
1043        (16, 4) => Ok(BLOCK_16X4),
1044        (8, 32) => Ok(BLOCK_8X32),
1045        (32, 8) => Ok(BLOCK_32X8),
1046        (16, 64) => Ok(BLOCK_16X64),
1047        (64, 16) => Ok(BLOCK_64X16),
1048        _ => Err(InvalidBlockSize),
1049      }
1050    }
1051
1052    for w in 0..256 {
1053      for h in 0..256 {
1054        let a = BlockSize::from_width_and_height_opt(w, h);
1055        let b = from_wh_opt_naive(w, h);
1056
1057        assert_eq!(a, b);
1058      }
1059    }
1060  }
1061}
rav1e/partition.rs

rav1e/
partition.rs
