encoder.rs 74.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::*;
29
use crate::header::*;
30

31
use bitstream_io::{BitWriter, BigEndian};
32
use std;
33
use std::{fmt, io, mem};
Josh Holmer's avatar
Josh Holmer committed
34
use std::io::Write;
Raphaël Zumer's avatar
Raphaël Zumer committed
35
use std::sync::Arc;
36
37

extern {
38
39
  pub fn av1_rtcd();
  pub fn aom_dsp_rtcd();
40
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
160
  pub cdfs: CDFContext,
  pub frame_mvs: Vec<Vec<MotionVector>>,
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
182
}

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;

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

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

197
#[derive(Copy, Clone, Debug)]
198
pub struct Sequence {
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
254
255
  // 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,
256
257
258
}

impl Sequence {
259
260
261
  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();
262
263
    assert!(width_bits <= 16);
    assert!(height_bits <= 16);
264

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

274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
    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,
289
290
291
292
293
294
295
296
297
      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,
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
      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,
      enable_cdef: true,
317
318
      enable_restoration: config.chroma_sampling != ChromaSampling::Cs422 &&
        config.chroma_sampling != ChromaSampling::Cs444, // FIXME: not working yet
319
320
321
322
323
324
325
326
      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,
327
    }
328
  }
fbossen's avatar
fbossen committed
329

330
331
332
333
334
  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
335

336
  pub fn get_skip_mode_allowed<T: Pixel>(&self, fi: &FrameInvariants<T>, reference_select: bool) -> bool {
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
    if fi.intra_only || !reference_select || !self.enable_order_hint {
      false
    } else {
      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 {
            if 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 {
fbossen's avatar
fbossen committed
361
        false
362
363
364
      } else if backward_idx >= 0 {
        // set skip_mode_frame
        true
fbossen's avatar
fbossen committed
365
      } else {
366
367
        let mut second_forward_idx: isize = -1;
        let mut second_forward_hint = 0;
fbossen's avatar
fbossen committed
368
369
370
        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;
371
372
373
374
            if self.get_relative_dist(ref_hint, forward_hint) < 0 {
              if 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;
fbossen's avatar
fbossen committed
375
376
377
378
              }
            }
          }
        }
379
380
381

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

#[derive(Debug)]
388
389
390
391
392
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>,
393
394
395
396
397
  pub qc: QuantizationContext,
  pub cdfs: CDFContext,
  pub deblock: DeblockState,
  pub segmentation: SegmentationState,
  pub restoration: RestorationState,
398
  pub frame_mvs: Vec<Vec<MotionVector>>,
399
400
}

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

408
  pub fn new_with_frame(fi: &FrameInvariants<T>, frame: Arc<Frame<T>>) -> Self {
409
    let rs = RestorationState::new(fi, &frame);
Raphaël Zumer's avatar
Raphaël Zumer committed
410
411
412
413
414
    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;

415
    Self {
416
      input: frame,
Raphaël Zumer's avatar
Raphaël Zumer committed
417
418
      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),
419
      rec: Frame::new(luma_width, luma_height, fi.sequence.chroma_sampling),
420
421
422
423
424
      qc: Default::default(),
      cdfs: CDFContext::new(0),
      deblock: Default::default(),
      segmentation: Default::default(),
      restoration: rs,
425
      frame_mvs: vec![vec![MotionVector{row: 0, col: 0}; fi.w_in_b * fi.h_in_b]; REF_FRAMES]
426
    }
427
  }
428
429
}

430
431
#[derive(Copy, Clone, Debug)]
pub struct DeblockState {
432
433
434
435
436
437
438
439
440
  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,
441
442
443
}

impl Default for DeblockState {
444
445
446
447
448
449
450
451
452
453
454
  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
455
    }
456
  }
457
458
}

459
460
#[derive(Copy, Clone, Debug)]
pub struct SegmentationState {
461
462
463
464
465
466
467
  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],
468
469
470
}

impl Default for SegmentationState {
471
472
473
474
475
476
477
478
479
  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],
480
    }
481
  }
482
483
}

484
485
// Frame Invariants are invariant inside a frame
#[allow(dead_code)]
Josh Holmer's avatar
Josh Holmer committed
486
#[derive(Debug, Clone)]
487
pub struct FrameInvariants<T: Pixel> {
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
  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,
  pub globalmv_transformation_type: [GlobalMVMode; ALTREF_FRAME + 1],
  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],
533
  pub rec_buffer: ReferenceFramesSet<T>,
534
535
536
  pub base_q_idx: u8,
  pub dc_delta_q: [i8; 3],
  pub ac_delta_q: [i8; 3],
537
538
  pub lambda: f64,
  pub me_lambda: f64,
539
540
541
542
  pub me_range_scale: u8,
  pub use_tx_domain_distortion: bool,
  pub inter_cfg: Option<InterPropsConfig>,
  pub enable_early_exit: bool,
543
544
}

545
pub(crate) fn pos_to_lvl(pos: u64, pyramid_depth: u64) -> u64 {
546
547
548
549
550
551
552
553
554
555
  // 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
}

556
557
558
impl<T: Pixel> FrameInvariants<T> {
  pub fn new(config: EncoderConfig, sequence: Sequence) -> Self {
    assert!(sequence.bit_depth <= mem::size_of::<T>() * 8, "bit depth cannot fit into u8");
559
560
561
    // 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.
562
    let min_partition_size = config.speed_settings.min_block_size;
563
564
565
    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;

566
567
568
    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

569
    Self {
570
      sequence,
571
572
573
574
      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),
575
576
      w_in_b,
      h_in_b,
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
      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,
      globalmv_transformation_type: [GlobalMVMode::IDENTITY; ALTREF_FRAME + 1],
      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],
616
617
      lambda: 0.0,
      me_lambda: 0.0,
618
619
620
621
      me_range_scale: 1,
      use_tx_domain_distortion,
      inter_cfg: None,
      enable_early_exit: true,
622
      config,
Josh Holmer's avatar
Josh Holmer committed
623
    }
624
  }
Josh Holmer's avatar
Josh Holmer committed
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668

  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;
    }
    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.
669
670
671
672
  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
673
674
675
676
677
678
    let mut fi = previous_fi.clone();
    fi.frame_type = FrameType::INTER;
    fi.intra_only = false;
    fi.apply_inter_props_cfg(idx_in_segment);
    let inter_cfg = fi.inter_cfg.unwrap();

679
680
681
682
683
684
685
    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
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
    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 {
      NONE_FRAME
    } 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
729
730
731
732
733
    fi.primary_ref_frame = if lvl > 0 {
      PRIMARY_REF_NONE
    } else {
      (ref_in_previous_group - LAST_FRAME) as u32
    };
Josh Holmer's avatar
Josh Holmer committed
734
735
736
737
738
739
740

    for i in 0..INTER_REFS_PER_FRAME {
      fi.ref_frames[i] = if lvl == 0 {
        if i == second_ref_frame - LAST_FRAME {
          (slot_idx + 4 - 2) as u8 % 4
        } else {
          (slot_idx + 4 - 1) as u8 % 4
741
        }
Josh Holmer's avatar
Josh Holmer committed
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
      } else {
        if i == second_ref_frame - LAST_FRAME {
          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 - LAST_FRAME {
          if lvl == 0 {
            (slot_idx + 4 - 1) as u8 % 4
          } else {
            slot_idx as u8
          }
        } else {
          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
          } else {
            3 + lvl1 as u8
          }
        }
      }
767
    }
Josh Holmer's avatar
Josh Holmer committed
768
769
770
771
772
773
774
775
776
777

    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)
  }
778

779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
  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;
800
801
802
803
804
805
    // 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;
    }
806
807
808
809
810
811
812
813
814
815
    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 =
      qps.lambda * ((1 << 2 * (self.sequence.bit_depth - 8)) as f64);
816
817
    self.me_lambda = self.lambda.sqrt();
  }
818
819
}

820
impl<T: Pixel> fmt::Display for FrameInvariants<T> {
821
822
823
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Frame {} - {}", self.number, self.frame_type)
  }
824
825
}

Josh Holmer's avatar
Josh Holmer committed
826
827
828
829
830
831
832
833
834
835
836
#[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,
}

837
838
839
840
841
842
843
pub fn write_temporal_delimiter(
  packet: &mut dyn io::Write
) -> io::Result<()> {
  packet.write(&TEMPORAL_DELIMITER)?;
  Ok(())
}

844
845
fn write_obus<T: Pixel>(
  packet: &mut dyn io::Write, fi: &mut FrameInvariants<T>, fs: &FrameState<T>
846
847
) -> io::Result<()> {
  let obu_extension = 0 as u32;
848

849
  let mut buf1 = Vec::new();
850

851
852
  // write sequence header obu if KEY_FRAME, preceded by 4-byte size
  if fi.frame_type == FrameType::KEY {
853
854
    let mut buf2 = Vec::new();
    {
855
856
      let mut bw2 = BitWriter::endian(&mut buf2, BigEndian);
      bw2.write_sequence_header_obu(fi)?;
Raphaël Zumer's avatar
Raphaël Zumer committed
857
      bw2.write_bit(true)?; // trailing bit
858
      bw2.byte_align()?;
859
860
861
    }

    {
862
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
863
      bw1.write_obu_header(ObuType::OBU_SEQUENCE_HEADER, obu_extension)?;
864
    }
865
    packet.write_all(&buf1).unwrap();
866
867
868
    buf1.clear();

    {
869
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
870
      bw1.write_uleb128(buf2.len() as u64)?;
871
    }
872
    packet.write_all(&buf1).unwrap();
873
874
    buf1.clear();

875
    packet.write_all(&buf2).unwrap();
876
    buf2.clear();
877
878

    if fi.sequence.content_light.is_some() {
879
880
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
      bw1.write_metadata_obu(ObuMetaType::OBU_META_HDR_CLL, fi.sequence)?;
881
882
883
884
885
      packet.write_all(&buf1).unwrap();
      buf1.clear();
    }

    if fi.sequence.mastering_display.is_some() {
886
887
      let mut bw1 = BitWriter::endian(&mut buf1, BigEndian);
      bw1.write_metadata_obu(ObuMetaType::OBU_META_HDR_MDCV, fi.sequence)?;
888
889
890
      packet.write_all(&buf1).unwrap();
      buf1.clear();
    }
891
  }
892

893
894
895
896
897
898
899
900
  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);
901
    bw1.write_obu_header(ObuType::OBU_FRAME_HEADER, obu_extension)?;
902
903
904
905
906
907
  }
  packet.write_all(&buf1).unwrap();
  buf1.clear();

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

911
912
913
914
915
916
917
  packet.write_all(&buf1).unwrap();
  buf1.clear();

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

  Ok(())
918
919
920
}

/// Write into `dst` the difference between the blocks at `src1` and `src2`
921
fn diff<T: Pixel>(dst: &mut [i16], src1: &PlaneSlice<'_, T>, src2: &PlaneSlice<'_, T>, width: usize, height: usize) {
fbossen's avatar
fbossen committed
922
923
924
925
  let src1_stride = src1.plane.cfg.stride;
  let src2_stride = src2.plane.cfg.stride;

  for ((l, s1), s2) in dst.chunks_mut(width).take(height)
926
927
928
    .zip(src1.as_slice().chunks(src1_stride))
    .zip(src2.as_slice().chunks(src2_stride)) {
      for ((r, v1), v2) in l.iter_mut().zip(s1).zip(s2) {
929
        *r = i16::cast_from(*v1) - i16::cast_from(*v2);
930
      }
931
932
933
    }
}

934
fn get_qidx<T: Pixel>(fi: &FrameInvariants<T>, fs: &FrameState<T>, cw: &ContextWriter, bo: &BlockOffset) -> u8 {
935
936
937
938
939
940
941
  let mut qidx = fi.base_q_idx;
  let sidx = cw.bc.at(bo).segmentation_idx as usize;
  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
942
943
}

944
945
946
// For a transform block,
// predict, transform, quantize, write coefficients to a bitstream,
// dequantize, inverse-transform.
947
948
pub fn encode_tx_block<T: Pixel>(
  fi: &FrameInvariants<T>, fs: &mut FrameState<T>, cw: &mut ContextWriter,
949
950
  w: &mut dyn Writer, p: usize, bo: &BlockOffset, mode: PredictionMode,
  tx_size: TxSize, tx_type: TxType, plane_bsize: BlockSize, po: &PlaneOffset,
951
  skip: bool, ac: &[i16], alpha: i16, rdo_type: RDOType, for_rdo_use: bool
952
) -> (bool, i64) {
953
954
955
  let qidx = get_qidx(fi, fs, cw, bo);
  let rec = &mut fs.rec.planes[p];
  let PlaneConfig { stride, xdec, ydec, .. } = fs.input.planes[p].cfg;
956

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

959
960
961
962
963
  if mode.is_intra() {
    let bit_depth = fi.sequence.bit_depth;
    let edge_buf = get_intra_edges(&rec.slice(po), tx_size, bit_depth, &fs.input.planes[p].cfg, fi.w_in_b, fi.h_in_b, Some(mode));
    mode.predict_intra(&mut rec.mut_slice(po), tx_size, bit_depth, &ac, alpha, &edge_buf);
  }
964

965
  if skip { return (false, -1); }
966

967
968
969
970
971
972
973
974
  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()];
975

976
977
978
979
980
  diff(residual,
       &fs.input.planes[p].slice(po),
       &rec.slice(po),
       tx_size.width(),
       tx_size.height());
981

982
  forward_transform(residual, coeffs, tx_size.width(), tx_size, tx_type, fi.sequence.bit_depth);
983

984
985
  let coded_tx_size = av1_get_coded_tx_size(tx_size).area();
  fs.qc.quantize(coeffs, qcoeffs, coded_tx_size);
986

987
988
  let has_coeff = cw.write_coeffs_lv_map(w, p, bo, &qcoeffs, mode, tx_size, tx_type, plane_bsize, xdec, ydec,
                                         fi.use_reduced_tx_set);
989

990
991
  // Reconstruct
  dequantize(qidx, qcoeffs, rcoeffs, tx_size, fi.sequence.bit_depth, fi.dc_delta_q[p], fi.ac_delta_q[p]);
992

993
  let mut tx_dist: i64 = -1;
994

995
996
  if !fi.use_tx_domain_distortion || !for_rdo_use {
    inverse_transform_add(rcoeffs, &mut rec.mut_slice(po).as_mut_slice(), stride, tx_size, tx_type, fi.sequence.bit_depth);
997
998
  }
  if rdo_type.needs_tx_dist() {
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
    // 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;
  }
  (has_coeff, tx_dist)
1013
1014
}

1015
1016
1017
1018
1019
pub fn motion_compensate<T: Pixel>(
  fi: &FrameInvariants<T>, fs: &mut FrameState<T>, cw: &mut ContextWriter,
  luma_mode: PredictionMode, ref_frames: [usize; 2], mvs: [MotionVector; 2],
  bsize: BlockSize, bo: &BlockOffset, luma_only: bool
) {
1020
  debug_assert!(!luma_mode.is_intra());
1021
1022
1023
1024
1025

  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.
1026
  let num_planes = 1 + if !luma_only && has_chroma(bo, bsize, xdec, ydec) { 2 } else { 0 };
1027

1028
1029
1030
1031
1032
1033
  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];
1034
1035
1036
1037

    if p > 0 && bsize < BlockSize::BLOCK_8X8 {
      let mut some_use_intra = false;
      if bsize == BlockSize::BLOCK_4X4 || bsize == BlockSize::BLOCK_4X8 {
1038
        some_use_intra |= cw.bc.at(&bo.with_offset(-1,0)).mode.is_intra(); };
1039
      if !some_use_intra && bsize == BlockSize::BLOCK_4X4 || bsize == BlockSize::BLOCK_8X4 {
1040
        some_use_intra |= cw.bc.at(&bo.with_offset(0,-1)).mode.is_intra(); };
1041
      if !some_use_intra && bsize == BlockSize::BLOCK_4X4 {
1042
        some_use_intra |= cw.bc.at(&bo.with_offset(-1,-1)).mode.is_intra(); };
1043
1044
1045

      if some_use_intra {
        luma_mode.predict_inter(fi, p, &po, &mut rec.mut_slice(&po), plane_bsize.width(),
1046
                                plane_bsize.height(), ref_frames, mvs);
1047
      } else {
1048
        assert!(xdec == 1 && ydec == 1);
1049
1050
        // TODO: these are absolutely only valid for 4:2:0
        if bsize == BlockSize::BLOCK_4X4 {
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
          let mv0 = cw.bc.at(&bo.with_offset(-1,-1)).mv;
          let rf0 = cw.bc.at(&bo.with_offset(-1,-1)).ref_frames;
          let mv1 = cw.bc.at(&bo.with_offset(0,-1)).mv;
          let rf1 = cw.bc.at(&bo.with_offset(0,-1)).ref_frames;
          let po1 = PlaneOffset { x: po.x+2, y: po.y };
          let mv2 = cw.bc.at(&bo.with_offset(-1,0)).mv;
          let rf2 = cw.bc.at(&bo.with_offset(-1,0)).ref_frames;
          let po2 = PlaneOffset { x: po.x, y: po.y+2 };
          let po3 = PlaneOffset { x: po.x+2, y: po.y+2 };
          luma_mode.predict_inter(fi, p, &po, &mut rec.mut_slice(&po), 2, 2, rf0, mv0);
          luma_mode.predict_inter(fi, p, &po1, &mut rec.mut_slice(&po1), 2, 2, rf1, mv1);
          luma_mode.predict_inter(fi, p, &po2, &mut rec.mut_slice(&po2), 2, 2, rf2, mv2);
          luma_mode.predict_inter(fi, p, &po3, &mut rec.mut_slice(&po3), 2, 2, ref_frames, mvs);
1064
1065
        }
        if bsize == BlockSize::BLOCK_8X4 {
1066
1067
1068
1069
1070
          let mv1 = cw.bc.at(&bo.with_offset(0,-1)).mv;
          let rf1 = cw.bc.at(&bo.with_offset(0,-1)).ref_frames;
          luma_mode.predict_inter(fi, p, &po, &mut rec.mut_slice(&po), 4, 2, rf1, mv1);
          let po3 = PlaneOffset { x: po.x, y: po.y+2 };
          luma_mode.predict_inter(fi, p, &po3, &mut rec.mut_slice(&po3), 4, 2, ref_frames, mvs);
1071
1072
        }
        if bsize == BlockSize::BLOCK_4X8 {
1073
1074
1075
1076
1077
          let mv2 = cw.bc.at(&bo.with_offset(-1,0)).mv;
          let rf2 = cw.bc.at(&bo.with_offset(-1,0)).ref_frames;
          luma_mode.predict_inter(fi, p, &po, &mut rec.mut_slice(&po), 2, 4, rf2, mv2);
          let po3 = PlaneOffset { x: po.x+2, y: po.y };
          luma_mode.predict_inter(fi, p, &po3, &mut rec.mut_slice(&po3), 2, 4, ref_frames, mvs);
1078
        }
1079
1080
1081
      }
    } else {
      luma_mode.predict_inter(fi, p, &po, &mut rec.mut_slice(&po), plane_bsize.width(),
1082
                              plane_bsize.height(), ref_frames, mvs);
1083
1084
1085
1086
    }
  }
}

1087
1088
1089
1090
1091
pub fn encode_block_a<T: Pixel>(
  seq: &Sequence, fs: &FrameState<T>,
  cw: &mut ContextWriter, w: &mut dyn Writer,
  bsize: BlockSize, bo: &BlockOffset, skip: bool
) -> bool {
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
  cw.bc.set_skip(bo, bsize, skip);
  if fs.segmentation.enabled && fs.segmentation.update_map && fs.segmentation.preskip {
    cw.write_segmentation(w, bo, bsize, false, fs.segmentation.last_active_segid);
  }
  cw.write_skip(w, bo, skip);
  if fs.segmentation.enabled && fs.segmentation.update_map && !fs.segmentation.preskip {
    cw.write_segmentation(w, bo, bsize, skip, fs.segmentation.last_active_segid);
  }
  if !skip && seq.enable_cdef {
    cw.bc.cdef_coded = true;
  }
  cw.bc.cdef_coded
1104
1105
}

1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
pub fn encode_block_b<T: Pixel>(
  fi: &FrameInvariants<T>, fs: &mut FrameState<T>,
  cw: &mut ContextWriter, w: &mut dyn Writer,
  luma_mode: PredictionMode, chroma_mode: PredictionMode,
  ref_frames: [usize; 2], mvs: [MotionVector; 2],
  bsize: BlockSize, bo: &BlockOffset, skip: bool,
  cfl: CFLParams, tx_size: TxSize, tx_type: TxType,
  mode_context: usize, mv_stack: &[CandidateMV],
  rdo_type: RDOType, for_rdo_use: bool
) -> i64 {
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
  let is_inter = !luma_mode.is_intra();
  if is_inter { assert!(luma_mode == chroma_mode); };
  let sb_size = if fi.sequence.use_128x128_superblock {
    BlockSize::BLOCK_128X128
  } else {
    BlockSize::BLOCK_64X64
  };
  let PlaneConfig { xdec, ydec, .. } = fs.input.planes[1].cfg;
  if skip {
    cw.bc.reset_skip_context(bo, bsize, xdec, ydec);
  }
  cw.bc.set_block_size(bo, bsize);
  cw.bc.set_mode(bo, bsize, luma_mode);
  cw.bc.set_ref_frames(bo, bsize, ref_frames);
  cw.bc.set_motion_vectors(bo, bsize, mvs);

  //write_q_deltas();
  if cw.bc.code_deltas && fs.deblock.block_deltas_enabled && (bsize < sb_size || !skip) {
    cw.write_block_deblock_deltas(w, bo, fs.deblock.block_delta_multi);
  }
  cw.bc.code_deltas = false;
1137

1138
1139
1140
1141
1142
  if fi.frame_type == FrameType::INTER {
    cw.write_is_inter(w, bo, is_inter);
    if is_inter {
      cw.fill_neighbours_ref_counts(bo);
      cw.write_ref_frames(w, fi, bo);
1143

1144
1145
1146
1147
1148
      if luma_mode >= PredictionMode::NEAREST_NEARESTMV {
        cw.write_compound_mode(w, luma_mode, mode_context);
      } else {
        cw.write_inter_mode(w, luma_mode, mode_context);
      }
fbossen's avatar
fbossen committed
1149

1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
      let ref_mv_idx = 0;
      let num_mv_found = mv_stack.len();

      if luma_mode == PredictionMode::NEWMV || luma_mode == PredictionMode::NEW_NEWMV {
        if luma_mode == PredictionMode::NEW_NEWMV { assert!(num_mv_found >= 2); }
        for idx in 0..2 {
          if num_mv_found > idx + 1 {
            let drl_mode = ref_mv_idx > idx;
            let ctx: usize = (mv_stack[idx].weight < REF_CAT_LEVEL) as usize
              + (mv_stack[idx + 1].weight < REF_CAT_LEVEL) as usize;
            cw.write_drl_mode(w, drl_mode, ctx);
            if !drl_mode { break; }
          }
        }
      }
fbossen's avatar
fbossen committed
1165

1166
1167
1168
1169
1170
      let ref_mvs = if num_mv_found > 0 {
        [mv_stack[ref_mv_idx].this_mv, mv_stack[ref_mv_idx].comp_mv]
      } else {
        [MotionVector{ row: 0, col: 0 }; 2]
      };
fbossen's avatar
fbossen committed
1171

1172
1173
1174
1175
1176
1177
1178
      let mv_precision = if fi.force_integer_mv != 0 {
        MvSubpelPrecision::MV_SUBPEL_NONE
      } else if fi.allow_high_precision_mv {
        MvSubpelPrecision::MV_SUBPEL_HIGH_PRECISION
      } else {
        MvSubpelPrecision::MV_SUBPEL_LOW_PRECISION
      };
fbossen's avatar
fbossen committed
1179

1180
1181
1182
1183
1184
1185
1186
1187
      if luma_mode == PredictionMode::NEWMV ||
        luma_mode == PredictionMode::NEW_NEWMV ||
        luma_mode == PredictionMode::NEW_NEARESTMV {
          cw.write_mv(w, mvs[0], ref_mvs[0], mv_precision);
        }
      if luma_mode == PredictionMode::NEW_NEWMV ||
        luma_mode == PredictionMode::NEAREST_NEWMV {
          cw.write_mv(w, mvs[1], ref_mvs[1], mv_precision);
1188
1189
        }

1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
      if luma_mode >= PredictionMode::NEAR0MV && luma_mode <= PredictionMode::NEAR2MV {
        let ref_mv_idx = luma_mode as usize - PredictionMode::NEAR0MV as usize + 1;
        if luma_mode != PredictionMode::NEAR0MV { assert!(num_mv_found > ref_mv_idx); }

        for idx in 1..3 {
          if num_mv_found > idx + 1 {
            let drl_mode = ref_mv_idx > idx;
            let ctx: usize = (mv_stack[idx].weight < REF_CAT_LEVEL) as usize
              + (mv_stack[idx + 1].weight < REF_CAT_LEVEL) as usize;

            cw.write_drl_mode(w, drl_mode, ctx);
            if !drl_mode { break; }
          }
1203
        }
1204
1205
1206
1207
1208
1209
        if mv_stack.len() > 1 {
          assert!(mv_stack[ref_mv_idx].this_mv.row == mvs[0].row);
          assert!(mv_stack[ref_mv_idx].this_mv.col == mvs[0].col);
        } else {
          assert!(0 == mvs[0].row);
          assert!(0 == mvs[0].col);
1210
        }
1211
1212
1213
1214
1215
1216
1217
      } else if luma_mode == PredictionMode::NEARESTMV {
        if mv_stack.len() > 0 {
          assert!(mv_stack[0].this_mv.row == mvs[0].row);
          assert!(mv_stack[0].this_mv.col == mvs[0].col);
        } else {
          assert!(0 == mvs[0].row);
          assert!(0 == mvs[0].col);
1218
        }
1219
1220
1221
      }
    } else {
      cw.write_intra_mode(w, bsize, luma_mode);
1222
    }
1223
1224
1225
  } else {
    cw.write_intra_mode_kf(w, bo, luma_mode);
  }
1226

1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
  if !is_inter {
    if luma_mode.is_directional() && bsize >= BlockSize::BLOCK_8X8 {
      cw.write_angle_delta(w, 0, luma_mode);
    }
    if has_chroma(bo, bsize, xdec, ydec) {
      cw.write_intra_uv_mode(w, chroma_mode, luma_mode, bsize);
      if chroma_mode.is_cfl() {
        assert!(bsize.cfl_allowed());
        cw.write_cfl_alphas(w, cfl);
      }
      if chroma_mode.is_directional() && bsize >= BlockSize::BLOCK_8X8 {
        cw.write_angle_delta(w, 0, chroma_mode);
      }
1240
    }
1241
1242
1243
1244
1245
1246
1247
1248
    // TODO: Extra condition related to palette mode, see `read_filter_intra_mode_info` in decodemv.c
    if luma_mode == PredictionMode::DC_PRED && bsize.width() <= 32 && bsize.height() <= 32 {
      cw.write_use_filter_intra(w,false, bsize); // Always turn off FILTER_INTRA
    }
  }

  if is_inter {
    motion_compensate(fi, fs, cw, luma_mode, ref_frames, mvs, bsize, bo, false);
1249
    write_tx_tree(fi, fs, cw, w, luma_mode, bo, bsize, tx_size, tx_type, skip, false, rdo_type, for_rdo_use)
1250
  } else {
1251
    write_tx_blocks(fi, fs, cw, w, luma_mode, chroma_mode, bo, bsize, tx_size, tx_type, skip, cfl, false, rdo_type, for_rdo_use)
1252
  }
1253
1254
}

1255
1256
pub fn luma_ac<T: Pixel>(
  ac: &mut [i16], fs: &mut FrameState<T>, bo: &BlockOffset, bsize: BlockSize
1257
1258
1259
) {
  let PlaneConfig { xdec, ydec, .. } = fs.input.planes[1].cfg;
  let plane_bsize = get_plane_block_size(bsize, xdec, ydec);
1260
  let po = if bsize.is_sub8x8() {
1261
1262
    let offset = bsize.sub8x8_offset();
    bo.with_offset(offset.0, offset.1).plane_offset(&fs.input.planes[0].cfg)