encoder.rs 80.4 KB
Newer Older
rzumer's avatar
rzumer committed
1 2 3 4 5 6 7 8 9
// Copyright (c) 2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

Raphaël Zumer's avatar
Raphaël Zumer committed
10 11 12 13 14 15 16 17 18 19 20
use crate::api::*;
use crate::cdef::*;
use crate::context::*;
use crate::deblock::*;
use crate::ec::*;
use crate::lrf::*;
use crate::mc::*;
use crate::me::*;
use crate::partition::*;
use crate::plane::*;
use crate::quantize::*;
21 22 23
use crate::rate::QuantizerParameters;
use crate::rate::FRAME_SUBTYPE_I;
use crate::rate::FRAME_SUBTYPE_P;
Raphaël Zumer's avatar
Raphaël Zumer committed
24 25 26 27 28
use crate::rdo::*;
use crate::segmentation::*;
use crate::transform::*;
use crate::util::*;
use crate::partition::PartitionType::*;
Thomas Daede's avatar
Thomas Daede committed
29
use crate::partition::RefType::*;
30
use crate::header::*;
31

32
use arg_enum_proc_macro::ArgEnum;
33
use bitstream_io::{BitWriter, BigEndian};
34
use bincode::{serialize, deserialize};
35
use std;
36
use std::{fmt, io, mem};
Josh Holmer's avatar
Josh Holmer committed
37
use std::io::Write;
38
use std::io::Read;
Raphaël Zumer's avatar
Raphaël Zumer committed
39
use std::sync::Arc;
40
use std::fs::File;
41 42

#[derive(Debug, Clone)]
43 44
pub struct Frame<T: Pixel> {
  pub planes: [Plane<T>; 3]
45 46
}

47 48
pub static TEMPORAL_DELIMITER: [u8; 2] = [0x12, 0x00];

49 50
const FRAME_MARGIN: usize = 16 + SUBPEL_FILTER_SIZE;

51 52
impl<T: Pixel> Frame<T> {
  pub fn new(width: usize, height: usize, chroma_sampling: ChromaSampling) -> Self {
53 54 55 56 57 58 59 60 61 62 63 64
    let luma_width = width.align_power_of_two(3);
    let luma_height = height.align_power_of_two(3);
    let luma_padding = MAX_SB_SIZE + FRAME_MARGIN;

    let (chroma_sampling_period_x, chroma_sampling_period_y) =
      chroma_sampling.sampling_period();
    let chroma_width = luma_width / chroma_sampling_period_x;
    let chroma_height = luma_height / chroma_sampling_period_y;
    let chroma_padding_x = luma_padding / chroma_sampling_period_x;
    let chroma_padding_y = luma_padding / chroma_sampling_period_y;
    let chroma_decimation_x = chroma_sampling_period_x - 1;
    let chroma_decimation_y = chroma_sampling_period_y - 1;
65

66 67 68
    Frame {
      planes: [
        Plane::new(
69
          luma_width, luma_height,
70
          0, 0,
71
          luma_padding, luma_padding
72 73 74
        ),
        Plane::new(
          chroma_width, chroma_height,
75 76
          chroma_decimation_x, chroma_decimation_y,
          chroma_padding_x, chroma_padding_y
77 78 79
        ),
        Plane::new(
          chroma_width, chroma_height,
80 81
          chroma_decimation_x, chroma_decimation_y,
          chroma_padding_x, chroma_padding_y
82 83
        )
      ]
84
    }
85
  }
fbossen's avatar
fbossen committed
86

87 88 89
  pub fn pad(&mut self, w: usize, h: usize) {
    for p in self.planes.iter_mut() {
      p.pad(w, h);
fbossen's avatar
fbossen committed
90
    }
91
  }
92

93 94 95 96 97 98 99 100
  /// Returns a `PixelIter` containing the data of this frame's planes in YUV format.
  /// Each point in the `PixelIter` is a triple consisting of a Y, U, and V component.
  /// The `PixelIter` is laid out as contiguous rows, e.g. to get a given 0-indexed row
  /// you could use `data.skip(width * row_idx).take(width)`.
  ///
  /// This data retains any padding, e.g. it uses the width and height specifed in
  /// the Y-plane's `cfg` struct, and not the display width and height specied in
  /// `FrameInvariants`.
101
  pub fn iter(&self) -> PixelIter<'_, T> {
102 103
    PixelIter::new(&self.planes)
  }
Josh Holmer's avatar
Josh Holmer committed
104 105 106
}

#[derive(Debug)]
107 108
pub struct PixelIter<'a, T: Pixel> {
  planes: &'a [Plane<T>; 3],
Josh Holmer's avatar
Josh Holmer committed
109 110 111 112
  y: usize,
  x: usize,
}

113 114
impl<'a, T: Pixel> PixelIter<'a, T> {
  pub fn new(planes: &'a [Plane<T>; 3]) -> Self {
Josh Holmer's avatar
Josh Holmer committed
115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130
    PixelIter {
      planes,
      y: 0,
      x: 0,
    }
  }

  fn width(&self) -> usize {
    self.planes[0].cfg.width
  }

  fn height(&self) -> usize {
    self.planes[0].cfg.height
  }
}

131 132
impl<'a, T: Pixel> Iterator for PixelIter<'a, T> {
  type Item = (T, T, T);
Josh Holmer's avatar
Josh Holmer committed
133 134 135 136 137 138 139

  fn next(&mut self) -> Option<<Self as Iterator>::Item> {
    if self.y == self.height() - 1 && self.x == self.width() - 1 {
      return None;
    }
    let pixel = (
      self.planes[0].p(self.x, self.y),
140 141
      self.planes[1].p(self.x >> self.planes[1].cfg.xdec, self.y >> self.planes[1].cfg.ydec),
      self.planes[2].p(self.x >> self.planes[2].cfg.xdec, self.y >> self.planes[2].cfg.ydec),
Josh Holmer's avatar
Josh Holmer committed
142 143 144 145 146 147 148 149 150
    );
    if self.x == self.width() - 1 {
      self.x = 0;
      self.y += 1;
    } else {
      self.x += 1;
    }
    Some(pixel)
  }
151 152
}

153
#[derive(Debug, Clone)]
154
pub struct ReferenceFrame<T: Pixel> {
fbossen's avatar
fbossen committed
155
  pub order_hint: u32,
156 157 158
  pub frame: Frame<T>,
  pub input_hres: Plane<T>,
  pub input_qres: Plane<T>,
159
  pub cdfs: CDFContext,
Romain Vimont's avatar
Romain Vimont committed
160
  pub frame_mvs: Vec<FrameMotionVectors>,
161 162
}

163
#[derive(Debug, Clone, Default)]
164 165
pub struct ReferenceFramesSet<T: Pixel> {
  pub frames: [Option<Arc<ReferenceFrame<T>>>; (REF_FRAMES as usize)],
166
  pub deblock: [DeblockState; (REF_FRAMES as usize)]
167 168
}

169 170 171
impl<T: Pixel> ReferenceFramesSet<T> {
  pub fn new() -> Self {
    Self {
172 173
      frames: Default::default(),
      deblock: Default::default()
174
    }
175
  }
176 177 178 179 180 181
}

const MAX_NUM_TEMPORAL_LAYERS: usize = 8;
const MAX_NUM_SPATIAL_LAYERS: usize = 4;
const MAX_NUM_OPERATING_POINTS: usize = MAX_NUM_TEMPORAL_LAYERS * MAX_NUM_SPATIAL_LAYERS;

182 183 184 185 186
#[derive(ArgEnum, Copy, Clone, Debug, PartialEq)]
#[repr(C)]
pub enum Tune {
  Psnr,
  Psychovisual
187 188 189
}

impl Default for Tune {
190
  fn default() -> Self {
191
    Tune::Psychovisual
192
  }
193 194
}

195
#[derive(Copy, Clone, Debug)]
196
pub struct Sequence {
197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253
  // OBU Sequence header of AV1
  pub profile: u8,
  pub num_bits_width: u32,
  pub num_bits_height: u32,
  pub bit_depth: usize,
  pub chroma_sampling: ChromaSampling,
  pub chroma_sample_position: ChromaSamplePosition,
  pub pixel_range: PixelRange,
  pub color_description: Option<ColorDescription>,
  pub mastering_display: Option<MasteringDisplay>,
  pub content_light: Option<ContentLight>,
  pub max_frame_width: u32,
  pub max_frame_height: u32,
  pub frame_id_numbers_present_flag: bool,
  pub frame_id_length: u32,
  pub delta_frame_id_length: u32,
  pub use_128x128_superblock: bool,
  pub order_hint_bits_minus_1: u32,
  pub force_screen_content_tools: u32,  // 0 - force off
  // 1 - force on
  // 2 - adaptive
  pub force_integer_mv: u32,      // 0 - Not to force. MV can be in 1/4 or 1/8
  // 1 - force to integer
  // 2 - adaptive
  pub still_picture: bool,               // Video is a single frame still picture
  pub reduced_still_picture_hdr: bool,   // Use reduced header for still picture
  pub enable_intra_edge_filter: bool,    // enables/disables corner/edge/upsampling
  pub enable_interintra_compound: bool,  // enables/disables interintra_compound
  pub enable_masked_compound: bool,      // enables/disables masked compound
  pub enable_dual_filter: bool,         // 0 - disable dual interpolation filter
  // 1 - enable vert/horiz filter selection
  pub enable_order_hint: bool,     // 0 - disable order hint, and related tools
  // jnt_comp, ref_frame_mvs, frame_sign_bias
  // if 0, enable_jnt_comp and
  // enable_ref_frame_mvs must be set zs 0.
  pub enable_jnt_comp: bool,        // 0 - disable joint compound modes
  // 1 - enable it
  pub enable_ref_frame_mvs: bool,  // 0 - disable ref frame mvs
  // 1 - enable it
  pub enable_warped_motion: bool,   // 0 - disable warped motion for sequence
  // 1 - enable it for the sequence
  pub enable_superres: bool,// 0 - Disable superres for the sequence, and disable
  //     transmitting per-frame superres enabled flag.
  // 1 - Enable superres for the sequence, and also
  //     enable per-frame flag to denote if superres is
  //     enabled for that frame.
  pub enable_cdef: bool,         // To turn on/off CDEF
  pub enable_restoration: bool,  // To turn on/off loop restoration
  pub operating_points_cnt_minus_1: usize,
  pub operating_point_idc: [u16; MAX_NUM_OPERATING_POINTS],
  pub display_model_info_present_flag: bool,
  pub decoder_model_info_present_flag: bool,
  pub level: [[usize; 2]; MAX_NUM_OPERATING_POINTS],	// minor, major
  pub tier: [usize; MAX_NUM_OPERATING_POINTS],  // seq_tier in the spec. One bit: 0
  // or 1.
  pub film_grain_params_present: bool,
  pub separate_uv_delta_q: bool,
254 255 256
}

impl Sequence {
257 258 259
  pub fn new(config: &EncoderConfig) -> Sequence {
    let width_bits = 32 - (config.width as u32).leading_zeros();
    let height_bits = 32 - (config.height as u32).leading_zeros();
260 261
    assert!(width_bits <= 16);
    assert!(height_bits <= 16);
262

263 264
    let profile = if config.bit_depth == 12 ||
      config.chroma_sampling == ChromaSampling::Cs422 {
265
      2
266
    } else if config.chroma_sampling == ChromaSampling::Cs444 {
267 268 269 270
      1
    } else {
      0
    };
271

272 273 274 275 276 277 278 279 280 281 282 283 284 285 286
    let mut operating_point_idc = [0 as u16; MAX_NUM_OPERATING_POINTS];
    let mut level = [[1, 2 as usize]; MAX_NUM_OPERATING_POINTS];
    let mut tier = [0 as usize; MAX_NUM_OPERATING_POINTS];

    for i in 0..MAX_NUM_OPERATING_POINTS {
      operating_point_idc[i] = 0;
      level[i][0] = 1;    // minor
      level[i][1] = 2;    // major
      tier[i] = 0;
    }

    Sequence {
      profile,
      num_bits_width: width_bits,
      num_bits_height: height_bits,
287 288 289 290 291 292 293 294 295
      bit_depth: config.bit_depth,
      chroma_sampling: config.chroma_sampling,
      chroma_sample_position: config.chroma_sample_position,
      pixel_range: config.pixel_range,
      color_description: config.color_description,
      mastering_display: config.mastering_display,
      content_light: config.content_light,
      max_frame_width: config.width as u32,
      max_frame_height: config.height as u32,
296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
      frame_id_numbers_present_flag: false,
      frame_id_length: 0,
      delta_frame_id_length: 0,
      use_128x128_superblock: false,
      order_hint_bits_minus_1: 5,
      force_screen_content_tools: 0,
      force_integer_mv: 2,
      still_picture: false,
      reduced_still_picture_hdr: false,
      enable_intra_edge_filter: false,
      enable_interintra_compound: false,
      enable_masked_compound: false,
      enable_dual_filter: false,
      enable_order_hint: true,
      enable_jnt_comp: false,
      enable_ref_frame_mvs: false,
      enable_warped_motion: false,
      enable_superres: false,
314
      enable_cdef: config.speed_settings.cdef,
315 316
      enable_restoration: config.chroma_sampling != ChromaSampling::Cs422 &&
        config.chroma_sampling != ChromaSampling::Cs444, // FIXME: not working yet
317 318 319 320 321 322 323 324
      operating_points_cnt_minus_1: 0,
      operating_point_idc,
      display_model_info_present_flag: false,
      decoder_model_info_present_flag: false,
      level,
      tier,
      film_grain_params_present: false,
      separate_uv_delta_q: false,
325
    }
326
  }
fbossen's avatar
fbossen committed
327

328 329 330 331 332
  pub fn get_relative_dist(&self, a: u32, b: u32) -> i32 {
    let diff = a as i32 - b as i32;
    let m = 1 << self.order_hint_bits_minus_1;
    (diff & (m - 1)) - (diff & m)
  }
fbossen's avatar
fbossen committed
333

334
  pub fn get_skip_mode_allowed<T: Pixel>(&self, fi: &FrameInvariants<T>, reference_select: bool) -> bool {
335
    if fi.intra_only || !reference_select || !self.enable_order_hint {
336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361
      return false;
    }

    let mut forward_idx: isize = -1;
    let mut backward_idx: isize = -1;
    let mut forward_hint = 0;
    let mut backward_hint = 0;

    for i in 0..INTER_REFS_PER_FRAME {
      if let Some(ref rec) = fi.rec_buffer.frames[fi.ref_frames[i] as usize] {
        let ref_hint = rec.order_hint;

        if self.get_relative_dist(ref_hint, fi.order_hint) < 0 {
          if forward_idx < 0 || self.get_relative_dist(ref_hint, forward_hint) > 0 {
            forward_idx = i as isize;
            forward_hint = ref_hint;
          }
        } else if self.get_relative_dist(ref_hint, fi.order_hint) > 0 &&
          (backward_idx < 0 || self.get_relative_dist(ref_hint, backward_hint) > 0) {
          backward_idx = i as isize;
          backward_hint = ref_hint;
        }
      }
    }

    if forward_idx < 0 {
362
      false
363 364 365
    } else if backward_idx >= 0 {
      // set skip_mode_frame
      true
366
    } else {
367 368 369
      let mut second_forward_idx: isize = -1;
      let mut second_forward_hint = 0;

370 371 372
      for i in 0..INTER_REFS_PER_FRAME {
        if let Some(ref rec) = fi.rec_buffer.frames[fi.ref_frames[i] as usize] {
          let ref_hint = rec.order_hint;
373 374 375 376 377

          if self.get_relative_dist(ref_hint, forward_hint) < 0 &&
            (second_forward_idx < 0 || self.get_relative_dist(ref_hint, second_forward_hint) > 0) {
              second_forward_idx = i as isize;
              second_forward_hint = ref_hint;
378 379 380
          }
        }
      }
381

382 383
      // TODO: Set skip_mode_frame, when second_forward_idx is not less than 0.
      second_forward_idx >= 0
fbossen's avatar
fbossen committed
384
    }
385
  }
386 387 388
}

#[derive(Debug)]
389 390 391 392 393
pub struct FrameState<T: Pixel> {
  pub input: Arc<Frame<T>>,
  pub input_hres: Plane<T>, // half-resolution version of input luma
  pub input_qres: Plane<T>, // quarter-resolution version of input luma
  pub rec: Frame<T>,
394 395 396 397 398
  pub qc: QuantizationContext,
  pub cdfs: CDFContext,
  pub deblock: DeblockState,
  pub segmentation: SegmentationState,
  pub restoration: RestorationState,
Romain Vimont's avatar
Romain Vimont committed
399
  pub frame_mvs: Vec<FrameMotionVectors>,
400
  pub t: RDOTracker,
401 402
}

403 404
impl<T: Pixel> FrameState<T> {
  pub fn new(fi: &FrameInvariants<T>) -> Self {
405
    // TODO(negge): Use fi.cfg.chroma_sampling when we store VideoDetails in FrameInvariants
406
    FrameState::new_with_frame(fi, Arc::new(Frame::new(
407
      fi.width, fi.height, fi.sequence.chroma_sampling)))
408 409
  }

410
  pub fn new_with_frame(fi: &FrameInvariants<T>, frame: Arc<Frame<T>>) -> Self {
411
    let rs = RestorationState::new(fi, &frame);
Raphaël Zumer's avatar
Raphaël Zumer committed
412 413 414 415 416
    let luma_width = frame.planes[0].cfg.width;
    let luma_height = frame.planes[0].cfg.height;
    let luma_padding_x = frame.planes[0].cfg.xpad;
    let luma_padding_y = frame.planes[0].cfg.ypad;

417
    Self {
418
      input: frame,
Raphaël Zumer's avatar
Raphaël Zumer committed
419 420
      input_hres: Plane::new(luma_width / 2, luma_height / 2, 1, 1, luma_padding_x / 2, luma_padding_y / 2),
      input_qres: Plane::new(luma_width / 4, luma_height / 4, 2, 2, luma_padding_x / 4, luma_padding_y / 4),
421
      rec: Frame::new(luma_width, luma_height, fi.sequence.chroma_sampling),
422 423 424 425 426
      qc: Default::default(),
      cdfs: CDFContext::new(0),
      deblock: Default::default(),
      segmentation: Default::default(),
      restoration: rs,
Romain Vimont's avatar
Romain Vimont committed
427 428 429 430 431 432 433
      frame_mvs: {
        let mut vec = Vec::with_capacity(REF_FRAMES);
        for _ in 0..REF_FRAMES {
          vec.push(FrameMotionVectors::new(fi.w_in_b, fi.h_in_b));
        }
        vec
      },
434
      t: RDOTracker::new()
435
    }
436
  }
437 438
}

439 440
#[derive(Copy, Clone, Debug)]
pub struct DeblockState {
441 442 443 444 445 446 447 448 449
  pub levels: [u8; PLANES+1],  // Y vertical edges, Y horizontal, U, V
  pub sharpness: u8,
  pub deltas_enabled: bool,
  pub delta_updates_enabled: bool,
  pub ref_deltas: [i8; REF_FRAMES],
  pub mode_deltas: [i8; 2],
  pub block_deltas_enabled: bool,
  pub block_delta_shift: u8,
  pub block_delta_multi: bool,
450 451 452
}

impl Default for DeblockState {
453 454 455 456 457 458 459 460 461 462 463
  fn default() -> Self {
    DeblockState {
      levels: [8,8,4,4],
      sharpness: 0,
      deltas_enabled: false, // requires delta_q_enabled
      delta_updates_enabled: false,
      ref_deltas: [1, 0, 0, 0, 0, -1, -1, -1],
      mode_deltas: [0, 0],
      block_deltas_enabled: false,
      block_delta_shift: 0,
      block_delta_multi: false
464
    }
465
  }
466 467
}

468 469
#[derive(Copy, Clone, Debug)]
pub struct SegmentationState {
470 471 472 473 474 475 476
  pub enabled: bool,
  pub update_data: bool,
  pub update_map: bool,
  pub preskip: bool,
  pub last_active_segid: u8,
  pub features: [[bool; SegLvl::SEG_LVL_MAX as usize]; 8],
  pub data: [[i16; SegLvl::SEG_LVL_MAX as usize]; 8],
477 478 479
}

impl Default for SegmentationState {
480 481 482 483 484 485 486 487 488
  fn default() -> Self {
    SegmentationState {
      enabled: false,
      update_data: false,
      update_map: false,
      preskip: true,
      last_active_segid: 0,
      features: [[false; SegLvl::SEG_LVL_MAX as usize]; 8],
      data: [[0; SegLvl::SEG_LVL_MAX as usize]; 8],
489
    }
490
  }
491 492
}

493 494
// Frame Invariants are invariant inside a frame
#[allow(dead_code)]
Josh Holmer's avatar
Josh Holmer committed
495
#[derive(Debug, Clone)]
496
pub struct FrameInvariants<T: Pixel> {
497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517
  pub sequence: Sequence,
  pub width: usize,
  pub height: usize,
  pub sb_width: usize,
  pub sb_height: usize,
  pub w_in_b: usize,
  pub h_in_b: usize,
  pub number: u64,
  pub order_hint: u32,
  pub show_frame: bool,
  pub showable_frame: bool,
  pub error_resilient: bool,
  pub intra_only: bool,
  pub allow_high_precision_mv: bool,
  pub frame_type: FrameType,
  pub show_existing_frame: bool,
  pub frame_to_show_map_idx: u32,
  pub use_reduced_tx_set: bool,
  pub reference_mode: ReferenceMode,
  pub use_prev_frame_mvs: bool,
  pub min_partition_size: BlockSize,
Thomas Daede's avatar
Thomas Daede committed
518
  pub globalmv_transformation_type: [GlobalMVMode; INTER_REFS_PER_FRAME],
519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541
  pub num_tg: usize,
  pub large_scale_tile: bool,
  pub disable_cdf_update: bool,
  pub allow_screen_content_tools: u32,
  pub force_integer_mv: u32,
  pub primary_ref_frame: u32,
  pub refresh_frame_flags: u32,  // a bitmask that specifies which
  // reference frame slots will be updated with the current frame
  // after it is decoded.
  pub allow_intrabc: bool,
  pub use_ref_frame_mvs: bool,
  pub is_filter_switchable: bool,
  pub is_motion_mode_switchable: bool,
  pub disable_frame_end_update_cdf: bool,
  pub allow_warped_motion: bool,
  pub cdef_damping: u8,
  pub cdef_bits: u8,
  pub cdef_y_strengths: [u8; 8],
  pub cdef_uv_strengths: [u8; 8],
  pub delta_q_present: bool,
  pub config: EncoderConfig,
  pub ref_frames: [u8; INTER_REFS_PER_FRAME],
  pub ref_frame_sign_bias: [bool; INTER_REFS_PER_FRAME],
542
  pub rec_buffer: ReferenceFramesSet<T>,
543 544 545
  pub base_q_idx: u8,
  pub dc_delta_q: [i8; 3],
  pub ac_delta_q: [i8; 3],
546 547
  pub lambda: f64,
  pub me_lambda: f64,
548 549
  pub me_range_scale: u8,
  pub use_tx_domain_distortion: bool,
550
  pub use_tx_domain_rate: bool,
551 552
  pub inter_cfg: Option<InterPropsConfig>,
  pub enable_early_exit: bool,
553
  pub tx_mode_select: bool,
554 555
}

556
pub(crate) fn pos_to_lvl(pos: u64, pyramid_depth: u64) -> u64 {
557 558 559 560 561 562 563 564 565 566
  // Derive level within pyramid for a frame with a given coding order position
  // For example, with a pyramid of depth 2, the 2 least significant bits of the
  // position determine the level:
  // 00 -> 0
  // 01 -> 2
  // 10 -> 1
  // 11 -> 2
  pyramid_depth - (pos | (1 << pyramid_depth)).trailing_zeros() as u64
}

567
impl<T: Pixel> FrameInvariants<T> {
568
  #[allow(clippy::erasing_op, clippy::identity_op)]
569 570
  pub fn new(config: EncoderConfig, sequence: Sequence) -> Self {
    assert!(sequence.bit_depth <= mem::size_of::<T>() * 8, "bit depth cannot fit into u8");
571 572 573
    // Speed level decides the minimum partition size, i.e. higher speed --> larger min partition size,
    // with exception that SBs on right or bottom frame borders split down to BLOCK_4X4.
    // At speed = 0, RDO search is exhaustive.
574
    let min_partition_size = config.speed_settings.min_block_size;
575 576
    let use_reduced_tx_set = config.speed_settings.reduced_tx_set;
    let use_tx_domain_distortion = config.tune == Tune::Psnr && config.speed_settings.tx_domain_distortion;
577
    let use_tx_domain_rate = config.speed_settings.tx_domain_rate;
578

579 580 581
    let w_in_b = 2 * config.width.align_power_of_two_and_shift(3); // MiCols, ((width+7)/8)<<3 >> MI_SIZE_LOG2
    let h_in_b = 2 * config.height.align_power_of_two_and_shift(3); // MiRows, ((height+7)/8)<<3 >> MI_SIZE_LOG2

582
    Self {
583
      sequence,
584 585 586 587
      width: config.width,
      height: config.height,
      sb_width: config.width.align_power_of_two_and_shift(6),
      sb_height: config.height.align_power_of_two_and_shift(6),
588 589
      w_in_b,
      h_in_b,
590 591 592 593 594 595 596 597 598 599 600 601 602 603
      number: 0,
      order_hint: 0,
      show_frame: true,
      showable_frame: true,
      error_resilient: false,
      intra_only: false,
      allow_high_precision_mv: false,
      frame_type: FrameType::KEY,
      show_existing_frame: false,
      frame_to_show_map_idx: 0,
      use_reduced_tx_set,
      reference_mode: ReferenceMode::SINGLE,
      use_prev_frame_mvs: false,
      min_partition_size,
Thomas Daede's avatar
Thomas Daede committed
604
      globalmv_transformation_type: [GlobalMVMode::IDENTITY; INTER_REFS_PER_FRAME],
605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628
      num_tg: 1,
      large_scale_tile: false,
      disable_cdf_update: false,
      allow_screen_content_tools: 0,
      force_integer_mv: 0,
      primary_ref_frame: PRIMARY_REF_NONE,
      refresh_frame_flags: 0,
      allow_intrabc: false,
      use_ref_frame_mvs: false,
      is_filter_switchable: false,
      is_motion_mode_switchable: false, // 0: only the SIMPLE motion mode will be used.
      disable_frame_end_update_cdf: false,
      allow_warped_motion: false,
      cdef_damping: 3,
      cdef_bits: 3,
      cdef_y_strengths: [0*4+0, 1*4+0, 2*4+1, 3*4+1, 5*4+2, 7*4+3, 10*4+3, 13*4+3],
      cdef_uv_strengths: [0*4+0, 1*4+0, 2*4+1, 3*4+1, 5*4+2, 7*4+3, 10*4+3, 13*4+3],
      delta_q_present: false,
      ref_frames: [0; INTER_REFS_PER_FRAME],
      ref_frame_sign_bias: [false; INTER_REFS_PER_FRAME],
      rec_buffer: ReferenceFramesSet::new(),
      base_q_idx: config.quantizer as u8,
      dc_delta_q: [0; 3],
      ac_delta_q: [0; 3],
629 630
      lambda: 0.0,
      me_lambda: 0.0,
631 632
      me_range_scale: 1,
      use_tx_domain_distortion,
633
      use_tx_domain_rate,
634 635
      inter_cfg: None,
      enable_early_exit: true,
636
      config,
637
      tx_mode_select : false,
Josh Holmer's avatar
Josh Holmer committed
638
    }
639
  }
Josh Holmer's avatar
Josh Holmer committed
640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655

  pub fn new_key_frame(previous_fi: &Self, segment_start_frame: u64) -> Self {
    let mut fi = previous_fi.clone();
    fi.frame_type = FrameType::KEY;
    fi.intra_only = true;
    fi.inter_cfg = None;
    fi.order_hint = 0;
    fi.refresh_frame_flags = ALL_REF_FRAMES_MASK;
    fi.show_frame = true;
    fi.show_existing_frame = false;
    fi.frame_to_show_map_idx = 0;
    fi.primary_ref_frame = PRIMARY_REF_NONE;
    fi.number = segment_start_frame;
    for i in 0..INTER_REFS_PER_FRAME {
      fi.ref_frames[i] = 0;
    }
656

657
    fi.tx_mode_select = fi.config.speed_settings.rdo_tx_decision;
658
    // FIXME: tx partition for intra not supported for chroma 422
659 660
    if fi.tx_mode_select && fi.sequence.chroma_sampling == ChromaSampling::Cs422 {
      fi.tx_mode_select = false;
661 662
    }

Josh Holmer's avatar
Josh Holmer committed
663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690
    fi
  }

  fn apply_inter_props_cfg(&mut self, idx_in_segment: u64) {
    let reorder = !self.config.low_latency;
    let multiref = reorder || self.config.speed_settings.multiref;

    let pyramid_depth = if reorder { 2 } else { 0 };
    let group_src_len = 1 << pyramid_depth;
    let group_len = group_src_len + if reorder { pyramid_depth } else { 0 };

    let idx_in_group = (idx_in_segment - 1) % group_len;
    let group_idx = (idx_in_segment - 1) / group_len;

    self.inter_cfg = Some(InterPropsConfig {
      reorder,
      multiref,
      pyramid_depth,
      group_src_len,
      group_len,
      idx_in_group,
      group_idx,
    })
  }

  /// Returns the created FrameInvariants along with a bool indicating success.
  /// This interface provides simpler usage, because we always need the produced
  /// FrameInvariants regardless of success or failure.
691 692 693 694
  pub fn new_inter_frame(
    previous_fi: &Self, segment_start_frame: u64, idx_in_segment: u64,
    next_keyframe: u64
  ) -> (Self, bool) {
Josh Holmer's avatar
Josh Holmer committed
695 696 697 698
    let mut fi = previous_fi.clone();
    fi.frame_type = FrameType::INTER;
    fi.intra_only = false;
    fi.apply_inter_props_cfg(idx_in_segment);
699
    fi.tx_mode_select = false;
Josh Holmer's avatar
Josh Holmer committed
700 701
    let inter_cfg = fi.inter_cfg.unwrap();

702 703 704 705 706 707 708
    fi.order_hint =
      (inter_cfg.group_src_len * inter_cfg.group_idx +
       if inter_cfg.reorder && inter_cfg.idx_in_group < inter_cfg.pyramid_depth {
         inter_cfg.group_src_len >> inter_cfg.idx_in_group
       } else {
         inter_cfg.idx_in_group - inter_cfg.pyramid_depth + 1
       }) as u32;
Josh Holmer's avatar
Josh Holmer committed
709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742
    let number = segment_start_frame + fi.order_hint as u64;
    if number >= next_keyframe {
      fi.show_existing_frame = false;
      fi.show_frame = false;
      return (fi, false);
    }

    let lvl = if !inter_cfg.reorder {
      0
    } else if inter_cfg.idx_in_group < inter_cfg.pyramid_depth {
      inter_cfg.idx_in_group
    } else {
      pos_to_lvl(inter_cfg.idx_in_group - inter_cfg.pyramid_depth + 1, inter_cfg.pyramid_depth)
    };

    // Frames with lvl == 0 are stored in slots 0..4 and frames with higher values
    // of lvl in slots 4..8
    let slot_idx = if lvl == 0 {
      (fi.order_hint >> inter_cfg.pyramid_depth) % 4 as u32
    } else {
      3 + lvl as u32
    };
    fi.show_frame = !inter_cfg.reorder || inter_cfg.idx_in_group >= inter_cfg.pyramid_depth;
    fi.show_existing_frame = fi.show_frame && inter_cfg.reorder &&
      (inter_cfg.idx_in_group - inter_cfg.pyramid_depth + 1).count_ones() == 1 &&
      inter_cfg.idx_in_group != inter_cfg.pyramid_depth;
    fi.frame_to_show_map_idx = slot_idx;
    fi.refresh_frame_flags = if fi.show_existing_frame {
      0
    } else {
      1 << slot_idx
    };

    let second_ref_frame = if !inter_cfg.multiref {
Thomas Daede's avatar
Thomas Daede committed
743
      LAST_FRAME // make second_ref_frame match first
Josh Holmer's avatar
Josh Holmer committed
744 745 746 747 748 749 750 751
    } else if !inter_cfg.reorder || inter_cfg.idx_in_group == 0 {
      LAST2_FRAME
    } else {
      ALTREF_FRAME
    };
    let ref_in_previous_group = LAST3_FRAME;

    // reuse probability estimates from previous frames only in top level frames
752 753 754
    fi.primary_ref_frame = if lvl > 0 {
      PRIMARY_REF_NONE
    } else {
Thomas Daede's avatar
Thomas Daede committed
755
      (ref_in_previous_group.to_index()) as u32
756
    };
Josh Holmer's avatar
Josh Holmer committed
757 758 759

    for i in 0..INTER_REFS_PER_FRAME {
      fi.ref_frames[i] = if lvl == 0 {
Thomas Daede's avatar
Thomas Daede committed
760
        if i == second_ref_frame.to_index() {
Josh Holmer's avatar
Josh Holmer committed
761 762 763
          (slot_idx + 4 - 2) as u8 % 4
        } else {
          (slot_idx + 4 - 1) as u8 % 4
764
        }
765 766 767 768 769 770 771 772 773 774 775 776 777 778
      } else if i == second_ref_frame.to_index() {
        let oh = fi.order_hint + (inter_cfg.group_src_len as u32 >> lvl);
        let lvl2 = pos_to_lvl(oh as u64, inter_cfg.pyramid_depth);
        if lvl2 == 0 {
          ((oh >> inter_cfg.pyramid_depth) % 4) as u8
        } else {
          3 + lvl2 as u8
        }
      } else if i == ref_in_previous_group.to_index() {
        if lvl == 0 {
          (slot_idx + 4 - 1) as u8 % 4
        } else {
          slot_idx as u8
        }
Josh Holmer's avatar
Josh Holmer committed
779
      } else {
780 781 782 783
        let oh = fi.order_hint - (inter_cfg.group_src_len as u32 >> lvl);
        let lvl1 = pos_to_lvl(oh as u64, inter_cfg.pyramid_depth);
        if lvl1 == 0 {
          ((oh >> inter_cfg.pyramid_depth) % 4) as u8
Josh Holmer's avatar
Josh Holmer committed
784
        } else {
785
          3 + lvl1 as u8
Josh Holmer's avatar
Josh Holmer committed
786 787
        }
      }
788
    }
Josh Holmer's avatar
Josh Holmer committed
789 790 791 792 793 794 795 796 797 798

    fi.reference_mode = if inter_cfg.multiref && inter_cfg.reorder && inter_cfg.idx_in_group != 0 {
      ReferenceMode::SELECT
    } else {
      ReferenceMode::SINGLE
    };
    fi.number = number;
    fi.me_range_scale = (inter_cfg.group_src_len >> lvl) as u8;
    (fi, true)
  }
799

800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820
  pub fn get_frame_subtype(&self) -> usize {
    if self.frame_type == FrameType::KEY {
      FRAME_SUBTYPE_I
    } else {
      let inter_cfg = self.inter_cfg.unwrap();
      let lvl = if !inter_cfg.reorder {
        0
      } else if inter_cfg.idx_in_group < inter_cfg.pyramid_depth {
        inter_cfg.idx_in_group
      } else {
        pos_to_lvl(
          inter_cfg.idx_in_group - inter_cfg.pyramid_depth + 1,
          inter_cfg.pyramid_depth
        )
      };
      FRAME_SUBTYPE_P + (lvl as usize)
    }
  }

  pub fn set_quantizers(&mut self, qps: &QuantizerParameters) {
    self.base_q_idx = qps.ac_qi;
821 822 823 824 825 826
    // TODO: Separate qi values for each color plane.
    if self.frame_type != FrameType::KEY {
      self.cdef_bits = 3 - ((self.base_q_idx.max(128) - 128) >> 5);
    } else {
      self.cdef_bits = 3;
    }
827 828 829 830 831 832 833 834 835
    self.base_q_idx = qps.ac_qi;
    // TODO: Separate qi values for each color plane.
    debug_assert!(qps.dc_qi as i32 - qps.ac_qi as i32 >= -128);
    debug_assert!((qps.dc_qi as i32 - qps.ac_qi as i32) < 128);
    for pi in 0..3 {
      self.dc_delta_q[pi] = (qps.dc_qi as i32 - qps.ac_qi as i32) as i8;
      self.ac_delta_q[pi] = 0;
    }
    self.lambda =
836
      qps.lambda * ((1 << (2 * (self.sequence.bit_depth - 8))) as f64);
837 838
    self.me_lambda = self.lambda.sqrt();
  }
839 840
}

841
impl<T: Pixel> fmt::Display for FrameInvariants<T> {
842 843 844
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Frame {} - {}", self.number, self.frame_type)
  }
845 846
}

Josh Holmer's avatar
Josh Holmer committed
847 848 849 850 851 852 853 854 855 856 857
#[derive(Debug, Clone, Copy)]
pub struct InterPropsConfig {
  pub reorder: bool,
  pub multiref: bool,
  pub pyramid_depth: u64,
  pub group_src_len: u64,
  pub group_len: u64,
  pub idx_in_group: u64,
  pub group_idx: u64,
}

858 859 860
pub fn write_temporal_delimiter(
  packet: &mut dyn io::Write
) -> io::Result<()> {
861
  packet.write_all(&TEMPORAL_DELIMITER)?;
862 863 864
  Ok(())
}

865 866
fn write_obus<T: Pixel>(
  packet: &mut dyn io::Write, fi: &mut FrameInvariants<T>, fs: &FrameState<T>
867 868
) -> io::Result<()> {
  let obu_extension = 0 as u32;
869

870
  let mut buf1 = Vec::new();
871

872 873
  // write sequence header obu if KEY_FRAME, preceded by 4-byte size
  if fi.frame_type == FrameType::KEY {
874 875
    let mut buf2 = Vec::new();
    {
876 877
      let mut bw2 = BitWriter::endian(&mut buf2, BigEndian);
      bw2.write_sequence_header_obu(fi)?;
Raphaël Zumer's avatar
Raphaël Zumer committed
878
      bw2.write_bit(true)?; // trailing bit
879
      bw2.byte_align()?;
880 881 882
    }

    {
883
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
884
      bw1.write_obu_header(ObuType::OBU_SEQUENCE_HEADER, obu_extension)?;
885
    }
886
    packet.write_all(&buf1).unwrap();
887 888 889
    buf1.clear();

    {
890
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
891
      bw1.write_uleb128(buf2.len() as u64)?;
892
    }
893
    packet.write_all(&buf1).unwrap();
894 895
    buf1.clear();

896
    packet.write_all(&buf2).unwrap();
897
    buf2.clear();
898 899

    if fi.sequence.content_light.is_some() {
900 901
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
      bw1.write_metadata_obu(ObuMetaType::OBU_META_HDR_CLL, fi.sequence)?;
902 903 904 905 906
      packet.write_all(&buf1).unwrap();
      buf1.clear();
    }

    if fi.sequence.mastering_display.is_some() {
907 908
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
      bw1.write_metadata_obu(ObuMetaType::OBU_META_HDR_MDCV, fi.sequence)?;
909 910 911
      packet.write_all(&buf1).unwrap();
      buf1.clear();
    }
912
  }
913

914 915 916 917 918 919 920 921
  let mut buf2 = Vec::new();
  {
    let mut bw2 = BitWriter::endian(&mut buf2, BigEndian);
    bw2.write_frame_header_obu(fi, fs)?;
  }

  {
    let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
922
    bw1.write_obu_header(ObuType::OBU_FRAME_HEADER, obu_extension)?;
923 924 925 926 927 928
  }
  packet.write_all(&buf1).unwrap();
  buf1.clear();

  {
    let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
929
    bw1.write_uleb128(buf2.len() as u64)?;
930
  }
Raphaël Zumer's avatar
Raphaël Zumer committed
931

932 933 934 935 936 937 938
  packet.write_all(&buf1).unwrap();
  buf1.clear();

  packet.write_all(&buf2).unwrap();
  buf2.clear();

  Ok(())
939 940 941
}

/// Write into `dst` the difference between the blocks at `src1` and `src2`
942
fn diff<T: Pixel>(dst: &mut [i16], src1: &PlaneSlice<'_, T>, src2: &PlaneSlice<'_, T>, width: usize, height: usize) {
fbossen's avatar
fbossen committed
943
  for ((l, s1), s2) in dst.chunks_mut(width).take(height)
944 945
    .zip(src1.rows_iter())
    .zip(src2.rows_iter()) {
946
      for ((r, v1), v2) in l.iter_mut().zip(s1).zip(s2) {
947
        *r = i16::cast_from(*v1) - i16::cast_from(*v2);
948
      }
949 950 951
    }
}

Romain Vimont's avatar
Romain Vimont committed
952
fn get_qidx<T: Pixel>(fi: &FrameInvariants<T>, fs: &FrameState<T>, cw: &ContextWriter, bo: BlockOffset) -> u8 {
953
  let mut qidx = fi.base_q_idx;
954
  let sidx = cw.bc.blocks[bo].segmentation_idx as usize;
955 956 957 958 959
  if fs.segmentation.features[sidx][SegLvl::SEG_LVL_ALT_Q as usize] {
    let delta = fs.segmentation.data[sidx][SegLvl::SEG_LVL_ALT_Q as usize];
    qidx = clamp((qidx as i16) + delta, 0, 255) as u8;
  }
  qidx
960 961
}

962 963 964
// For a transform block,
// predict, transform, quantize, write coefficients to a bitstream,
// dequantize, inverse-transform.
965 966
pub fn encode_tx_block<T: Pixel>(
  fi: &FrameInvariants<T>, fs: &mut FrameState<T>, cw: &mut ContextWriter,
Romain Vimont's avatar
Romain Vimont committed
967
  w: &mut dyn Writer, p: usize, bo: BlockOffset, mode: PredictionMode,
Romain Vimont's avatar
Romain Vimont committed
968
  tx_size: TxSize, tx_type: TxType, plane_bsize: BlockSize, po: PlaneOffset,
969
  skip: bool, ac: &[i16], alpha: i16, rdo_type: RDOType, for_rdo_use: bool
970
) -> (bool, i64) {
971 972
  let qidx = get_qidx(fi, fs, cw, bo);
  let rec = &mut fs.rec.planes[p];
973
  let PlaneConfig { xdec, ydec, .. } = fs.input.planes[p].cfg;
974

975
  assert!(tx_size.sqr() <= TxSize::TX_32X32 || tx_type == TxType::DCT_DCT);
976

977 978
  if mode.is_intra() {
    let bit_depth = fi.sequence.bit_depth;
979
    let edge_buf = get_intra_edges(rec, po, tx_size, bit_depth, Some(mode));
980 981
    mode.predict_intra(&mut rec.mut_slice(po), tx_size, bit_depth, &ac, alpha, &edge_buf);
  }
982

983
  if skip { return (false, -1); }
984

985 986 987 988 989 990 991 992
  let mut residual_storage: AlignedArray<[i16; 64 * 64]> = UninitializedAlignedArray();
  let mut coeffs_storage: AlignedArray<[i32; 64 * 64]> = UninitializedAlignedArray();
  let mut qcoeffs_storage: AlignedArray<[i32; 64 * 64]> = UninitializedAlignedArray();
  let mut rcoeffs_storage: AlignedArray<[i32; 64 * 64]> = UninitializedAlignedArray();
  let residual = &mut residual_storage.array[..tx_size.area()];
  let coeffs = &mut coeffs_storage.array[..tx_size.area()];
  let qcoeffs = &mut qcoeffs_storage.array[..tx_size.area()];
  let rcoeffs = &mut rcoeffs_storage.array[..tx_size.area()];
993

994 995 996 997 998
  diff(residual,
       &fs.input.planes[p].slice(po),
       &rec.slice(po),
       tx_size.width(),
       tx_size.height());
999

1000
  forward_transform(residual, coeffs, tx_size.width(), tx_size, tx_type, fi.sequence.bit_depth);
1001

1002 1003
  let coded_tx_size = av1_get_coded_tx_size(tx_size).area();
  fs.qc.quantize(coeffs, qcoeffs, coded_tx_size);
1004

1005 1006 1007 1008 1009 1010 1011 1012
  let tell_coeffs = w.tell_frac();
  let has_coeff = if !for_rdo_use || rdo_type.needs_coeff_rate() {
    cw.write_coeffs_lv_map(w, p, bo, &qcoeffs, mode, tx_size, tx_type, plane_bsize, xdec, ydec,
                           fi.use_reduced_tx_set)
  } else {
    true
  };
  let cost_coeffs = w.tell_frac() - tell_coeffs;
1013 1014
  // Reconstruct
  dequantize(qidx, qcoeffs, rcoeffs, tx_size, fi.sequence.bit_depth, fi.dc_delta_q[p], fi.ac_delta_q[p]);
1015

1016
  let mut tx_dist: i64 = -1;
1017

1018
  if !fi.use_tx_domain_distortion || !for_rdo_use {
1019
    inverse_transform_add(rcoeffs, &mut rec.mut_slice(po), tx_size, tx_type, fi.sequence.bit_depth);
1020 1021
  }
  if rdo_type.needs_tx_dist() {
1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034
    // Store tx-domain distortion of this block
    tx_dist = coeffs
      .iter()
      .zip(rcoeffs)
      .map(|(a, b)| {
        let c = *a as i32 - *b as i32;
        (c * c) as u64
      }).sum::<u64>() as i64;

    let tx_dist_scale_bits = 2*(3 - get_log_tx_scale(tx_size));
    let tx_dist_scale_rounding_offset = 1 << (tx_dist_scale_bits - 1);
    tx_dist = (tx_dist + tx_dist_scale_rounding_offset) >> tx_dist_scale_bits;
  }
1035 1036 1037 1038 1039 1040 1041 1042 1043
  if fi.config.train_rdo {
    fs.t.add_rate(fi.base_q_idx, tx_size, tx_dist as u64, cost_coeffs as u64);
  }

  if rdo_type == RDOType::TxDistEstRate {
    // look up rate and distortion in table
    let estimated_rate = estimate_rate(fi.base_q_idx, tx_size, tx_dist as u64);
    w.add_bits_frac(estimated_rate as u32);
  }
1044
  (has_coeff, tx_dist)
1045 1046
}

1047 1048
pub fn motion_compensate<T: Pixel>(
  fi: &FrameInvariants<T>, fs: &mut FrameState<T>, cw: &mut ContextWriter,
Thomas Daede's avatar
Thomas Daede committed
1049
  luma_mode: PredictionMode, ref_frames: [RefType; 2], mvs: [MotionVector; 2],
Romain Vimont's avatar
Romain Vimont committed
1050
  bsize: BlockSize, bo: BlockOffset, luma_only: bool
1051
) {
1052
  debug_assert!(!luma_mode.is_intra());
1053 1054 1055 1056 1057

  let PlaneConfig { xdec, ydec, .. } = fs.input.planes[1].cfg;

  // Inter mode prediction can take place once for a whole partition,
  // instead of each tx-block.
1058
  let num_planes = 1 + if !luma_only && has_chroma(bo, bsize, xdec, ydec) { 2 } else { 0 };
1059

1060 1061 1062 1063 1064 1065
  for p in 0..num_planes {
    let plane_bsize = if p == 0 { bsize }
    else { get_plane_block_size(bsize, xdec, ydec) };

    let po = bo.plane_offset(&fs.input.planes[p].cfg);
    let rec = &mut fs.rec.planes[p];
1066 1067 1068 1069

    if p > 0 && bsize < BlockSize::BLOCK_8X8 {
      let mut some_use_intra = false;
      if bsize == BlockSize::BLOCK_4X4 || bsize == BlockSize::BLOCK_4X8 {
1070
        some_use_intra |= cw.bc.blocks[bo.with_offset(-1,0)].mode.is_intra(); };
1071
      if !some_use_intra && bsize == BlockSize::BLOCK_4X4 || bsize == BlockSize::BLOCK_8X4 {
1072
        some_use_intra |= cw.bc.blocks[bo.with_offset(0,-1)].mode.is_intra(); };
1073
      if !some_use_intra && bsize == BlockSize::BLOCK_4X4 {
1074
        some_use_intra |= cw.bc.blocks[bo.with_offset(-1,-1)].mode.is_intra(); };
1075 1076

      if some_use_intra {
Romain Vimont's avatar
Romain Vimont committed
1077
        luma_mode.predict_inter(fi, p, po, &mut rec.mut_slice(po), plane_bsize.width(),
1078
                                plane_bsize.height(), ref_frames, mvs);
1079
      } else {
1080
        assert!(xdec == 1 && ydec == 1);
1081 1082
        // TODO: these are absolutely only valid for 4:2:0
        if bsize == BlockSize::BLOCK_4X4 {
1083 1084 1085 1086
          let mv0 = cw.bc.blocks[bo.with_offset(-1,-1)].mv;
          let rf0 = cw.bc.blocks[bo.with_offset(-1,-1)].ref_frames;
          let mv1 = cw.bc.blocks[bo.with_offset(0,-1)].mv;
          let rf1 = cw.bc.blocks[bo.with_offset(0,-1)].ref_frames;
1087
          let po1 = PlaneOffset { x: po.x+2, y: po.y };
1088 1089
          let mv2 = cw