#![allow(dead_code)]
#[doc(hidden)] pub mod bitstream;
#[doc(hidden)] pub mod cabac_engine;
#[doc(hidden)] pub mod cabac_path;
#[doc(hidden)] pub mod cabac_syntax;
#[doc(hidden)] pub mod cavlc;
#[doc(hidden)] pub mod deblock;
pub mod deep;
pub mod field;
#[doc(hidden)] pub mod intra4x4;
#[doc(hidden)] pub mod intra_pred;
#[doc(hidden)] pub mod macroblock;
pub mod mbaff;
pub mod mbaff_cabac;
#[doc(hidden)] pub mod me;
pub mod multiref;
#[doc(hidden)] pub mod nal;
#[doc(hidden)] pub mod pps;
pub mod rate_control;
#[doc(hidden)] pub mod rdo;
pub mod scp;
pub mod sei;
pub mod session;
#[doc(hidden)] pub mod slice;
pub mod slices;
pub mod sp;
#[doc(hidden)] pub mod sps;
#[doc(hidden)] pub mod transform;
use crate::cavlc::CoeffTokenContext;
#[derive(Debug, Clone, Copy)]
struct MbQpTracker {
cur: i32,
mbaff_flag_pending: Option<bool>,
mbaff_field_mb: bool,
wp_l0: Option<WeightedPredL0>,
ref_idx_l0: i32,
num_ref_idx_l0_active_minus1_slice: u32,
intra_nb: Option<(bool, bool, bool)>,
}
#[derive(Debug, Clone, Copy)]
struct WeightedPredL0 {
log2_wd: u32,
weight: i32,
offset: i32,
}
impl WeightedPredL0 {
fn apply(&self, pred: &mut [i32]) {
for p in pred.iter_mut() {
let v = if self.log2_wd >= 1 {
((*p * self.weight + (1i32 << (self.log2_wd - 1))) >> self.log2_wd) + self.offset
} else {
*p * self.weight + self.offset
};
*p = v.clamp(0, 255);
}
}
}
impl MbQpTracker {
fn new(cur: i32) -> Self {
Self {
cur,
mbaff_flag_pending: None,
mbaff_field_mb: false,
wp_l0: None,
ref_idx_l0: 0,
num_ref_idx_l0_active_minus1_slice: 0,
intra_nb: None,
}
}
fn write_pending_mbaff_flag(&mut self, sw: &mut BitWriter) {
if let Some(f) = self.mbaff_flag_pending.take() {
sw.u(1, u32::from(f));
}
}
fn num_ref_idx_l0_active_minus1(&self) -> u32 {
let n = self.num_ref_idx_l0_active_minus1_slice;
if self.mbaff_field_mb {
2 * (n + 1) - 1
} else {
n
}
}
}
use crate::encoder::bitstream::BitWriter;
use crate::encoder::cavlc::encode_residual_block_cavlc;
use crate::encoder::deblock::{
chroma_nz_mask_from_blocks, luma_nz_mask_from_blocks, MbDeblockInfo,
};
use crate::encoder::intra4x4::{
availability as i4x4_avail, predict_4x4, Intra4x4Mode, LUMA_4X4_XY as I4X4_XY,
};
use crate::encoder::intra_pred::{
predict_16x16, predict_16x16_avail, predict_chroma_8x8_avail, sad_8x8, I16x16Mode,
IntraChromaMode,
};
use crate::encoder::macroblock::{
write_b_16x16_mb, write_b_8x8_mixed_mb, write_i8x8_mb, write_i_nxn_mb, write_intra16x16_mb,
write_intra16x16_mb_in_inter_slice, write_p_8x8_all_pl08x8_mb, B16x16McbConfig,
B8x8MixedMcbConfig, BPartPred, BPred16x16, BSubMbCell, I16x16McbConfig, I8x8McbConfig,
INxNMcbConfig, P8x8AllPL08x8McbConfig, PL016x16McbConfig,
};
use crate::encoder::me::{search_quarter_pel_16x16, search_quarter_pel_8x8};
use crate::encoder::nal::build_nal_unit;
use crate::encoder::pps::{build_baseline_pps_rbsp, BaselinePpsConfig};
use crate::encoder::rdo::{cost_combined, lambda_ssd, ssd_16x16, ssd_4x4};
use crate::encoder::slice::{
write_b_slice_header, write_idr_i_slice_header, write_p_slice_header, BSliceHeaderConfig,
ExplicitBipredWeightTable, ExplicitPredWeightTableL0, FieldPicSignal, IdrSliceHeaderConfig,
PSliceHeaderConfig,
};
use crate::encoder::sps::{build_baseline_sps_rbsp, BaselineSpsConfig};
use crate::encoder::transform::{
deinterleave_8x8_to_4x4, forward_core_4x4, forward_core_8x8, forward_hadamard_2x2,
forward_hadamard_4x4, quantize_4x4_ac_w, quantize_4x4_w, quantize_8x8_w, quantize_chroma_dc_w,
quantize_luma_dc_w, scan_8x8_for_cavlc, zigzag_scan_4x4, zigzag_scan_4x4_ac,
};
use crate::inter_pred::{interpolate_chroma, interpolate_luma};
use crate::intra_pred::{
filter_samples_8x8, predict_8x8, Intra8x8Mode, Neighbour4x4Availability, Samples8x8,
};
use crate::macroblock_layer::{derive_nc_luma, CavlcNcGrid, LumaNcKind};
use crate::mv_deriv::{Mv, MvpredInputs, NeighbourMv};
use crate::nal::NalUnitType;
use crate::transform::{
inverse_hadamard_chroma_dc_420, inverse_hadamard_luma_dc_16x16, inverse_scan_4x4_zigzag_ac,
inverse_transform_4x4, inverse_transform_4x4_dc_preserved, inverse_transform_8x8, qp_bd_offset,
qp_y_to_qp_c_with_bd_offset, FLAT_4X4_16,
};
const LUMA_4X4_BLK: [(usize, usize); 16] = [
(0, 0),
(1, 0),
(0, 1),
(1, 1),
(2, 0),
(3, 0),
(2, 1),
(3, 1),
(0, 2),
(1, 2),
(0, 3),
(1, 3),
(2, 2),
(3, 2),
(2, 3),
(3, 3),
];
pub struct YuvFrame<'a> {
pub width: u32,
pub height: u32,
pub y: &'a [u8],
pub u: &'a [u8],
pub v: &'a [u8],
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ScalingMatrixMode {
#[default]
Flat,
SeqDefault,
PicDefault,
SeqCustom(CustomScalingLists),
PicCustom(CustomScalingLists),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CustomScalingLists {
pub intra4: [i32; 16],
pub inter4: [i32; 16],
pub intra8: [i32; 64],
pub inter8: [i32; 64],
}
impl CustomScalingLists {
pub fn is_valid(&self) -> bool {
self.intra4.iter().all(|v| (1..=255).contains(v))
&& self.inter4.iter().all(|v| (1..=255).contains(v))
&& self.intra8.iter().all(|v| (1..=255).contains(v))
&& self.inter8.iter().all(|v| (1..=255).contains(v))
}
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScalingListsSpec {
Default,
Custom(CustomScalingLists),
}
pub(crate) struct WeightPair {
pub(crate) w4: [i32; 16],
pub(crate) w8: [i32; 64],
}
impl ScalingMatrixMode {
pub(crate) fn is_flat(self) -> bool {
matches!(self, ScalingMatrixMode::Flat)
}
pub(crate) fn seq_spec(self) -> Option<ScalingListsSpec> {
match self {
ScalingMatrixMode::SeqDefault => Some(ScalingListsSpec::Default),
ScalingMatrixMode::SeqCustom(l) => Some(ScalingListsSpec::Custom(l)),
_ => None,
}
}
pub(crate) fn pic_spec(self) -> Option<ScalingListsSpec> {
match self {
ScalingMatrixMode::PicDefault => Some(ScalingListsSpec::Default),
ScalingMatrixMode::PicCustom(l) => Some(ScalingListsSpec::Custom(l)),
_ => None,
}
}
pub(crate) fn intra_weights(self) -> WeightPair {
use crate::transform::{
weight_scale_4x4_from_scan, weight_scale_8x8_from_scan, DEFAULT_4X4_INTRA,
DEFAULT_8X8_INTRA, FLAT_8X8_16,
};
match self {
ScalingMatrixMode::Flat => WeightPair {
w4: FLAT_4X4_16,
w8: FLAT_8X8_16,
},
ScalingMatrixMode::SeqDefault | ScalingMatrixMode::PicDefault => WeightPair {
w4: weight_scale_4x4_from_scan(&DEFAULT_4X4_INTRA),
w8: weight_scale_8x8_from_scan(&DEFAULT_8X8_INTRA),
},
ScalingMatrixMode::SeqCustom(l) | ScalingMatrixMode::PicCustom(l) => WeightPair {
w4: weight_scale_4x4_from_scan(&l.intra4),
w8: weight_scale_8x8_from_scan(&l.intra8),
},
}
}
pub(crate) fn inter_weights(self) -> WeightPair {
use crate::transform::{
weight_scale_4x4_from_scan, weight_scale_8x8_from_scan, DEFAULT_4X4_INTER,
DEFAULT_8X8_INTER, FLAT_8X8_16,
};
match self {
ScalingMatrixMode::Flat => WeightPair {
w4: FLAT_4X4_16,
w8: FLAT_8X8_16,
},
ScalingMatrixMode::SeqDefault | ScalingMatrixMode::PicDefault => WeightPair {
w4: weight_scale_4x4_from_scan(&DEFAULT_4X4_INTER),
w8: weight_scale_8x8_from_scan(&DEFAULT_8X8_INTER),
},
ScalingMatrixMode::SeqCustom(l) | ScalingMatrixMode::PicCustom(l) => WeightPair {
w4: weight_scale_4x4_from_scan(&l.inter4),
w8: weight_scale_8x8_from_scan(&l.inter8),
},
}
}
}
#[derive(Debug, Clone)]
pub struct EncoderConfig {
pub width: u32,
pub height: u32,
pub qp: i32,
pub profile_idc: u8,
pub max_num_ref_frames: u32,
pub direct_temporal_mv_pred: bool,
pub b_l0_bi_only: bool,
pub explicit_weighted_bipred: bool,
pub explicit_weighted_pred: bool,
pub intra_in_inter: bool,
pub chroma_format_idc: u32,
pub colour_plane_id: Option<u8>,
pub cabac: bool,
pub vui: Option<crate::vui::VuiParameters>,
pub level_idc: u8,
pub bit_depth_chroma_minus8: u32,
pub trellis_quant: bool,
pub trellis_full: bool,
pub trellis_quant_intra: bool,
pub trellis_quant_intra_chroma: bool,
pub transform_8x8: bool,
pub scaling_matrix: ScalingMatrixMode,
pub table_8_10_cy_offset: i32,
}
impl EncoderConfig {
pub fn new(width: u32, height: u32) -> Self {
let width_in_mbs = width.div_ceil(16);
let height_in_mbs = height.div_ceil(16);
Self {
width,
height,
qp: 26,
profile_idc: 66,
max_num_ref_frames: 1,
direct_temporal_mv_pred: false,
b_l0_bi_only: false,
explicit_weighted_bipred: false,
explicit_weighted_pred: false,
intra_in_inter: true,
chroma_format_idc: 1,
colour_plane_id: None,
cabac: false,
vui: None,
level_idc: min_level_idc_for_picture_size(width_in_mbs, height_in_mbs),
bit_depth_chroma_minus8: 0,
trellis_quant: true,
trellis_full: false,
trellis_quant_intra: false,
trellis_quant_intra_chroma: false,
transform_8x8: false,
scaling_matrix: ScalingMatrixMode::default(),
table_8_10_cy_offset: 0,
}
}
}
pub(crate) fn min_level_idc_for_picture_size(width_in_mbs: u32, height_in_mbs: u32) -> u8 {
let pic_size = width_in_mbs.saturating_mul(height_in_mbs);
const LEVELS: &[(u8, u32)] = &[
(10, 99), (11, 396), (12, 396), (13, 396), (20, 396), (21, 792), (22, 1_620), (30, 1_620), (31, 3_600), (32, 5_120), (40, 8_192), (41, 8_192), (42, 8_704), (50, 22_080), (51, 36_864), (52, 36_864), (60, 139_264), (61, 139_264), (62, 139_264), ];
for &(level, max_fs) in LEVELS {
if pic_size <= max_fs {
return level;
}
}
62
}
#[derive(Debug, Clone, Default)]
pub(crate) struct IntraGridSlot {
available: bool,
is_i_nxn: bool,
intra_4x4_pred_modes: [u8; 16],
is_i8x8: bool,
intra_8x8_pred_modes: [u8; 4],
}
#[derive(Debug, Clone)]
pub(crate) struct IntraGrid {
width_mbs: usize,
slots: Vec<IntraGridSlot>,
}
impl IntraGrid {
fn new(width_mbs: usize, height_mbs: usize) -> Self {
Self {
width_mbs,
slots: vec![IntraGridSlot::default(); width_mbs * height_mbs],
}
}
fn slot(&self, mb_x: usize, mb_y: usize) -> &IntraGridSlot {
&self.slots[mb_y * self.width_mbs + mb_x]
}
fn slot_mut(&mut self, mb_x: usize, mb_y: usize) -> &mut IntraGridSlot {
&mut self.slots[mb_y * self.width_mbs + mb_x]
}
}
fn predicted_intra4x4_mode(
grid: &IntraGrid,
mb_x: usize,
mb_y: usize,
block_idx: usize,
pic_w_mbs: usize,
) -> u8 {
use crate::encoder::intra4x4::LUMA_4X4_XY as XY;
let (bx, by) = XY[block_idx];
let neigh = |xd: i32, yd: i32| -> Option<(usize, usize, usize)> {
let xn = bx as i32 + xd;
let yn = by as i32 + yd;
let (nmb_x, nmb_y) = if xn < 0 && (0..16).contains(&yn) {
if mb_x == 0 {
return None;
}
(mb_x - 1, mb_y)
} else if (0..16).contains(&xn) && yn < 0 {
if mb_y == 0 {
return None;
}
(mb_x, mb_y - 1)
} else if (0..16).contains(&xn) && (0..16).contains(&yn) {
(mb_x, mb_y)
} else {
return None;
};
let xw = ((xn + 16) % 16) as usize;
let yw = ((yn + 16) % 16) as usize;
let nblk = 8 * (yw / 8) + 4 * (xw / 8) + 2 * ((yw % 8) / 4) + ((xw % 8) / 4);
Some((nmb_x, nmb_y, nblk))
};
let na = neigh(-1, 0);
let nb = neigh(0, -1);
let mb_a_avail = na
.and_then(|(nx, ny, _)| Some(grid.slot(nx, ny)).filter(|s| s.available))
.is_some();
let mb_b_avail = nb
.and_then(|(nx, ny, _)| Some(grid.slot(nx, ny)).filter(|s| s.available))
.is_some();
let dc_pred_flag = !mb_a_avail || !mb_b_avail;
let mode_for = |neigh: Option<(usize, usize, usize)>| -> u8 {
if dc_pred_flag {
return 2;
}
let Some((nx, ny, nblk)) = neigh else {
return 2;
};
let s = grid.slot(nx, ny);
if !s.available {
return 2;
}
if s.is_i_nxn {
s.intra_4x4_pred_modes[nblk]
} else if s.is_i8x8 {
s.intra_8x8_pred_modes[nblk >> 2]
} else {
2
}
};
let mode_a = mode_for(na);
let mode_b = mode_for(nb);
let _ = pic_w_mbs;
mode_a.min(mode_b)
}
fn predicted_intra8x8_mode(grid: &IntraGrid, mb_x: usize, mb_y: usize, blk8: usize) -> u8 {
let neigh = |xd: i32, yd: i32| -> Option<(usize, usize, usize)> {
let xn = ((blk8 as i32) % 2) * 8 + xd;
let yn = ((blk8 as i32) / 2) * 8 + yd;
let (nmb_x, nmb_y) = if xn < 0 && (0..16).contains(&yn) {
if mb_x == 0 {
return None;
}
(mb_x - 1, mb_y)
} else if (0..16).contains(&xn) && yn < 0 {
if mb_y == 0 {
return None;
}
(mb_x, mb_y - 1)
} else if (0..16).contains(&xn) && (0..16).contains(&yn) {
(mb_x, mb_y)
} else {
return None;
};
let xw = ((xn + 16) % 16) as usize;
let yw = ((yn + 16) % 16) as usize;
let nblk8 = 2 * (yw / 8) + (xw / 8);
Some((nmb_x, nmb_y, nblk8))
};
let na = neigh(-1, 0);
let nb = neigh(0, -1);
let avail = |n: Option<(usize, usize, usize)>| -> bool {
n.map(|(nx, ny, _)| grid.slot(nx, ny).available)
.unwrap_or(false)
};
let dc_pred_flag = !avail(na) || !avail(nb);
let mode_for = |n: Option<(usize, usize, usize)>, n_for_4x4: usize| -> u8 {
if dc_pred_flag {
return 2;
}
let Some((nx, ny, nblk8)) = n else {
return 2;
};
let s = grid.slot(nx, ny);
if !s.available {
return 2;
}
if s.is_i8x8 {
s.intra_8x8_pred_modes[nblk8]
} else if s.is_i_nxn {
s.intra_4x4_pred_modes[nblk8 * 4 + n_for_4x4]
} else {
2
}
};
mode_for(na, 1).min(mode_for(nb, 2))
}
fn gather_8x8_from_recon(
recon_y: &[u8],
width: usize,
mb_x: usize,
mb_y: usize,
blk8: usize,
) -> Samples8x8 {
let bx = (mb_x * 16 + (blk8 % 2) * 8) as i32;
let by = (mb_y * 16 + (blk8 / 2) * 8) as i32;
let left_avail = bx > 0;
let top_avail = by > 0;
let tl_avail = bx > 0 && by > 0;
let tr_avail = by > 0 && blk8 != 3 && (bx + 8) < width as i32;
let at = |x: i32, y: i32| -> i32 { recon_y[(y as usize) * width + (x as usize)] as i32 };
let top_left = if tl_avail { at(bx - 1, by - 1) } else { 0 };
let mut top = [0i32; 8];
let mut top_right = [0i32; 8];
for x in 0..8i32 {
top[x as usize] = if top_avail { at(bx + x, by - 1) } else { 0 };
top_right[x as usize] = if tr_avail { at(bx + 8 + x, by - 1) } else { 0 };
}
let mut left = [0i32; 8];
for y in 0..8i32 {
left[y as usize] = if left_avail { at(bx - 1, by + y) } else { 0 };
}
Samples8x8 {
top_left,
top,
top_right,
left,
availability: Neighbour4x4Availability {
top_left: tl_avail,
top: top_avail,
top_right: tr_avail,
left: left_avail,
},
}
}
fn intra8x8_mode_available(mode: Intra8x8Mode, avail: &Neighbour4x4Availability) -> bool {
use Intra8x8Mode::*;
match mode {
Dc => true,
Vertical | DiagonalDownLeft | VerticalLeft => avail.top,
Horizontal | HorizontalUp => avail.left,
DiagonalDownRight | VerticalRight | HorizontalDown => {
avail.top && avail.left && avail.top_left
}
}
}
struct MbTrial {
luma_cost: u64,
deblock: MbDeblockInfo,
}
struct MbStateSnapshot {
recon_y_mb: [u8; 256],
recon_u_mb: [u8; 256],
recon_v_mb: [u8; 256],
chroma_tile: (usize, usize),
sw: BitWriter,
nc_slot: crate::macroblock_layer::CavlcMbNc,
intra_slot: IntraGridSlot,
}
fn chroma_tile_dims(chroma_array_type: u32) -> (usize, usize) {
match chroma_array_type {
0 => (0, 0),
2 => (8, 16),
3 => (16, 16),
_ => (8, 8),
}
}
impl MbStateSnapshot {
#[allow(clippy::too_many_arguments)]
fn capture(
recon_y: &[u8],
recon_u: &[u8],
recon_v: &[u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &BitWriter,
nc_grid: &CavlcNcGrid,
intra_grid: &IntraGrid,
mb_addr: usize,
chroma_array_type: u32,
) -> Self {
let (cw, ch) = chroma_tile_dims(chroma_array_type);
let mut snap = Self {
recon_y_mb: [0u8; 256],
recon_u_mb: [0u8; 256],
recon_v_mb: [0u8; 256],
chroma_tile: (cw, ch),
sw: sw.clone(),
nc_slot: nc_grid.mbs[mb_addr],
intra_slot: intra_grid.slot(mb_x, mb_y).clone(),
};
for j in 0..16usize {
let src = (mb_y * 16 + j) * width + mb_x * 16;
snap.recon_y_mb[j * 16..j * 16 + 16].copy_from_slice(&recon_y[src..src + 16]);
}
for j in 0..ch {
let src = (mb_y * ch + j) * chroma_width + mb_x * cw;
snap.recon_u_mb[j * cw..j * cw + cw].copy_from_slice(&recon_u[src..src + cw]);
snap.recon_v_mb[j * cw..j * cw + cw].copy_from_slice(&recon_v[src..src + cw]);
}
snap
}
#[allow(clippy::too_many_arguments)]
fn restore(
&self,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_addr: usize,
) {
for j in 0..16usize {
let dst = (mb_y * 16 + j) * width + mb_x * 16;
recon_y[dst..dst + 16].copy_from_slice(&self.recon_y_mb[j * 16..j * 16 + 16]);
}
let (cw, ch) = self.chroma_tile;
for j in 0..ch {
let dst = (mb_y * ch + j) * chroma_width + mb_x * cw;
recon_u[dst..dst + cw].copy_from_slice(&self.recon_u_mb[j * cw..j * cw + cw]);
recon_v[dst..dst + cw].copy_from_slice(&self.recon_v_mb[j * cw..j * cw + cw]);
}
*sw = self.sw.clone();
nc_grid.mbs[mb_addr] = self.nc_slot;
*intra_grid.slot_mut(mb_x, mb_y) = self.intra_slot.clone();
}
#[allow(clippy::too_many_arguments)]
fn install(
&self,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_addr: usize,
) {
self.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr,
);
}
}
struct InterMbStateSnapshot {
inner: MbStateSnapshot,
mv_l0: MvGridSlot,
mv_l1: Option<MvGridSlot>,
pending_skip: u32,
}
impl InterMbStateSnapshot {
#[allow(clippy::too_many_arguments)]
fn capture_p(
recon_y: &[u8],
recon_u: &[u8],
recon_v: &[u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &BitWriter,
nc_grid: &CavlcNcGrid,
intra_grid: &IntraGrid,
mv_grid: &MvGrid,
mb_addr: usize,
_width_mbs: usize,
pending_skip: u32,
) -> Self {
Self {
inner: MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr,
1,
),
mv_l0: *mv_grid.slot(mb_x, mb_y),
mv_l1: None,
pending_skip,
}
}
#[allow(clippy::too_many_arguments)]
fn capture_b(
recon_y: &[u8],
recon_u: &[u8],
recon_v: &[u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &BitWriter,
nc_grid: &CavlcNcGrid,
intra_grid: &IntraGrid,
mv_grid_l0: &MvGrid,
mv_grid_l1: &MvGrid,
mb_addr: usize,
_width_mbs: usize,
pending_skip: u32,
) -> Self {
Self {
inner: MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr,
1,
),
mv_l0: *mv_grid_l0.slot(mb_x, mb_y),
mv_l1: Some(*mv_grid_l1.slot(mb_x, mb_y)),
pending_skip,
}
}
#[allow(clippy::too_many_arguments)]
fn restore_p(
&self,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid: &mut MvGrid,
mb_addr: usize,
pending_skip: &mut u32,
) {
self.inner.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr,
);
*mv_grid.slot_mut(mb_x, mb_y) = self.mv_l0;
*pending_skip = self.pending_skip;
}
#[allow(clippy::too_many_arguments)]
fn restore_b(
&self,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
width: usize,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid_l0: &mut MvGrid,
mv_grid_l1: &mut MvGrid,
mb_addr: usize,
pending_skip: &mut u32,
) {
self.inner.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr,
);
*mv_grid_l0.slot_mut(mb_x, mb_y) = self.mv_l0;
if let Some(slot) = self.mv_l1 {
*mv_grid_l1.slot_mut(mb_x, mb_y) = slot;
}
*pending_skip = self.pending_skip;
}
}
impl Encoder {
pub fn parameter_sets(&self, cabac: bool) -> (Vec<u8>, Vec<u8>) {
let width_mbs = self.cfg.width / 16;
let height_mbs = self.cfg.height / 16;
let profile_idc = if (self.cfg.transform_8x8 || !self.cfg.scaling_matrix.is_flat())
&& self.cfg.profile_idc < 100
{
100
} else {
self.cfg.profile_idc
};
let sps_cfg = BaselineSpsConfig {
seq_parameter_set_id: 0,
level_idc: self.cfg.level_idc,
width_in_mbs: width_mbs,
height_in_mbs: height_mbs,
log2_max_frame_num_minus4: 4,
log2_max_poc_lsb_minus4: 4,
max_num_ref_frames: self.cfg.max_num_ref_frames,
profile_idc,
chroma_format_idc: self.cfg.chroma_format_idc,
separate_colour_plane: self.cfg.colour_plane_id.is_some(),
seq_scaling_lists: self.cfg.scaling_matrix.seq_spec(),
bit_depth_luma_minus8: 0,
bit_depth_chroma_minus8: 0,
interlaced_fields: false,
mbaff: false,
vui: self.cfg.vui.clone(),
};
let pps_cfg = BaselinePpsConfig {
redundant_pic_cnt_present_flag: false,
slice_groups: None,
constrained_intra_pred_flag: false,
pic_scaling_lists: self.cfg.scaling_matrix.pic_spec(),
chroma_format_idc: self.cfg.chroma_format_idc,
pic_parameter_set_id: 0,
seq_parameter_set_id: 0,
pic_init_qp_minus26: self.cfg.qp - 26,
chroma_qp_index_offset: 0,
weighted_pred_flag: !cabac && self.cfg.explicit_weighted_pred,
weighted_bipred_idc: if !cabac && self.cfg.explicit_weighted_bipred {
1
} else {
0
},
entropy_coding_mode_flag: cabac,
transform_8x8_mode_flag: self.cfg.transform_8x8,
};
(
build_baseline_sps_rbsp(&sps_cfg),
build_baseline_pps_rbsp(&pps_cfg),
)
}
pub fn parameter_sets_annex_b(&self, cabac: bool) -> Vec<u8> {
let (sps, pps) = self.parameter_sets(cabac);
let mut out = build_nal_unit(3, NalUnitType::Sps, &sps);
out.extend_from_slice(&build_nal_unit(3, NalUnitType::Pps, &pps));
out
}
}
pub struct Encoder {
cfg: EncoderConfig,
idr_pic_id: u32,
}
impl Encoder {
pub fn config(&self) -> &EncoderConfig {
&self.cfg
}
pub fn idr_pic_id(&self) -> u32 {
self.idr_pic_id
}
pub fn set_idr_pic_id(&mut self, idr_pic_id: u32) {
assert!(idr_pic_id <= 65535, "idr_pic_id out of 0..=65535 (§7.4.3)");
self.idr_pic_id = idr_pic_id;
}
pub fn with_idr_pic_id(mut self, idr_pic_id: u32) -> Self {
self.set_idr_pic_id(idr_pic_id);
self
}
}
impl Encoder {
pub fn new(cfg: EncoderConfig) -> Self {
assert!(
cfg.width % 16 == 0 && cfg.height % 16 == 0,
"encoder requires 16-pixel-aligned dimensions",
);
assert!((0..=51).contains(&cfg.qp));
assert!(
matches!(cfg.chroma_format_idc, 0..=3),
"encoder only supports 4:0:0 (chroma_format_idc=0), 4:2:0 \
(chroma_format_idc=1), 4:2:2 (chroma_format_idc=2), and \
4:4:4 (chroma_format_idc=3); got {}",
cfg.chroma_format_idc,
);
if cfg.chroma_format_idc == 0 {
assert!(
matches!(cfg.profile_idc, 100 | 110 | 122 | 244),
"4:0:0 (chroma_format_idc=0) requires a High-family profile \
(100 / 110 / 122 / 244); got {}",
cfg.profile_idc,
);
}
if let Some(plane) = cfg.colour_plane_id {
assert!(plane <= 2, "colour_plane_id must be 0..=2 (§7.4.3)");
assert_eq!(
cfg.chroma_format_idc, 0,
"a separate-colour-plane per-plane encoder codes monochrome \
pictures (ChromaArrayType 0): set chroma_format_idc = 0",
);
assert_eq!(
cfg.profile_idc, 244,
"separate_colour_plane_flag lives in the 4:4:4 SPS syntax \
(§7.3.2.1.1): profile_idc must be 244",
);
assert!(
cfg.scaling_matrix.is_flat(),
"separate-colour-plane encoding currently requires flat scaling lists",
);
}
if cfg.chroma_format_idc == 2 {
assert_eq!(
cfg.profile_idc, 122,
"4:2:2 (chroma_format_idc=2) requires profile_idc=122 (High 4:2:2)",
);
}
if cfg.chroma_format_idc == 3 {
assert_eq!(
cfg.profile_idc, 244,
"4:4:4 (chroma_format_idc=3) requires profile_idc=244 (High 4:4:4 Predictive)",
);
}
Self { cfg, idr_pic_id: 0 }
}
pub fn encode_idr(&self, frame: &YuvFrame<'_>) -> EncodedIdr {
self.encode_idr_with_qp(frame, self.cfg.qp)
}
pub fn encode_idr_with_qp(&self, frame: &YuvFrame<'_>, frame_qp: i32) -> EncodedIdr {
self.encode_idr_impl(frame, frame_qp, None)
}
pub fn encode_idr_rate_adaptive(
&self,
frame: &YuvFrame<'_>,
frame_qp: i32,
frame_budget_bits: u64,
) -> EncodedIdr {
assert_eq!(
self.cfg.chroma_format_idc, 1,
"encode_idr_rate_adaptive requires 4:2:0"
);
self.encode_idr_impl(frame, frame_qp, Some(frame_budget_bits))
}
fn encode_idr_impl(
&self,
frame: &YuvFrame<'_>,
frame_qp: i32,
row_budget_bits: Option<u64>,
) -> EncodedIdr {
assert!(
(0..=51).contains(&frame_qp),
"frame_qp {frame_qp} out of 0..=51"
);
assert_eq!(frame.width, self.cfg.width);
assert_eq!(frame.height, self.cfg.height);
let width_mbs = self.cfg.width / 16;
let height_mbs = self.cfg.height / 16;
let chroma_width = match self.cfg.chroma_format_idc {
0 => 0usize,
3 => self.cfg.width as usize,
_ => (self.cfg.width / 2) as usize,
};
let chroma_height = match self.cfg.chroma_format_idc {
0 => 0usize,
2 => self.cfg.height as usize,
3 => self.cfg.height as usize,
_ => (self.cfg.height / 2) as usize,
};
debug_assert_eq!(
frame.u.len(),
chroma_width * chroma_height,
"Cb plane length {} does not match {}x{} for chroma_format_idc={}",
frame.u.len(),
chroma_width,
chroma_height,
self.cfg.chroma_format_idc,
);
debug_assert_eq!(frame.v.len(), chroma_width * chroma_height);
if !self.cfg.scaling_matrix.is_flat() {
assert!(
!self.cfg.trellis_quant_intra && !self.cfg.trellis_quant_intra_chroma,
"trellis intra refinement assumes flat scaling lists",
);
}
if let ScalingMatrixMode::SeqCustom(l) | ScalingMatrixMode::PicCustom(l) =
self.cfg.scaling_matrix
{
assert!(
l.is_valid(),
"custom scaling-list values must all be in 1..=255",
);
}
let (sps_rbsp, pps_rbsp) = self.parameter_sets(false);
let mut stream: Vec<u8> = Vec::new();
stream.extend_from_slice(&build_nal_unit(3, NalUnitType::Sps, &sps_rbsp));
stream.extend_from_slice(&build_nal_unit(3, NalUnitType::Pps, &pps_rbsp));
let mut sw = BitWriter::new();
write_idr_i_slice_header(
&mut sw,
&IdrSliceHeaderConfig {
first_mb_in_slice: 0,
slice_type_raw: 7,
pic_parameter_set_id: 0,
colour_plane_id: self.cfg.colour_plane_id,
frame_num: 0,
frame_num_bits: 8,
idr_pic_id: self.idr_pic_id,
pic_order_cnt_lsb: 0,
poc_lsb_bits: 8,
slice_qp_delta: frame_qp - self.cfg.qp,
disable_deblocking_filter_idc: if self.cfg.chroma_format_idc == 3 {
1
} else {
0
},
slice_alpha_c0_offset_div2: 0,
slice_beta_offset_div2: 0,
field: FieldPicSignal::FrameMbsOnly,
idr: true,
nal_ref_idc: 3,
long_term_reference_flag: false,
mmco: &[],
redundant_pic_cnt: None,
slice_group_change_cycle: None,
},
);
let mut recon_y = vec![0u8; (self.cfg.width * self.cfg.height) as usize];
let mut recon_u = vec![0u8; chroma_width * chroma_height];
let mut recon_v = vec![0u8; chroma_width * chroma_height];
let mut qp_tracker = MbQpTracker::new(frame_qp);
let mut want_qp = frame_qp;
let qp_bd_offset_c = qp_bd_offset(self.cfg.bit_depth_chroma_minus8);
let mut nc_grid = CavlcNcGrid::new(width_mbs, height_mbs);
let mut intra_grid = IntraGrid::new(width_mbs as usize, height_mbs as usize);
let mut mb_deblock_infos: Vec<MbDeblockInfo> =
vec![MbDeblockInfo::default(); (width_mbs * height_mbs) as usize];
if self.cfg.transform_8x8 {
assert!(
matches!(self.cfg.chroma_format_idc, 0..=3),
"transform_8x8 encode requires chroma_format_idc in {{0, 1, 2, 3}}"
);
}
let mut i8x8_mb_count = 0u32;
let mut i4x4_mb_count = 0u32;
for mb_y in 0..height_mbs {
if let Some(budget) = row_budget_bits {
if mb_y > 0 && budget > 0 {
let frac = f64::from(mb_y) / f64::from(height_mbs);
let spent = sw.bits_emitted() as f64;
let expected = (budget as f64 * frac).max(1.0);
let ratio = spent.max(1.0) / expected;
let adj = if ratio >= 1.0 {
(6.0 * ratio.log2()).round() as i32
} else {
(3.0 * ratio.log2()).round() as i32
};
let mut new_want = frame_qp + adj;
if new_want < frame_qp && frac < 0.6 {
new_want = frame_qp;
}
want_qp = new_want
.clamp(want_qp - 1, want_qp + 2)
.clamp(frame_qp - 2, frame_qp + 6)
.clamp(0, 51);
}
}
let qp_y = want_qp;
let qp_c = qp_y_to_qp_c_with_bd_offset(
qp_y,
0,
qp_bd_offset_c,
);
for mb_x in 0..width_mbs {
let dbl = self.encode_mb(
frame,
mb_x as usize,
mb_y as usize,
qp_y,
qp_c,
chroma_width,
chroma_height,
&mut recon_y,
&mut recon_u,
&mut recon_v,
&mut sw,
&mut nc_grid,
&mut intra_grid,
&mut qp_tracker,
);
if dbl.transform_size_8x8_flag {
i8x8_mb_count += 1;
}
if dbl.is_intra_4x4 {
i4x4_mb_count += 1;
}
let mb_addr = (mb_y * width_mbs + mb_x) as usize;
mb_deblock_infos[mb_addr] = dbl;
}
}
sw.rbsp_trailing_bits();
let slice_rbsp = sw.into_bytes();
stream.extend_from_slice(&build_nal_unit(3, NalUnitType::SliceIdr, &slice_rbsp));
if self.cfg.chroma_format_idc != 3 {
crate::encoder::deblock::deblock_recon_with_chroma_array_type(
self.cfg.width,
self.cfg.height,
chroma_width as u32,
chroma_height as u32,
&mut recon_y,
&mut recon_u,
&mut recon_v,
&mb_deblock_infos,
0,
width_mbs,
height_mbs,
self.cfg.chroma_format_idc,
);
}
let n_parts = (width_mbs as usize) * (height_mbs as usize) * 4;
let partition_mvs = vec![
FrameRefPartitionMv {
mv_l0: (0, 0),
ref_idx_l0: -1,
is_intra: true,
};
n_parts
];
EncodedIdr {
annex_b: stream,
recon_y,
recon_u,
recon_v,
recon_width: self.cfg.width,
recon_height: self.cfg.height,
partition_mvs,
i8x8_mb_count,
i4x4_mb_count,
}
}
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_arguments)]
fn encode_mb(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
qp_tracker: &mut MbQpTracker,
) -> MbDeblockInfo {
let mb_qp_delta = qp_y - qp_tracker.cur;
let dbl = self.encode_mb_inner(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
);
qp_tracker.cur = dbl.qp_y;
dbl
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_inner(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_qp_delta: i32,
) -> MbDeblockInfo {
let width = self.cfg.width as usize;
let mb_addr = (mb_y as u32) * self.cfg.width / 16 + (mb_x as u32);
if self.cfg.chroma_format_idc == 0 {
let lambda = lambda_ssd(qp_y, true);
let snap = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
0,
);
let bits_before = sw.bits_emitted();
let mb_ssd_luma = |recon_y: &[u8]| -> u64 {
let mut d = 0u64;
for j in 0..16usize {
for i in 0..16usize {
let s = frame.y[(mb_y * 16 + j) * width + mb_x * 16 + i] as i64;
let r = recon_y[(mb_y * 16 + j) * width + mb_x * 16 + i] as i64;
d += ((s - r) * (s - r)) as u64;
}
}
d
};
let trial_a = self.encode_mb_intra16x16(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
None,
);
let d_a = mb_ssd_luma(recon_y);
let r_a = sw.bits_emitted() - bits_before;
let cost_a = cost_combined(d_a, r_a as u64, lambda);
let post_a = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
0,
);
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
let dbl_b = self.encode_mb_intra4x4_mono(
frame,
mb_x,
mb_y,
qp_y,
recon_y,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
);
let d_b = mb_ssd_luma(recon_y);
let r_b = sw.bits_emitted() - bits_before;
let cost_b = cost_combined(d_b, r_b as u64, lambda);
if cost_a <= cost_b {
post_a.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
return trial_a.deblock;
}
return dbl_b;
}
if self.cfg.chroma_format_idc == 3 {
let lambda = lambda_ssd(qp_y, true);
let snap = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
3,
);
let bits_before = sw.bits_emitted();
let trial_a = self.encode_mb_intra16x16(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
None,
);
let d_a = mb_ssd_3planes_444(frame, recon_y, recon_u, recon_v, width, mb_x, mb_y);
let r_a = sw.bits_emitted() - bits_before;
let cost_a = cost_combined(d_a, r_a as u64, lambda);
let post_a = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
3,
);
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
let dbl_b = self.encode_mb_intra4x4_444(
frame, mb_x, mb_y, qp_y, qp_c, recon_y, recon_u, recon_v, sw, nc_grid, intra_grid,
);
let d_b = mb_ssd_3planes_444(frame, recon_y, recon_u, recon_v, width, mb_x, mb_y);
let r_b = sw.bits_emitted() - bits_before;
let cost_b = cost_combined(d_b, r_b as u64, lambda);
if !self.cfg.transform_8x8 {
if cost_a <= cost_b {
post_a.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
return trial_a.deblock;
}
return dbl_b;
}
let post_b = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
3,
);
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
let dbl_c = self.encode_mb_intra8x8_444(
frame, mb_x, mb_y, qp_y, qp_c, recon_y, recon_u, recon_v, sw, nc_grid, intra_grid,
);
let d_c = mb_ssd_3planes_444(frame, recon_y, recon_u, recon_v, width, mb_x, mb_y);
let r_c = sw.bits_emitted() - bits_before;
let cost_c = cost_combined(d_c, r_c as u64, lambda);
if cost_a <= cost_b && cost_a <= cost_c {
post_a.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
return trial_a.deblock;
}
if cost_b <= cost_c {
post_b.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
return dbl_b;
}
return dbl_c;
}
let snap = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
1,
);
let trial_a = self.encode_mb_intra16x16(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
None,
);
let snap_a = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
1,
);
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
let trial_b = self.encode_mb_intra4x4(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
);
let trial_c = if self.cfg.transform_8x8 {
let snap_b = MbStateSnapshot::capture(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
1,
);
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
let trial_c = self.encode_mb_intra8x8(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mb_qp_delta,
);
Some((trial_c, snap_b))
} else {
None
};
if let Some((trial_c, snap_b)) = trial_c {
if trial_a.luma_cost <= trial_b.luma_cost && trial_a.luma_cost <= trial_c.luma_cost {
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
snap_a.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
trial_a.deblock
} else if trial_b.luma_cost <= trial_c.luma_cost {
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
snap_b.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
trial_b.deblock
} else {
trial_c.deblock
}
} else if trial_a.luma_cost <= trial_b.luma_cost {
snap.restore(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
snap_a.install(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mb_addr as usize,
);
trial_a.deblock
} else {
trial_b.deblock
}
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_intra16x16(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_qp_delta: i32,
nb: Option<(bool, bool, bool)>,
) -> MbTrial {
debug_assert!(
matches!(self.cfg.chroma_format_idc, 0 | 1) || mb_qp_delta == 0,
"round-430 row modulation (non-zero mb_qp_delta) is 4:2:0/4:0:0-only; \
the 4:2:2 / 4:4:4 writer branches still hard-code delta 0",
);
let width = self.cfg.width as usize;
let height = self.cfg.height as usize;
let lambda = lambda_ssd(qp_y, true);
let wq = self.cfg.scaling_matrix.intra_weights();
let mut best_luma_mode = I16x16Mode::Dc;
let mut best_luma_pred = [0i32; 256];
let mut best_luma_cost = u64::MAX;
for &mode in &[
I16x16Mode::Vertical,
I16x16Mode::Horizontal,
I16x16Mode::Dc,
I16x16Mode::Plane,
] {
let Some(pred) = predict_16x16_avail(mode, recon_y, width, height, mb_x, mb_y, nb)
else {
continue;
};
let cost = trial_intra16x16_cost(frame, &pred, mb_x, mb_y, qp_y, width, lambda, &wq.w4);
if cost < best_luma_cost {
best_luma_cost = cost;
best_luma_mode = mode;
best_luma_pred = pred;
}
}
let mut residual = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
let s = frame.y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
residual[j * 16 + i] = s - best_luma_pred[j * 16 + i];
}
}
let mut coeffs_4x4 = [[0i32; 16]; 16];
for (blk, slot) in coeffs_4x4.iter_mut().enumerate() {
let bx = blk % 4;
let by = blk / 4;
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by * 4 + j) * 16 + bx * 4 + i];
}
}
*slot = forward_core_4x4(&block);
}
let mut dc_coeffs = [0i32; 16];
for (blk, coeffs) in coeffs_4x4.iter().enumerate() {
let bx = blk % 4;
let by = blk / 4;
dc_coeffs[by * 4 + bx] = coeffs[0];
}
let dc_hadamard = forward_hadamard_4x4(&dc_coeffs);
let dc_levels = quantize_luma_dc_w(&dc_hadamard, qp_y, true, wq.w4[0]);
let mut luma_ac_levels = [[0i32; 16]; 16];
let mut luma_ac_quant_raster = [[0i32; 16]; 16]; let mut any_luma_ac_nz = false;
for blkz in 0..16usize {
let (bx, by) = LUMA_4X4_BLK[blkz];
let raster = by * 4 + bx;
let z = quantize_4x4_ac_w(&coeffs_4x4[raster], qp_y, true, &wq.w4);
luma_ac_quant_raster[raster] = z;
let scan_ac = zigzag_scan_4x4_ac(&z);
luma_ac_levels[blkz] = scan_ac;
if scan_ac.iter().any(|&v| v != 0) {
any_luma_ac_nz = true;
}
}
#[allow(unused_mut)]
let mut cbp_luma: u8 = if any_luma_ac_nz { 15 } else { 0 };
let inv_dc =
inverse_hadamard_luma_dc_16x16(&dc_levels, qp_y, &wq.w4, 8).expect("luma DC inverse");
#[allow(clippy::needless_range_loop)] for blk in 0..16usize {
let bx = blk % 4;
let by = blk / 4;
let mut coeffs_in = if cbp_luma == 15 {
luma_ac_quant_raster[blk]
} else {
[0i32; 16]
};
coeffs_in[0] = inv_dc[by * 4 + bx];
let r = inverse_transform_4x4_dc_preserved(&coeffs_in, qp_y, &wq.w4, 8)
.expect("inverse 4x4");
for j in 0..4 {
for i in 0..4 {
let px = mb_x * 16 + bx * 4 + i;
let py = mb_y * 16 + by * 4 + j;
let p = best_luma_pred[(by * 4 + j) * 16 + bx * 4 + i];
let raw = p + r[j * 4 + i];
recon_y[py * width + px] = raw.clamp(0, 255) as u8;
}
}
}
let chroma_array_type = self.cfg.chroma_format_idc;
let mut chroma_block_for_mask: Option<ChromaBlock> = None;
let mut cb_dc_444 = [0i32; 16];
let mut cr_dc_444 = [0i32; 16];
let mut cb_ac_scan_444 = [[0i32; 16]; 16];
let mut cr_ac_scan_444 = [[0i32; 16]; 16];
let mut cb_ac_quant_444 = [[0i32; 16]; 16];
let mut cr_ac_quant_444 = [[0i32; 16]; 16];
let (cbp_chroma, best_chroma_mode_u8): (u8, u8) = if chroma_array_type == 3 {
let (cb_dc, cb_ac_scan, cb_ac_quant, cb_nz) = encode_intra16x16_chroma_plane_luma_like(
best_luma_mode,
frame.u,
recon_u,
chroma_width,
chroma_height,
mb_x,
mb_y,
qp_c,
&wq.w4,
);
let (cr_dc, cr_ac_scan, cr_ac_quant, cr_nz) = encode_intra16x16_chroma_plane_luma_like(
best_luma_mode,
frame.v,
recon_v,
chroma_width,
chroma_height,
mb_x,
mb_y,
qp_c,
&wq.w4,
);
cb_dc_444 = cb_dc;
cr_dc_444 = cr_dc;
cb_ac_scan_444 = cb_ac_scan;
cr_ac_scan_444 = cr_ac_scan;
cb_ac_quant_444 = cb_ac_quant;
cr_ac_quant_444 = cr_ac_quant;
if cb_nz || cr_nz {
cbp_luma = 15;
}
(0u8, 0u8)
} else if chroma_array_type == 0 {
(0u8, 0u8)
} else {
let chroma_block = encode_chroma_block(
chroma_array_type,
frame.u,
frame.v,
recon_u,
recon_v,
chroma_width,
chroma_height,
mb_x,
mb_y,
qp_c,
&wq.w4,
nb,
);
let cm = chroma_block.pred_mode.as_u8();
let ccbp = chroma_block.cbp_chroma;
install_chroma_recon(
chroma_array_type,
chroma_width,
mb_x,
mb_y,
&chroma_block,
recon_u,
recon_v,
);
chroma_block_for_mask = Some(chroma_block);
(ccbp, cm)
};
let n_chroma_blk = ChromaBlock::cb_block_count(chroma_array_type);
let mut u_dc_levels = [0i32; 4];
let mut v_dc_levels = [0i32; 4];
let mut u_ac_levels = [[0i32; 16]; 4];
let mut v_ac_levels = [[0i32; 16]; 4];
if chroma_array_type == 1 {
let cb = chroma_block_for_mask.as_ref().unwrap();
u_dc_levels.copy_from_slice(&cb.dc_cb[..4]);
v_dc_levels.copy_from_slice(&cb.dc_cr[..4]);
u_ac_levels.copy_from_slice(&cb.ac_cb[..4]);
v_ac_levels.copy_from_slice(&cb.ac_cr[..4]);
}
let _ = chroma_height;
let width_mbs = self.cfg.width / 16;
let mb_addr = (mb_y as u32) * width_mbs + (mb_x as u32);
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16]; }
let nc_dc = derive_nc_luma(nc_grid, mb_addr, 0, LumaNcKind::Intra16x16Dc, true, false);
let mut luma_ac_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blkz in 0..16u8 {
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr, blkz, LumaNcKind::Ac, true, false);
luma_ac_nc[blkz as usize] = nc;
if cbp_luma == 15 {
let tc = luma_ac_levels[blkz as usize]
.iter()
.filter(|&&v| v != 0)
.count() as u8;
own_totals[blkz as usize] = tc;
}
}
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc_dc_ctx = CoeffTokenContext::Numeric(nc_dc);
if chroma_array_type == 1 {
let chroma_block = chroma_block_for_mask.as_ref().unwrap();
let (nc_cb, nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr,
true,
cbp_chroma,
&chroma_block.ac_cb[..4],
&chroma_block.ac_cr[..4],
1,
);
let mb_cfg = I16x16McbConfig {
pred_mode: best_luma_mode.as_u8(),
intra_chroma_pred_mode: best_chroma_mode_u8,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_dc_levels_raster: dc_levels,
luma_ac_levels,
luma_ac_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: nc_cb,
chroma_ac_nc_cr: nc_cr,
};
write_intra16x16_mb(sw, &mb_cfg, nc_dc_ctx).expect("write Intra16x16 mb");
} else if chroma_array_type == 0 {
crate::encoder::macroblock::write_intra16x16_mb_chroma(
sw,
best_luma_mode.as_u8(),
0,
cbp_luma,
cbp_chroma,
mb_qp_delta,
&dc_levels,
&luma_ac_levels,
&luma_ac_nc,
crate::encoder::macroblock::ChromaWriteKind::Mono,
nc_dc_ctx,
)
.expect("write Intra16x16 mb (4:0:0)");
} else if chroma_array_type == 2 {
let chroma_block = chroma_block_for_mask.as_ref().unwrap();
let (nc_cb, nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr,
true,
cbp_chroma,
&chroma_block.ac_cb,
&chroma_block.ac_cr,
2,
);
crate::encoder::macroblock::write_intra16x16_mb_chroma(
sw,
best_luma_mode.as_u8(),
best_chroma_mode_u8,
cbp_luma,
cbp_chroma,
0,
&dc_levels,
&luma_ac_levels,
&luma_ac_nc,
crate::encoder::macroblock::ChromaWriteKind::Yuv422 {
chroma_dc_cb: &chroma_block.dc_cb,
chroma_dc_cr: &chroma_block.dc_cr,
chroma_ac_cb: &chroma_block.ac_cb,
chroma_ac_cr: &chroma_block.ac_cr,
cb_ac_nc: &nc_cb,
cr_ac_nc: &nc_cr,
},
nc_dc_ctx,
)
.expect("write Intra16x16 mb (4:2:2)");
} else {
let (cb_dc_nc, cb_ac_nc) =
derive_plane_nc_444_i16x16(nc_grid, mb_addr, false, cbp_luma, &cb_ac_scan_444);
let (cr_dc_nc, cr_ac_nc) =
derive_plane_nc_444_i16x16(nc_grid, mb_addr, true, cbp_luma, &cr_ac_scan_444);
crate::encoder::macroblock::write_intra16x16_mb_chroma(
sw,
best_luma_mode.as_u8(),
0,
cbp_luma,
cbp_chroma,
0,
&dc_levels,
&luma_ac_levels,
&luma_ac_nc,
crate::encoder::macroblock::ChromaWriteKind::Yuv444 {
cb_dc_raster: &cb_dc_444,
cr_dc_raster: &cr_dc_444,
cb_ac_levels: &cb_ac_scan_444,
cr_ac_levels: &cr_ac_scan_444,
cb_ac_nc: &cb_ac_nc,
cr_ac_nc: &cr_ac_nc,
cb_dc_nc,
cr_dc_nc,
cbp_luma_for_ac_gate: cbp_luma,
},
nc_dc_ctx,
)
.expect("write Intra16x16 mb (4:4:4)");
let _ = (cb_ac_quant_444, cr_ac_quant_444);
}
let slot = intra_grid.slot_mut(mb_x, mb_y);
slot.available = true;
slot.is_i_nxn = false;
let luma_nonzero_4x4 = if cbp_luma == 15 {
let mut nz = [false; 16];
for (blkz, scan_ac) in luma_ac_levels.iter().enumerate() {
nz[blkz] = scan_ac.iter().any(|&v| v != 0);
}
luma_nz_mask_from_blocks(&nz)
} else {
0
};
let chroma_nz_mask = chroma_block_for_mask
.as_ref()
.map(|cb| chroma_nz_mask_from_chroma_block(cb, n_chroma_blk))
.unwrap_or(0);
MbTrial {
luma_cost: best_luma_cost,
deblock: MbDeblockInfo {
is_intra: true,
qp_y,
luma_nonzero_4x4,
chroma_nonzero_4x4: chroma_nz_mask,
..Default::default()
},
}
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_intra4x4(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_qp_delta: i32,
) -> MbTrial {
debug_assert!(
self.cfg.chroma_format_idc == 1 || mb_qp_delta == 0,
"round-430 row modulation (non-zero mb_qp_delta) is 4:2:0-only; \
the 4:2:2 / 4:4:4 writer branches still hard-code delta 0",
);
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let lambda = lambda_ssd(qp_y, true);
{
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = true;
s.intra_4x4_pred_modes = [0u8; 16];
}
let mut chosen_mode = [0u8; 16]; let mut prev_flag = [false; 16];
let mut rem = [0u8; 16];
let mut luma_4x4_levels_scan = [[0i32; 16]; 16]; let wq = self.cfg.scaling_matrix.intra_weights();
let mut luma_4x4_quant_raster = [[0i32; 16]; 16]; let mut blk_has_nz = [false; 16];
let mb_overhead_bits: u64 = 9 + 16;
let mut mb_luma_cost: u64 = cost_combined(0, mb_overhead_bits, lambda);
for blk in 0..16usize {
let (bx, by) = I4X4_XY[blk];
let avail = i4x4_avail(mb_x, mb_y, bx, by, width_mbs);
let predicted = predicted_intra4x4_mode(intra_grid, mb_x, mb_y, blk, width_mbs);
let mut best_mode = Intra4x4Mode::Dc;
let mut best_cost = u64::MAX;
let mut best_pred = [0i32; 16];
let mut best_z_raster = [0i32; 16];
let mut best_z_scan = [0i32; 16];
let mut best_recon = [0i32; 16];
for &mode in &Intra4x4Mode::ALL {
let Some(pred) = predict_4x4(mode, recon_y, width, mb_x, mb_y, bx, by, avail)
else {
continue;
};
let candidate = mode.as_u8();
let mut block_res = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
let s = frame.y[(mb_y * 16 + by + j) * width + (mb_x * 16 + bx + i)] as i32;
block_res[j * 4 + i] = s - pred[j * 4 + i];
}
}
let w_coeffs = forward_core_4x4(&block_res);
let z_raster = quantize_4x4_w(&w_coeffs, qp_y, true, &wq.w4);
let z_scan = zigzag_scan_4x4(&z_raster);
let r = inverse_transform_4x4(&z_raster, qp_y, &wq.w4, 8).expect("inverse 4x4");
let mut recon_block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
recon_block[j * 4 + i] = pred[j * 4 + i] + r[j * 4 + i];
}
}
let d = ssd_4x4(frame.y, width, mb_x * 16 + bx, mb_y * 16 + by, &recon_block);
let mut bits: u64 = 1;
if candidate != predicted {
bits += 3;
}
let mut trial_w = BitWriter::new();
if encode_residual_block_cavlc(
&mut trial_w,
CoeffTokenContext::Numeric(0),
16,
&z_scan,
)
.is_ok()
{
bits += trial_w.bits_emitted() as u64;
}
let cost = cost_combined(d, bits, lambda);
if cost < best_cost {
best_cost = cost;
best_mode = mode;
best_pred = pred;
best_z_raster = z_raster;
best_z_scan = z_scan;
best_recon = recon_block;
}
}
let chosen = best_mode.as_u8();
chosen_mode[blk] = chosen;
mb_luma_cost = mb_luma_cost.saturating_add(best_cost);
if chosen == predicted {
prev_flag[blk] = true;
} else {
prev_flag[blk] = false;
let r = if chosen < predicted {
chosen
} else {
chosen - 1
};
rem[blk] = r;
}
intra_grid.slot_mut(mb_x, mb_y).intra_4x4_pred_modes[blk] = chosen;
luma_4x4_quant_raster[blk] = best_z_raster;
luma_4x4_levels_scan[blk] = best_z_scan;
blk_has_nz[blk] = best_z_raster.iter().any(|&v| v != 0);
let _ = best_pred;
for j in 0..4 {
for i in 0..4 {
let px = mb_x * 16 + bx + i;
let py = mb_y * 16 + by + j;
recon_y[py * width + px] = best_recon[j * 4 + i].clamp(0, 255) as u8;
}
}
}
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
let any_nz = (0..4).any(|sub| blk_has_nz[blk8 * 4 + sub]);
if any_nz {
cbp_luma |= 1u8 << blk8;
}
}
for blk8 in 0..4usize {
if (cbp_luma >> blk8) & 1 == 1 {
continue;
}
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = I4X4_XY[blk];
let mode = Intra4x4Mode::from_u8(chosen_mode[blk]).expect("valid mode");
let avail = i4x4_avail(mb_x, mb_y, bx, by, width_mbs);
let pred = predict_4x4(mode, recon_y, width, mb_x, mb_y, bx, by, avail)
.expect("predictor exists for chosen mode");
for j in 0..4 {
for i in 0..4 {
let px = mb_x * 16 + bx + i;
let py = mb_y * 16 + by + j;
recon_y[py * width + px] = pred[j * 4 + i].clamp(0, 255) as u8;
}
}
}
}
let chroma_array_type = self.cfg.chroma_format_idc;
let chroma_block = encode_chroma_block(
chroma_array_type,
frame.u,
frame.v,
recon_u,
recon_v,
chroma_width,
chroma_height,
mb_x,
mb_y,
qp_c,
&wq.w4,
None,
);
let best_chroma_mode = chroma_block.pred_mode;
let cbp_chroma = chroma_block.cbp_chroma;
let n_chroma_blk = ChromaBlock::cb_block_count(chroma_array_type);
let mut u_dc_levels = [0i32; 4];
let mut v_dc_levels = [0i32; 4];
let mut u_ac_levels = [[0i32; 16]; 4];
let mut v_ac_levels = [[0i32; 16]; 4];
if chroma_array_type == 1 {
u_dc_levels.copy_from_slice(&chroma_block.dc_cb[..4]);
v_dc_levels.copy_from_slice(&chroma_block.dc_cr[..4]);
u_ac_levels.copy_from_slice(&chroma_block.ac_cb[..4]);
v_ac_levels.copy_from_slice(&chroma_block.ac_cr[..4]);
}
install_chroma_recon(
chroma_array_type,
chroma_width,
mb_x,
mb_y,
&chroma_block,
recon_u,
recon_v,
);
let _ = chroma_height;
let pic_w_mbs = self.cfg.width / 16;
let mb_addr = (mb_y as u32) * pic_w_mbs + (mb_x as u32);
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr, blk as u8, LumaNcKind::Ac, true, false);
luma_4x4_nc[blk] = nc;
if (cbp_luma >> blk8) & 1 == 1 {
let tc = luma_4x4_levels_scan[blk]
.iter()
.filter(|&&v| v != 0)
.count() as u8;
own_totals[blk] = tc;
} else {
own_totals[blk] = 0;
}
}
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let (chroma_nc_cb, chroma_nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr,
true,
cbp_chroma,
&chroma_block.ac_cb[..ChromaBlock::cb_block_count(chroma_array_type)],
&chroma_block.ac_cr[..ChromaBlock::cb_block_count(chroma_array_type)],
chroma_array_type,
);
if chroma_array_type == 1 {
let mb_cfg = INxNMcbConfig {
emit_transform_size_8x8_zero: self.cfg.transform_8x8,
prev_intra4x4_pred_mode_flag: prev_flag,
rem_intra4x4_pred_mode: rem,
intra_chroma_pred_mode: best_chroma_mode.as_u8(),
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_i_nxn_mb(sw, &mb_cfg).expect("write I_NxN mb");
} else {
crate::encoder::macroblock::write_i_nxn_mb_chroma(
sw,
self.cfg.transform_8x8,
&prev_flag,
&rem,
best_chroma_mode.as_u8(),
cbp_luma,
cbp_chroma,
0,
&luma_4x4_levels_scan,
&luma_4x4_nc,
crate::encoder::macroblock::ChromaWriteKind::Yuv422 {
chroma_dc_cb: &chroma_block.dc_cb,
chroma_dc_cr: &chroma_block.dc_cr,
chroma_ac_cb: &chroma_block.ac_cb,
chroma_ac_cr: &chroma_block.ac_cr,
cb_ac_nc: &chroma_nc_cb,
cr_ac_nc: &chroma_nc_cr,
},
)
.expect("write I_NxN mb (4:2:2)");
}
let luma_nonzero_4x4 = luma_nz_mask_from_blocks(&blk_has_nz);
let chroma_nz_mask = chroma_nz_mask_from_chroma_block(&chroma_block, n_chroma_blk);
MbTrial {
luma_cost: mb_luma_cost,
deblock: MbDeblockInfo {
is_intra: true,
qp_y: if cbp_luma > 0 || cbp_chroma > 0 {
qp_y
} else {
qp_y - mb_qp_delta
},
luma_nonzero_4x4,
chroma_nonzero_4x4: chroma_nz_mask,
is_intra_4x4: true,
..Default::default()
},
}
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_intra8x8(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_qp_delta: i32,
) -> MbTrial {
debug_assert!(
self.cfg.chroma_format_idc == 1 || mb_qp_delta == 0,
"round-430 row modulation (non-zero mb_qp_delta) is 4:2:0-only; \
the 4:2:2 / 4:4:4 writer branches still hard-code delta 0",
);
let width = self.cfg.width as usize;
let lambda = lambda_ssd(qp_y, true);
let wq = self.cfg.scaling_matrix.intra_weights();
let sl8 = wq.w8;
{
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
s.is_i8x8 = true;
s.intra_8x8_pred_modes = [2u8; 4];
}
let mut prev_flag = [false; 4];
let mut rem = [0u8; 4];
let mut luma_4x4_levels = [[0i32; 16]; 16];
let mut blk8_has_nz = [false; 4];
let mb_overhead_bits: u64 = 10 + 16;
let mut mb_luma_cost: u64 = cost_combined(0, mb_overhead_bits, lambda);
for blk8 in 0..4usize {
let bx = mb_x * 16 + (blk8 % 2) * 8;
let by = mb_y * 16 + (blk8 / 2) * 8;
let predicted = predicted_intra8x8_mode(intra_grid, mb_x, mb_y, blk8);
let raw = gather_8x8_from_recon(recon_y, width, mb_x, mb_y, blk8);
let filtered = filter_samples_8x8(&raw, 8);
let mut best_mode = Intra8x8Mode::Dc;
let mut best_cost = u64::MAX;
let mut best_z_raster = [0i32; 64];
let mut best_recon = [0i32; 64];
for mode_idx in 0..9u8 {
let mode = Intra8x8Mode::from_index(mode_idx).expect("mode index in range");
if !intra8x8_mode_available(mode, &raw.availability) {
continue;
}
let mut pred = [0i32; 64];
predict_8x8(mode, &filtered, 8, &mut pred);
let mut res = [0i32; 64];
for j in 0..8 {
for i in 0..8 {
let s = frame.y[(by + j) * width + bx + i] as i32;
res[j * 8 + i] = s - pred[j * 8 + i];
}
}
let w_coeffs = forward_core_8x8(&res);
let z_raster = quantize_8x8_w(&w_coeffs, qp_y, true, &wq.w8);
let r = inverse_transform_8x8(&z_raster, qp_y, &sl8, 8).expect("inverse 8x8");
let mut recon_blk = [0i32; 64];
let mut d: u64 = 0;
for j in 0..8 {
for i in 0..8 {
let v = (pred[j * 8 + i] + r[j * 8 + i]).clamp(0, 255);
recon_blk[j * 8 + i] = v;
let s = frame.y[(by + j) * width + bx + i] as i32;
let e = (s - v) as i64;
d += (e * e) as u64;
}
}
let mut bits: u64 = 1;
if mode_idx != predicted {
bits += 3;
}
let scan = scan_8x8_for_cavlc(&z_raster, sw.field_scan());
let subs = deinterleave_8x8_to_4x4(&scan);
let mut trial_w = BitWriter::new();
for sub in &subs {
if encode_residual_block_cavlc(
&mut trial_w,
CoeffTokenContext::Numeric(0),
16,
sub,
)
.is_err()
{
break;
}
}
bits += trial_w.bits_emitted() as u64;
let cost = cost_combined(d, bits, lambda);
if cost < best_cost {
best_cost = cost;
best_mode = mode;
best_z_raster = z_raster;
best_recon = recon_blk;
}
}
let chosen = best_mode.as_index();
mb_luma_cost = mb_luma_cost.saturating_add(best_cost);
if chosen == predicted {
prev_flag[blk8] = true;
} else {
prev_flag[blk8] = false;
rem[blk8] = if chosen < predicted {
chosen
} else {
chosen - 1
};
}
intra_grid.slot_mut(mb_x, mb_y).intra_8x8_pred_modes[blk8] = chosen;
let scan = scan_8x8_for_cavlc(&best_z_raster, sw.field_scan());
let subs = deinterleave_8x8_to_4x4(&scan);
for (i4x4, sub) in subs.iter().enumerate() {
luma_4x4_levels[blk8 * 4 + i4x4] = *sub;
}
blk8_has_nz[blk8] = best_z_raster.iter().any(|&v| v != 0);
for j in 0..8 {
for i in 0..8 {
recon_y[(by + j) * width + bx + i] = best_recon[j * 8 + i] as u8;
}
}
}
let mut cbp_luma: u8 = 0;
for (blk8, &has_nz) in blk8_has_nz.iter().enumerate() {
if has_nz {
cbp_luma |= 1u8 << blk8;
}
}
let chroma_array_type = self.cfg.chroma_format_idc;
debug_assert!(
matches!(chroma_array_type, 1 | 2),
"Intra_8x8 encode supports 4:2:0 and 4:2:2 chroma"
);
let chroma_block = encode_chroma_block(
chroma_array_type,
frame.u,
frame.v,
recon_u,
recon_v,
chroma_width,
chroma_height,
mb_x,
mb_y,
qp_c,
&wq.w4,
None,
);
let cbp_chroma = chroma_block.cbp_chroma;
let n_chroma_blk = ChromaBlock::cb_block_count(chroma_array_type);
install_chroma_recon(
chroma_array_type,
chroma_width,
mb_x,
mb_y,
&chroma_block,
recon_u,
recon_v,
);
let _ = chroma_height;
let pic_w_mbs = self.cfg.width / 16;
let mb_addr = (mb_y as u32) * pic_w_mbs + (mb_x as u32);
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr, blk as u8, LumaNcKind::Ac, true, false);
luma_4x4_nc[blk] = nc;
if (cbp_luma >> blk8) & 1 == 1 {
let tc = luma_4x4_levels[blk].iter().filter(|&&v| v != 0).count() as u8;
own_totals[blk] = tc;
} else {
own_totals[blk] = 0;
}
}
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let (chroma_nc_cb, chroma_nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr,
true,
cbp_chroma,
&chroma_block.ac_cb[..ChromaBlock::cb_block_count(chroma_array_type)],
&chroma_block.ac_cr[..ChromaBlock::cb_block_count(chroma_array_type)],
chroma_array_type,
);
if chroma_array_type == 1 {
let mut u_dc_levels = [0i32; 4];
let mut v_dc_levels = [0i32; 4];
let mut u_ac_levels = [[0i32; 16]; 4];
let mut v_ac_levels = [[0i32; 16]; 4];
u_dc_levels.copy_from_slice(&chroma_block.dc_cb[..4]);
v_dc_levels.copy_from_slice(&chroma_block.dc_cr[..4]);
u_ac_levels.copy_from_slice(&chroma_block.ac_cb[..4]);
v_ac_levels.copy_from_slice(&chroma_block.ac_cr[..4]);
let mb_cfg = I8x8McbConfig {
prev_intra8x8_pred_mode_flag: prev_flag,
rem_intra8x8_pred_mode: rem,
intra_chroma_pred_mode: chroma_block.pred_mode.as_u8(),
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_i8x8_mb(sw, &mb_cfg).expect("write I_8x8 mb");
} else {
crate::encoder::macroblock::write_i8x8_mb_chroma(
sw,
&prev_flag,
&rem,
chroma_block.pred_mode.as_u8(),
cbp_luma,
cbp_chroma,
0,
&luma_4x4_levels,
&luma_4x4_nc,
crate::encoder::macroblock::ChromaWriteKind::Yuv422 {
chroma_dc_cb: &chroma_block.dc_cb,
chroma_dc_cr: &chroma_block.dc_cr,
chroma_ac_cb: &chroma_block.ac_cb,
chroma_ac_cr: &chroma_block.ac_cr,
cb_ac_nc: &chroma_nc_cb,
cr_ac_nc: &chroma_nc_cr,
},
)
.expect("write I_8x8 mb (4:2:2)");
}
let mut blk_has_nz = [false; 16];
for blk8 in 0..4usize {
for sub in 0..4usize {
blk_has_nz[blk8 * 4 + sub] = blk8_has_nz[blk8];
}
}
let luma_nonzero_4x4 = luma_nz_mask_from_blocks(&blk_has_nz);
let chroma_nz_mask = chroma_nz_mask_from_chroma_block(&chroma_block, n_chroma_blk);
MbTrial {
luma_cost: mb_luma_cost,
deblock: MbDeblockInfo {
is_intra: true,
qp_y: if cbp_luma > 0 || cbp_chroma > 0 {
qp_y
} else {
qp_y - mb_qp_delta
},
luma_nonzero_4x4,
chroma_nonzero_4x4: chroma_nz_mask,
transform_size_8x8_flag: true,
..Default::default()
},
}
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_intra4x4_444(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
) -> MbDeblockInfo {
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let lambda = lambda_ssd(qp_y, true);
let wq = self.cfg.scaling_matrix.intra_weights();
{
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = true;
s.is_i8x8 = false;
s.intra_4x4_pred_modes = [0u8; 16];
}
let mut chosen_mode = [0u8; 16];
let mut prev_flag = [false; 16];
let mut rem = [0u8; 16];
let mut luma_4x4_levels = [[0i32; 16]; 16];
let mut cb_4x4_levels = [[0i32; 16]; 16];
let mut cr_4x4_levels = [[0i32; 16]; 16];
let mut blk_luma_nz = [false; 16];
let mut blk_cb_nz = [false; 16];
let mut blk_cr_nz = [false; 16];
for blk in 0..16usize {
let (bx, by) = I4X4_XY[blk];
let avail = i4x4_avail(mb_x, mb_y, bx, by, width_mbs);
let predicted = predicted_intra4x4_mode(intra_grid, mb_x, mb_y, blk, width_mbs);
let mut best_mode = Intra4x4Mode::Dc;
let mut best_cost = u64::MAX;
let mut best_z_scan = [0i32; 16];
let mut best_recon = [0i32; 16];
for &mode in &Intra4x4Mode::ALL {
let Some(pred) = predict_4x4(mode, recon_y, width, mb_x, mb_y, bx, by, avail)
else {
continue;
};
let candidate = mode.as_u8();
let mut block_res = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
let s = frame.y[(mb_y * 16 + by + j) * width + (mb_x * 16 + bx + i)] as i32;
block_res[j * 4 + i] = s - pred[j * 4 + i];
}
}
let w_coeffs = forward_core_4x4(&block_res);
let z_raster = quantize_4x4_w(&w_coeffs, qp_y, true, &wq.w4);
let z_scan = zigzag_scan_4x4(&z_raster);
let r = inverse_transform_4x4(&z_raster, qp_y, &wq.w4, 8).expect("inverse 4x4");
let mut recon_block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
recon_block[j * 4 + i] = pred[j * 4 + i] + r[j * 4 + i];
}
}
let d = ssd_4x4(frame.y, width, mb_x * 16 + bx, mb_y * 16 + by, &recon_block);
let mut bits: u64 = 1;
if candidate != predicted {
bits += 3;
}
let mut trial_w = BitWriter::new();
if encode_residual_block_cavlc(
&mut trial_w,
CoeffTokenContext::Numeric(0),
16,
&z_scan,
)
.is_ok()
{
bits += trial_w.bits_emitted() as u64;
}
let cost = cost_combined(d, bits, lambda);
if cost < best_cost {
best_cost = cost;
best_mode = mode;
best_z_scan = z_scan;
best_recon = recon_block;
}
}
let chosen = best_mode.as_u8();
chosen_mode[blk] = chosen;
if chosen == predicted {
prev_flag[blk] = true;
} else {
prev_flag[blk] = false;
rem[blk] = if chosen < predicted {
chosen
} else {
chosen - 1
};
}
intra_grid.slot_mut(mb_x, mb_y).intra_4x4_pred_modes[blk] = chosen;
luma_4x4_levels[blk] = best_z_scan;
blk_luma_nz[blk] = best_z_scan.iter().any(|&v| v != 0);
for j in 0..4 {
for i in 0..4 {
recon_y[(mb_y * 16 + by + j) * width + mb_x * 16 + bx + i] =
best_recon[j * 4 + i].clamp(0, 255) as u8;
}
}
for plane_is_cr in [false, true] {
let (plane_src, plane_recon): (&[u8], &mut [u8]) = if plane_is_cr {
(frame.v, &mut *recon_v)
} else {
(frame.u, &mut *recon_u)
};
let pred = predict_4x4(best_mode, plane_recon, width, mb_x, mb_y, bx, by, avail)
.expect("chosen mode is available on the chroma plane (same geometry)");
let mut block_res = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
let s =
plane_src[(mb_y * 16 + by + j) * width + (mb_x * 16 + bx + i)] as i32;
block_res[j * 4 + i] = s - pred[j * 4 + i];
}
}
let w_coeffs = forward_core_4x4(&block_res);
let z_raster = quantize_4x4_w(&w_coeffs, qp_c, true, &wq.w4);
let z_scan = zigzag_scan_4x4(&z_raster);
let r = inverse_transform_4x4(&z_raster, qp_c, &wq.w4, 8).expect("inverse 4x4");
for j in 0..4 {
for i in 0..4 {
let v = (pred[j * 4 + i] + r[j * 4 + i]).clamp(0, 255);
plane_recon[(mb_y * 16 + by + j) * width + mb_x * 16 + bx + i] = v as u8;
}
}
let any_nz = z_scan.iter().any(|&v| v != 0);
if plane_is_cr {
cr_4x4_levels[blk] = z_scan;
blk_cr_nz[blk] = any_nz;
} else {
cb_4x4_levels[blk] = z_scan;
blk_cb_nz[blk] = any_nz;
}
}
}
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
if (0..4).any(|sub| {
let blk = blk8 * 4 + sub;
blk_luma_nz[blk] || blk_cb_nz[blk] || blk_cr_nz[blk]
}) {
cbp_luma |= 1u8 << blk8;
}
}
let pic_w_mbs = self.cfg.width / 16;
let mb_addr = (mb_y as u32) * pic_w_mbs + (mb_x as u32);
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
cur.cb_luma_total_coeff = [0u8; 16];
cur.cr_luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr, blk as u8, LumaNcKind::Ac, true, false);
luma_4x4_nc[blk] = nc;
own_totals[blk] = if (cbp_luma >> blk8) & 1 == 1 {
luma_4x4_levels[blk].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let mut cb_4x4_nc = [0i32; 16];
let mut cr_4x4_nc = [0i32; 16];
for plane_is_cr in [false, true] {
let levels = if plane_is_cr {
&cr_4x4_levels
} else {
&cb_4x4_levels
};
let mut own_plane = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
if plane_is_cr {
cur.cr_luma_total_coeff = own_plane;
} else {
cur.cb_luma_total_coeff = own_plane;
}
}
let nc = crate::macroblock_layer::derive_nc_plane_luma_like(
nc_grid,
mb_addr,
blk as u8,
plane_is_cr,
true,
false,
);
if plane_is_cr {
cr_4x4_nc[blk] = nc;
} else {
cb_4x4_nc[blk] = nc;
}
own_plane[blk] = if (cbp_luma >> blk8) & 1 == 1 {
levels[blk].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
let cur = &mut nc_grid.mbs[mb_addr as usize];
if plane_is_cr {
cur.cr_luma_total_coeff = own_plane;
} else {
cur.cb_luma_total_coeff = own_plane;
}
}
crate::encoder::macroblock::write_i4x4_mb_444(
sw,
self.cfg.transform_8x8,
&prev_flag,
&rem,
cbp_luma,
0,
&luma_4x4_levels,
&luma_4x4_nc,
&cb_4x4_levels,
&cb_4x4_nc,
&cr_4x4_levels,
&cr_4x4_nc,
)
.expect("write I_4x4 mb (4:4:4)");
MbDeblockInfo {
is_intra: true,
qp_y,
is_intra_4x4: true,
..Default::default()
}
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_intra4x4_mono(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
recon_y: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_qp_delta: i32,
) -> MbDeblockInfo {
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let lambda = lambda_ssd(qp_y, true);
let wq = self.cfg.scaling_matrix.intra_weights();
{
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = true;
s.is_i8x8 = false;
s.intra_4x4_pred_modes = [0u8; 16];
}
let mut prev_flag = [false; 16];
let mut rem = [0u8; 16];
let mut luma_4x4_levels = [[0i32; 16]; 16];
let mut blk_luma_nz = [false; 16];
for blk in 0..16usize {
let (bx, by) = I4X4_XY[blk];
let avail = i4x4_avail(mb_x, mb_y, bx, by, width_mbs);
let predicted = predicted_intra4x4_mode(intra_grid, mb_x, mb_y, blk, width_mbs);
let mut best_mode = Intra4x4Mode::Dc;
let mut best_cost = u64::MAX;
let mut best_z_scan = [0i32; 16];
let mut best_recon = [0i32; 16];
for &mode in &Intra4x4Mode::ALL {
let Some(pred) = predict_4x4(mode, recon_y, width, mb_x, mb_y, bx, by, avail)
else {
continue;
};
let candidate = mode.as_u8();
let mut block_res = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
let s = frame.y[(mb_y * 16 + by + j) * width + (mb_x * 16 + bx + i)] as i32;
block_res[j * 4 + i] = s - pred[j * 4 + i];
}
}
let w_coeffs = forward_core_4x4(&block_res);
let z_raster = quantize_4x4_w(&w_coeffs, qp_y, true, &wq.w4);
let z_scan = zigzag_scan_4x4(&z_raster);
let r = inverse_transform_4x4(&z_raster, qp_y, &wq.w4, 8).expect("inverse 4x4");
let mut recon_block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
recon_block[j * 4 + i] = pred[j * 4 + i] + r[j * 4 + i];
}
}
let d = ssd_4x4(frame.y, width, mb_x * 16 + bx, mb_y * 16 + by, &recon_block);
let mut bits: u64 = 1;
if candidate != predicted {
bits += 3;
}
let mut trial_w = BitWriter::new();
if encode_residual_block_cavlc(
&mut trial_w,
CoeffTokenContext::Numeric(0),
16,
&z_scan,
)
.is_ok()
{
bits += trial_w.bits_emitted() as u64;
}
let cost = cost_combined(d, bits, lambda);
if cost < best_cost {
best_cost = cost;
best_mode = mode;
best_z_scan = z_scan;
best_recon = recon_block;
}
}
let chosen = best_mode.as_u8();
if chosen == predicted {
prev_flag[blk] = true;
} else {
prev_flag[blk] = false;
rem[blk] = if chosen < predicted {
chosen
} else {
chosen - 1
};
}
intra_grid.slot_mut(mb_x, mb_y).intra_4x4_pred_modes[blk] = chosen;
luma_4x4_levels[blk] = best_z_scan;
blk_luma_nz[blk] = best_z_scan.iter().any(|&v| v != 0);
for j in 0..4 {
for i in 0..4 {
recon_y[(mb_y * 16 + by + j) * width + mb_x * 16 + bx + i] =
best_recon[j * 4 + i].clamp(0, 255) as u8;
}
}
}
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
if (0..4).any(|sub| blk_luma_nz[blk8 * 4 + sub]) {
cbp_luma |= 1u8 << blk8;
}
}
let pic_w_mbs = self.cfg.width / 16;
let mb_addr = (mb_y as u32) * pic_w_mbs + (mb_x as u32);
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr, blk as u8, LumaNcKind::Ac, true, false);
luma_4x4_nc[blk] = nc;
own_totals[blk] = if (cbp_luma >> blk8) & 1 == 1 {
luma_4x4_levels[blk].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
crate::encoder::macroblock::write_i4x4_mb_mono(
sw,
self.cfg.transform_8x8,
&prev_flag,
&rem,
cbp_luma,
mb_qp_delta,
&luma_4x4_levels,
&luma_4x4_nc,
)
.expect("write I_4x4 mb (4:0:0)");
let mut nz = [false; 16];
for (blk, &b) in blk_luma_nz.iter().enumerate() {
nz[blk] = b && (cbp_luma >> (blk / 4)) & 1 == 1;
}
MbDeblockInfo {
is_intra: true,
qp_y: if cbp_luma > 0 {
qp_y
} else {
qp_y - mb_qp_delta
},
luma_nonzero_4x4: luma_nz_mask_from_blocks(&nz),
is_intra_4x4: true,
..Default::default()
}
}
#[allow(clippy::too_many_arguments)]
fn encode_mb_intra8x8_444(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
) -> MbDeblockInfo {
let width = self.cfg.width as usize;
let lambda = lambda_ssd(qp_y, true);
let wq = self.cfg.scaling_matrix.intra_weights();
let sl8 = wq.w8;
{
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
s.is_i8x8 = true;
s.intra_8x8_pred_modes = [2u8; 4];
}
let mut prev_flag = [false; 4];
let mut rem = [0u8; 4];
let mut luma_4x4_levels = [[0i32; 16]; 16];
let mut cb_4x4_levels = [[0i32; 16]; 16];
let mut cr_4x4_levels = [[0i32; 16]; 16];
let mut blk8_luma_nz = [false; 4];
let mut blk8_cb_nz = [false; 4];
let mut blk8_cr_nz = [false; 4];
for blk8 in 0..4usize {
let bx = mb_x * 16 + (blk8 % 2) * 8;
let by = mb_y * 16 + (blk8 / 2) * 8;
let predicted = predicted_intra8x8_mode(intra_grid, mb_x, mb_y, blk8);
let raw = gather_8x8_from_recon(recon_y, width, mb_x, mb_y, blk8);
let filtered = filter_samples_8x8(&raw, 8);
let mut best_mode = Intra8x8Mode::Dc;
let mut best_cost = u64::MAX;
let mut best_z_raster = [0i32; 64];
let mut best_recon = [0i32; 64];
for mode_idx in 0..9u8 {
let mode = Intra8x8Mode::from_index(mode_idx).expect("mode index in range");
if !intra8x8_mode_available(mode, &raw.availability) {
continue;
}
let mut pred = [0i32; 64];
predict_8x8(mode, &filtered, 8, &mut pred);
let mut res = [0i32; 64];
for j in 0..8 {
for i in 0..8 {
let s = frame.y[(by + j) * width + bx + i] as i32;
res[j * 8 + i] = s - pred[j * 8 + i];
}
}
let w_coeffs = forward_core_8x8(&res);
let z_raster = quantize_8x8_w(&w_coeffs, qp_y, true, &sl8);
let r = inverse_transform_8x8(&z_raster, qp_y, &sl8, 8).expect("inverse 8x8");
let mut recon_blk = [0i32; 64];
let mut d: u64 = 0;
for j in 0..8 {
for i in 0..8 {
let v = (pred[j * 8 + i] + r[j * 8 + i]).clamp(0, 255);
recon_blk[j * 8 + i] = v;
let s = frame.y[(by + j) * width + bx + i] as i32;
let e = (s - v) as i64;
d += (e * e) as u64;
}
}
let mut bits: u64 = 1;
if mode_idx != predicted {
bits += 3;
}
let scan = scan_8x8_for_cavlc(&z_raster, sw.field_scan());
let subs = deinterleave_8x8_to_4x4(&scan);
let mut trial_w = BitWriter::new();
for sub in &subs {
if encode_residual_block_cavlc(
&mut trial_w,
CoeffTokenContext::Numeric(0),
16,
sub,
)
.is_err()
{
break;
}
}
bits += trial_w.bits_emitted() as u64;
let cost = cost_combined(d, bits, lambda);
if cost < best_cost {
best_cost = cost;
best_mode = mode;
best_z_raster = z_raster;
best_recon = recon_blk;
}
}
let chosen = best_mode.as_index();
if chosen == predicted {
prev_flag[blk8] = true;
} else {
prev_flag[blk8] = false;
rem[blk8] = if chosen < predicted {
chosen
} else {
chosen - 1
};
}
intra_grid.slot_mut(mb_x, mb_y).intra_8x8_pred_modes[blk8] = chosen;
let scan = scan_8x8_for_cavlc(&best_z_raster, sw.field_scan());
let subs = deinterleave_8x8_to_4x4(&scan);
for (i4x4, sub) in subs.iter().enumerate() {
luma_4x4_levels[blk8 * 4 + i4x4] = *sub;
}
blk8_luma_nz[blk8] = best_z_raster.iter().any(|&v| v != 0);
for j in 0..8 {
for i in 0..8 {
recon_y[(by + j) * width + bx + i] = best_recon[j * 8 + i] as u8;
}
}
for plane_is_cr in [false, true] {
let (plane_src, plane_recon): (&[u8], &mut [u8]) = if plane_is_cr {
(frame.v, &mut *recon_v)
} else {
(frame.u, &mut *recon_u)
};
let raw_c = gather_8x8_from_recon(plane_recon, width, mb_x, mb_y, blk8);
let filtered_c = filter_samples_8x8(&raw_c, 8);
let mut pred = [0i32; 64];
predict_8x8(best_mode, &filtered_c, 8, &mut pred);
let mut res = [0i32; 64];
for j in 0..8 {
for i in 0..8 {
let s = plane_src[(by + j) * width + bx + i] as i32;
res[j * 8 + i] = s - pred[j * 8 + i];
}
}
let w_coeffs = forward_core_8x8(&res);
let z_raster = quantize_8x8_w(&w_coeffs, qp_c, true, &sl8);
let r = inverse_transform_8x8(&z_raster, qp_c, &sl8, 8).expect("inverse 8x8");
for j in 0..8 {
for i in 0..8 {
let v = (pred[j * 8 + i] + r[j * 8 + i]).clamp(0, 255);
plane_recon[(by + j) * width + bx + i] = v as u8;
}
}
let scan = scan_8x8_for_cavlc(&z_raster, sw.field_scan());
let subs = deinterleave_8x8_to_4x4(&scan);
let any_nz = z_raster.iter().any(|&v| v != 0);
if plane_is_cr {
for (i4x4, sub) in subs.iter().enumerate() {
cr_4x4_levels[blk8 * 4 + i4x4] = *sub;
}
blk8_cr_nz[blk8] = any_nz;
} else {
for (i4x4, sub) in subs.iter().enumerate() {
cb_4x4_levels[blk8 * 4 + i4x4] = *sub;
}
blk8_cb_nz[blk8] = any_nz;
}
}
}
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
if blk8_luma_nz[blk8] || blk8_cb_nz[blk8] || blk8_cr_nz[blk8] {
cbp_luma |= 1u8 << blk8;
}
}
let pic_w_mbs = self.cfg.width / 16;
let mb_addr = (mb_y as u32) * pic_w_mbs + (mb_x as u32);
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
cur.cb_luma_total_coeff = [0u8; 16];
cur.cr_luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr, blk as u8, LumaNcKind::Ac, true, false);
luma_4x4_nc[blk] = nc;
own_totals[blk] = if (cbp_luma >> blk8) & 1 == 1 {
luma_4x4_levels[blk].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
nc_grid.mbs[mb_addr as usize].luma_total_coeff = own_totals;
let mut cb_4x4_nc = [0i32; 16];
let mut cr_4x4_nc = [0i32; 16];
for plane_is_cr in [false, true] {
let levels = if plane_is_cr {
&cr_4x4_levels
} else {
&cb_4x4_levels
};
let mut own_plane = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
if plane_is_cr {
cur.cr_luma_total_coeff = own_plane;
} else {
cur.cb_luma_total_coeff = own_plane;
}
}
let nc = crate::macroblock_layer::derive_nc_plane_luma_like(
nc_grid,
mb_addr,
blk as u8,
plane_is_cr,
true,
false,
);
if plane_is_cr {
cr_4x4_nc[blk] = nc;
} else {
cb_4x4_nc[blk] = nc;
}
own_plane[blk] = if (cbp_luma >> blk8) & 1 == 1 {
levels[blk].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
let cur = &mut nc_grid.mbs[mb_addr as usize];
if plane_is_cr {
cur.cr_luma_total_coeff = own_plane;
} else {
cur.cb_luma_total_coeff = own_plane;
}
}
crate::encoder::macroblock::write_i8x8_mb_444(
sw,
&prev_flag,
&rem,
cbp_luma,
0,
&luma_4x4_levels,
&luma_4x4_nc,
&cb_4x4_levels,
&cb_4x4_nc,
&cr_4x4_levels,
&cr_4x4_nc,
)
.expect("write I_8x8 mb (4:4:4)");
MbDeblockInfo {
is_intra: true,
qp_y,
transform_size_8x8_flag: true,
..Default::default()
}
}
}
pub struct EncodedIdr {
pub annex_b: Vec<u8>,
pub recon_y: Vec<u8>,
pub recon_u: Vec<u8>,
pub recon_v: Vec<u8>,
pub recon_width: u32,
pub recon_height: u32,
pub partition_mvs: Vec<FrameRefPartitionMv>,
pub i8x8_mb_count: u32,
pub i4x4_mb_count: u32,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct FrameRefPartitionMv {
pub mv_l0: (i16, i16),
pub ref_idx_l0: i8,
pub is_intra: bool,
}
pub struct EncodedP {
pub annex_b: Vec<u8>,
pub recon_y: Vec<u8>,
pub recon_u: Vec<u8>,
pub recon_v: Vec<u8>,
pub recon_width: u32,
pub recon_height: u32,
pub frame_num: u32,
pub pic_order_cnt_lsb: u32,
pub partition_mvs: Vec<FrameRefPartitionMv>,
pub i8x8_mb_count: u32,
pub pred_weight_table: Option<ExplicitPredWeightTableL0>,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct MvGridSlot {
available: bool,
is_intra: bool,
ref_idx_l0_8x8: [i32; 4],
mv_l0_8x8: [Mv; 4],
}
impl MvGridSlot {
fn ref_idx_l0(&self) -> i32 {
self.ref_idx_l0_8x8[0]
}
fn mv_l0(&self) -> Mv {
self.mv_l0_8x8[0]
}
}
impl Default for MvGridSlot {
fn default() -> Self {
Self {
available: false,
is_intra: false,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct MvGrid {
width_mbs: usize,
slots: Vec<MvGridSlot>,
}
impl MvGrid {
fn new(width_mbs: usize, height_mbs: usize) -> Self {
Self {
width_mbs,
slots: vec![MvGridSlot::default(); width_mbs * height_mbs],
}
}
fn slot(&self, mb_x: usize, mb_y: usize) -> &MvGridSlot {
&self.slots[mb_y * self.width_mbs + mb_x]
}
fn slot_mut(&mut self, mb_x: usize, mb_y: usize) -> &mut MvGridSlot {
&mut self.slots[mb_y * self.width_mbs + mb_x]
}
}
fn neighbour_mvs_16x16(
grid: &MvGrid,
mb_x: usize,
mb_y: usize,
) -> (NeighbourMv, NeighbourMv, NeighbourMv, NeighbourMv) {
let make = |maybe: Option<&MvGridSlot>, part_idx: usize| -> NeighbourMv {
match maybe {
Some(s) if s.available => {
if s.is_intra {
NeighbourMv::intra_but_mb_available()
} else {
NeighbourMv {
available: s.ref_idx_l0_8x8[part_idx] >= 0,
mb_available: true,
partition_available: true,
ref_idx: s.ref_idx_l0_8x8[part_idx],
mv: s.mv_l0_8x8[part_idx],
}
}
}
_ => NeighbourMv::UNAVAILABLE,
}
};
let a = if mb_x > 0 {
Some(grid.slot(mb_x - 1, mb_y))
} else {
None
};
let b = if mb_y > 0 {
Some(grid.slot(mb_x, mb_y - 1))
} else {
None
};
let c = if mb_y > 0 && mb_x + 1 < grid.width_mbs {
Some(grid.slot(mb_x + 1, mb_y - 1))
} else {
None
};
let d = if mb_x > 0 && mb_y > 0 {
Some(grid.slot(mb_x - 1, mb_y - 1))
} else {
None
};
(
make(a, 1), make(b, 2), make(c, 2), make(d, 3), )
}
pub(crate) fn mvp_for_16x16(grid: &MvGrid, mb_x: usize, mb_y: usize, ref_idx: i32) -> Mv {
let (a, b, c, d) = neighbour_mvs_16x16(grid, mb_x, mb_y);
let inputs = MvpredInputs::from_abc(a, b, c, ref_idx, Default::default());
crate::mv_deriv::derive_mvpred_with_d(&inputs, d)
}
pub(crate) fn p_skip_mv(grid: &MvGrid, mb_x: usize, mb_y: usize) -> (i32, Mv) {
let (a, b, c, d) = neighbour_mvs_16x16(grid, mb_x, mb_y);
crate::mv_deriv::derive_p_skip_mv_with_d(a, b, c, d)
}
#[allow(clippy::too_many_arguments)]
fn neighbour_8x8_partition_mv(
grid: &MvGrid,
mb_x: usize,
mb_y: usize,
bx_4x4: i32,
by_4x4: i32,
inflight_avail: [bool; 4],
mv_inflight: [Mv; 4],
ref_inflight: [i32; 4],
) -> NeighbourMv {
let (nmb_x, nmb_y, lx, ly) = if bx_4x4 < 0 && (0..4).contains(&by_4x4) {
if mb_x == 0 {
return NeighbourMv::UNAVAILABLE;
}
(mb_x - 1, mb_y, 3, by_4x4)
} else if (0..4).contains(&bx_4x4) && by_4x4 < 0 {
if mb_y == 0 {
return NeighbourMv::UNAVAILABLE;
}
(mb_x, mb_y - 1, bx_4x4, 3)
} else if bx_4x4 < 0 && by_4x4 < 0 {
if mb_x == 0 || mb_y == 0 {
return NeighbourMv::UNAVAILABLE;
}
(mb_x - 1, mb_y - 1, 3, 3)
} else if (0..4).contains(&bx_4x4) && (0..4).contains(&by_4x4) {
let part = 2 * (by_4x4 as usize / 2) + (bx_4x4 as usize / 2);
if inflight_avail[part] {
return NeighbourMv {
available: ref_inflight[part] >= 0,
mb_available: true,
partition_available: true,
ref_idx: ref_inflight[part],
mv: mv_inflight[part],
};
} else {
return NeighbourMv::partition_not_yet_decoded_same_mb();
}
} else if (4..).contains(&bx_4x4) && by_4x4 < 0 {
if mb_y == 0 || mb_x + 1 >= grid.width_mbs {
return NeighbourMv::UNAVAILABLE;
}
(mb_x + 1, mb_y - 1, bx_4x4 - 4, 3)
} else {
return NeighbourMv::UNAVAILABLE;
};
let s = grid.slot(nmb_x, nmb_y);
if !s.available {
return NeighbourMv::UNAVAILABLE;
}
if s.is_intra {
return NeighbourMv::intra_but_mb_available();
}
let part = 2 * (ly as usize / 2) + (lx as usize / 2);
NeighbourMv {
available: s.ref_idx_l0_8x8[part] >= 0,
mb_available: true,
partition_available: true,
ref_idx: s.ref_idx_l0_8x8[part],
mv: s.mv_l0_8x8[part],
}
}
#[allow(clippy::too_many_arguments)]
fn mvp_for_p_8x8_partition(
grid: &MvGrid,
mb_x: usize,
mb_y: usize,
sub_idx: usize,
ref_idx: i32,
inflight_avail: [bool; 4],
mv_inflight: [Mv; 4],
ref_inflight: [i32; 4],
) -> Mv {
let sub_x = (sub_idx % 2) as i32;
let sub_y = (sub_idx / 2) as i32;
let bx = sub_x * 2;
let by = sub_y * 2;
let a = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx - 1,
by,
inflight_avail,
mv_inflight,
ref_inflight,
);
let b = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx,
by - 1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let c = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx + 2,
by - 1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let d = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx - 1,
by - 1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let inputs = MvpredInputs::from_abc(a, b, c, ref_idx, Default::default());
crate::mv_deriv::derive_mvpred_with_d(&inputs, d)
}
fn build_inter_pred_luma(
ref_y: &[u8],
ref_w: u32,
ref_h: u32,
mb_x: usize,
mb_y: usize,
mv: Mv,
) -> [i32; 256] {
let stride = ref_w as usize;
let h = ref_h as usize;
let ref_i32: Vec<i32> = ref_y.iter().take(stride * h).map(|&v| v as i32).collect();
let mv_int_x = mv.x >> 2;
let mv_int_y = mv.y >> 2;
let x_frac = (mv.x & 3) as u8;
let y_frac = (mv.y & 3) as u8;
let int_x = (mb_x * 16) as i32 + mv_int_x;
let int_y = (mb_y * 16) as i32 + mv_int_y;
let mut dst = [0i32; 256];
interpolate_luma(
&ref_i32, stride, stride, h, int_x, int_y, x_frac, y_frac, 16, 16, 8, &mut dst, 16,
)
.expect("luma interpolation");
dst
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_inter_pred_luma_8x8(
ref_y: &[u8],
ref_w: u32,
ref_h: u32,
mb_x: usize,
mb_y: usize,
sub_x: usize,
sub_y: usize,
mv: Mv,
) -> [i32; 64] {
let stride = ref_w as usize;
let h = ref_h as usize;
let ref_i32: Vec<i32> = ref_y.iter().take(stride * h).map(|&v| v as i32).collect();
let mv_int_x = mv.x >> 2;
let mv_int_y = mv.y >> 2;
let x_frac = (mv.x & 3) as u8;
let y_frac = (mv.y & 3) as u8;
let int_x = (mb_x * 16 + sub_x * 8) as i32 + mv_int_x;
let int_y = (mb_y * 16 + sub_y * 8) as i32 + mv_int_y;
let mut dst = [0i32; 64];
interpolate_luma(
&ref_i32, stride, stride, h, int_x, int_y, x_frac, y_frac, 8, 8, 8, &mut dst, 8,
)
.expect("luma interpolation 8x8");
dst
}
#[allow(clippy::too_many_arguments)]
pub(super) fn build_chroma_rect_pred_fmt(
plane: &[u8],
plane_w: u32,
plane_h: u32,
cx0: usize,
cy0: usize,
rw: usize,
rh: usize,
mv: Mv,
fmt: u32,
) -> Vec<i32> {
if fmt == 0 || rw == 0 || rh == 0 {
return Vec::new();
}
let stride = plane_w as usize;
let h = plane_h as usize;
let ref_i32: Vec<i32> = plane.iter().take(stride * h).map(|&v| v as i32).collect();
let mut dst = vec![0i32; rw * rh];
if fmt == 3 {
let int_x = cx0 as i32 + (mv.x >> 2);
let int_y = cy0 as i32 + (mv.y >> 2);
let x_frac = (mv.x & 3) as u8;
let y_frac = (mv.y & 3) as u8;
interpolate_luma(
&ref_i32, stride, stride, h, int_x, int_y, x_frac, y_frac, rw as u32, rh as u32, 8,
&mut dst, rw,
)
.expect("4:4:4 luma-like chroma interpolation");
} else {
let mv_cx = mv.x;
let mv_cy = if fmt == 2 { mv.y * 2 } else { mv.y };
let int_x = cx0 as i32 + (mv_cx >> 3);
let int_y = cy0 as i32 + (mv_cy >> 3);
let x_frac = (mv_cx & 7) as u8;
let y_frac = (mv_cy & 7) as u8;
interpolate_chroma(
&ref_i32, stride, stride, h, int_x, int_y, x_frac, y_frac, rw as u32, rh as u32, 8,
&mut dst, rw,
)
.expect("chroma rect interpolation");
}
dst
}
#[allow(clippy::too_many_arguments)]
pub(super) fn build_b_partition_chroma_rect_fmt(
pred: BPartPred,
plane_l0: &[u8],
plane_l1: &[u8],
plane_w: u32,
plane_h: u32,
cx0: usize,
cy0: usize,
rw: usize,
rh: usize,
mv_l0: Mv,
mv_l1: Mv,
fmt: u32,
) -> Vec<i32> {
match pred {
BPartPred::L0 => {
build_chroma_rect_pred_fmt(plane_l0, plane_w, plane_h, cx0, cy0, rw, rh, mv_l0, fmt)
}
BPartPred::L1 => {
build_chroma_rect_pred_fmt(plane_l1, plane_w, plane_h, cx0, cy0, rw, rh, mv_l1, fmt)
}
BPartPred::Bi => {
let a = build_chroma_rect_pred_fmt(
plane_l0, plane_w, plane_h, cx0, cy0, rw, rh, mv_l0, fmt,
);
let b = build_chroma_rect_pred_fmt(
plane_l1, plane_w, plane_h, cx0, cy0, rw, rh, mv_l1, fmt,
);
a.iter().zip(&b).map(|(&x, &y)| (x + y + 1) >> 1).collect()
}
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_inter_pred_chroma_4x4(
ref_c: &[u8],
ref_cw: u32,
ref_ch: u32,
mb_x: usize,
mb_y: usize,
sub_x: usize,
sub_y: usize,
mv: Mv,
) -> [i32; 16] {
let stride = ref_cw as usize;
let h = ref_ch as usize;
let ref_i32: Vec<i32> = ref_c.iter().take(stride * h).map(|&v| v as i32).collect();
let mv_cx = mv.x;
let mv_cy = mv.y;
let int_x = (mb_x * 8 + sub_x * 4) as i32 + (mv_cx >> 3);
let int_y = (mb_y * 8 + sub_y * 4) as i32 + (mv_cy >> 3);
let x_frac = (mv_cx & 7) as u8;
let y_frac = (mv_cy & 7) as u8;
let mut dst = [0i32; 16];
interpolate_chroma(
&ref_i32, stride, stride, h, int_x, int_y, x_frac, y_frac, 4, 4, 8, &mut dst, 4,
)
.expect("chroma interpolation 4x4");
dst
}
fn build_inter_pred_chroma(
ref_c: &[u8],
ref_cw: u32,
ref_ch: u32,
mb_x: usize,
mb_y: usize,
mv: Mv,
) -> [i32; 64] {
let stride = ref_cw as usize;
let h = ref_ch as usize;
let ref_i32: Vec<i32> = ref_c.iter().take(stride * h).map(|&v| v as i32).collect();
let mv_cx = mv.x;
let mv_cy = mv.y;
let int_x = (mb_x * 8) as i32 + (mv_cx >> 3);
let int_y = (mb_y * 8) as i32 + (mv_cy >> 3);
let x_frac = (mv_cx & 7) as u8;
let y_frac = (mv_cy & 7) as u8;
let mut dst = [0i32; 64];
interpolate_chroma(
&ref_i32, stride, stride, h, int_x, int_y, x_frac, y_frac, 8, 8, 8, &mut dst, 8,
)
.expect("chroma interpolation");
dst
}
fn encode_chroma_residual_inter(
src_plane: &[u8],
chroma_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 64],
qp_c: i32,
w4: &[i32; 16],
) -> ([i32; 4], [[i32; 16]; 4], [i32; 64]) {
encode_chroma_residual(src_plane, chroma_stride, mb_x, mb_y, pred, qp_c, w4)
}
fn mb_ssd_3planes_444(
frame: &YuvFrame<'_>,
recon_y: &[u8],
recon_u: &[u8],
recon_v: &[u8],
width: usize,
mb_x: usize,
mb_y: usize,
) -> u64 {
let mut d: u64 = 0;
for (src, rec) in [(frame.y, recon_y), (frame.u, recon_u), (frame.v, recon_v)] {
for j in 0..16usize {
let row = (mb_y * 16 + j) * width + mb_x * 16;
for i in 0..16usize {
let e = src[row + i] as i64 - rec[row + i] as i64;
d += (e * e) as u64;
}
}
}
d
}
pub(crate) fn derive_chroma_ac_nc_and_commit_totals(
nc_grid: &mut CavlcNcGrid,
mb_addr: u32,
is_intra: bool,
cbp_chroma: u8,
ac_cb: &[[i32; 16]],
ac_cr: &[[i32; 16]],
chroma_array_type: u32,
) -> ([i32; 8], [i32; 8]) {
let n_blk = if chroma_array_type == 2 { 8usize } else { 4 };
let mut nc_cb = [0i32; 8];
let mut nc_cr = [0i32; 8];
for is_cr in [false, true] {
let (ncs, acs): (&mut [i32; 8], &[[i32; 16]]) = if is_cr {
(&mut nc_cr, ac_cr)
} else {
(&mut nc_cb, ac_cb)
};
let mut own = [0u8; 8];
for blk in 0..n_blk {
{
let slot = &mut nc_grid.mbs[mb_addr as usize];
if is_cr {
slot.cr_total_coeff = own;
} else {
slot.cb_total_coeff = own;
}
}
ncs[blk] = crate::macroblock_layer::derive_nc_chroma_ac(
nc_grid,
mb_addr,
blk as u8,
is_cr,
chroma_array_type,
is_intra,
false,
);
own[blk] = if cbp_chroma == 2 {
acs[blk][..15].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
let slot = &mut nc_grid.mbs[mb_addr as usize];
if is_cr {
slot.cr_total_coeff = own;
} else {
slot.cb_total_coeff = own;
}
}
(nc_cb, nc_cr)
}
pub(crate) fn derive_plane_nc_444_i16x16(
nc_grid: &mut CavlcNcGrid,
mb_addr: u32,
is_cr: bool,
cbp_luma: u8,
ac_scan: &[[i32; 16]; 16],
) -> (i32, [i32; 16]) {
let dc_nc =
crate::macroblock_layer::derive_nc_plane_luma_like(nc_grid, mb_addr, 0, is_cr, true, false);
let mut ac_nc = [0i32; 16];
let mut own = [0u8; 16];
for blk in 0..16usize {
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
if is_cr {
cur.cr_luma_total_coeff = own;
} else {
cur.cb_luma_total_coeff = own;
}
}
ac_nc[blk] = crate::macroblock_layer::derive_nc_plane_luma_like(
nc_grid, mb_addr, blk as u8, is_cr, true, false,
);
own[blk] = if cbp_luma == 15 {
ac_scan[blk][..15].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
let cur = &mut nc_grid.mbs[mb_addr as usize];
if is_cr {
cur.cr_luma_total_coeff = own;
} else {
cur.cb_luma_total_coeff = own;
}
(dc_nc, ac_nc)
}
pub(crate) fn derive_plane_nc_444_inter(
nc_grid: &mut CavlcNcGrid,
mb_addr: u32,
is_cr: bool,
cbp_luma: u8,
levels_scan: &[[i32; 16]; 16],
) -> [i32; 16] {
let mut ac_nc = [0i32; 16];
let mut own = [0u8; 16];
for blk in 0..16usize {
{
let cur = &mut nc_grid.mbs[mb_addr as usize];
if is_cr {
cur.cr_luma_total_coeff = own;
} else {
cur.cb_luma_total_coeff = own;
}
}
ac_nc[blk] = crate::macroblock_layer::derive_nc_plane_luma_like(
nc_grid, mb_addr, blk as u8, is_cr, false, false,
);
own[blk] = if (cbp_luma >> (blk / 4)) & 1 == 1 {
levels_scan[blk].iter().filter(|&&v| v != 0).count() as u8
} else {
0
};
}
let cur = &mut nc_grid.mbs[mb_addr as usize];
if is_cr {
cur.cr_luma_total_coeff = own;
} else {
cur.cb_luma_total_coeff = own;
}
ac_nc
}
struct ChromaBlock {
pred_mode: IntraChromaMode,
cbp_chroma: u8,
dc_cb: [i32; 8],
dc_cr: [i32; 8],
ac_cb: [[i32; 16]; 8],
ac_cr: [[i32; 16]; 8],
recon_cb: Vec<i32>,
recon_cr: Vec<i32>,
pred_cb: Vec<i32>,
pred_cr: Vec<i32>,
}
impl ChromaBlock {
fn cb_block_count(chroma_array_type: u32) -> usize {
match chroma_array_type {
1 => 4,
2 => 8,
_ => 0,
}
}
}
#[allow(clippy::too_many_arguments)]
fn encode_chroma_block(
chroma_array_type: u32,
frame_u: &[u8],
frame_v: &[u8],
recon_u: &[u8],
recon_v: &[u8],
chroma_width: usize,
chroma_height: usize,
mb_x: usize,
mb_y: usize,
qp_c: i32,
w4: &[i32; 16],
nb: Option<(bool, bool, bool)>,
) -> ChromaBlock {
use crate::encoder::intra_pred::{predict_chroma_8x16, sad_8x16};
let (h_chroma, n_blk) = match chroma_array_type {
1 => (8usize, 4usize),
2 => (16usize, 8usize),
_ => panic!("encode_chroma_block: unsupported chroma_array_type {chroma_array_type}"),
};
let mut best_mode = IntraChromaMode::Dc;
let mut best_pred_u_64 = [0i32; 64];
let mut best_pred_v_64 = [0i32; 64];
let mut best_pred_u_128 = [0i32; 128];
let mut best_pred_v_128 = [0i32; 128];
let mut best_sad = u64::MAX;
for &mode in &[
IntraChromaMode::Dc,
IntraChromaMode::Horizontal,
IntraChromaMode::Vertical,
IntraChromaMode::Plane,
] {
if chroma_array_type == 1 {
let (Some(pu), Some(pv)) = (
predict_chroma_8x8_avail(
mode,
recon_u,
chroma_width,
chroma_height,
mb_x,
mb_y,
nb,
),
predict_chroma_8x8_avail(
mode,
recon_v,
chroma_width,
chroma_height,
mb_x,
mb_y,
nb,
),
) else {
continue;
};
let s = sad_8x8(frame_u, chroma_width, mb_x * 8, mb_y * 8, &pu)
+ sad_8x8(frame_v, chroma_width, mb_x * 8, mb_y * 8, &pv);
if s < best_sad {
best_sad = s;
best_mode = mode;
best_pred_u_64 = pu;
best_pred_v_64 = pv;
}
} else {
let (Some(pu), Some(pv)) = (
predict_chroma_8x16(mode, recon_u, chroma_width, chroma_height, mb_x, mb_y),
predict_chroma_8x16(mode, recon_v, chroma_width, chroma_height, mb_x, mb_y),
) else {
continue;
};
let s = sad_8x16(frame_u, chroma_width, mb_x * 8, mb_y * 16, &pu)
+ sad_8x16(frame_v, chroma_width, mb_x * 8, mb_y * 16, &pv);
if s < best_sad {
best_sad = s;
best_mode = mode;
best_pred_u_128 = pu;
best_pred_v_128 = pv;
}
}
}
let mut dc_cb = [0i32; 8];
let mut dc_cr = [0i32; 8];
let mut ac_cb = [[0i32; 16]; 8];
let mut ac_cr = [[0i32; 16]; 8];
let mut recon_cb = vec![0i32; 8 * h_chroma];
let mut recon_cr = vec![0i32; 8 * h_chroma];
if chroma_array_type == 1 {
let (u_dc, u_ac, u_recon) =
encode_chroma_residual(frame_u, chroma_width, mb_x, mb_y, &best_pred_u_64, qp_c, w4);
let (v_dc, v_ac, v_recon) =
encode_chroma_residual(frame_v, chroma_width, mb_x, mb_y, &best_pred_v_64, qp_c, w4);
dc_cb[..4].copy_from_slice(&u_dc);
dc_cr[..4].copy_from_slice(&v_dc);
ac_cb[..4].copy_from_slice(&u_ac);
ac_cr[..4].copy_from_slice(&v_ac);
recon_cb[..64].copy_from_slice(&u_recon);
recon_cr[..64].copy_from_slice(&v_recon);
} else {
let (u_dc, u_ac, u_recon) = encode_chroma_residual_422(
frame_u,
chroma_width,
mb_x,
mb_y,
&best_pred_u_128,
qp_c,
w4,
);
let (v_dc, v_ac, v_recon) = encode_chroma_residual_422(
frame_v,
chroma_width,
mb_x,
mb_y,
&best_pred_v_128,
qp_c,
w4,
);
dc_cb = u_dc;
dc_cr = v_dc;
ac_cb = u_ac;
ac_cr = v_ac;
recon_cb[..128].copy_from_slice(&u_recon);
recon_cr[..128].copy_from_slice(&v_recon);
}
let any_dc_nz =
dc_cb[..n_blk].iter().any(|&v| v != 0) || dc_cr[..n_blk].iter().any(|&v| v != 0);
let any_ac_nz = ac_cb[..n_blk].iter().any(|blk| blk.iter().any(|&v| v != 0))
|| ac_cr[..n_blk].iter().any(|blk| blk.iter().any(|&v| v != 0));
let cbp_chroma: u8 = if any_ac_nz {
2
} else if any_dc_nz {
1
} else {
0
};
let pred_cb: Vec<i32> = if chroma_array_type == 1 {
best_pred_u_64.to_vec()
} else {
best_pred_u_128.to_vec()
};
let pred_cr: Vec<i32> = if chroma_array_type == 1 {
best_pred_v_64.to_vec()
} else {
best_pred_v_128.to_vec()
};
ChromaBlock {
pred_mode: best_mode,
cbp_chroma,
dc_cb,
dc_cr,
ac_cb,
ac_cr,
recon_cb,
recon_cr,
pred_cb,
pred_cr,
}
}
fn chroma_nz_mask_from_chroma_block(block: &ChromaBlock, n_blk: usize) -> u16 {
let mut mask = 0u16;
if block.cbp_chroma != 2 {
return mask;
}
for k in 0..n_blk {
if block.ac_cb[k].iter().any(|&v| v != 0) {
mask |= 1u16 << k;
}
if block.ac_cr[k].iter().any(|&v| v != 0) {
mask |= 1u16 << (8 + k);
}
}
mask
}
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
fn encode_intra16x16_chroma_plane_luma_like(
mode: I16x16Mode,
src_plane: &[u8],
recon_chroma: &mut [u8],
chroma_width: usize,
chroma_height: usize,
mb_x: usize,
mb_y: usize,
qp_c: i32,
w4: &[i32; 16],
) -> ([i32; 16], [[i32; 16]; 16], [[i32; 16]; 16], bool) {
let pred = predict_16x16(mode, recon_chroma, chroma_width, chroma_height, mb_x, mb_y)
.unwrap_or([128i32; 256]);
let mut residual = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
let s = src_plane[(mb_y * 16 + j) * chroma_width + (mb_x * 16 + i)] as i32;
residual[j * 16 + i] = s - pred[j * 16 + i];
}
}
let mut coeffs_4x4 = [[0i32; 16]; 16];
for (blk, slot) in coeffs_4x4.iter_mut().enumerate() {
let bx = blk % 4;
let by = blk / 4;
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by * 4 + j) * 16 + bx * 4 + i];
}
}
*slot = forward_core_4x4(&block);
}
let mut dc_coeffs = [0i32; 16];
for (blk, c) in coeffs_4x4.iter().enumerate() {
let bx = blk % 4;
let by = blk / 4;
dc_coeffs[by * 4 + bx] = c[0];
}
let dc_t = forward_hadamard_4x4(&dc_coeffs);
let dc_levels = quantize_luma_dc_w(&dc_t, qp_c, true, w4[0]);
let mut ac_quant_raster = [[0i32; 16]; 16];
let mut ac_scan = [[0i32; 16]; 16];
let mut any_ac_nz = false;
for blkz in 0..16usize {
let (bx, by) = LUMA_4X4_BLK[blkz];
let raster = by * 4 + bx;
let z = quantize_4x4_ac_w(&coeffs_4x4[raster], qp_c, true, w4);
ac_quant_raster[raster] = z;
ac_scan[blkz] = zigzag_scan_4x4_ac(&z);
if ac_scan[blkz].iter().any(|&v| v != 0) {
any_ac_nz = true;
}
}
let inv_dc =
inverse_hadamard_luma_dc_16x16(&dc_levels, qp_c, w4, 8).expect("4:4:4 chroma DC inverse");
#[allow(clippy::needless_range_loop)] for blk in 0..16usize {
let bx = blk % 4;
let by = blk / 4;
let mut coeffs_in = if any_ac_nz {
ac_quant_raster[blk]
} else {
[0i32; 16]
};
coeffs_in[0] = inv_dc[by * 4 + bx];
let r = inverse_transform_4x4_dc_preserved(&coeffs_in, qp_c, w4, 8)
.expect("4:4:4 chroma 4x4 inverse");
for j in 0..4 {
for i in 0..4 {
let px = mb_x * 16 + bx * 4 + i;
let py = mb_y * 16 + by * 4 + j;
let p = pred[(by * 4 + j) * 16 + bx * 4 + i];
let raw = p + r[j * 4 + i];
recon_chroma[py * chroma_width + px] = raw.clamp(0, 255) as u8;
}
}
}
(dc_levels, ac_scan, ac_quant_raster, any_ac_nz)
}
fn install_chroma_recon(
chroma_array_type: u32,
chroma_width: usize,
mb_x: usize,
mb_y: usize,
block: &ChromaBlock,
recon_u: &mut [u8],
recon_v: &mut [u8],
) {
let (h_chroma, base_y) = match chroma_array_type {
1 => (8usize, mb_y * 8),
2 => (16usize, mb_y * 16),
_ => return,
};
for j in 0..h_chroma {
for i in 0..8usize {
let cx = mb_x * 8 + i;
let cy = base_y + j;
let pred_u = block.pred_cb[j * 8 + i];
let pred_v = block.pred_cr[j * 8 + i];
let res_u = if block.cbp_chroma == 0 {
0
} else {
block.recon_cb[j * 8 + i]
};
let res_v = if block.cbp_chroma == 0 {
0
} else {
block.recon_cr[j * 8 + i]
};
recon_u[cy * chroma_width + cx] = (pred_u + res_u).clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = (pred_v + res_v).clamp(0, 255) as u8;
}
}
}
fn encode_chroma_residual_422(
src_plane: &[u8],
chroma_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 128],
qp_c: i32,
w4: &[i32; 16],
) -> ([i32; 8], [[i32; 16]; 8], [i32; 128]) {
use crate::encoder::transform::{
forward_hadamard_4x2, quantize_4x4_ac_w, quantize_chroma_dc_422_w, scan_chroma_dc_422,
};
use crate::transform::inverse_hadamard_chroma_dc_422;
let mut residual = [0i32; 128];
for j in 0..16usize {
for i in 0..8usize {
let s = src_plane[(mb_y * 16 + j) * chroma_stride + mb_x * 8 + i] as i32;
residual[j * 8 + i] = s - pred[j * 8 + i];
}
}
let blocks_xy: [(usize, usize); 8] = [
(0, 0),
(4, 0),
(0, 4),
(4, 4),
(0, 8),
(4, 8),
(0, 12),
(4, 12),
];
let mut sub_coeffs = [[0i32; 16]; 8];
for (k, &(bx, by)) in blocks_xy.iter().enumerate() {
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by + j) * 8 + bx + i];
}
}
sub_coeffs[k] = forward_core_4x4(&block);
}
let dc_in: [i32; 8] = [
sub_coeffs[0][0], sub_coeffs[1][0], sub_coeffs[2][0], sub_coeffs[3][0], sub_coeffs[4][0], sub_coeffs[5][0], sub_coeffs[6][0], sub_coeffs[7][0], ];
let dc_t = forward_hadamard_4x2(&dc_in);
let dc_q = quantize_chroma_dc_422_w(&dc_t, qp_c, true, w4[0]);
let dc_levels = scan_chroma_dc_422(&dc_q);
let mut ac_levels = [[0i32; 16]; 8];
for (k, coeffs) in sub_coeffs.iter().enumerate() {
let z = quantize_4x4_ac_w(coeffs, qp_c, true, w4);
ac_levels[k] = zigzag_scan_4x4_ac(&z);
}
let dc_back =
inverse_hadamard_chroma_dc_422(&dc_levels, qp_c, w4, 8).expect("chroma DC 422 inverse");
let dc_block_index_to_dc_back: [usize; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
let mut recon_residual = [0i32; 128];
for (k, &(bx, by)) in blocks_xy.iter().enumerate() {
let mut coeffs = inverse_scan_4x4_zigzag_ac(&ac_levels[k]);
coeffs[0] = dc_back[dc_block_index_to_dc_back[k]];
let r = inverse_transform_4x4_dc_preserved(&coeffs, qp_c, w4, 8)
.expect("chroma 4x4 422 inverse");
for j in 0..4 {
for i in 0..4 {
recon_residual[(by + j) * 8 + bx + i] = r[j * 4 + i];
}
}
}
(dc_levels, ac_levels, recon_residual)
}
fn encode_chroma_residual(
src_plane: &[u8],
chroma_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 64],
qp_c: i32,
w4: &[i32; 16],
) -> ([i32; 4], [[i32; 16]; 4], [i32; 64]) {
let mut residual = [0i32; 64];
for j in 0..8usize {
for i in 0..8usize {
let s = src_plane[(mb_y * 8 + j) * chroma_stride + mb_x * 8 + i] as i32;
residual[j * 8 + i] = s - pred[j * 8 + i];
}
}
let blocks_xy = [(0usize, 0usize), (4, 0), (0, 4), (4, 4)];
let mut sub_coeffs = [[0i32; 16]; 4]; for (k, &(bx, by)) in blocks_xy.iter().enumerate() {
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by + j) * 8 + bx + i];
}
}
sub_coeffs[k] = forward_core_4x4(&block);
}
let dc_in: [i32; 4] = [
sub_coeffs[0][0],
sub_coeffs[1][0],
sub_coeffs[2][0],
sub_coeffs[3][0],
];
let dc_t = forward_hadamard_2x2(&dc_in);
let dc_levels = quantize_chroma_dc_w(&dc_t, qp_c, true, w4[0]);
let mut ac_levels = [[0i32; 16]; 4];
for (k, coeffs) in sub_coeffs.iter().enumerate() {
let z = quantize_4x4_ac_w(coeffs, qp_c, true, w4);
ac_levels[k] = zigzag_scan_4x4_ac(&z);
}
let dc_back =
inverse_hadamard_chroma_dc_420(&dc_levels, qp_c, w4, 8).expect("chroma DC inverse");
let mut recon_residual = [0i32; 64];
for (k, &(bx, by)) in blocks_xy.iter().enumerate() {
let mut coeffs = inverse_scan_4x4_zigzag_ac(&ac_levels[k]);
coeffs[0] = dc_back[k];
let r =
inverse_transform_4x4_dc_preserved(&coeffs, qp_c, w4, 8).expect("chroma 4x4 inverse");
for j in 0..4 {
for i in 0..4 {
recon_residual[(by + j) * 8 + bx + i] = r[j * 4 + i];
}
}
}
(dc_levels, ac_levels, recon_residual)
}
#[allow(clippy::too_many_arguments)]
fn trial_intra16x16_cost(
frame: &YuvFrame<'_>,
pred: &[i32; 256],
mb_x: usize,
mb_y: usize,
qp_y: i32,
width: usize,
lambda: f64,
w4: &[i32; 16],
) -> u64 {
let mut residual = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
let s = frame.y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
residual[j * 16 + i] = s - pred[j * 16 + i];
}
}
let mut coeffs_4x4 = [[0i32; 16]; 16];
for (blk, slot) in coeffs_4x4.iter_mut().enumerate() {
let bx = blk % 4;
let by = blk / 4;
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by * 4 + j) * 16 + bx * 4 + i];
}
}
*slot = forward_core_4x4(&block);
}
let mut dc_coeffs = [0i32; 16];
for (blk, c) in coeffs_4x4.iter().enumerate() {
let bx = blk % 4;
let by = blk / 4;
dc_coeffs[by * 4 + bx] = c[0];
}
let dc_t = forward_hadamard_4x4(&dc_coeffs);
let dc_levels = quantize_luma_dc_w(&dc_t, qp_y, true, w4[0]);
let mut ac_quant_raster = [[0i32; 16]; 16];
let mut ac_scan = [[0i32; 16]; 16];
let mut any_ac_nz = false;
for blkz in 0..16usize {
let (bx, by) = LUMA_4X4_BLK[blkz];
let raster = by * 4 + bx;
let z = quantize_4x4_ac_w(&coeffs_4x4[raster], qp_y, true, w4);
ac_quant_raster[raster] = z;
ac_scan[blkz] = zigzag_scan_4x4_ac(&z);
if ac_scan[blkz].iter().any(|&v| v != 0) {
any_ac_nz = true;
}
}
let inv_dc = inverse_hadamard_luma_dc_16x16(&dc_levels, qp_y, w4, 8).expect("luma DC inverse");
let mut recon = [0i32; 256];
#[allow(clippy::needless_range_loop)]
for blk in 0..16usize {
let bx = blk % 4;
let by = blk / 4;
let mut coeffs_in = if any_ac_nz {
ac_quant_raster[blk]
} else {
[0i32; 16]
};
coeffs_in[0] = inv_dc[by * 4 + bx];
let r = inverse_transform_4x4_dc_preserved(&coeffs_in, qp_y, w4, 8).expect("inv 4x4");
for j in 0..4 {
for i in 0..4 {
let p = pred[(by * 4 + j) * 16 + bx * 4 + i];
recon[(by * 4 + j) * 16 + bx * 4 + i] = p + r[j * 4 + i];
}
}
}
let d = ssd_16x16(frame.y, width, mb_x * 16, mb_y * 16, &recon);
let mut bits: u64 = 0;
bits += 7;
bits += 1;
let dc_scan = zigzag_scan_4x4(&dc_levels);
let mut trial_w = BitWriter::new();
if encode_residual_block_cavlc(&mut trial_w, CoeffTokenContext::Numeric(0), 16, &dc_scan)
.is_ok()
{
bits += trial_w.bits_emitted() as u64;
}
if any_ac_nz {
for scan in ac_scan.iter() {
let mut tw = BitWriter::new();
if encode_residual_block_cavlc(&mut tw, CoeffTokenContext::Numeric(0), 15, &scan[..15])
.is_ok()
{
bits += tw.bits_emitted() as u64;
}
}
}
cost_combined(d, bits, lambda)
}
impl Encoder {
#[allow(clippy::too_many_arguments)]
pub fn encode_p(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
) -> EncodedP {
self.encode_p_with_qp(frame, prev, frame_num, pic_order_cnt_lsb, self.cfg.qp)
}
pub fn encode_p_with_qp(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
frame_qp: i32,
) -> EncodedP {
self.encode_p_impl(frame, prev, frame_num, pic_order_cnt_lsb, frame_qp, None)
}
pub fn encode_p_rate_adaptive(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
frame_qp: i32,
frame_budget_bits: u64,
) -> EncodedP {
assert_eq!(
self.cfg.chroma_format_idc, 1,
"encode_p_rate_adaptive requires 4:2:0"
);
self.encode_p_impl(
frame,
prev,
frame_num,
pic_order_cnt_lsb,
frame_qp,
Some(frame_budget_bits),
)
}
fn encode_p_impl(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
frame_qp: i32,
row_budget_bits: Option<u64>,
) -> EncodedP {
assert!(
(0..=51).contains(&frame_qp),
"frame_qp {frame_qp} out of 0..=51"
);
assert_eq!(frame.width, self.cfg.width);
assert_eq!(frame.height, self.cfg.height);
assert_eq!(prev.width, self.cfg.width);
assert_eq!(prev.height, self.cfg.height);
let width_mbs = self.cfg.width / 16;
let height_mbs = self.cfg.height / 16;
let chroma_width = match self.cfg.chroma_format_idc {
0 => 0usize,
3 => self.cfg.width as usize,
_ => (self.cfg.width / 2) as usize,
};
let chroma_height = match self.cfg.chroma_format_idc {
0 => 0usize,
2 | 3 => self.cfg.height as usize,
_ => (self.cfg.height / 2) as usize,
};
let mut stream: Vec<u8> = Vec::new();
let log2_max_frame_num_minus4: u32 = 4;
let log2_max_poc_lsb_minus4: u32 = 4;
let chroma_qp_index_offset: i32 = 0;
let p_wp_table: Option<ExplicitPredWeightTableL0> = if self.cfg.explicit_weighted_pred {
Some(compute_explicit_p_weight(
frame.y,
prev.recon_y,
self.cfg.width as usize,
self.cfg.height as usize,
))
} else {
None
};
let wp_l0 = p_wp_table.map(|t| {
let (weight, offset) = t.effective();
WeightedPredL0 {
log2_wd: t.luma_log2_weight_denom,
weight,
offset,
}
});
let me_plane: Vec<u8> = match wp_l0 {
Some(wp) => {
let mut v: Vec<i32> = prev.recon_y.iter().map(|&p| p as i32).collect();
wp.apply(&mut v);
v.iter().map(|&p| p as u8).collect()
}
None => Vec::new(),
};
let me_luma: &[u8] = if me_plane.is_empty() {
prev.recon_y
} else {
&me_plane
};
let mut sw = BitWriter::new();
write_p_slice_header(
&mut sw,
&PSliceHeaderConfig {
first_mb_in_slice: 0,
slice_type_raw: 5, pic_parameter_set_id: 0,
colour_plane_id: self.cfg.colour_plane_id,
frame_num,
frame_num_bits: log2_max_frame_num_minus4 + 4,
pic_order_cnt_lsb,
poc_lsb_bits: log2_max_poc_lsb_minus4 + 4,
slice_qp_delta: frame_qp - self.cfg.qp,
disable_deblocking_filter_idc: if self.cfg.chroma_format_idc == 3 {
1
} else {
0
},
slice_alpha_c0_offset_div2: 0,
slice_beta_offset_div2: 0,
nal_ref_idc: 2,
cabac: None,
field: FieldPicSignal::FrameMbsOnly,
rplm_l0: &[],
mmco: &[],
pred_weight_table: p_wp_table,
num_ref_idx_l0_active_minus1: None,
redundant_pic_cnt: None,
slice_group_change_cycle: None,
},
);
let mut recon_y = vec![0u8; (self.cfg.width * self.cfg.height) as usize];
let mut recon_u = vec![0u8; chroma_width * chroma_height];
let mut recon_v = vec![0u8; chroma_width * chroma_height];
let qp_bd_offset_c = qp_bd_offset(self.cfg.bit_depth_chroma_minus8);
let mut nc_grid = CavlcNcGrid::new(width_mbs, height_mbs);
let mut intra_grid = IntraGrid::new(width_mbs as usize, height_mbs as usize);
let mut mv_grid = MvGrid::new(width_mbs as usize, height_mbs as usize);
let mut mb_deblock_infos: Vec<MbDeblockInfo> =
vec![MbDeblockInfo::default(); (width_mbs * height_mbs) as usize];
let mut pending_skip: u32 = 0;
let mut i8x8_mb_count = 0u32;
let mut qp_tracker = MbQpTracker::new(frame_qp);
qp_tracker.wp_l0 = wp_l0;
let mut want_qp = frame_qp;
for mb_y in 0..height_mbs as usize {
if let Some(budget) = row_budget_bits {
if mb_y > 0 && budget > 0 {
let frac = mb_y as f64 / height_mbs as f64;
let spent = sw.bits_emitted() as f64;
let expected = (budget as f64 * frac).max(1.0);
let ratio = spent.max(1.0) / expected;
let adj = if ratio >= 1.0 {
(6.0 * ratio.log2()).round() as i32
} else {
(3.0 * ratio.log2()).round() as i32
};
let mut new_want = frame_qp + adj;
if new_want < frame_qp && frac < 0.6 {
new_want = frame_qp;
}
want_qp = new_want
.clamp(want_qp - 1, want_qp + 2)
.clamp(frame_qp - 2, frame_qp + 6)
.clamp(0, 51);
}
}
let qp_y = want_qp;
let qp_c = qp_y_to_qp_c_with_bd_offset(qp_y, chroma_qp_index_offset, qp_bd_offset_c);
for mb_x in 0..width_mbs as usize {
let mb_addr = mb_y * width_mbs as usize + mb_x;
let dbl = self.encode_p_mb_with_intra_fallback(
frame,
prev,
me_luma,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
&mut recon_y,
&mut recon_u,
&mut recon_v,
&mut sw,
&mut nc_grid,
&mut intra_grid,
&mut mv_grid,
&mut pending_skip,
&mut qp_tracker,
);
if dbl.transform_size_8x8_flag {
i8x8_mb_count += 1;
}
mb_deblock_infos[mb_addr] = dbl;
}
}
if pending_skip > 0 {
sw.ue(pending_skip);
pending_skip = 0;
}
let _ = pending_skip;
sw.rbsp_trailing_bits();
let slice_rbsp = sw.into_bytes();
stream.extend_from_slice(&build_nal_unit(2, NalUnitType::SliceNonIdr, &slice_rbsp));
if self.cfg.chroma_format_idc != 3 {
crate::encoder::deblock::deblock_recon_with_chroma_array_type(
self.cfg.width,
self.cfg.height,
chroma_width as u32,
chroma_height as u32,
&mut recon_y,
&mut recon_u,
&mut recon_v,
&mb_deblock_infos,
chroma_qp_index_offset,
width_mbs,
height_mbs,
self.cfg.chroma_format_idc,
);
}
let n_parts = (width_mbs as usize) * (height_mbs as usize) * 4;
let mut partition_mvs = Vec::with_capacity(n_parts);
for slot in &mv_grid.slots {
for part in 0..4usize {
let (mv, ref_idx, is_intra) = if !slot.available || slot.is_intra {
(Mv::ZERO, -1, true)
} else {
(slot.mv_l0_8x8[part], slot.ref_idx_l0_8x8[part], false)
};
partition_mvs.push(FrameRefPartitionMv {
mv_l0: (
mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
),
ref_idx_l0: ref_idx.clamp(-1, 127) as i8,
is_intra,
});
}
}
EncodedP {
annex_b: stream,
recon_y,
recon_u,
recon_v,
recon_width: self.cfg.width,
recon_height: self.cfg.height,
frame_num,
pic_order_cnt_lsb,
partition_mvs,
i8x8_mb_count,
pred_weight_table: p_wp_table,
}
}
#[allow(clippy::too_many_arguments)]
#[doc(hidden)]
fn encode_p_mb(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
me_luma: &[u8],
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid: &mut MvGrid,
pending_skip: &mut u32,
qp_tracker: &mut MbQpTracker,
) -> MbDeblockInfo {
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let me = search_quarter_pel_16x16(
frame.y,
width,
self.cfg.width,
self.cfg.height,
me_luma,
prev.width as usize,
prev.width,
prev.height,
mb_x,
mb_y,
16,
16,
);
let chosen_mv = Mv::new(me.mv_x, me.mv_y);
let me_16x16_sad = me.sad;
let ref_idx = qp_tracker.ref_idx_l0;
let ref_key = prev.pic_order_cnt.wrapping_mul(4);
let mvp = mvp_for_16x16(mv_grid, mb_x, mb_y, ref_idx);
let (skip_ref, skip_mv) = p_skip_mv(mv_grid, mb_x, mb_y);
let chroma_array_type = self.cfg.chroma_format_idc;
let mut pred_y =
build_inter_pred_luma(prev.recon_y, prev.width, prev.height, mb_x, mb_y, chosen_mv);
if let Some(wp) = qp_tracker.wp_l0 {
wp.apply(&mut pred_y);
}
let chroma_w = chroma_width as u32;
let chroma_h = chroma_height as u32;
let (pred_u, pred_v): (Vec<i32>, Vec<i32>) = match chroma_array_type {
0 => (Vec::new(), Vec::new()),
3 => (
build_inter_pred_luma(prev.recon_u, chroma_w, chroma_h, mb_x, mb_y, chosen_mv)
.to_vec(),
build_inter_pred_luma(prev.recon_v, chroma_w, chroma_h, mb_x, mb_y, chosen_mv)
.to_vec(),
),
2 => (
cabac_path::build_inter_pred_chroma_local_422(
prev.recon_u,
chroma_w,
chroma_h,
mb_x,
mb_y,
chosen_mv,
)
.to_vec(),
cabac_path::build_inter_pred_chroma_local_422(
prev.recon_v,
chroma_w,
chroma_h,
mb_x,
mb_y,
chosen_mv,
)
.to_vec(),
),
_ => {
let cmv = Mv::new(chosen_mv.x, chosen_mv.y + self.cfg.table_8_10_cy_offset);
(
build_inter_pred_chroma(prev.recon_u, chroma_w, chroma_h, mb_x, mb_y, cmv)
.to_vec(),
build_inter_pred_chroma(prev.recon_v, chroma_w, chroma_h, mb_x, mb_y, cmv)
.to_vec(),
)
}
};
let me_4mv: [crate::encoder::me::MeResult; 4] = std::array::from_fn(|i| {
let sub_x = i % 2;
let sub_y = i / 2;
search_quarter_pel_8x8(
frame.y,
width,
self.cfg.width,
self.cfg.height,
me_luma,
prev.width as usize,
prev.width,
prev.height,
mb_x,
mb_y,
sub_x,
sub_y,
16,
16,
)
});
let me_4mv_sad_sum: u32 = me_4mv.iter().map(|m| m.sad).sum();
let try_4mv = chroma_array_type == 1
&& me_16x16_sad >= 1024
&& (me_4mv_sad_sum as u64) * 10 < (me_16x16_sad as u64) * 7;
if try_4mv {
return self.encode_p_mb_4mv(
frame,
prev,
me_luma,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid,
pending_skip,
qp_tracker,
me_4mv,
);
}
let wq = self.cfg.scaling_matrix.inter_weights();
let inter_luma = forward_inter_luma(frame.y, width, mb_x, mb_y, &pred_y, qp_y, &wq.w4);
let mut luma_4x4_levels_scan = inter_luma.levels_scan;
let luma_4x4_quant_raster = inter_luma.quant_raster;
let mut blk_has_nz = inter_luma.blk_has_nz;
let mut recon_block_residual = inter_luma.recon_residual;
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
let any_nz = (0..4).any(|sub| blk_has_nz[blk8 * 4 + sub]);
if any_nz {
cbp_luma |= 1u8 << blk8;
}
}
let mut transform_size_8x8 = false;
if self.cfg.transform_8x8 {
let inter_luma8 = forward_inter_luma_8x8(
frame.y,
width,
mb_x,
mb_y,
&pred_y,
qp_y,
&wq.w8,
sw.field_scan(),
);
let mut cbp_luma8: u8 = 0;
for blk8 in 0..4usize {
if inter_luma8.blk_has_nz[blk8 * 4] {
cbp_luma8 |= 1u8 << blk8;
}
}
let lambda = lambda_ssd(qp_y, false);
let j4 = inter_luma_transform_cost(
frame.y,
width,
mb_x,
mb_y,
&pred_y,
&inter_luma,
cbp_luma,
lambda,
);
let j8 = inter_luma_transform_cost(
frame.y,
width,
mb_x,
mb_y,
&pred_y,
&inter_luma8,
cbp_luma8,
lambda,
);
if j8 < j4 {
transform_size_8x8 = true;
luma_4x4_levels_scan = inter_luma8.levels_scan;
blk_has_nz = inter_luma8.blk_has_nz;
recon_block_residual = inter_luma8.recon_residual;
cbp_luma = cbp_luma8;
}
}
let mut u_dc_levels = [0i32; 4];
let mut v_dc_levels = [0i32; 4];
let mut u_ac_levels = [[0i32; 16]; 4];
let mut v_ac_levels = [[0i32; 16]; 4];
let mut u_recon_residual = [0i32; 64];
let mut v_recon_residual = [0i32; 64];
let mut u_dc8 = [0i32; 8];
let mut v_dc8 = [0i32; 8];
let mut u_ac8 = [[0i32; 16]; 8];
let mut v_ac8 = [[0i32; 16]; 8];
let mut u_recon_res_422 = [0i32; 128];
let mut v_recon_res_422 = [0i32; 128];
let mut cb444 = InterLumaForward::default();
let mut cr444 = InterLumaForward::default();
let mut cbp_chroma: u8 = 0;
match chroma_array_type {
0 => {}
3 => {
let pu: &[i32; 256] = pred_u[..256].try_into().expect("pred tile");
let pv: &[i32; 256] = pred_v[..256].try_into().expect("pred tile");
if transform_size_8x8 {
cb444 = forward_inter_luma_8x8(
frame.u,
chroma_width,
mb_x,
mb_y,
pu,
qp_c,
&wq.w8,
sw.field_scan(),
);
cr444 = forward_inter_luma_8x8(
frame.v,
chroma_width,
mb_x,
mb_y,
pv,
qp_c,
&wq.w8,
sw.field_scan(),
);
} else {
cb444 = forward_inter_luma(frame.u, chroma_width, mb_x, mb_y, pu, qp_c, &wq.w4);
cr444 = forward_inter_luma(frame.v, chroma_width, mb_x, mb_y, pv, qp_c, &wq.w4);
}
for blk8 in 0..4usize {
let cb_nz = (0..4).any(|sub| cb444.blk_has_nz[blk8 * 4 + sub]);
let cr_nz = (0..4).any(|sub| cr444.blk_has_nz[blk8 * 4 + sub]);
if cb_nz || cr_nz {
cbp_luma |= 1u8 << blk8;
}
}
if cbp_luma == 0 {
transform_size_8x8 = false;
}
}
2 => {
let pu: &[i32; 128] = pred_u[..128].try_into().expect("pred tile");
let pv: &[i32; 128] = pred_v[..128].try_into().expect("pred tile");
let (dcu, acu, ru) =
encode_chroma_residual_422(frame.u, chroma_width, mb_x, mb_y, pu, qp_c, &wq.w4);
let (dcv, acv, rv) =
encode_chroma_residual_422(frame.v, chroma_width, mb_x, mb_y, pv, qp_c, &wq.w4);
u_dc8 = dcu;
v_dc8 = dcv;
u_ac8 = acu;
v_ac8 = acv;
u_recon_res_422 = ru;
v_recon_res_422 = rv;
let any_dc_nz = u_dc8.iter().any(|&v| v != 0) || v_dc8.iter().any(|&v| v != 0);
let any_ac_nz = u_ac8.iter().any(|b| b.iter().any(|&v| v != 0))
|| v_ac8.iter().any(|b| b.iter().any(|&v| v != 0));
cbp_chroma = if any_ac_nz {
2
} else if any_dc_nz {
1
} else {
0
};
}
_ => {
let pu: &[i32; 64] = pred_u[..64].try_into().expect("pred tile");
let pv: &[i32; 64] = pred_v[..64].try_into().expect("pred tile");
let (dcu, acu, ru) = encode_chroma_residual_inter(
frame.u,
chroma_width,
mb_x,
mb_y,
pu,
qp_c,
&wq.w4,
);
let (dcv, acv, rv) = encode_chroma_residual_inter(
frame.v,
chroma_width,
mb_x,
mb_y,
pv,
qp_c,
&wq.w4,
);
u_dc_levels = dcu;
v_dc_levels = dcv;
u_ac_levels = acu;
v_ac_levels = acv;
u_recon_residual = ru;
v_recon_residual = rv;
let any_dc_nz =
u_dc_levels.iter().any(|&v| v != 0) || v_dc_levels.iter().any(|&v| v != 0);
let any_ac_nz = u_ac_levels.iter().any(|blk| blk.iter().any(|&v| v != 0))
|| v_ac_levels.iter().any(|blk| blk.iter().any(|&v| v != 0));
cbp_chroma = if any_ac_nz {
2
} else if any_dc_nz {
1
} else {
0
};
}
}
let _ = skip_ref;
let is_skip = ref_idx == 0 && chosen_mv == skip_mv && cbp_luma == 0 && cbp_chroma == 0;
if is_skip {
*pending_skip += 1;
*mv_grid.slot_mut(mb_x, mb_y) = MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: [0; 4],
mv_l0_8x8: [skip_mv; 4],
};
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
{
let cur = &mut nc_grid.mbs[mb_addr];
cur.is_available = true;
cur.is_intra = false;
cur.is_skip = true;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
cur.cb_total_coeff = [0u8; 8];
cur.cr_total_coeff = [0u8; 8];
}
for j in 0..16usize {
for i in 0..16usize {
let py = mb_y * 16 + j;
let px = mb_x * 16 + i;
recon_y[py * width + px] = pred_y[j * 16 + i].clamp(0, 255) as u8;
}
}
let (ct_w, ct_h) = chroma_tile_dims(chroma_array_type);
for j in 0..ct_h {
for i in 0..ct_w {
let cy = mb_y * ct_h + j;
let cx = mb_x * ct_w + i;
recon_u[cy * chroma_width + cx] = pred_u[j * ct_w + i].clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = pred_v[j * ct_w + i].clamp(0, 255) as u8;
}
}
let _ = chroma_height;
let mv_t = (
skip_mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
skip_mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
return MbDeblockInfo {
is_intra: false,
qp_y: qp_tracker.cur,
luma_nonzero_4x4: 0,
chroma_nonzero_4x4: 0,
mv_l0: [mv_t; 16],
ref_idx_l0: [0; 4],
ref_poc_l0: [ref_key; 4],
..Default::default()
};
}
sw.ue(*pending_skip);
*pending_skip = 0;
qp_tracker.write_pending_mbaff_flag(sw);
let mvd_x = chosen_mv.x - mvp.x;
let mvd_y = chosen_mv.y - mvp.y;
for blk8 in 0..4usize {
let send = (cbp_luma >> blk8) & 1 == 1;
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = LUMA_4X4_BLK[blk];
let bx_pix = bx * 4;
let by_pix = by * 4;
for j in 0..4usize {
for i in 0..4usize {
let p = pred_y[(by_pix + j) * 16 + bx_pix + i];
let r = if send {
recon_block_residual[blk][j * 4 + i]
} else {
0
};
let py = mb_y * 16 + by_pix + j;
let px = mb_x * 16 + bx_pix + i;
recon_y[py * width + px] = (p + r).clamp(0, 255) as u8;
}
}
}
}
match chroma_array_type {
0 => {}
3 => {
for (pred, fwd, recon) in [
(&pred_u, &cb444, &mut *recon_u),
(&pred_v, &cr444, &mut *recon_v),
] {
for blk8 in 0..4usize {
let send = (cbp_luma >> blk8) & 1 == 1;
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = LUMA_4X4_BLK[blk];
for j in 0..4usize {
for i in 0..4usize {
let p = pred[(by * 4 + j) * 16 + bx * 4 + i];
let r = if send {
fwd.recon_residual[blk][j * 4 + i]
} else {
0
};
let cy = mb_y * 16 + by * 4 + j;
let cx = mb_x * 16 + bx * 4 + i;
recon[cy * chroma_width + cx] = (p + r).clamp(0, 255) as u8;
}
}
}
}
}
}
2 => {
for j in 0..16usize {
for i in 0..8usize {
let cy = mb_y * 16 + j;
let cx = mb_x * 8 + i;
let pu = pred_u[j * 8 + i];
let pv = pred_v[j * 8 + i];
let res_u = if cbp_chroma == 0 {
0
} else {
u_recon_res_422[j * 8 + i]
};
let res_v = if cbp_chroma == 0 {
0
} else {
v_recon_res_422[j * 8 + i]
};
recon_u[cy * chroma_width + cx] = (pu + res_u).clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = (pv + res_v).clamp(0, 255) as u8;
}
}
}
_ => {
for j in 0..8usize {
for i in 0..8usize {
let cy = mb_y * 8 + j;
let cx = mb_x * 8 + i;
let pu = pred_u[j * 8 + i];
let pv = pred_v[j * 8 + i];
let res_u = if cbp_chroma == 0 {
0
} else {
u_recon_residual[j * 8 + i]
};
let res_v = if cbp_chroma == 0 {
0
} else {
v_recon_residual[j * 8 + i]
};
recon_u[cy * chroma_width + cx] = (pu + res_u).clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = (pv + res_v).clamp(0, 255) as u8;
}
}
}
}
let _ = chroma_height;
{
let cur = &mut nc_grid.mbs[mb_addr];
cur.is_available = true;
cur.is_intra = false;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let nc = derive_nc_luma(
nc_grid,
mb_addr as u32,
blk as u8,
LumaNcKind::Ac,
true,
false,
);
luma_4x4_nc[blk] = nc;
if (cbp_luma >> blk8) & 1 == 1 {
let tc = luma_4x4_levels_scan[blk]
.iter()
.filter(|&&v| v != 0)
.count() as u8;
own_totals[blk] = tc;
} else {
own_totals[blk] = 0;
}
}
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let mb_qp_delta = if cbp_luma > 0 || cbp_chroma > 0 {
let d = qp_y - qp_tracker.cur;
qp_tracker.cur = qp_y;
d
} else {
0
};
let mb_cfg = PL016x16McbConfig {
ref_idx_l0: ref_idx as u32,
transform_size_8x8_flag: if self.cfg.transform_8x8 {
Some(transform_size_8x8)
} else {
None
},
mvd_l0_x: mvd_x,
mvd_l0_y: mvd_y,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: [0; 4],
chroma_dc_cr: [0; 4],
chroma_ac_cb: [[0; 16]; 4],
chroma_ac_cr: [[0; 16]; 4],
chroma_ac_nc_cb: [0; 8],
chroma_ac_nc_cr: [0; 8],
};
let _ = luma_4x4_quant_raster; match chroma_array_type {
0 => {
crate::encoder::macroblock::write_p_l0_16x16_mb_chroma(
sw,
&mb_cfg,
qp_tracker.num_ref_idx_l0_active_minus1(),
0,
crate::encoder::macroblock::ChromaWriteKind::Mono,
)
.expect("write P_L0_16x16 mb (4:0:0)");
}
3 => {
let cb_nc = derive_plane_nc_444_inter(
nc_grid,
mb_addr as u32,
false,
cbp_luma,
&cb444.levels_scan,
);
let cr_nc = derive_plane_nc_444_inter(
nc_grid,
mb_addr as u32,
true,
cbp_luma,
&cr444.levels_scan,
);
crate::encoder::macroblock::write_p_l0_16x16_mb_chroma(
sw,
&mb_cfg,
qp_tracker.num_ref_idx_l0_active_minus1(),
3,
crate::encoder::macroblock::ChromaWriteKind::Yuv444Inter {
cb_levels: &cb444.levels_scan,
cr_levels: &cr444.levels_scan,
cb_nc: &cb_nc,
cr_nc: &cr_nc,
cbp_luma,
},
)
.expect("write P_L0_16x16 mb (4:4:4)");
}
2 => {
let (nc_cb, nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr as u32,
false,
cbp_chroma,
&u_ac8,
&v_ac8,
2,
);
crate::encoder::macroblock::write_p_l0_16x16_mb_chroma(
sw,
&mb_cfg,
qp_tracker.num_ref_idx_l0_active_minus1(),
2,
crate::encoder::macroblock::ChromaWriteKind::Yuv422 {
chroma_dc_cb: &u_dc8,
chroma_dc_cr: &v_dc8,
chroma_ac_cb: &u_ac8,
chroma_ac_cr: &v_ac8,
cb_ac_nc: &nc_cb,
cr_ac_nc: &nc_cr,
},
)
.expect("write P_L0_16x16 mb (4:2:2)");
}
_ => {
let (chroma_nc_cb, chroma_nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr as u32,
false,
cbp_chroma,
&u_ac_levels,
&v_ac_levels,
1,
);
crate::encoder::macroblock::write_p_l0_16x16_mb_chroma(
sw,
&mb_cfg,
qp_tracker.num_ref_idx_l0_active_minus1(),
1,
crate::encoder::macroblock::ChromaWriteKind::Yuv420 {
chroma_dc_cb: &u_dc_levels,
chroma_dc_cr: &v_dc_levels,
chroma_ac_cb: &u_ac_levels,
chroma_ac_cr: &v_ac_levels,
cb_ac_nc: &chroma_nc_cb,
cr_ac_nc: &chroma_nc_cr,
},
)
.expect("write P_L0_16x16 mb");
}
}
*mv_grid.slot_mut(mb_x, mb_y) = MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: [ref_idx; 4],
mv_l0_8x8: [chosen_mv; 4],
};
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
let luma_nonzero_4x4 = luma_nz_mask_from_blocks(&blk_has_nz);
let chroma_nonzero_4x4: u16 = match chroma_array_type {
0 | 3 => 0,
2 => {
let mut m = 0u16;
for blk in 0..8usize {
if cbp_chroma == 2 && u_ac8[blk].iter().any(|&v| v != 0) {
m |= 1 << blk;
}
if cbp_chroma == 2 && v_ac8[blk].iter().any(|&v| v != 0) {
m |= 1 << (8 + blk);
}
}
m
}
_ => {
let cb_nz: [bool; 4] = std::array::from_fn(|i| {
cbp_chroma == 2 && u_ac_levels[i].iter().any(|&v| v != 0)
});
let cr_nz: [bool; 4] = std::array::from_fn(|i| {
cbp_chroma == 2 && v_ac_levels[i].iter().any(|&v| v != 0)
});
chroma_nz_mask_from_blocks(&cb_nz, &cr_nz)
}
};
let mv_t = (
chosen_mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
chosen_mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
MbDeblockInfo {
is_intra: false,
qp_y: qp_tracker.cur,
luma_nonzero_4x4,
chroma_nonzero_4x4,
transform_size_8x8_flag: transform_size_8x8,
mv_l0: [mv_t; 16],
ref_idx_l0: [ref_idx as i8; 4],
ref_poc_l0: [ref_key; 4],
..Default::default()
}
}
#[allow(clippy::too_many_arguments)]
fn encode_p_mb_4mv(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
_me_luma: &[u8],
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid: &mut MvGrid,
pending_skip: &mut u32,
qp_tracker: &mut MbQpTracker,
me_4mv: [crate::encoder::me::MeResult; 4],
) -> MbDeblockInfo {
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let chroma_w = prev.width / 2;
let chroma_h = prev.height / 2;
let _ = chroma_height;
let ref_idx = qp_tracker.ref_idx_l0;
let ref_key = prev.pic_order_cnt.wrapping_mul(4);
let mvs: [Mv; 4] = std::array::from_fn(|i| Mv::new(me_4mv[i].mv_x, me_4mv[i].mv_y));
let mut mvds: [(i32, i32); 4] = [(0, 0); 4];
let mut inflight_avail = [false; 4];
let mut mv_inflight = [Mv::ZERO; 4];
let mut ref_inflight = [-1i32; 4];
for sub_idx in 0..4usize {
let mvp = mvp_for_p_8x8_partition(
mv_grid,
mb_x,
mb_y,
sub_idx,
ref_idx,
inflight_avail,
mv_inflight,
ref_inflight,
);
mvds[sub_idx] = (mvs[sub_idx].x - mvp.x, mvs[sub_idx].y - mvp.y);
inflight_avail[sub_idx] = true;
mv_inflight[sub_idx] = mvs[sub_idx];
ref_inflight[sub_idx] = ref_idx;
}
let mut pred_y = [0i32; 256];
for (sub_idx, &mv) in mvs.iter().enumerate() {
let sub_x = sub_idx % 2;
let sub_y = sub_idx / 2;
let p = build_inter_pred_luma_8x8(
prev.recon_y,
prev.width,
prev.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv,
);
for j in 0..8usize {
for i in 0..8usize {
pred_y[(sub_y * 8 + j) * 16 + sub_x * 8 + i] = p[j * 8 + i];
}
}
}
if let Some(wp) = qp_tracker.wp_l0 {
wp.apply(&mut pred_y);
}
let mut pred_u = [0i32; 64];
let mut pred_v = [0i32; 64];
for (sub_idx, &mv) in mvs.iter().enumerate() {
let sub_x = sub_idx % 2;
let sub_y = sub_idx / 2;
let cmv = Mv::new(mv.x, mv.y + self.cfg.table_8_10_cy_offset);
let pu = build_inter_pred_chroma_4x4(
prev.recon_u,
chroma_w,
chroma_h,
mb_x,
mb_y,
sub_x,
sub_y,
cmv,
);
let pv = build_inter_pred_chroma_4x4(
prev.recon_v,
chroma_w,
chroma_h,
mb_x,
mb_y,
sub_x,
sub_y,
cmv,
);
for j in 0..4usize {
for i in 0..4usize {
pred_u[(sub_y * 4 + j) * 8 + sub_x * 4 + i] = pu[j * 4 + i];
pred_v[(sub_y * 4 + j) * 8 + sub_x * 4 + i] = pv[j * 4 + i];
}
}
}
let wq = self.cfg.scaling_matrix.inter_weights();
let inter_luma = forward_inter_luma(frame.y, width, mb_x, mb_y, &pred_y, qp_y, &wq.w4);
let luma_4x4_levels_scan = inter_luma.levels_scan;
let blk_has_nz = inter_luma.blk_has_nz;
let recon_block_residual = inter_luma.recon_residual;
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
let any_nz = (0..4).any(|sub| blk_has_nz[blk8 * 4 + sub]);
if any_nz {
cbp_luma |= 1u8 << blk8;
}
}
let (u_dc_levels, u_ac_levels, u_recon_residual) =
encode_chroma_residual_inter(frame.u, chroma_width, mb_x, mb_y, &pred_u, qp_c, &wq.w4);
let (v_dc_levels, v_ac_levels, v_recon_residual) =
encode_chroma_residual_inter(frame.v, chroma_width, mb_x, mb_y, &pred_v, qp_c, &wq.w4);
let any_dc_nz = u_dc_levels.iter().any(|&v| v != 0) || v_dc_levels.iter().any(|&v| v != 0);
let any_ac_nz = u_ac_levels.iter().any(|blk| blk.iter().any(|&v| v != 0))
|| v_ac_levels.iter().any(|blk| blk.iter().any(|&v| v != 0));
let cbp_chroma: u8 = if any_ac_nz {
2
} else if any_dc_nz {
1
} else {
0
};
sw.ue(*pending_skip);
*pending_skip = 0;
qp_tracker.write_pending_mbaff_flag(sw);
for blk8 in 0..4usize {
let send = (cbp_luma >> blk8) & 1 == 1;
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = LUMA_4X4_BLK[blk];
let bx_pix = bx * 4;
let by_pix = by * 4;
for j in 0..4usize {
for i in 0..4usize {
let p = pred_y[(by_pix + j) * 16 + bx_pix + i];
let r = if send {
recon_block_residual[blk][j * 4 + i]
} else {
0
};
let py = mb_y * 16 + by_pix + j;
let px = mb_x * 16 + bx_pix + i;
recon_y[py * width + px] = (p + r).clamp(0, 255) as u8;
}
}
}
}
for j in 0..8usize {
for i in 0..8usize {
let cy = mb_y * 8 + j;
let cx = mb_x * 8 + i;
let pu = pred_u[j * 8 + i];
let pv = pred_v[j * 8 + i];
let res_u = if cbp_chroma == 0 {
0
} else {
u_recon_residual[j * 8 + i]
};
let res_v = if cbp_chroma == 0 {
0
} else {
v_recon_residual[j * 8 + i]
};
recon_u[cy * chroma_width + cx] = (pu + res_u).clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = (pv + res_v).clamp(0, 255) as u8;
}
}
{
let cur = &mut nc_grid.mbs[mb_addr];
cur.is_available = true;
cur.is_intra = false;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let nc = derive_nc_luma(
nc_grid,
mb_addr as u32,
blk as u8,
LumaNcKind::Ac,
true,
false,
);
luma_4x4_nc[blk] = nc;
if (cbp_luma >> blk8) & 1 == 1 {
let tc = luma_4x4_levels_scan[blk]
.iter()
.filter(|&&v| v != 0)
.count() as u8;
own_totals[blk] = tc;
} else {
own_totals[blk] = 0;
}
}
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let (chroma_nc_cb, chroma_nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr as u32,
false,
cbp_chroma,
&u_ac_levels,
&v_ac_levels,
1,
);
let mb_cfg = P8x8AllPL08x8McbConfig {
ref_idx_l0: ref_idx as u32,
emit_transform_size_8x8_zero: self.cfg.transform_8x8,
mvd_l0: mvds,
cbp_luma,
cbp_chroma,
mb_qp_delta: {
if cbp_luma > 0 || cbp_chroma > 0 {
let d = qp_y - qp_tracker.cur;
qp_tracker.cur = qp_y;
d
} else {
0
}
},
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_p_8x8_all_pl08x8_mb(sw, &mb_cfg, qp_tracker.num_ref_idx_l0_active_minus1())
.expect("write P_8x8 mb");
*mv_grid.slot_mut(mb_x, mb_y) = MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: [ref_idx; 4],
mv_l0_8x8: mvs,
};
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
let luma_nonzero_4x4 = luma_nz_mask_from_blocks(&blk_has_nz);
let cb_nz: [bool; 4] =
std::array::from_fn(|i| cbp_chroma == 2 && u_ac_levels[i].iter().any(|&v| v != 0));
let cr_nz: [bool; 4] =
std::array::from_fn(|i| cbp_chroma == 2 && v_ac_levels[i].iter().any(|&v| v != 0));
let mut mv_l0 = [(0i16, 0i16); 16];
for blk in 0..16usize {
let (bx, by) = LUMA_4X4_BLK[blk];
let sub_idx = 2 * (by / 2) + (bx / 2);
let mv = mvs[sub_idx];
mv_l0[blk] = (
mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
}
MbDeblockInfo {
is_intra: false,
qp_y: qp_tracker.cur,
luma_nonzero_4x4,
chroma_nonzero_4x4: chroma_nz_mask_from_blocks(&cb_nz, &cr_nz),
mv_l0,
ref_idx_l0: [ref_idx as i8; 4],
ref_poc_l0: [ref_key; 4],
..Default::default()
}
}
#[allow(clippy::too_many_arguments)]
fn encode_p_mb_with_intra_fallback(
&self,
frame: &YuvFrame<'_>,
prev: &EncodedFrameRef<'_>,
me_luma: &[u8],
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid: &mut MvGrid,
pending_skip: &mut u32,
qp_tracker: &mut MbQpTracker,
) -> MbDeblockInfo {
if !self.cfg.intra_in_inter || self.cfg.chroma_format_idc != 1 {
return self.encode_p_mb(
frame,
prev,
me_luma,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid,
pending_skip,
qp_tracker,
);
}
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let snap_pre = InterMbStateSnapshot::capture_p(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid,
mb_addr,
width_mbs,
*pending_skip,
);
let pre_bits = sw.bits_emitted();
let qp_tracker_pre = *qp_tracker;
let inter_dbl = self.encode_p_mb(
frame,
prev,
me_luma,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid,
pending_skip,
qp_tracker,
);
let inter_bits = sw.bits_emitted() - pre_bits;
let inter_d = ssd_16x16(frame.y, width, mb_x * 16, mb_y * 16, &{
let mut buf = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
buf[j * 16 + i] = recon_y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
}
}
buf
});
let lambda = lambda_ssd(qp_y, false);
let j_inter = cost_combined(inter_d, inter_bits as u64, lambda);
let snap_inter = InterMbStateSnapshot::capture_p(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid,
mb_addr,
width_mbs,
*pending_skip,
);
let qp_tracker_inter = *qp_tracker;
snap_pre.restore_p(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid,
mb_addr,
pending_skip,
);
*qp_tracker = qp_tracker_pre;
let intra_dbl = self.encode_p_mb_intra16x16(
qp_tracker,
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid,
pending_skip,
);
let intra_bits = sw.bits_emitted() - pre_bits;
let intra_d = ssd_16x16(frame.y, width, mb_x * 16, mb_y * 16, &{
let mut buf = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
buf[j * 16 + i] = recon_y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
}
}
buf
});
let j_intra = cost_combined(intra_d, intra_bits as u64, lambda);
if j_inter <= j_intra {
snap_pre.restore_p(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid,
mb_addr,
pending_skip,
);
snap_inter.restore_p(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid,
mb_addr,
pending_skip,
);
*qp_tracker = qp_tracker_inter;
inter_dbl
} else {
intra_dbl
}
}
#[allow(clippy::too_many_arguments)]
fn encode_p_mb_intra16x16(
&self,
qp_tracker: &mut MbQpTracker,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid: &mut MvGrid,
pending_skip: &mut u32,
) -> MbDeblockInfo {
sw.ue(*pending_skip);
*pending_skip = 0;
qp_tracker.write_pending_mbaff_flag(sw);
let mb_qp_delta = qp_y - qp_tracker.cur;
qp_tracker.cur = qp_y;
self.encode_intra16x16_in_inter_slice(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
5,
mb_qp_delta,
qp_tracker.intra_nb,
);
*mv_grid.slot_mut(mb_x, mb_y) = MvGridSlot {
available: true,
is_intra: true,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
};
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let any_nz = nc_grid.mbs[mb_addr]
.luma_total_coeff
.iter()
.any(|&v| v != 0);
let luma_nonzero_4x4: u16 = if any_nz {
let mut nz = [false; 16];
for (blkz, slot) in nz.iter_mut().enumerate() {
*slot = nc_grid.mbs[mb_addr].luma_total_coeff[blkz] != 0;
}
luma_nz_mask_from_blocks(&nz)
} else {
0
};
MbDeblockInfo {
is_intra: true,
qp_y: qp_tracker.cur,
luma_nonzero_4x4,
chroma_nonzero_4x4: 0,
..Default::default()
}
}
#[allow(clippy::too_many_arguments)]
fn encode_intra16x16_in_inter_slice(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mb_type_offset: u32,
mb_qp_delta: i32,
nb: Option<(bool, bool, bool)>,
) {
let width = self.cfg.width as usize;
let height = self.cfg.height as usize;
let lambda = lambda_ssd(qp_y, true);
let wq = self.cfg.scaling_matrix.intra_weights();
let mut best_luma_mode = I16x16Mode::Dc;
let mut best_luma_pred = [0i32; 256];
let mut best_luma_cost = u64::MAX;
for &mode in &[
I16x16Mode::Vertical,
I16x16Mode::Horizontal,
I16x16Mode::Dc,
I16x16Mode::Plane,
] {
let Some(pred) = predict_16x16_avail(mode, recon_y, width, height, mb_x, mb_y, nb)
else {
continue;
};
let cost = trial_intra16x16_cost(frame, &pred, mb_x, mb_y, qp_y, width, lambda, &wq.w4);
if cost < best_luma_cost {
best_luma_cost = cost;
best_luma_mode = mode;
best_luma_pred = pred;
}
}
let mut residual = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
let s = frame.y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
residual[j * 16 + i] = s - best_luma_pred[j * 16 + i];
}
}
let mut coeffs_4x4 = [[0i32; 16]; 16];
for (blk, slot) in coeffs_4x4.iter_mut().enumerate() {
let bx = blk % 4;
let by = blk / 4;
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by * 4 + j) * 16 + bx * 4 + i];
}
}
*slot = forward_core_4x4(&block);
}
let mut dc_coeffs = [0i32; 16];
for (blk, c) in coeffs_4x4.iter().enumerate() {
let bx = blk % 4;
let by = blk / 4;
dc_coeffs[by * 4 + bx] = c[0];
}
let dc_t = forward_hadamard_4x4(&dc_coeffs);
let dc_levels = quantize_luma_dc_w(&dc_t, qp_y, true, wq.w4[0]);
let mut luma_ac_levels = [[0i32; 16]; 16];
let mut luma_ac_quant_raster = [[0i32; 16]; 16];
let mut any_luma_ac_nz = false;
for blkz in 0..16usize {
let (bx, by) = LUMA_4X4_BLK[blkz];
let raster = by * 4 + bx;
let z = quantize_4x4_ac_w(&coeffs_4x4[raster], qp_y, true, &wq.w4);
luma_ac_quant_raster[raster] = z;
let scan_ac = zigzag_scan_4x4_ac(&z);
luma_ac_levels[blkz] = scan_ac;
if scan_ac.iter().any(|&v| v != 0) {
any_luma_ac_nz = true;
}
}
let cbp_luma: u8 = if any_luma_ac_nz { 15 } else { 0 };
let inv_dc =
inverse_hadamard_luma_dc_16x16(&dc_levels, qp_y, &wq.w4, 8).expect("luma DC inverse");
#[allow(clippy::needless_range_loop)]
for blk in 0..16usize {
let bx = blk % 4;
let by = blk / 4;
let mut coeffs_in = if cbp_luma == 15 {
luma_ac_quant_raster[blk]
} else {
[0i32; 16]
};
coeffs_in[0] = inv_dc[by * 4 + bx];
let r = inverse_transform_4x4_dc_preserved(&coeffs_in, qp_y, &wq.w4, 8)
.expect("inverse 4x4");
for j in 0..4 {
for i in 0..4 {
let px = mb_x * 16 + bx * 4 + i;
let py = mb_y * 16 + by * 4 + j;
let p = best_luma_pred[(by * 4 + j) * 16 + bx * 4 + i];
let raw = p + r[j * 4 + i];
recon_y[py * width + px] = raw.clamp(0, 255) as u8;
}
}
}
let chroma_block = encode_chroma_block(
self.cfg.chroma_format_idc,
frame.u,
frame.v,
recon_u,
recon_v,
chroma_width,
chroma_height,
mb_x,
mb_y,
qp_c,
&wq.w4,
nb,
);
let cbp_chroma = chroma_block.cbp_chroma;
let best_chroma_mode_u8 = chroma_block.pred_mode.as_u8();
install_chroma_recon(
self.cfg.chroma_format_idc,
chroma_width,
mb_x,
mb_y,
&chroma_block,
recon_u,
recon_v,
);
let _ = chroma_height;
let mut u_dc_levels = [0i32; 4];
let mut v_dc_levels = [0i32; 4];
let mut u_ac_levels = [[0i32; 16]; 4];
let mut v_ac_levels = [[0i32; 16]; 4];
u_dc_levels.copy_from_slice(&chroma_block.dc_cb[..4]);
v_dc_levels.copy_from_slice(&chroma_block.dc_cr[..4]);
u_ac_levels.copy_from_slice(&chroma_block.ac_cb[..4]);
v_ac_levels.copy_from_slice(&chroma_block.ac_cr[..4]);
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
{
let cur = &mut nc_grid.mbs[mb_addr];
cur.is_available = true;
cur.is_intra = true;
cur.is_skip = false;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let nc_dc = derive_nc_luma(
nc_grid,
mb_addr as u32,
0,
LumaNcKind::Intra16x16Dc,
true,
false,
);
let mut luma_ac_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blkz in 0..16u8 {
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let nc = derive_nc_luma(nc_grid, mb_addr as u32, blkz, LumaNcKind::Ac, true, false);
luma_ac_nc[blkz as usize] = nc;
if cbp_luma == 15 {
let tc = luma_ac_levels[blkz as usize]
.iter()
.filter(|&&v| v != 0)
.count() as u8;
own_totals[blkz as usize] = tc;
}
}
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let (chroma_nc_cb, chroma_nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr as u32,
true,
cbp_chroma,
&u_ac_levels,
&v_ac_levels,
1,
);
let mb_cfg = I16x16McbConfig {
pred_mode: best_luma_mode.as_u8(),
intra_chroma_pred_mode: best_chroma_mode_u8,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_dc_levels_raster: dc_levels,
luma_ac_levels,
luma_ac_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_intra16x16_mb_in_inter_slice(
sw,
&mb_cfg,
mb_type_offset,
CoeffTokenContext::Numeric(nc_dc),
)
.expect("write Intra_16x16 in inter slice");
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
let _ = best_luma_cost;
}
}
pub struct EncodedFrameRef<'a> {
pub width: u32,
pub height: u32,
pub recon_y: &'a [u8],
pub recon_u: &'a [u8],
pub recon_v: &'a [u8],
pub partition_mvs: &'a [FrameRefPartitionMv],
pub pic_order_cnt: i32,
}
impl<'a> From<&'a EncodedIdr> for EncodedFrameRef<'a> {
fn from(e: &'a EncodedIdr) -> Self {
Self {
width: e.recon_width,
height: e.recon_height,
recon_y: &e.recon_y,
recon_u: &e.recon_u,
recon_v: &e.recon_v,
partition_mvs: &e.partition_mvs,
pic_order_cnt: 0, }
}
}
impl<'a> From<&'a EncodedP> for EncodedFrameRef<'a> {
fn from(e: &'a EncodedP) -> Self {
Self {
width: e.recon_width,
height: e.recon_height,
recon_y: &e.recon_y,
recon_u: &e.recon_u,
recon_v: &e.recon_v,
partition_mvs: &e.partition_mvs,
pic_order_cnt: e.pic_order_cnt_lsb as i32,
}
}
}
#[derive(Default)]
struct InterLumaForward {
levels_scan: [[i32; 16]; 16],
quant_raster: [[i32; 16]; 16],
blk_has_nz: [bool; 16],
recon_residual: [[i32; 16]; 16],
}
#[allow(clippy::too_many_arguments)]
fn forward_inter_luma_8x8(
src_y: &[u8],
src_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 256],
qp_y: i32,
w8: &[i32; 64],
field_scan: bool,
) -> InterLumaForward {
let mut levels_scan = [[0i32; 16]; 16];
let mut blk_has_nz = [false; 16];
let mut recon_residual = [[0i32; 16]; 16];
for blk8 in 0..4usize {
let qx = (blk8 % 2) * 8;
let qy = (blk8 / 2) * 8;
let mut res = [0i32; 64];
for j in 0..8 {
for i in 0..8 {
let s = src_y[(mb_y * 16 + qy + j) * src_stride + (mb_x * 16 + qx + i)] as i32;
res[j * 8 + i] = s - pred[(qy + j) * 16 + qx + i];
}
}
let w = forward_core_8x8(&res);
let z = quantize_8x8_w(&w, qp_y, false, w8);
let any_nz = z.iter().any(|&v| v != 0);
let scan = scan_8x8_for_cavlc(&z, field_scan);
let subs = deinterleave_8x8_to_4x4(&scan);
for (i4x4, sub) in subs.iter().enumerate() {
levels_scan[blk8 * 4 + i4x4] = *sub;
blk_has_nz[blk8 * 4 + i4x4] = any_nz;
}
let r = inverse_transform_8x8(&z, qp_y, w8, 8).expect("inv 8x8");
for sub in 0..4usize {
let sx = (sub % 2) * 4;
let sy = (sub / 2) * 4;
let mut tile = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
tile[j * 4 + i] = r[(sy + j) * 8 + sx + i];
}
}
recon_residual[blk8 * 4 + sub] = tile;
}
}
InterLumaForward {
levels_scan,
quant_raster: [[0i32; 16]; 16],
blk_has_nz,
recon_residual,
}
}
#[allow(clippy::too_many_arguments)]
fn inter_luma_transform_cost(
src_y: &[u8],
src_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 256],
fwd: &InterLumaForward,
cbp_luma: u8,
lambda: f64,
) -> u64 {
let mut d: u64 = 0;
for (blk, &(bx, by)) in LUMA_4X4_BLK.iter().enumerate() {
let send = (cbp_luma >> (blk / 4)) & 1 == 1;
for j in 0..4usize {
for i in 0..4usize {
let p = pred[(by * 4 + j) * 16 + bx * 4 + i];
let r = if send {
fwd.recon_residual[blk][j * 4 + i]
} else {
0
};
let recon = (p + r).clamp(0, 255);
let s =
src_y[(mb_y * 16 + by * 4 + j) * src_stride + (mb_x * 16 + bx * 4 + i)] as i32;
let e = (s - recon) as i64;
d += (e * e) as u64;
}
}
}
let mut trial_w = BitWriter::new();
for blk in 0..16usize {
if (cbp_luma >> (blk / 4)) & 1 == 1 {
let _ = encode_residual_block_cavlc(
&mut trial_w,
CoeffTokenContext::Numeric(0),
16,
&fwd.levels_scan[blk],
);
}
}
cost_combined(d, trial_w.bits_emitted() as u64, lambda)
}
fn forward_inter_luma(
src_y: &[u8],
src_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 256],
qp_y: i32,
w4: &[i32; 16],
) -> InterLumaForward {
let mut residual = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
let s = src_y[(mb_y * 16 + j) * src_stride + (mb_x * 16 + i)] as i32;
residual[j * 16 + i] = s - pred[j * 16 + i];
}
}
let mut levels_scan = [[0i32; 16]; 16];
let mut quant_raster = [[0i32; 16]; 16];
let mut blk_has_nz = [false; 16];
let mut recon_residual = [[0i32; 16]; 16];
for blkz in 0..16usize {
let (bx, by) = LUMA_4X4_BLK[blkz];
let mut block = [0i32; 16];
for j in 0..4 {
for i in 0..4 {
block[j * 4 + i] = residual[(by * 4 + j) * 16 + bx * 4 + i];
}
}
let w = forward_core_4x4(&block);
let z = quantize_4x4_w(&w, qp_y, false, w4);
let scan = zigzag_scan_4x4(&z);
levels_scan[blkz] = scan;
quant_raster[blkz] = z;
blk_has_nz[blkz] = z.iter().any(|&v| v != 0);
let r = inverse_transform_4x4(&z, qp_y, w4, 8).expect("inv 4x4");
recon_residual[blkz] = r;
}
InterLumaForward {
levels_scan,
quant_raster,
blk_has_nz,
recon_residual,
}
}
#[derive(Debug, Clone, Copy)]
struct BDirectDerivation {
ref_idx_l0: i32,
ref_idx_l1: i32,
mv_l0_per_8x8: [Mv; 4],
mv_l1_per_8x8: [Mv; 4],
direct_zero: bool,
}
impl BDirectDerivation {
fn is_uniform(&self) -> bool {
let l0_uniform = self
.mv_l0_per_8x8
.iter()
.all(|m| *m == self.mv_l0_per_8x8[0]);
let l1_uniform = self
.mv_l1_per_8x8
.iter()
.all(|m| *m == self.mv_l1_per_8x8[0]);
l0_uniform && l1_uniform
}
fn mv_l0(&self) -> Mv {
self.mv_l0_per_8x8[0]
}
fn mv_l1(&self) -> Mv {
self.mv_l1_per_8x8[0]
}
}
fn col_zero_flag(
l1_partition_mvs: &[FrameRefPartitionMv],
width_mbs: usize,
mb_x: usize,
mb_y: usize,
part_8x8: usize,
) -> bool {
let idx = (mb_y * width_mbs + mb_x) * 4 + part_8x8;
let Some(slot) = l1_partition_mvs.get(idx) else {
return false;
};
if slot.is_intra {
return false;
}
if slot.ref_idx_l0 != 0 {
return false;
}
slot.mv_l0.0.abs() <= 1 && slot.mv_l0.1.abs() <= 1
}
fn b_spatial_direct_derive(
mv_grid_l0: &MvGrid,
mv_grid_l1: &MvGrid,
l1_partition_mvs: &[FrameRefPartitionMv],
width_mbs: usize,
mb_x: usize,
mb_y: usize,
) -> BDirectDerivation {
let (a_l0, b_l0, c_l0, d_l0) = neighbour_mvs_16x16(mv_grid_l0, mb_x, mb_y);
let (a_l1, b_l1, c_l1, d_l1) = neighbour_mvs_16x16(mv_grid_l1, mb_x, mb_y);
let (ref_idx_l0_helper, mv_l0_mb, ref_idx_l1_helper, mv_l1_mb) =
crate::mv_deriv::derive_b_spatial_direct_with_d(
a_l0, a_l1, b_l0, b_l1, c_l0, c_l1, d_l0, d_l1,
);
let c_l0_eff = if !c_l0.partition_available && d_l0.partition_available {
d_l0
} else {
c_l0
};
let c_l1_eff = if !c_l1.partition_available && d_l1.partition_available {
d_l1
} else {
c_l1
};
let min_pos = |x: i32, y: i32| -> i32 {
if x >= 0 && y >= 0 {
x.min(y)
} else {
x.max(y)
}
};
let r0_minp = min_pos(a_l0.ref_idx, min_pos(b_l0.ref_idx, c_l0_eff.ref_idx));
let r1_minp = min_pos(a_l1.ref_idx, min_pos(b_l1.ref_idx, c_l1_eff.ref_idx));
let direct_zero = r0_minp < 0 && r1_minp < 0;
let ref_idx_l0_final: i32 = if direct_zero {
0
} else if r0_minp < 0 {
-1
} else {
ref_idx_l0_helper
};
let ref_idx_l1_final: i32 = if direct_zero {
0
} else if r1_minp < 0 {
-1
} else {
ref_idx_l1_helper
};
let mut mv_l0_per_8x8 = [Mv::ZERO; 4];
let mut mv_l1_per_8x8 = [Mv::ZERO; 4];
for part in 0..4usize {
let cz = col_zero_flag(l1_partition_mvs, width_mbs, mb_x, mb_y, part);
let l0_force_zero = direct_zero || ref_idx_l0_final < 0 || (ref_idx_l0_final == 0 && cz);
let l1_force_zero = direct_zero || ref_idx_l1_final < 0 || (ref_idx_l1_final == 0 && cz);
mv_l0_per_8x8[part] = if l0_force_zero { Mv::ZERO } else { mv_l0_mb };
mv_l1_per_8x8[part] = if l1_force_zero { Mv::ZERO } else { mv_l1_mb };
}
BDirectDerivation {
ref_idx_l0: ref_idx_l0_final,
ref_idx_l1: ref_idx_l1_final,
mv_l0_per_8x8,
mv_l1_per_8x8,
direct_zero,
}
}
fn b_temporal_direct_derive(
l1_partition_mvs: &[FrameRefPartitionMv],
width_mbs: usize,
mb_x: usize,
mb_y: usize,
curr_poc: i32,
pic0_poc: i32,
pic1_poc: i32,
) -> BDirectDerivation {
let mut mv_l0_per_8x8 = [Mv::ZERO; 4];
let mut mv_l1_per_8x8 = [Mv::ZERO; 4];
for part in 0..4usize {
let idx = (mb_y * width_mbs + mb_x) * 4 + part;
let (mv_col, is_intra) = match l1_partition_mvs.get(idx) {
Some(slot) if !slot.is_intra => {
(Mv::new(slot.mv_l0.0 as i32, slot.mv_l0.1 as i32), false)
}
_ => (Mv::ZERO, true),
};
if is_intra {
mv_l0_per_8x8[part] = Mv::ZERO;
mv_l1_per_8x8[part] = Mv::ZERO;
continue;
}
let (mv_l0, mv_l1) =
crate::mv_deriv::derive_b_temporal_direct(mv_col, pic1_poc, curr_poc, pic0_poc);
mv_l0_per_8x8[part] = mv_l0;
mv_l1_per_8x8[part] = mv_l1;
}
BDirectDerivation {
ref_idx_l0: 0,
ref_idx_l1: 0,
mv_l0_per_8x8,
mv_l1_per_8x8,
direct_zero: false,
}
}
pub struct EncodedB {
pub annex_b: Vec<u8>,
pub recon_y: Vec<u8>,
pub recon_u: Vec<u8>,
pub recon_v: Vec<u8>,
pub recon_width: u32,
pub recon_height: u32,
pub frame_num: u32,
pub pic_order_cnt_lsb: u32,
pub partition_mvs: Vec<FrameRefPartitionMv>,
pub i8x8_mb_count: u32,
}
impl<'a> From<&'a EncodedB> for EncodedFrameRef<'a> {
fn from(e: &'a EncodedB) -> Self {
Self {
width: e.recon_width,
height: e.recon_height,
recon_y: &e.recon_y,
recon_u: &e.recon_u,
recon_v: &e.recon_v,
partition_mvs: &e.partition_mvs,
pic_order_cnt: e.pic_order_cnt_lsb as i32,
}
}
}
#[derive(Debug, Clone, Copy)]
struct BMbDecision {
pred: BPred16x16,
mv_l0: Mv,
mv_l1: Mv,
mvd_l0_x: i32,
mvd_l0_y: i32,
mvd_l1_x: i32,
mvd_l1_y: i32,
pred_y: [i32; 256],
pred_u: [i32; 64],
pred_v: [i32; 64],
}
#[derive(Debug, Clone, Copy)]
struct WeightedBipredLuma {
log2_wd: u32,
weight_l0: i32,
offset_l0: i32,
weight_l1: i32,
offset_l1: i32,
apply_chroma: bool,
}
#[inline]
fn weighted_bipred_merge(a: i32, b: i32, wp: &WeightedBipredLuma) -> i32 {
let round = 1i32 << wp.log2_wd;
let off_avg = (wp.offset_l0 + wp.offset_l1 + 1) >> 1;
let v = (a * wp.weight_l0 + b * wp.weight_l1 + round) >> (wp.log2_wd + 1);
(v + off_avg).clamp(0, 255)
}
#[allow(clippy::too_many_arguments)]
fn compute_explicit_bipred_weights(
src_y: &[u8],
ref_l0_y: &[u8],
ref_l1_y: &[u8],
width: usize,
height: usize,
ref_l0_stride: usize,
ref_l1_stride: usize,
) -> ExplicitBipredWeightTable {
const LOG2_WD: u32 = 5; const DENOM: i32 = 1 << (LOG2_WD + 1);
let default_table = ExplicitBipredWeightTable {
luma_log2_weight_denom: LOG2_WD,
luma_weight_l0: 32,
luma_offset_l0: 0,
luma_weight_l1: 32,
luma_offset_l1: 0,
};
let sse_for = |w0: i32, w1: i32, o: i32| -> u64 {
let mut sse: u64 = 0;
let logwd_plus1 = LOG2_WD + 1;
let round = 1i32 << LOG2_WD;
let off_avg = (o + o + 1) >> 1;
for j in 0..height {
for i in 0..width {
let p0 = ref_l0_y[j * ref_l0_stride + i] as i32;
let p1 = ref_l1_y[j * ref_l1_stride + i] as i32;
let v = ((p0 * w0 + p1 * w1 + round) >> logwd_plus1) + off_avg;
let pred = v.clamp(0, 255);
let s = src_y[j * width + i] as i32;
let d = s - pred;
sse = sse.saturating_add((d * d) as u64);
}
}
sse
};
let total = (width * height) as f64;
let mut best_sse = u64::MAX;
let mut best_w0 = 32i32;
let mut best_off = 0i32;
for w0 in (0..=DENOM).step_by(4) {
let w1 = DENOM - w0;
let mut sum_resid: f64 = 0.0;
let logwd_plus1 = LOG2_WD + 1;
let round = 1i32 << LOG2_WD;
for j in 0..height {
for i in 0..width {
let p0 = ref_l0_y[j * ref_l0_stride + i] as i32;
let p1 = ref_l1_y[j * ref_l1_stride + i] as i32;
let v = (p0 * w0 + p1 * w1 + round) >> logwd_plus1;
let s = src_y[j * width + i] as i32;
sum_resid += (s - v) as f64;
}
}
let o_raw = (sum_resid / total).round() as i32;
let o = o_raw.clamp(-127, 127);
let sse = sse_for(w0, w1, o);
if sse < best_sse {
best_sse = sse;
best_w0 = w0;
best_off = o;
}
}
let mut default_sse: u64 = 0;
for j in 0..height {
for i in 0..width {
let p0 = ref_l0_y[j * ref_l0_stride + i] as i32;
let p1 = ref_l1_y[j * ref_l1_stride + i] as i32;
let pred = (p0 + p1 + 1) >> 1;
let s = src_y[j * width + i] as i32;
let d = s - pred;
default_sse = default_sse.saturating_add((d * d) as u64);
}
}
let margin = default_sse / 100;
if best_sse + margin >= default_sse {
return default_table;
}
let w1 = DENOM - best_w0;
ExplicitBipredWeightTable {
luma_log2_weight_denom: LOG2_WD,
luma_weight_l0: best_w0,
luma_offset_l0: best_off,
luma_weight_l1: w1,
luma_offset_l1: best_off,
}
}
fn compute_explicit_p_weight(
src_y: &[u8],
ref_y: &[u8],
width: usize,
height: usize,
) -> ExplicitPredWeightTableL0 {
const LOG2_WD: u32 = 5;
let identity = ExplicitPredWeightTableL0 {
luma_log2_weight_denom: LOG2_WD,
luma: None,
};
let n = (width * height) as f64;
if n == 0.0 {
return identity;
}
let (mut sx, mut sr, mut srr, mut sxr) = (0f64, 0f64, 0f64, 0f64);
for (&x, &r) in src_y.iter().zip(ref_y.iter()) {
let (x, r) = (x as f64, r as f64);
sx += x;
sr += r;
srr += r * r;
sxr += x * r;
}
let var_r = srr - sr * sr / n;
let a = if var_r.abs() < 1e-9 {
1.0
} else {
(sxr - sx * sr / n) / var_r
};
let b = (sx - a * sr) / n;
let w = ((a * (1u32 << LOG2_WD) as f64).round() as i32).clamp(-128, 127);
let o = (b.round() as i32).clamp(-128, 127);
if w == (1i32 << LOG2_WD) && o == 0 {
return identity;
}
let cand = WeightedPredL0 {
log2_wd: LOG2_WD,
weight: w,
offset: o,
};
let (mut ssd_id, mut ssd_w) = (0u64, 0u64);
let mut buf = [0i32; 1];
for (&x, &r) in src_y.iter().zip(ref_y.iter()) {
buf[0] = r as i32;
cand.apply(&mut buf);
let d_id = x as i64 - r as i64;
let d_w = x as i64 - buf[0] as i64;
ssd_id += (d_id * d_id) as u64;
ssd_w += (d_w * d_w) as u64;
}
if ssd_w * 100 <= ssd_id * 99 {
ExplicitPredWeightTableL0 {
luma_log2_weight_denom: LOG2_WD,
luma: Some((w, o)),
}
} else {
identity
}
}
#[allow(clippy::too_many_arguments)]
fn build_b_predictors(
pred: BPred16x16,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
mv_l0: Mv,
mv_l1: Mv,
weighted: Option<WeightedBipredLuma>,
) -> ([i32; 256], [i32; 64], [i32; 64]) {
let chroma_w_l0 = ref_l0.width / 2;
let chroma_h_l0 = ref_l0.height / 2;
let chroma_w_l1 = ref_l1.width / 2;
let chroma_h_l1 = ref_l1.height / 2;
let p_y_l0 = if pred.uses_l0() {
Some(build_inter_pred_luma(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
mv_l0,
))
} else {
None
};
let p_y_l1 = if pred.uses_l1() {
Some(build_inter_pred_luma(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
mv_l1,
))
} else {
None
};
let p_u_l0 = if pred.uses_l0() {
Some(build_inter_pred_chroma(
ref_l0.recon_u,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
mv_l0,
))
} else {
None
};
let p_u_l1 = if pred.uses_l1() {
Some(build_inter_pred_chroma(
ref_l1.recon_u,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
mv_l1,
))
} else {
None
};
let p_v_l0 = if pred.uses_l0() {
Some(build_inter_pred_chroma(
ref_l0.recon_v,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
mv_l0,
))
} else {
None
};
let p_v_l1 = if pred.uses_l1() {
Some(build_inter_pred_chroma(
ref_l1.recon_v,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
mv_l1,
))
} else {
None
};
let mut pred_y = [0i32; 256];
let mut pred_u = [0i32; 64];
let mut pred_v = [0i32; 64];
match pred {
BPred16x16::L0 => {
pred_y = p_y_l0.unwrap();
pred_u = p_u_l0.unwrap();
pred_v = p_v_l0.unwrap();
}
BPred16x16::L1 => {
pred_y = p_y_l1.unwrap();
pred_u = p_u_l1.unwrap();
pred_v = p_v_l1.unwrap();
}
BPred16x16::Bi => {
let l0_y = p_y_l0.unwrap();
let l1_y = p_y_l1.unwrap();
if let Some(wp) = weighted {
let logwd = wp.log2_wd;
let off_avg = (wp.offset_l0 + wp.offset_l1 + 1) >> 1;
let round = 1i32 << logwd;
for i in 0..256 {
let v =
(l0_y[i] * wp.weight_l0 + l1_y[i] * wp.weight_l1 + round) >> (logwd + 1);
pred_y[i] = (v + off_avg).clamp(0, 255);
}
} else {
for i in 0..256 {
pred_y[i] = (l0_y[i] + l1_y[i] + 1) >> 1;
}
}
let l0_u = p_u_l0.unwrap();
let l1_u = p_u_l1.unwrap();
let l0_v = p_v_l0.unwrap();
let l1_v = p_v_l1.unwrap();
match weighted {
Some(wp) if wp.apply_chroma => {
for i in 0..64 {
pred_u[i] = weighted_bipred_merge(l0_u[i], l1_u[i], &wp);
pred_v[i] = weighted_bipred_merge(l0_v[i], l1_v[i], &wp);
}
}
_ => {
for i in 0..64 {
pred_u[i] = (l0_u[i] + l1_u[i] + 1) >> 1;
pred_v[i] = (l0_v[i] + l1_v[i] + 1) >> 1;
}
}
}
}
}
(pred_y, pred_u, pred_v)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DirectPart8x8 {
L0,
L1,
Bi,
None,
}
impl DirectPart8x8 {
fn uses_l0(self) -> bool {
matches!(self, Self::L0 | Self::Bi)
}
fn uses_l1(self) -> bool {
matches!(self, Self::L1 | Self::Bi)
}
}
fn direct_per_8x8_kinds(d: &BDirectDerivation) -> [DirectPart8x8; 4] {
let l0_used = d.ref_idx_l0 >= 0;
let l1_used = d.ref_idx_l1 >= 0;
let kind = match (l0_used, l1_used) {
(true, true) => DirectPart8x8::Bi,
(true, false) => DirectPart8x8::L0,
(false, true) => DirectPart8x8::L1,
(false, false) => DirectPart8x8::None,
};
[kind; 4]
}
#[allow(clippy::too_many_arguments)]
fn build_b_direct_8x8_predictors(
weighted: Option<WeightedBipredLuma>,
kinds: [DirectPart8x8; 4],
mv_l0_per_8x8: [Mv; 4],
mv_l1_per_8x8: [Mv; 4],
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
) -> ([i32; 256], [i32; 64], [i32; 64]) {
let mut pred_y = [0i32; 256];
let mut pred_u = [0i32; 64];
let mut pred_v = [0i32; 64];
let chroma_w_l0 = ref_l0.width / 2;
let chroma_h_l0 = ref_l0.height / 2;
let chroma_w_l1 = ref_l1.width / 2;
let chroma_h_l1 = ref_l1.height / 2;
for q in 0..4usize {
let sub_x = q % 2;
let sub_y = q / 2;
let kind = kinds[q];
let mv0 = mv_l0_per_8x8[q];
let mv1 = mv_l1_per_8x8[q];
let l0_y = if kind.uses_l0() {
Some(build_inter_pred_luma_8x8(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv0,
))
} else {
None
};
let l1_y = if kind.uses_l1() {
Some(build_inter_pred_luma_8x8(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv1,
))
} else {
None
};
let q_y: [i32; 64] = match kind {
DirectPart8x8::L0 => l0_y.unwrap(),
DirectPart8x8::L1 => l1_y.unwrap(),
DirectPart8x8::Bi => {
let l0 = l0_y.unwrap();
let l1 = l1_y.unwrap();
let mut out = [0i32; 64];
match weighted {
Some(wp) => {
for i in 0..64 {
out[i] = weighted_bipred_merge(l0[i], l1[i], &wp);
}
}
None => {
for i in 0..64 {
out[i] = (l0[i] + l1[i] + 1) >> 1;
}
}
}
out
}
DirectPart8x8::None => {
build_inter_pred_luma_8x8(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sub_x,
sub_y,
Mv::ZERO,
)
}
};
let l0_u = if kind.uses_l0() {
Some(build_inter_pred_chroma_4x4(
ref_l0.recon_u,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
sub_x,
sub_y,
mv0,
))
} else {
None
};
let l0_v = if kind.uses_l0() {
Some(build_inter_pred_chroma_4x4(
ref_l0.recon_v,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
sub_x,
sub_y,
mv0,
))
} else {
None
};
let l1_u = if kind.uses_l1() {
Some(build_inter_pred_chroma_4x4(
ref_l1.recon_u,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
sub_x,
sub_y,
mv1,
))
} else {
None
};
let l1_v = if kind.uses_l1() {
Some(build_inter_pred_chroma_4x4(
ref_l1.recon_v,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
sub_x,
sub_y,
mv1,
))
} else {
None
};
let merge_chroma = |c0: Option<[i32; 16]>, c1: Option<[i32; 16]>| -> [i32; 16] {
match kind {
DirectPart8x8::L0 => c0.unwrap(),
DirectPart8x8::L1 => c1.unwrap(),
DirectPart8x8::Bi => {
let a = c0.unwrap();
let b = c1.unwrap();
let mut out = [0i32; 16];
match weighted {
Some(wp) if wp.apply_chroma => {
for i in 0..16 {
out[i] = weighted_bipred_merge(a[i], b[i], &wp);
}
}
_ => {
for i in 0..16 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
}
}
out
}
DirectPart8x8::None => c0.unwrap_or([0; 16]),
}
};
let q_u = merge_chroma(l0_u, l1_u);
let q_v = merge_chroma(l0_v, l1_v);
for j in 0..8usize {
for i in 0..8usize {
let dst_y = (sub_y * 8 + j) * 16 + sub_x * 8 + i;
pred_y[dst_y] = q_y[j * 8 + i];
}
}
for j in 0..4usize {
for i in 0..4usize {
let dst_c = (sub_y * 4 + j) * 8 + sub_x * 4 + i;
pred_u[dst_c] = q_u[j * 4 + i];
pred_v[dst_c] = q_v[j * 4 + i];
}
}
}
(pred_y, pred_u, pred_v)
}
pub(crate) fn sad_8x8_pred(
src_y: &[u8],
src_stride: usize,
mb_x: usize,
mb_y: usize,
off_x: usize,
off_y: usize,
pred: &[i32; 64],
) -> u32 {
let mut s = 0u32;
for j in 0..8usize {
for i in 0..8usize {
let src = src_y[(mb_y * 16 + off_y + j) * src_stride + mb_x * 16 + off_x + i] as i32;
let p = pred[j * 8 + i].clamp(0, 255);
s = s.saturating_add((src - p).unsigned_abs());
}
}
s
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_b_explicit_8x8_cell_luma(
cell: BSubMbCell,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
sub_x: usize,
sub_y: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> [i32; 64] {
match cell {
BSubMbCell::L0 => build_inter_pred_luma_8x8(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l0,
),
BSubMbCell::L1 => build_inter_pred_luma_8x8(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l1,
),
BSubMbCell::Bi => {
let a = build_inter_pred_luma_8x8(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l0,
);
let b = build_inter_pred_luma_8x8(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l1,
);
let mut out = [0i32; 64];
for i in 0..64 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
out
}
BSubMbCell::Direct => {
build_inter_pred_luma_8x8(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sub_x,
sub_y,
Mv::ZERO,
)
}
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_b_explicit_4x4_cell_chroma(
cell: BSubMbCell,
ref_plane_l0: &[u8],
ref_plane_l1: &[u8],
chroma_w_l0: u32,
chroma_h_l0: u32,
chroma_w_l1: u32,
chroma_h_l1: u32,
mb_x: usize,
mb_y: usize,
sub_x: usize,
sub_y: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> [i32; 16] {
match cell {
BSubMbCell::L0 => build_inter_pred_chroma_4x4(
ref_plane_l0,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l0,
),
BSubMbCell::L1 => build_inter_pred_chroma_4x4(
ref_plane_l1,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l1,
),
BSubMbCell::Bi => {
let a = build_inter_pred_chroma_4x4(
ref_plane_l0,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l0,
);
let b = build_inter_pred_chroma_4x4(
ref_plane_l1,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
sub_x,
sub_y,
mv_l1,
);
let mut out = [0i32; 16];
for i in 0..16 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
out
}
BSubMbCell::Direct => build_inter_pred_chroma_4x4(
ref_plane_l0,
chroma_w_l0,
chroma_h_l0,
mb_x,
mb_y,
sub_x,
sub_y,
Mv::ZERO,
),
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn pick_best_b8x8_cell(
src_y: &[u8],
src_stride: usize,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
sub_x: usize,
sub_y: usize,
direct_pred: &[i32; 64],
me_l0: Mv,
me_l1: Mv,
qp_y: i32,
) -> (BSubMbCell, Mv, Mv, [i32; 64], u32) {
let off_x = sub_x * 8;
let off_y = sub_y * 8;
let sad_direct = sad_8x8_pred(src_y, src_stride, mb_x, mb_y, off_x, off_y, direct_pred);
let pred_l0 = build_b_explicit_8x8_cell_luma(
BSubMbCell::L0,
ref_l0,
ref_l1,
mb_x,
mb_y,
sub_x,
sub_y,
me_l0,
me_l1,
);
let sad_l0 = sad_8x8_pred(src_y, src_stride, mb_x, mb_y, off_x, off_y, &pred_l0);
let pred_l1 = build_b_explicit_8x8_cell_luma(
BSubMbCell::L1,
ref_l0,
ref_l1,
mb_x,
mb_y,
sub_x,
sub_y,
me_l0,
me_l1,
);
let sad_l1 = sad_8x8_pred(src_y, src_stride, mb_x, mb_y, off_x, off_y, &pred_l1);
let pred_bi = build_b_explicit_8x8_cell_luma(
BSubMbCell::Bi,
ref_l0,
ref_l1,
mb_x,
mb_y,
sub_x,
sub_y,
me_l0,
me_l1,
);
let sad_bi = sad_8x8_pred(src_y, src_stride, mb_x, mb_y, off_x, off_y, &pred_bi);
let lambda = (qp_y / 6).clamp(0, 8) as u32;
let cost_direct = 0u32;
let cost_l0 = (lambda + 1) * 6;
let cost_l1 = (lambda + 1) * 6;
let cost_bi = (lambda + 1) * 12;
let j_direct = sad_direct.saturating_add(cost_direct);
let j_l0 = sad_l0.saturating_add(cost_l0);
let j_l1 = sad_l1.saturating_add(cost_l1);
let j_bi = sad_bi.saturating_add(cost_bi);
let mut best = (
BSubMbCell::Direct,
Mv::ZERO,
Mv::ZERO,
*direct_pred,
j_direct,
sad_direct,
);
if j_l0 < best.4 {
best = (BSubMbCell::L0, me_l0, Mv::ZERO, pred_l0, j_l0, sad_l0);
}
if j_l1 < best.4 {
best = (BSubMbCell::L1, Mv::ZERO, me_l1, pred_l1, j_l1, sad_l1);
}
if j_bi < best.4 {
best = (BSubMbCell::Bi, me_l0, me_l1, pred_bi, j_bi, sad_bi);
}
(best.0, best.1, best.2, best.3, best.5)
}
fn sad_16x16_pred(
src_y: &[u8],
src_stride: usize,
mb_x: usize,
mb_y: usize,
pred: &[i32; 256],
) -> u32 {
let mut s = 0u32;
for j in 0..16usize {
for i in 0..16usize {
let src = src_y[(mb_y * 16 + j) * src_stride + mb_x * 16 + i] as i32;
let p = pred[j * 16 + i].clamp(0, 255);
s = s.saturating_add((src - p).unsigned_abs());
}
}
s
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_b_partition_pred_luma_16x8(
pred: BPartPred,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
partition_y_off: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> [i32; 128] {
let mut out = [0i32; 128];
let l0 = if pred.uses_l0() {
Some(build_inter_pred_luma_partition(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
0,
partition_y_off,
16,
8,
mv_l0,
))
} else {
None
};
let l1 = if pred.uses_l1() {
Some(build_inter_pred_luma_partition(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
0,
partition_y_off,
16,
8,
mv_l1,
))
} else {
None
};
match pred {
BPartPred::L0 => out.copy_from_slice(&l0.unwrap()),
BPartPred::L1 => out.copy_from_slice(&l1.unwrap()),
BPartPred::Bi => {
let a = l0.unwrap();
let b = l1.unwrap();
for i in 0..128 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
}
}
out
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_b_partition_pred_luma_8x16(
pred: BPartPred,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
partition_x_off: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> [i32; 128] {
let mut out = [0i32; 128];
let l0 = if pred.uses_l0() {
Some(build_inter_pred_luma_partition(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
partition_x_off,
0,
8,
16,
mv_l0,
))
} else {
None
};
let l1 = if pred.uses_l1() {
Some(build_inter_pred_luma_partition(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
partition_x_off,
0,
8,
16,
mv_l1,
))
} else {
None
};
match pred {
BPartPred::L0 => out.copy_from_slice(&l0.unwrap()),
BPartPred::L1 => out.copy_from_slice(&l1.unwrap()),
BPartPred::Bi => {
let a = l0.unwrap();
let b = l1.unwrap();
for i in 0..128 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
}
}
out
}
#[allow(clippy::too_many_arguments)]
fn build_inter_pred_luma_partition(
ref_y: &[u8],
ref_w: u32,
ref_h: u32,
mb_x: usize,
mb_y: usize,
off_x: usize,
off_y: usize,
w: u32,
h: u32,
mv: Mv,
) -> Vec<i32> {
let stride = ref_w as usize;
let h_us = ref_h as usize;
let ref_i32: Vec<i32> = ref_y
.iter()
.take(stride * h_us)
.map(|&v| v as i32)
.collect();
let mv_int_x = mv.x >> 2;
let mv_int_y = mv.y >> 2;
let x_frac = (mv.x & 3) as u8;
let y_frac = (mv.y & 3) as u8;
let int_x = (mb_x * 16 + off_x) as i32 + mv_int_x;
let int_y = (mb_y * 16 + off_y) as i32 + mv_int_y;
let mut dst = vec![0i32; (w * h) as usize];
interpolate_luma(
&ref_i32, stride, stride, h_us, int_x, int_y, x_frac, y_frac, w, h, 8, &mut dst, w as usize,
)
.expect("luma interpolation partition");
dst
}
#[allow(clippy::too_many_arguments)]
fn build_inter_pred_chroma_partition(
ref_c: &[u8],
ref_cw: u32,
ref_ch: u32,
mb_x: usize,
mb_y: usize,
off_x_luma: usize,
off_y_luma: usize,
w_chroma: u32,
h_chroma: u32,
mv: Mv,
) -> Vec<i32> {
let stride = ref_cw as usize;
let h_us = ref_ch as usize;
let ref_i32: Vec<i32> = ref_c
.iter()
.take(stride * h_us)
.map(|&v| v as i32)
.collect();
let int_x = (mb_x * 8 + off_x_luma / 2) as i32 + (mv.x >> 3);
let int_y = (mb_y * 8 + off_y_luma / 2) as i32 + (mv.y >> 3);
let x_frac = (mv.x & 7) as u8;
let y_frac = (mv.y & 7) as u8;
let mut dst = vec![0i32; (w_chroma * h_chroma) as usize];
interpolate_chroma(
&ref_i32,
stride,
stride,
h_us,
int_x,
int_y,
x_frac,
y_frac,
w_chroma,
h_chroma,
8,
&mut dst,
w_chroma as usize,
)
.expect("chroma interpolation partition");
dst
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_b_partition_pred_chroma_16x8(
pred: BPartPred,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
partition_y_off_luma: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> ([i32; 32], [i32; 32]) {
let cw_l0 = ref_l0.width / 2;
let ch_l0 = ref_l0.height / 2;
let cw_l1 = ref_l1.width / 2;
let ch_l1 = ref_l1.height / 2;
let make = |plane_l0: &[u8], plane_l1: &[u8]| -> [i32; 32] {
let mut out = [0i32; 32];
let a = if pred.uses_l0() {
Some(build_inter_pred_chroma_partition(
plane_l0,
cw_l0,
ch_l0,
mb_x,
mb_y,
0,
partition_y_off_luma,
8,
4,
mv_l0,
))
} else {
None
};
let b = if pred.uses_l1() {
Some(build_inter_pred_chroma_partition(
plane_l1,
cw_l1,
ch_l1,
mb_x,
mb_y,
0,
partition_y_off_luma,
8,
4,
mv_l1,
))
} else {
None
};
match pred {
BPartPred::L0 => out.copy_from_slice(&a.unwrap()),
BPartPred::L1 => out.copy_from_slice(&b.unwrap()),
BPartPred::Bi => {
let a = a.unwrap();
let b = b.unwrap();
for i in 0..32 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
}
}
out
};
(
make(ref_l0.recon_u, ref_l1.recon_u),
make(ref_l0.recon_v, ref_l1.recon_v),
)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_b_partition_pred_chroma_8x16(
pred: BPartPred,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
partition_x_off_luma: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> ([i32; 32], [i32; 32]) {
let cw_l0 = ref_l0.width / 2;
let ch_l0 = ref_l0.height / 2;
let cw_l1 = ref_l1.width / 2;
let ch_l1 = ref_l1.height / 2;
let make = |plane_l0: &[u8], plane_l1: &[u8]| -> [i32; 32] {
let mut out = [0i32; 32];
let a = if pred.uses_l0() {
Some(build_inter_pred_chroma_partition(
plane_l0,
cw_l0,
ch_l0,
mb_x,
mb_y,
partition_x_off_luma,
0,
4,
8,
mv_l0,
))
} else {
None
};
let b = if pred.uses_l1() {
Some(build_inter_pred_chroma_partition(
plane_l1,
cw_l1,
ch_l1,
mb_x,
mb_y,
partition_x_off_luma,
0,
4,
8,
mv_l1,
))
} else {
None
};
match pred {
BPartPred::L0 => out.copy_from_slice(&a.unwrap()),
BPartPred::L1 => out.copy_from_slice(&b.unwrap()),
BPartPred::Bi => {
let a = a.unwrap();
let b = b.unwrap();
for i in 0..32 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
}
}
out
};
(
make(ref_l0.recon_u, ref_l1.recon_u),
make(ref_l0.recon_v, ref_l1.recon_v),
)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn sad_partition_pred(
src_y: &[u8],
src_stride: usize,
mb_x: usize,
mb_y: usize,
off_x: usize,
off_y: usize,
w: usize,
h: usize,
pred: &[i32],
) -> u32 {
let mut s = 0u32;
for j in 0..h {
for i in 0..w {
let src = src_y[(mb_y * 16 + off_y + j) * src_stride + mb_x * 16 + off_x + i] as i32;
let p = pred[j * w + i].clamp(0, 255);
s = s.saturating_add((src - p).unsigned_abs());
}
}
s
}
#[allow(clippy::too_many_arguments)]
fn mvp_for_b_16x8_partition(
grid: &MvGrid,
mb_x: usize,
mb_y: usize,
is_top: bool,
ref_idx: i32,
inflight_avail: [bool; 4],
mv_inflight: [Mv; 4],
ref_inflight: [i32; 4],
) -> Mv {
let by = if is_top { 0i32 } else { 2 };
let a = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
-1,
by,
inflight_avail,
mv_inflight,
ref_inflight,
);
let b = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
0,
by - 1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let c = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
4,
by - 1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let d = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
-1,
by - 1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let shape = if is_top {
crate::mv_deriv::MvpredShape::Partition16x8Top
} else {
crate::mv_deriv::MvpredShape::Partition16x8Bottom
};
let inputs = MvpredInputs::from_abc(a, b, c, ref_idx, shape);
crate::mv_deriv::derive_mvpred_with_d(&inputs, d)
}
#[allow(clippy::too_many_arguments)]
fn mvp_for_b_8x16_partition(
grid: &MvGrid,
mb_x: usize,
mb_y: usize,
is_left: bool,
ref_idx: i32,
inflight_avail: [bool; 4],
mv_inflight: [Mv; 4],
ref_inflight: [i32; 4],
) -> Mv {
let bx = if is_left { 0i32 } else { 2 };
let a = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx - 1,
0,
inflight_avail,
mv_inflight,
ref_inflight,
);
let b = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx,
-1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let c = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx + 2,
-1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let d = neighbour_8x8_partition_mv(
grid,
mb_x,
mb_y,
bx - 1,
-1,
inflight_avail,
mv_inflight,
ref_inflight,
);
let shape = if is_left {
crate::mv_deriv::MvpredShape::Partition8x16Left
} else {
crate::mv_deriv::MvpredShape::Partition8x16Right
};
let inputs = MvpredInputs::from_abc(a, b, c, ref_idx, shape);
crate::mv_deriv::derive_mvpred_with_d(&inputs, d)
}
#[allow(clippy::too_many_arguments)]
fn best_partition_mv(
src_y: &[u8],
src_stride: usize,
src_w: u32,
src_h: u32,
ref_y: &[u8],
ref_stride: usize,
ref_w: u32,
ref_h: u32,
mb_x: usize,
mb_y: usize,
off_x: usize,
off_y: usize,
part_w: u32,
part_h: u32,
candidates: &[Mv],
) -> (Mv, u32) {
let mut best_mv = candidates[0];
let mut best_sad = u32::MAX;
for &mv in candidates {
let pred = build_inter_pred_luma_partition(
ref_y, ref_w, ref_h, mb_x, mb_y, off_x, off_y, part_w, part_h, mv,
);
let sad = sad_partition_pred(
src_y,
src_stride,
mb_x,
mb_y,
off_x,
off_y,
part_w as usize,
part_h as usize,
&pred,
);
let _ = (src_w, src_h, ref_stride);
if sad < best_sad {
best_sad = sad;
best_mv = mv;
}
}
(best_mv, best_sad)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn partition_sad_b_16x8(
pred: BPartPred,
src_y: &[u8],
src_stride: usize,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
off_y: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> u32 {
let p = build_b_partition_pred_luma_16x8(pred, ref_l0, ref_l1, mb_x, mb_y, off_y, mv_l0, mv_l1);
sad_partition_pred(src_y, src_stride, mb_x, mb_y, 0, off_y, 16, 8, &p)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn partition_sad_b_8x16(
pred: BPartPred,
src_y: &[u8],
src_stride: usize,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
off_x: usize,
mv_l0: Mv,
mv_l1: Mv,
) -> u32 {
let p = build_b_partition_pred_luma_8x16(pred, ref_l0, ref_l1, mb_x, mb_y, off_x, mv_l0, mv_l1);
sad_partition_pred(src_y, src_stride, mb_x, mb_y, off_x, 0, 8, 16, &p)
}
pub(crate) fn best_partition_mode_b(
candidates_l0: &[Mv],
candidates_l1: &[Mv],
eval_sad: impl Fn(BPartPred, Mv, Mv) -> u32,
) -> (BPartPred, Mv, Mv, u32) {
let mut best = (BPartPred::L0, Mv::ZERO, Mv::ZERO, u32::MAX);
for &mv0 in candidates_l0 {
let s = eval_sad(BPartPred::L0, mv0, Mv::ZERO);
if s < best.3 {
best = (BPartPred::L0, mv0, Mv::ZERO, s);
}
}
for &mv1 in candidates_l1 {
let s = eval_sad(BPartPred::L1, Mv::ZERO, mv1);
if s < best.3 {
best = (BPartPred::L1, Mv::ZERO, mv1, s);
}
}
for &mv0 in candidates_l0 {
for &mv1 in candidates_l1 {
let s = eval_sad(BPartPred::Bi, mv0, mv1);
if s < best.3 {
best = (BPartPred::Bi, mv0, mv1, s);
}
}
}
best
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PartitionShape {
P16x8,
P8x16,
}
#[derive(Debug, Clone, Copy)]
struct PartitionChoiceB {
shape: PartitionShape,
mode_a: BPartPred,
mode_b: BPartPred,
mv_l0_a: Mv,
mv_l1_a: Mv,
mv_l0_b: Mv,
mv_l1_b: Mv,
}
fn build_partition_composite_b(
pc: &PartitionChoiceB,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
) -> ([i32; 256], [i32; 64], [i32; 64]) {
let mut pred_y = [0i32; 256];
let mut pred_u = [0i32; 64];
let mut pred_v = [0i32; 64];
match pc.shape {
PartitionShape::P16x8 => {
let py0 = build_b_partition_pred_luma_16x8(
pc.mode_a, ref_l0, ref_l1, mb_x, mb_y, 0, pc.mv_l0_a, pc.mv_l1_a,
);
let py1 = build_b_partition_pred_luma_16x8(
pc.mode_b, ref_l0, ref_l1, mb_x, mb_y, 8, pc.mv_l0_b, pc.mv_l1_b,
);
for j in 0..8usize {
for i in 0..16usize {
pred_y[j * 16 + i] = py0[j * 16 + i];
pred_y[(8 + j) * 16 + i] = py1[j * 16 + i];
}
}
let (cb0, cr0) = build_b_partition_pred_chroma_16x8(
pc.mode_a, ref_l0, ref_l1, mb_x, mb_y, 0, pc.mv_l0_a, pc.mv_l1_a,
);
let (cb1, cr1) = build_b_partition_pred_chroma_16x8(
pc.mode_b, ref_l0, ref_l1, mb_x, mb_y, 8, pc.mv_l0_b, pc.mv_l1_b,
);
for j in 0..4usize {
for i in 0..8usize {
pred_u[j * 8 + i] = cb0[j * 8 + i];
pred_v[j * 8 + i] = cr0[j * 8 + i];
pred_u[(4 + j) * 8 + i] = cb1[j * 8 + i];
pred_v[(4 + j) * 8 + i] = cr1[j * 8 + i];
}
}
}
PartitionShape::P8x16 => {
let py0 = build_b_partition_pred_luma_8x16(
pc.mode_a, ref_l0, ref_l1, mb_x, mb_y, 0, pc.mv_l0_a, pc.mv_l1_a,
);
let py1 = build_b_partition_pred_luma_8x16(
pc.mode_b, ref_l0, ref_l1, mb_x, mb_y, 8, pc.mv_l0_b, pc.mv_l1_b,
);
for j in 0..16usize {
for i in 0..8usize {
pred_y[j * 16 + i] = py0[j * 8 + i];
pred_y[j * 16 + 8 + i] = py1[j * 8 + i];
}
}
let (cb0, cr0) = build_b_partition_pred_chroma_8x16(
pc.mode_a, ref_l0, ref_l1, mb_x, mb_y, 0, pc.mv_l0_a, pc.mv_l1_a,
);
let (cb1, cr1) = build_b_partition_pred_chroma_8x16(
pc.mode_b, ref_l0, ref_l1, mb_x, mb_y, 8, pc.mv_l0_b, pc.mv_l1_b,
);
for j in 0..8usize {
for i in 0..4usize {
pred_u[j * 8 + i] = cb0[j * 4 + i];
pred_v[j * 8 + i] = cr0[j * 4 + i];
pred_u[j * 8 + 4 + i] = cb1[j * 4 + i];
pred_v[j * 8 + 4 + i] = cr1[j * 4 + i];
}
}
}
}
(pred_y, pred_u, pred_v)
}
type PartitionMvds = [(i32, i32); 2];
fn partition_mvds_b(
pc: &PartitionChoiceB,
mv_grid_l0: &MvGrid,
mv_grid_l1: &MvGrid,
mb_x: usize,
mb_y: usize,
) -> (PartitionMvds, PartitionMvds) {
let mut mvd_l0 = [(0i32, 0i32); 2];
let mut mvd_l1 = [(0i32, 0i32); 2];
let inflight_avail = [false; 4];
let mv_inflight = [Mv::ZERO; 4];
let ref_inflight = [-1i32; 4];
match pc.shape {
PartitionShape::P16x8 => {
if pc.mode_a.uses_l0() {
let mvp = mvp_for_b_16x8_partition(
mv_grid_l0,
mb_x,
mb_y,
true,
0,
inflight_avail,
mv_inflight,
ref_inflight,
);
mvd_l0[0] = (pc.mv_l0_a.x - mvp.x, pc.mv_l0_a.y - mvp.y);
}
if pc.mode_a.uses_l1() {
let mvp = mvp_for_b_16x8_partition(
mv_grid_l1,
mb_x,
mb_y,
true,
0,
inflight_avail,
mv_inflight,
ref_inflight,
);
mvd_l1[0] = (pc.mv_l1_a.x - mvp.x, pc.mv_l1_a.y - mvp.y);
}
let inflight_avail_bot = [true, true, false, false];
let mv_inflight_bot_l0 = [pc.mv_l0_a, pc.mv_l0_a, Mv::ZERO, Mv::ZERO];
let mv_inflight_bot_l1 = [pc.mv_l1_a, pc.mv_l1_a, Mv::ZERO, Mv::ZERO];
let ref_inflight_bot_l0 = [
if pc.mode_a.uses_l0() { 0 } else { -1 },
if pc.mode_a.uses_l0() { 0 } else { -1 },
-1,
-1,
];
let ref_inflight_bot_l1 = [
if pc.mode_a.uses_l1() { 0 } else { -1 },
if pc.mode_a.uses_l1() { 0 } else { -1 },
-1,
-1,
];
if pc.mode_b.uses_l0() {
let mvp = mvp_for_b_16x8_partition(
mv_grid_l0,
mb_x,
mb_y,
false,
0,
inflight_avail_bot,
mv_inflight_bot_l0,
ref_inflight_bot_l0,
);
mvd_l0[1] = (pc.mv_l0_b.x - mvp.x, pc.mv_l0_b.y - mvp.y);
}
if pc.mode_b.uses_l1() {
let mvp = mvp_for_b_16x8_partition(
mv_grid_l1,
mb_x,
mb_y,
false,
0,
inflight_avail_bot,
mv_inflight_bot_l1,
ref_inflight_bot_l1,
);
mvd_l1[1] = (pc.mv_l1_b.x - mvp.x, pc.mv_l1_b.y - mvp.y);
}
}
PartitionShape::P8x16 => {
if pc.mode_a.uses_l0() {
let mvp = mvp_for_b_8x16_partition(
mv_grid_l0,
mb_x,
mb_y,
true,
0,
inflight_avail,
mv_inflight,
ref_inflight,
);
mvd_l0[0] = (pc.mv_l0_a.x - mvp.x, pc.mv_l0_a.y - mvp.y);
}
if pc.mode_a.uses_l1() {
let mvp = mvp_for_b_8x16_partition(
mv_grid_l1,
mb_x,
mb_y,
true,
0,
inflight_avail,
mv_inflight,
ref_inflight,
);
mvd_l1[0] = (pc.mv_l1_a.x - mvp.x, pc.mv_l1_a.y - mvp.y);
}
let inflight_avail_right = [true, false, true, false];
let mv_inflight_right_l0 = [pc.mv_l0_a, Mv::ZERO, pc.mv_l0_a, Mv::ZERO];
let mv_inflight_right_l1 = [pc.mv_l1_a, Mv::ZERO, pc.mv_l1_a, Mv::ZERO];
let ref_inflight_right_l0 = [
if pc.mode_a.uses_l0() { 0 } else { -1 },
-1,
if pc.mode_a.uses_l0() { 0 } else { -1 },
-1,
];
let ref_inflight_right_l1 = [
if pc.mode_a.uses_l1() { 0 } else { -1 },
-1,
if pc.mode_a.uses_l1() { 0 } else { -1 },
-1,
];
if pc.mode_b.uses_l0() {
let mvp = mvp_for_b_8x16_partition(
mv_grid_l0,
mb_x,
mb_y,
false,
0,
inflight_avail_right,
mv_inflight_right_l0,
ref_inflight_right_l0,
);
mvd_l0[1] = (pc.mv_l0_b.x - mvp.x, pc.mv_l0_b.y - mvp.y);
}
if pc.mode_b.uses_l1() {
let mvp = mvp_for_b_8x16_partition(
mv_grid_l1,
mb_x,
mb_y,
false,
0,
inflight_avail_right,
mv_inflight_right_l1,
ref_inflight_right_l1,
);
mvd_l1[1] = (pc.mv_l1_b.x - mvp.x, pc.mv_l1_b.y - mvp.y);
}
}
}
(mvd_l0, mvd_l1)
}
fn grid_state_b(
pc: Option<&PartitionChoiceB>,
is_l0: bool,
pred: BPred16x16,
eff_mv: Mv,
) -> ([Mv; 4], [i32; 4], bool) {
if let Some(pc) = pc {
let (mv_a, mv_b, used_a, used_b) = if is_l0 {
(
pc.mv_l0_a,
pc.mv_l0_b,
pc.mode_a.uses_l0(),
pc.mode_b.uses_l0(),
)
} else {
(
pc.mv_l1_a,
pc.mv_l1_b,
pc.mode_a.uses_l1(),
pc.mode_b.uses_l1(),
)
};
let any_used = used_a || used_b;
let mut mvs = [Mv::ZERO; 4];
let mut refs = [-1i32; 4];
match pc.shape {
PartitionShape::P16x8 => {
if used_a {
mvs[0] = mv_a;
mvs[1] = mv_a;
refs[0] = 0;
refs[1] = 0;
}
if used_b {
mvs[2] = mv_b;
mvs[3] = mv_b;
refs[2] = 0;
refs[3] = 0;
}
}
PartitionShape::P8x16 => {
if used_a {
mvs[0] = mv_a;
mvs[2] = mv_a;
refs[0] = 0;
refs[2] = 0;
}
if used_b {
mvs[1] = mv_b;
mvs[3] = mv_b;
refs[1] = 0;
refs[3] = 0;
}
}
}
(mvs, refs, any_used)
} else {
let used = if is_l0 {
pred.uses_l0()
} else {
pred.uses_l1()
};
if used {
([eff_mv; 4], [0; 4], true)
} else {
([Mv::ZERO; 4], [-1; 4], false)
}
}
}
fn deblock_mv_arr_b(
pc: Option<&PartitionChoiceB>,
is_l0: bool,
pred: BPred16x16,
eff_mv: Mv,
poc_sentinel: i32,
) -> ([(i16, i16); 16], [i8; 4], [i32; 4]) {
if let Some(pc) = pc {
let (mv_a, mv_b, used_a, used_b) = if is_l0 {
(
pc.mv_l0_a,
pc.mv_l0_b,
pc.mode_a.uses_l0(),
pc.mode_b.uses_l0(),
)
} else {
(
pc.mv_l1_a,
pc.mv_l1_b,
pc.mode_a.uses_l1(),
pc.mode_b.uses_l1(),
)
};
let mv_a_t = (
mv_a.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv_a.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
let mv_b_t = (
mv_b.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv_b.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
let mut mvs = [(0i16, 0i16); 16];
let mut refs = [-1i8; 4];
let mut pocs = [i32::MIN; 4];
match pc.shape {
PartitionShape::P16x8 => {
for (blk, slot) in mvs.iter_mut().enumerate() {
let (_bx, by) = LUMA_4X4_BLK[blk];
*slot = if by < 2 {
if used_a {
mv_a_t
} else {
(0, 0)
}
} else if used_b {
mv_b_t
} else {
(0, 0)
};
}
refs[0] = if used_a { 0 } else { -1 };
refs[1] = if used_a { 0 } else { -1 };
refs[2] = if used_b { 0 } else { -1 };
refs[3] = if used_b { 0 } else { -1 };
pocs[0] = if used_a { poc_sentinel } else { i32::MIN };
pocs[1] = if used_a { poc_sentinel } else { i32::MIN };
pocs[2] = if used_b { poc_sentinel } else { i32::MIN };
pocs[3] = if used_b { poc_sentinel } else { i32::MIN };
}
PartitionShape::P8x16 => {
for (blk, slot) in mvs.iter_mut().enumerate() {
let (bx, _by) = LUMA_4X4_BLK[blk];
*slot = if bx < 2 {
if used_a {
mv_a_t
} else {
(0, 0)
}
} else if used_b {
mv_b_t
} else {
(0, 0)
};
}
refs[0] = if used_a { 0 } else { -1 };
refs[2] = if used_a { 0 } else { -1 };
refs[1] = if used_b { 0 } else { -1 };
refs[3] = if used_b { 0 } else { -1 };
pocs[0] = if used_a { poc_sentinel } else { i32::MIN };
pocs[2] = if used_a { poc_sentinel } else { i32::MIN };
pocs[1] = if used_b { poc_sentinel } else { i32::MIN };
pocs[3] = if used_b { poc_sentinel } else { i32::MIN };
}
}
(mvs, refs, pocs)
} else {
let used = if is_l0 {
pred.uses_l0()
} else {
pred.uses_l1()
};
if used {
let mv_t = (
eff_mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
eff_mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
([mv_t; 16], [0; 4], [poc_sentinel; 4])
} else {
([(0, 0); 16], [-1; 4], [i32::MIN; 4])
}
}
}
fn deblock_arrays_per_8x8_direct(
mv_per_8x8: [Mv; 4],
ref_idx: i32,
poc_sentinel: i32,
) -> ([(i16, i16); 16], [i8; 4], [i32; 4]) {
let used = ref_idx >= 0;
if !used {
return ([(0, 0); 16], [-1; 4], [i32::MIN; 4]);
}
let mut mvs = [(0i16, 0i16); 16];
for (blk, slot) in mvs.iter_mut().enumerate() {
let (bx, by) = LUMA_4X4_BLK[blk];
let q = (by / 2) * 2 + (bx / 2);
let mv = mv_per_8x8[q];
*slot = (
mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
}
let refs = [ref_idx.clamp(i8::MIN as i32, i8::MAX as i32) as i8; 4];
let pocs = [poc_sentinel; 4];
(mvs, refs, pocs)
}
fn deblock_arrays_per_8x8_mixed(
mv_per_8x8: [Mv; 4],
ref_idx_per_8x8: [i32; 4],
poc_sentinel: i32,
) -> ([(i16, i16); 16], [i8; 4], [i32; 4]) {
let mut mvs = [(0i16, 0i16); 16];
for (blk, slot) in mvs.iter_mut().enumerate() {
let (bx, by) = LUMA_4X4_BLK[blk];
let q = (by / 2) * 2 + (bx / 2);
if ref_idx_per_8x8[q] >= 0 {
let mv = mv_per_8x8[q];
*slot = (
mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
);
} else {
*slot = (0, 0);
}
}
let mut refs = [-1i8; 4];
let mut pocs = [i32::MIN; 4];
for q in 0..4usize {
if ref_idx_per_8x8[q] >= 0 {
refs[q] = ref_idx_per_8x8[q].clamp(i8::MIN as i32, i8::MAX as i32) as i8;
pocs[q] = poc_sentinel;
}
}
(mvs, refs, pocs)
}
impl Encoder {
pub fn encode_b(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
) -> EncodedB {
self.encode_b_with_qp(
frame,
ref_l0,
ref_l1,
frame_num,
pic_order_cnt_lsb,
self.cfg.qp,
)
}
pub fn encode_b_with_qp(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
frame_qp: i32,
) -> EncodedB {
self.encode_b_impl(
frame,
ref_l0,
ref_l1,
frame_num,
pic_order_cnt_lsb,
frame_qp,
None,
)
}
#[allow(clippy::too_many_arguments)]
pub fn encode_b_rate_adaptive(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
frame_qp: i32,
frame_budget_bits: u64,
) -> EncodedB {
assert_eq!(
self.cfg.chroma_format_idc, 1,
"encode_b_rate_adaptive requires 4:2:0"
);
self.encode_b_impl(
frame,
ref_l0,
ref_l1,
frame_num,
pic_order_cnt_lsb,
frame_qp,
Some(frame_budget_bits),
)
}
#[allow(clippy::too_many_arguments)]
fn encode_b_impl(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
frame_num: u32,
pic_order_cnt_lsb: u32,
frame_qp: i32,
row_budget_bits: Option<u64>,
) -> EncodedB {
assert!(
(0..=51).contains(&frame_qp),
"frame_qp {frame_qp} out of 0..=51"
);
assert_eq!(frame.width, self.cfg.width);
assert_eq!(frame.height, self.cfg.height);
assert_eq!(ref_l0.width, self.cfg.width);
assert_eq!(ref_l0.height, self.cfg.height);
assert_eq!(ref_l1.width, self.cfg.width);
assert_eq!(ref_l1.height, self.cfg.height);
assert!(
self.cfg.profile_idc >= 77,
"B-slices require Main profile or higher (§A.2.2 — Baseline forbids B)"
);
let width = self.cfg.width as usize;
let width_mbs = self.cfg.width / 16;
let height_mbs = self.cfg.height / 16;
let chroma_width = match self.cfg.chroma_format_idc {
0 => 0usize,
3 => self.cfg.width as usize,
_ => (self.cfg.width / 2) as usize,
};
let chroma_height = match self.cfg.chroma_format_idc {
0 => 0usize,
2 | 3 => self.cfg.height as usize,
_ => (self.cfg.height / 2) as usize,
};
let mut stream: Vec<u8> = Vec::new();
let log2_max_frame_num_minus4: u32 = 4;
let log2_max_poc_lsb_minus4: u32 = 4;
let chroma_qp_index_offset: i32 = 0;
let direct_spatial = !self.cfg.direct_temporal_mv_pred;
let weighted_table = if self.cfg.explicit_weighted_bipred {
Some(compute_explicit_bipred_weights(
frame.y,
ref_l0.recon_y,
ref_l1.recon_y,
self.cfg.width as usize,
self.cfg.height as usize,
ref_l0.width as usize,
ref_l1.width as usize,
))
} else {
None
};
let weighted_luma_for_mb: Option<WeightedBipredLuma> =
weighted_table.map(|t| WeightedBipredLuma {
log2_wd: t.luma_log2_weight_denom,
weight_l0: t.luma_weight_l0,
offset_l0: t.luma_offset_l0,
weight_l1: t.luma_weight_l1,
offset_l1: t.luma_offset_l1,
apply_chroma: false,
});
let mut sw = BitWriter::new();
write_b_slice_header(
&mut sw,
&BSliceHeaderConfig {
first_mb_in_slice: 0,
slice_type_raw: 6, pic_parameter_set_id: 0,
colour_plane_id: self.cfg.colour_plane_id,
frame_num,
frame_num_bits: log2_max_frame_num_minus4 + 4,
pic_order_cnt_lsb,
poc_lsb_bits: log2_max_poc_lsb_minus4 + 4,
direct_spatial_mv_pred_flag: direct_spatial,
slice_qp_delta: frame_qp - self.cfg.qp,
disable_deblocking_filter_idc: if self.cfg.chroma_format_idc == 3 {
1
} else {
0
},
slice_alpha_c0_offset_div2: 0,
slice_beta_offset_div2: 0,
nal_ref_idc: 0, pred_weight_table: weighted_table,
cabac: None,
field: FieldPicSignal::FrameMbsOnly,
},
);
let mut recon_y = vec![0u8; (self.cfg.width * self.cfg.height) as usize];
let mut recon_u = vec![0u8; chroma_width * chroma_height];
let mut recon_v = vec![0u8; chroma_width * chroma_height];
let qp_bd_offset_c = qp_bd_offset(self.cfg.bit_depth_chroma_minus8);
let mut nc_grid = CavlcNcGrid::new(width_mbs, height_mbs);
let mut intra_grid = IntraGrid::new(width_mbs as usize, height_mbs as usize);
let mut mv_grid_l0 = MvGrid::new(width_mbs as usize, height_mbs as usize);
let mut mv_grid_l1 = MvGrid::new(width_mbs as usize, height_mbs as usize);
let mut mb_deblock_infos: Vec<MbDeblockInfo> =
vec![MbDeblockInfo::default(); (width_mbs * height_mbs) as usize];
let mut pending_skip: u32 = 0;
let mut i8x8_mb_count = 0u32;
let mut qp_tracker = MbQpTracker::new(frame_qp);
let mut want_qp = frame_qp;
for mb_y in 0..height_mbs as usize {
if let Some(budget) = row_budget_bits {
if mb_y > 0 && budget > 0 {
let frac = mb_y as f64 / height_mbs as f64;
let spent = sw.bits_emitted() as f64;
let expected = (budget as f64 * frac).max(1.0);
let ratio = spent.max(1.0) / expected;
let adj = if ratio >= 1.0 {
(6.0 * ratio.log2()).round() as i32
} else {
(3.0 * ratio.log2()).round() as i32
};
let mut new_want = frame_qp + adj;
if new_want < frame_qp && frac < 0.6 {
new_want = frame_qp;
}
want_qp = new_want
.clamp(want_qp - 1, want_qp + 2)
.clamp(frame_qp - 2, frame_qp + 6)
.clamp(0, 51);
}
}
let qp_y = want_qp;
let qp_c = qp_y_to_qp_c_with_bd_offset(qp_y, chroma_qp_index_offset, qp_bd_offset_c);
for mb_x in 0..width_mbs as usize {
let mb_addr = mb_y * (width_mbs as usize) + mb_x;
let dbl = self.encode_b_mb_with_intra_fallback(
frame,
ref_l0,
ref_l1,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
width,
&mut recon_y,
&mut recon_u,
&mut recon_v,
&mut sw,
&mut nc_grid,
&mut intra_grid,
&mut mv_grid_l0,
&mut mv_grid_l1,
&mut pending_skip,
pic_order_cnt_lsb as i32,
weighted_luma_for_mb,
&mut qp_tracker,
);
if dbl.transform_size_8x8_flag {
i8x8_mb_count += 1;
}
mb_deblock_infos[mb_addr] = dbl;
}
}
sw.ue(pending_skip);
sw.rbsp_trailing_bits();
let slice_rbsp = sw.into_bytes();
stream.extend_from_slice(&build_nal_unit(0, NalUnitType::SliceNonIdr, &slice_rbsp));
if self.cfg.chroma_format_idc != 3 {
crate::encoder::deblock::deblock_recon_with_chroma_array_type(
self.cfg.width,
self.cfg.height,
chroma_width as u32,
chroma_height as u32,
&mut recon_y,
&mut recon_u,
&mut recon_v,
&mb_deblock_infos,
chroma_qp_index_offset,
width_mbs,
height_mbs,
self.cfg.chroma_format_idc,
);
}
let n_parts = (width_mbs as usize) * (height_mbs as usize) * 4;
let mut partition_mvs = Vec::with_capacity(n_parts);
for slot in &mv_grid_l0.slots {
for part in 0..4usize {
let (mv, ref_idx, is_intra) = if !slot.available || slot.is_intra {
(Mv::ZERO, -1, true)
} else {
(slot.mv_l0_8x8[part], slot.ref_idx_l0_8x8[part], false)
};
partition_mvs.push(FrameRefPartitionMv {
mv_l0: (
mv.x.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
mv.y.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
),
ref_idx_l0: ref_idx.clamp(-1, 127) as i8,
is_intra,
});
}
}
EncodedB {
annex_b: stream,
recon_y,
recon_u,
recon_v,
recon_width: self.cfg.width,
recon_height: self.cfg.height,
frame_num,
pic_order_cnt_lsb,
partition_mvs,
i8x8_mb_count,
}
}
#[allow(clippy::too_many_arguments)]
fn encode_b_mb(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
_chroma_height: usize,
src_stride: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid_l0: &mut MvGrid,
mv_grid_l1: &mut MvGrid,
pending_skip: &mut u32,
curr_poc: i32,
weighted: Option<WeightedBipredLuma>,
qp_tracker: &mut MbQpTracker,
) -> MbDeblockInfo {
if self.cfg.chroma_format_idc != 1 {
debug_assert_eq!(
qp_tracker.cur, qp_y,
"row modulation (non-zero mb_qp_delta) is 4:2:0-only"
);
return self.encode_b_mb_multifmt(
frame,
ref_l0,
ref_l1,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
_chroma_height,
src_stride,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
pending_skip,
curr_poc,
weighted,
);
}
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let me_l0 = search_quarter_pel_16x16(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l0.recon_y,
ref_l0.width as usize,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
16,
16,
);
let me_l1 = search_quarter_pel_16x16(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l1.recon_y,
ref_l1.width as usize,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
16,
16,
);
let mv_l0 = Mv::new(me_l0.mv_x, me_l0.mv_y);
let mv_l1 = Mv::new(me_l1.mv_x, me_l1.mv_y);
let (pred_y_l0, pred_u_l0, pred_v_l0) = build_b_predictors(
BPred16x16::L0,
ref_l0,
ref_l1,
mb_x,
mb_y,
mv_l0,
mv_l1,
weighted,
);
let (pred_y_l1, pred_u_l1, pred_v_l1) = build_b_predictors(
BPred16x16::L1,
ref_l0,
ref_l1,
mb_x,
mb_y,
mv_l0,
mv_l1,
weighted,
);
let (pred_y_bi, pred_u_bi, pred_v_bi) = build_b_predictors(
BPred16x16::Bi,
ref_l0,
ref_l1,
mb_x,
mb_y,
mv_l0,
mv_l1,
weighted,
);
let sad_l0 = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &pred_y_l0);
let sad_l1 = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &pred_y_l1);
let sad_bi = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &pred_y_bi);
let direct = if self.cfg.direct_temporal_mv_pred {
b_temporal_direct_derive(
ref_l1.partition_mvs,
width_mbs,
mb_x,
mb_y,
curr_poc,
ref_l0.pic_order_cnt,
ref_l1.pic_order_cnt,
)
} else {
b_spatial_direct_derive(
mv_grid_l0,
mv_grid_l1,
ref_l1.partition_mvs,
width_mbs,
mb_x,
mb_y,
)
};
let direct_uniform = direct.is_uniform();
let direct_l0_ok = !(self.cfg.b_l0_bi_only && direct.ref_idx_l0 < 0);
let direct_pred_kind: Option<BPred16x16> =
match (direct.ref_idx_l0 >= 0, direct.ref_idx_l1 >= 0) {
(true, true) => Some(BPred16x16::Bi),
(true, false) => Some(BPred16x16::L0),
(false, true) => Some(BPred16x16::L1),
(false, false) => None,
};
let (direct_competitive, direct_pred_y, direct_pred_u, direct_pred_v, direct_sad) =
match (direct_uniform && direct_l0_ok, direct_pred_kind) {
(true, Some(dpk)) => {
let (dy, du, dv) = build_b_predictors(
dpk,
ref_l0,
ref_l1,
mb_x,
mb_y,
direct.mv_l0(),
direct.mv_l1(),
weighted,
);
let dsad = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &dy);
(true, dy, du, dv, dsad)
}
_ => (false, [0i32; 256], [0i32; 64], [0i32; 64], u32::MAX),
};
let direct_8x8_kinds = direct_per_8x8_kinds(&direct);
let (
direct_8x8_competitive,
direct_8x8_pred_y,
direct_8x8_pred_u,
direct_8x8_pred_v,
direct_8x8_sad,
) = match direct_pred_kind.filter(|_| direct_l0_ok) {
Some(_) => {
let (dy, du, dv) = build_b_direct_8x8_predictors(
weighted,
direct_8x8_kinds,
direct.mv_l0_per_8x8,
direct.mv_l1_per_8x8,
ref_l0,
ref_l1,
mb_x,
mb_y,
);
let dsad = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &dy);
(true, dy, du, dv, dsad)
}
None => (false, [0i32; 256], [0i32; 64], [0i32; 64], u32::MAX),
};
let mut mixed_me_l0_cells: [Mv; 4] = [Mv::ZERO; 4];
let mut mixed_me_l1_cells: [Mv; 4] = [Mv::ZERO; 4];
for cell in 0..4usize {
let sx = cell % 2;
let sy = cell / 2;
let r0 = search_quarter_pel_8x8(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l0.recon_y,
ref_l0.width as usize,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sx,
sy,
4,
4,
);
let r1 = search_quarter_pel_8x8(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l1.recon_y,
ref_l1.width as usize,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
sx,
sy,
4,
4,
);
mixed_me_l0_cells[cell] = Mv::new(r0.mv_x, r0.mv_y);
mixed_me_l1_cells[cell] = Mv::new(r1.mv_x, r1.mv_y);
}
let mut mixed_cells = [BSubMbCell::Direct; 4];
let mut mixed_cell_mv_l0 = [Mv::ZERO; 4];
let mut mixed_cell_mv_l1 = [Mv::ZERO; 4];
let mut mixed_cell_pred_y_per_cell = [[0i32; 64]; 4];
let mut mixed_total_sad = 0u32;
let mixed_competitive = direct_pred_kind.is_some();
if mixed_competitive {
for cell in 0..4usize {
let sub_x = cell % 2;
let sub_y = cell / 2;
let mut direct_cell_pred = [0i32; 64];
for j in 0..8usize {
for i in 0..8usize {
direct_cell_pred[j * 8 + i] =
direct_8x8_pred_y[(sub_y * 8 + j) * 16 + sub_x * 8 + i];
}
}
let (cell_kind, cell_mv0, cell_mv1, cell_pred, cell_sad) = pick_best_b8x8_cell(
frame.y,
src_stride,
ref_l0,
ref_l1,
mb_x,
mb_y,
sub_x,
sub_y,
&direct_cell_pred,
mixed_me_l0_cells[cell],
mixed_me_l1_cells[cell],
qp_y,
);
mixed_cells[cell] = cell_kind;
mixed_cell_mv_l0[cell] = cell_mv0;
mixed_cell_mv_l1[cell] = cell_mv1;
mixed_cell_pred_y_per_cell[cell] = cell_pred;
mixed_total_sad = mixed_total_sad.saturating_add(cell_sad);
}
} else {
mixed_total_sad = u32::MAX;
}
let mixed_has_non_direct = mixed_cells.iter().any(|&c| c != BSubMbCell::Direct);
let mixed_competitive_for_pick = mixed_competitive
&& mixed_has_non_direct
&& !weighted.map(|w| w.apply_chroma).unwrap_or(false)
&& !self.cfg.b_l0_bi_only;
let (mixed_pred_y, mixed_pred_u, mixed_pred_v) = if mixed_competitive_for_pick {
let mut py = [0i32; 256];
let mut pu = [0i32; 64];
let mut pv = [0i32; 64];
let chroma_w_l0 = ref_l0.width / 2;
let chroma_h_l0 = ref_l0.height / 2;
let chroma_w_l1 = ref_l1.width / 2;
let chroma_h_l1 = ref_l1.height / 2;
for cell in 0..4usize {
let sub_x = cell % 2;
let sub_y = cell / 2;
let cell_pred_y = mixed_cell_pred_y_per_cell[cell];
for j in 0..8usize {
for i in 0..8usize {
py[(sub_y * 8 + j) * 16 + sub_x * 8 + i] = cell_pred_y[j * 8 + i];
}
}
let (mv0, mv1) = if mixed_cells[cell] == BSubMbCell::Direct {
(direct.mv_l0_per_8x8[cell], direct.mv_l1_per_8x8[cell])
} else {
(mixed_cell_mv_l0[cell], mixed_cell_mv_l1[cell])
};
let chroma_kind = match mixed_cells[cell] {
BSubMbCell::Direct => {
match (direct.ref_idx_l0 >= 0, direct.ref_idx_l1 >= 0) {
(true, true) => BSubMbCell::Bi,
(true, false) => BSubMbCell::L0,
(false, true) => BSubMbCell::L1,
(false, false) => BSubMbCell::L0, }
}
other => other,
};
let cu = build_b_explicit_4x4_cell_chroma(
chroma_kind,
ref_l0.recon_u,
ref_l1.recon_u,
chroma_w_l0,
chroma_h_l0,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
sub_x,
sub_y,
mv0,
mv1,
);
let cv = build_b_explicit_4x4_cell_chroma(
chroma_kind,
ref_l0.recon_v,
ref_l1.recon_v,
chroma_w_l0,
chroma_h_l0,
chroma_w_l1,
chroma_h_l1,
mb_x,
mb_y,
sub_x,
sub_y,
mv0,
mv1,
);
for j in 0..4usize {
for i in 0..4usize {
pu[(sub_y * 4 + j) * 8 + sub_x * 4 + i] = cu[j * 4 + i];
pv[(sub_y * 4 + j) * 8 + sub_x * 4 + i] = cv[j * 4 + i];
}
}
}
(py, pu, pv)
} else {
([0i32; 256], [0i32; 64], [0i32; 64])
};
#[derive(Debug, Clone, Copy)]
enum BMode {
Direct,
DirectPer8x8,
Mixed8x8,
Explicit(BPred16x16),
}
let direct_bias_sad: u32 = {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 16
};
let direct_8x8_bias_sad: u32 = {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 12
};
let prefer_per_8x8 = direct_8x8_competitive
&& (!direct_competitive
|| direct_8x8_sad.saturating_add(direct_8x8_bias_sad) < direct_sad);
let (best_direct_competitive, best_direct_sad, best_direct_is_per_8x8) = if prefer_per_8x8 {
(direct_8x8_competitive, direct_8x8_sad, true)
} else {
(direct_competitive, direct_sad, false)
};
let mixed_mb_header_bias_sad: u32 = if best_direct_is_per_8x8 {
0
} else {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 12
};
let mixed_tiebreak_bias_sad: u32 = {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 8
};
let prefer_mixed = mixed_competitive_for_pick
&& best_direct_competitive
&& mixed_total_sad
.saturating_add(mixed_mb_header_bias_sad)
.saturating_add(mixed_tiebreak_bias_sad)
< best_direct_sad;
let chosen_mode: BMode = if prefer_mixed {
BMode::Mixed8x8
} else if best_direct_competitive
&& best_direct_sad <= sad_bi.saturating_add(direct_bias_sad)
&& best_direct_sad <= sad_l0.saturating_add(direct_bias_sad)
&& best_direct_sad <= sad_l1.saturating_add(direct_bias_sad)
{
if best_direct_is_per_8x8 {
BMode::DirectPer8x8
} else {
BMode::Direct
}
} else if sad_bi <= sad_l0 && (sad_bi <= sad_l1 || self.cfg.b_l0_bi_only) {
BMode::Explicit(BPred16x16::Bi)
} else if sad_l0 <= sad_l1 || self.cfg.b_l0_bi_only {
BMode::Explicit(BPred16x16::L0)
} else {
BMode::Explicit(BPred16x16::L1)
};
let (pred, pred_y, pred_u, pred_v, eff_mv_l0, eff_mv_l1) = match chosen_mode {
BMode::Direct => {
(
direct_pred_kind.expect("Direct chosen → direct_pred_kind set"),
direct_pred_y,
direct_pred_u,
direct_pred_v,
direct.mv_l0(),
direct.mv_l1(),
)
}
BMode::DirectPer8x8 => {
(
direct_pred_kind.expect("DirectPer8x8 chosen → direct_pred_kind set"),
direct_8x8_pred_y,
direct_8x8_pred_u,
direct_8x8_pred_v,
direct.mv_l0_per_8x8[0],
direct.mv_l1_per_8x8[0],
)
}
BMode::Mixed8x8 => {
let any_l0 = mixed_cells.iter().any(|&c| c.writes_l0()) || direct.ref_idx_l0 >= 0;
let any_l1 = mixed_cells.iter().any(|&c| c.writes_l1()) || direct.ref_idx_l1 >= 0;
let summary = match (any_l0, any_l1) {
(true, true) => BPred16x16::Bi,
(true, false) => BPred16x16::L0,
(false, true) => BPred16x16::L1,
(false, false) => BPred16x16::L0, };
(
summary,
mixed_pred_y,
mixed_pred_u,
mixed_pred_v,
mixed_cell_mv_l0[0],
mixed_cell_mv_l1[0],
)
}
BMode::Explicit(BPred16x16::Bi) => (
BPred16x16::Bi,
pred_y_bi,
pred_u_bi,
pred_v_bi,
mv_l0,
mv_l1,
),
BMode::Explicit(BPred16x16::L0) => (
BPred16x16::L0,
pred_y_l0,
pred_u_l0,
pred_v_l0,
mv_l0,
Mv::ZERO,
),
BMode::Explicit(BPred16x16::L1) => (
BPred16x16::L1,
pred_y_l1,
pred_u_l1,
pred_v_l1,
Mv::ZERO,
mv_l1,
),
};
let mut is_direct = matches!(chosen_mode, BMode::Direct);
let is_direct_per_8x8 = matches!(chosen_mode, BMode::DirectPer8x8);
let is_mixed_8x8 = matches!(chosen_mode, BMode::Mixed8x8);
let mut me_l0_cells: [Mv; 4] = [Mv::ZERO; 4];
let mut me_l1_cells: [Mv; 4] = [Mv::ZERO; 4];
for cell in 0..4usize {
let sx = cell % 2;
let sy = cell / 2;
let r0 = search_quarter_pel_8x8(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l0.recon_y,
ref_l0.width as usize,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sx,
sy,
4,
4,
);
let r1 = search_quarter_pel_8x8(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l1.recon_y,
ref_l1.width as usize,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
sx,
sy,
4,
4,
);
me_l0_cells[cell] = Mv::new(r0.mv_x, r0.mv_y);
me_l1_cells[cell] = Mv::new(r1.mv_x, r1.mv_y);
}
let best_16x16_sad = match chosen_mode {
BMode::Direct => direct_sad,
BMode::DirectPer8x8 => direct_8x8_sad,
BMode::Mixed8x8 => mixed_total_sad,
BMode::Explicit(BPred16x16::Bi) => sad_bi,
BMode::Explicit(BPred16x16::L0) => sad_l0,
BMode::Explicit(BPred16x16::L1) => sad_l1,
};
let cand_top_l0 = [mv_l0, me_l0_cells[0], me_l0_cells[1]];
let cand_top_l1 = [mv_l1, me_l1_cells[0], me_l1_cells[1]];
let cand_bot_l0 = [mv_l0, me_l0_cells[2], me_l0_cells[3]];
let cand_bot_l1 = [mv_l1, me_l1_cells[2], me_l1_cells[3]];
let (top_mode, top_mv_l0, top_mv_l1, sad_top) =
best_partition_mode_b(&cand_top_l0, &cand_top_l1, |pred, mv0, mv1| {
partition_sad_b_16x8(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 0, mv0, mv1,
)
});
let (bot_mode, bot_mv_l0, bot_mv_l1, sad_bot) =
best_partition_mode_b(&cand_bot_l0, &cand_bot_l1, |pred, mv0, mv1| {
partition_sad_b_16x8(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 8, mv0, mv1,
)
});
let sad_16x8_total = sad_top.saturating_add(sad_bot);
let cand_left_l0 = [mv_l0, me_l0_cells[0], me_l0_cells[2]];
let cand_left_l1 = [mv_l1, me_l1_cells[0], me_l1_cells[2]];
let cand_right_l0 = [mv_l0, me_l0_cells[1], me_l0_cells[3]];
let cand_right_l1 = [mv_l1, me_l1_cells[1], me_l1_cells[3]];
let (left_mode, left_mv_l0, left_mv_l1, sad_left) =
best_partition_mode_b(&cand_left_l0, &cand_left_l1, |pred, mv0, mv1| {
partition_sad_b_8x16(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 0, mv0, mv1,
)
});
let (right_mode, right_mv_l0, right_mv_l1, sad_right) =
best_partition_mode_b(&cand_right_l0, &cand_right_l1, |pred, mv0, mv1| {
partition_sad_b_8x16(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 8, mv0, mv1,
)
});
let sad_8x16_total = sad_left.saturating_add(sad_right);
let part_bias_sad: u32 = {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 8
};
let part_threshold = best_16x16_sad.saturating_sub(part_bias_sad);
let (part_winner, part_winner_sad) = if sad_16x8_total <= sad_8x16_total {
(PartitionShape::P16x8, sad_16x8_total)
} else {
(PartitionShape::P8x16, sad_8x16_total)
};
let mut partition_choice: Option<PartitionChoiceB> = None;
let partitions_allowed =
!weighted.map(|w| w.apply_chroma).unwrap_or(false) && !self.cfg.b_l0_bi_only;
if partitions_allowed
&& !is_direct
&& !is_direct_per_8x8
&& !is_mixed_8x8
&& part_winner_sad < part_threshold
{
partition_choice = Some(match part_winner {
PartitionShape::P16x8 => PartitionChoiceB {
shape: PartitionShape::P16x8,
mode_a: top_mode,
mode_b: bot_mode,
mv_l0_a: top_mv_l0,
mv_l1_a: top_mv_l1,
mv_l0_b: bot_mv_l0,
mv_l1_b: bot_mv_l1,
},
PartitionShape::P8x16 => PartitionChoiceB {
shape: PartitionShape::P8x16,
mode_a: left_mode,
mode_b: right_mode,
mv_l0_a: left_mv_l0,
mv_l1_a: left_mv_l1,
mv_l0_b: right_mv_l0,
mv_l1_b: right_mv_l1,
},
});
is_direct = false;
}
let (pred_y, pred_u, pred_v) = if let Some(pc) = partition_choice.as_ref() {
build_partition_composite_b(pc, ref_l0, ref_l1, mb_x, mb_y)
} else {
(pred_y, pred_u, pred_v)
};
let mvp_l0 = mvp_for_16x16(mv_grid_l0, mb_x, mb_y, 0);
let mvp_l1 = mvp_for_16x16(mv_grid_l1, mb_x, mb_y, 0);
let (mvd_l0_x, mvd_l0_y) = if !is_direct
&& !is_direct_per_8x8
&& !is_mixed_8x8
&& partition_choice.is_none()
&& pred.uses_l0()
{
(eff_mv_l0.x - mvp_l0.x, eff_mv_l0.y - mvp_l0.y)
} else {
(0, 0)
};
let (mvd_l1_x, mvd_l1_y) = if !is_direct
&& !is_direct_per_8x8
&& !is_mixed_8x8
&& partition_choice.is_none()
&& pred.uses_l1()
{
(eff_mv_l1.x - mvp_l1.x, eff_mv_l1.y - mvp_l1.y)
} else {
(0, 0)
};
let (mvd_l0_parts, mvd_l1_parts) = if let Some(pc) = partition_choice.as_ref() {
partition_mvds_b(pc, mv_grid_l0, mv_grid_l1, mb_x, mb_y)
} else {
([(0, 0); 2], [(0, 0); 2])
};
let (mixed_mvd_l0_per_cell, mixed_mvd_l1_per_cell) = if is_mixed_8x8 {
let mut mvd0 = [(0i32, 0i32); 4];
let mut mvd1 = [(0i32, 0i32); 4];
let mut inflight_avail = [false; 4];
let mut mv_inflight_l0 = [Mv::ZERO; 4];
let mut ref_inflight_l0 = [-1i32; 4];
let mut mv_inflight_l1 = [Mv::ZERO; 4];
let mut ref_inflight_l1 = [-1i32; 4];
for cell in 0..4usize {
let kind = mixed_cells[cell];
if kind.writes_l0() {
let mvp = mvp_for_p_8x8_partition(
mv_grid_l0,
mb_x,
mb_y,
cell,
0,
inflight_avail,
mv_inflight_l0,
ref_inflight_l0,
);
let mv = mixed_cell_mv_l0[cell];
mvd0[cell] = (mv.x - mvp.x, mv.y - mvp.y);
}
if kind.writes_l1() {
let mvp = mvp_for_p_8x8_partition(
mv_grid_l1,
mb_x,
mb_y,
cell,
0,
inflight_avail,
mv_inflight_l1,
ref_inflight_l1,
);
let mv = mixed_cell_mv_l1[cell];
mvd1[cell] = (mv.x - mvp.x, mv.y - mvp.y);
}
inflight_avail[cell] = true;
let (cell_mv0_ifl, cell_ref0_ifl) = match kind {
BSubMbCell::Direct => (direct.mv_l0_per_8x8[cell], direct.ref_idx_l0),
BSubMbCell::L0 | BSubMbCell::Bi => (mixed_cell_mv_l0[cell], 0),
BSubMbCell::L1 => (Mv::ZERO, -1),
};
let (cell_mv1_ifl, cell_ref1_ifl) = match kind {
BSubMbCell::Direct => (direct.mv_l1_per_8x8[cell], direct.ref_idx_l1),
BSubMbCell::L1 | BSubMbCell::Bi => (mixed_cell_mv_l1[cell], 0),
BSubMbCell::L0 => (Mv::ZERO, -1),
};
mv_inflight_l0[cell] = cell_mv0_ifl;
ref_inflight_l0[cell] = cell_ref0_ifl;
mv_inflight_l1[cell] = cell_mv1_ifl;
ref_inflight_l1[cell] = cell_ref1_ifl;
}
(mvd0, mvd1)
} else {
([(0, 0); 4], [(0, 0); 4])
};
let _ = BMbDecision {
pred,
mv_l0: eff_mv_l0,
mv_l1: eff_mv_l1,
mvd_l0_x,
mvd_l0_y,
mvd_l1_x,
mvd_l1_y,
pred_y,
pred_u,
pred_v,
};
let wq = self.cfg.scaling_matrix.inter_weights();
let inter_luma = forward_inter_luma(frame.y, src_stride, mb_x, mb_y, &pred_y, qp_y, &wq.w4);
let mut luma_4x4_levels_scan = inter_luma.levels_scan;
let mut blk_has_nz = inter_luma.blk_has_nz;
let mut recon_block_residual = inter_luma.recon_residual;
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
let any_nz = (0..4).any(|sub| blk_has_nz[blk8 * 4 + sub]);
if any_nz {
cbp_luma |= 1u8 << blk8;
}
}
let mut transform_size_8x8 = false;
if self.cfg.transform_8x8 {
let inter_luma8 = forward_inter_luma_8x8(
frame.y,
src_stride,
mb_x,
mb_y,
&pred_y,
qp_y,
&wq.w8,
sw.field_scan(),
);
let mut cbp_luma8: u8 = 0;
for blk8 in 0..4usize {
if inter_luma8.blk_has_nz[blk8 * 4] {
cbp_luma8 |= 1u8 << blk8;
}
}
let lambda = lambda_ssd(qp_y, false);
let j4 = inter_luma_transform_cost(
frame.y,
src_stride,
mb_x,
mb_y,
&pred_y,
&inter_luma,
cbp_luma,
lambda,
);
let j8 = inter_luma_transform_cost(
frame.y,
src_stride,
mb_x,
mb_y,
&pred_y,
&inter_luma8,
cbp_luma8,
lambda,
);
if j8 < j4 {
transform_size_8x8 = true;
luma_4x4_levels_scan = inter_luma8.levels_scan;
blk_has_nz = inter_luma8.blk_has_nz;
recon_block_residual = inter_luma8.recon_residual;
cbp_luma = cbp_luma8;
}
}
let b_t8x8_flag = if self.cfg.transform_8x8 {
Some(transform_size_8x8)
} else {
None
};
let (u_dc_levels, u_ac_levels, u_recon_residual) =
encode_chroma_residual_inter(frame.u, chroma_width, mb_x, mb_y, &pred_u, qp_c, &wq.w4);
let (v_dc_levels, v_ac_levels, v_recon_residual) =
encode_chroma_residual_inter(frame.v, chroma_width, mb_x, mb_y, &pred_v, qp_c, &wq.w4);
let any_dc_nz = u_dc_levels.iter().any(|&v| v != 0) || v_dc_levels.iter().any(|&v| v != 0);
let any_ac_nz = u_ac_levels.iter().any(|blk| blk.iter().any(|&v| v != 0))
|| v_ac_levels.iter().any(|blk| blk.iter().any(|&v| v != 0));
let cbp_chroma: u8 = if any_ac_nz {
2
} else if any_dc_nz {
1
} else {
0
};
let width = self.cfg.width as usize;
for blk8 in 0..4usize {
let send = (cbp_luma >> blk8) & 1 == 1;
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = LUMA_4X4_BLK[blk];
let bx_pix = bx * 4;
let by_pix = by * 4;
for j in 0..4usize {
for i in 0..4usize {
let p = pred_y[(by_pix + j) * 16 + bx_pix + i];
let r = if send {
recon_block_residual[blk][j * 4 + i]
} else {
0
};
let py = mb_y * 16 + by_pix + j;
let px = mb_x * 16 + bx_pix + i;
recon_y[py * width + px] = (p + r).clamp(0, 255) as u8;
}
}
}
}
for j in 0..8usize {
for i in 0..8usize {
let cy = mb_y * 8 + j;
let cx = mb_x * 8 + i;
let pu = pred_u[j * 8 + i];
let pv = pred_v[j * 8 + i];
let res_u = if cbp_chroma == 0 {
0
} else {
u_recon_residual[j * 8 + i]
};
let res_v = if cbp_chroma == 0 {
0
} else {
v_recon_residual[j * 8 + i]
};
recon_u[cy * chroma_width + cx] = (pu + res_u).clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = (pv + res_v).clamp(0, 255) as u8;
}
}
let mb_is_skip = is_direct && cbp_luma == 0 && cbp_chroma == 0;
{
let cur = &mut nc_grid.mbs[mb_addr];
cur.is_available = true;
cur.is_intra = false;
cur.is_skip = mb_is_skip;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let nc = derive_nc_luma(
nc_grid,
mb_addr as u32,
blk as u8,
LumaNcKind::Ac,
true,
false,
);
luma_4x4_nc[blk] = nc;
if (cbp_luma >> blk8) & 1 == 1 {
let tc = luma_4x4_levels_scan[blk]
.iter()
.filter(|&&v| v != 0)
.count() as u8;
own_totals[blk] = tc;
} else {
own_totals[blk] = 0;
}
}
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let (chroma_nc_cb, chroma_nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr as u32,
false,
cbp_chroma,
&u_ac_levels,
&v_ac_levels,
1,
);
let mb_qp_delta = if !mb_is_skip && (cbp_luma > 0 || cbp_chroma > 0) {
let d = qp_y - qp_tracker.cur;
qp_tracker.cur = qp_y;
d
} else {
0
};
if mb_is_skip {
*pending_skip += 1;
} else {
sw.ue(*pending_skip);
*pending_skip = 0;
if is_direct {
let dcfg = crate::encoder::macroblock::BDirect16x16McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
crate::encoder::macroblock::write_b_direct_16x16_mb(sw, &dcfg)
.expect("write B_Direct_16x16 mb");
} else if is_direct_per_8x8 {
let dcfg = crate::encoder::macroblock::B8x8AllDirectMcbConfig {
transform_size_8x8_flag: b_t8x8_flag,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
crate::encoder::macroblock::write_b_8x8_all_direct_mb(sw, &dcfg)
.expect("write B_8x8 + 4× B_Direct_8x8 mb");
} else if is_mixed_8x8 {
let mcfg = B8x8MixedMcbConfig {
transform_size_8x8_flag: b_t8x8_flag,
cells: mixed_cells,
mvd_l0: mixed_mvd_l0_per_cell,
mvd_l1: mixed_mvd_l1_per_cell,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_b_8x8_mixed_mb(sw, &mcfg, 0, 0).expect("write B_8x8 mixed mb");
} else if let Some(pc) = partition_choice.as_ref() {
use crate::encoder::macroblock::{
write_b_16x8_mb, write_b_8x16_mb, B16x8McbConfig, B8x16McbConfig,
};
match pc.shape {
PartitionShape::P16x8 => {
let cfg = B16x8McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
top: pc.mode_a,
bottom: pc.mode_b,
mvd_l0: mvd_l0_parts,
mvd_l1: mvd_l1_parts,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_b_16x8_mb(sw, &cfg, 0, 0).expect("write B_*_16x8 mb");
}
PartitionShape::P8x16 => {
let cfg = B8x16McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
left: pc.mode_a,
right: pc.mode_b,
mvd_l0: mvd_l0_parts,
mvd_l1: mvd_l1_parts,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_b_8x16_mb(sw, &cfg, 0, 0).expect("write B_*_8x16 mb");
}
}
} else {
let mb_cfg = B16x16McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
pred,
mvd_l0_x,
mvd_l0_y,
mvd_l1_x,
mvd_l1_y,
cbp_luma,
cbp_chroma,
mb_qp_delta,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: u_dc_levels,
chroma_dc_cr: v_dc_levels,
chroma_ac_cb: u_ac_levels,
chroma_ac_cr: v_ac_levels,
chroma_ac_nc_cb: chroma_nc_cb,
chroma_ac_nc_cr: chroma_nc_cr,
};
write_b_16x16_mb(sw, &mb_cfg, 0, 0).expect("write B_*_16x16 mb");
}
}
let (l0_mvs_8x8, l0_refs_8x8, l0_used) = if is_direct_per_8x8 {
let used = direct.ref_idx_l0 >= 0;
(
direct.mv_l0_per_8x8,
if used {
[direct.ref_idx_l0; 4]
} else {
[-1; 4]
},
used,
)
} else if is_mixed_8x8 {
let mut mvs = [Mv::ZERO; 4];
let mut refs = [-1i32; 4];
let mut any_used = false;
for cell in 0..4usize {
let kind = mixed_cells[cell];
match kind {
BSubMbCell::Direct => {
if direct.ref_idx_l0 >= 0 {
mvs[cell] = direct.mv_l0_per_8x8[cell];
refs[cell] = direct.ref_idx_l0;
any_used = true;
}
}
BSubMbCell::L0 | BSubMbCell::Bi => {
mvs[cell] = mixed_cell_mv_l0[cell];
refs[cell] = 0;
any_used = true;
}
BSubMbCell::L1 => {}
}
}
(mvs, refs, any_used)
} else {
grid_state_b(
partition_choice.as_ref(),
true, pred,
eff_mv_l0,
)
};
let (l1_mvs_8x8, l1_refs_8x8, l1_used) = if is_direct_per_8x8 {
let used = direct.ref_idx_l1 >= 0;
(
direct.mv_l1_per_8x8,
if used {
[direct.ref_idx_l1; 4]
} else {
[-1; 4]
},
used,
)
} else if is_mixed_8x8 {
let mut mvs = [Mv::ZERO; 4];
let mut refs = [-1i32; 4];
let mut any_used = false;
for cell in 0..4usize {
let kind = mixed_cells[cell];
match kind {
BSubMbCell::Direct => {
if direct.ref_idx_l1 >= 0 {
mvs[cell] = direct.mv_l1_per_8x8[cell];
refs[cell] = direct.ref_idx_l1;
any_used = true;
}
}
BSubMbCell::L1 | BSubMbCell::Bi => {
mvs[cell] = mixed_cell_mv_l1[cell];
refs[cell] = 0;
any_used = true;
}
BSubMbCell::L0 => {}
}
}
(mvs, refs, any_used)
} else {
grid_state_b(
partition_choice.as_ref(),
false, pred,
eff_mv_l1,
)
};
*mv_grid_l0.slot_mut(mb_x, mb_y) = if l0_used {
MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: l0_refs_8x8,
mv_l0_8x8: l0_mvs_8x8,
}
} else {
MvGridSlot {
available: true,
is_intra: true,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
}
};
*mv_grid_l1.slot_mut(mb_x, mb_y) = if l1_used {
MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: l1_refs_8x8,
mv_l0_8x8: l1_mvs_8x8,
}
} else {
MvGridSlot {
available: true,
is_intra: true,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
}
};
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
let luma_nonzero_4x4 = luma_nz_mask_from_blocks(&blk_has_nz);
let cb_nz: [bool; 4] =
std::array::from_fn(|i| cbp_chroma == 2 && u_ac_levels[i].iter().any(|&v| v != 0));
let cr_nz: [bool; 4] =
std::array::from_fn(|i| cbp_chroma == 2 && v_ac_levels[i].iter().any(|&v| v != 0));
let (mv_l0_arr, ref_idx_l0, ref_poc_l0) = if is_direct_per_8x8 {
deblock_arrays_per_8x8_direct(direct.mv_l0_per_8x8, direct.ref_idx_l0, 0i32)
} else if is_mixed_8x8 {
deblock_arrays_per_8x8_mixed(l0_mvs_8x8, l0_refs_8x8, 0i32)
} else {
deblock_mv_arr_b(partition_choice.as_ref(), true, pred, eff_mv_l0, 0i32)
};
let (mv_l1_arr, ref_idx_l1, ref_poc_l1) = if is_direct_per_8x8 {
deblock_arrays_per_8x8_direct(direct.mv_l1_per_8x8, direct.ref_idx_l1, 1000i32)
} else if is_mixed_8x8 {
deblock_arrays_per_8x8_mixed(l1_mvs_8x8, l1_refs_8x8, 1000i32)
} else {
deblock_mv_arr_b(partition_choice.as_ref(), false, pred, eff_mv_l1, 1000i32)
};
MbDeblockInfo {
is_intra: false,
qp_y: qp_tracker.cur,
luma_nonzero_4x4,
chroma_nonzero_4x4: chroma_nz_mask_from_blocks(&cb_nz, &cr_nz),
transform_size_8x8_flag: transform_size_8x8,
mv_l0: mv_l0_arr,
ref_idx_l0,
ref_poc_l0,
mv_l1: mv_l1_arr,
ref_idx_l1,
ref_poc_l1,
is_intra_4x4: false,
}
}
#[allow(clippy::too_many_arguments)]
fn encode_b_mb_multifmt(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
_chroma_height: usize,
src_stride: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid_l0: &mut MvGrid,
mv_grid_l1: &mut MvGrid,
pending_skip: &mut u32,
curr_poc: i32,
weighted: Option<WeightedBipredLuma>,
) -> MbDeblockInfo {
let chroma_array_type = self.cfg.chroma_format_idc;
debug_assert!(matches!(chroma_array_type, 0 | 2 | 3));
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let chroma_w = chroma_width as u32;
let chroma_h = _chroma_height as u32;
let me_l0 = search_quarter_pel_16x16(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l0.recon_y,
ref_l0.width as usize,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
16,
16,
);
let me_l1 = search_quarter_pel_16x16(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l1.recon_y,
ref_l1.width as usize,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
16,
16,
);
let mv_l0 = Mv::new(me_l0.mv_x, me_l0.mv_y);
let mv_l1 = Mv::new(me_l1.mv_x, me_l1.mv_y);
let chroma_plane_pred = |plane: &[u8], mv: Mv| -> Vec<i32> {
match chroma_array_type {
0 => Vec::new(),
3 => build_inter_pred_luma(plane, chroma_w, chroma_h, mb_x, mb_y, mv).to_vec(),
_ => cabac_path::build_inter_pred_chroma_local_422(
plane, chroma_w, chroma_h, mb_x, mb_y, mv,
)
.to_vec(),
}
};
let build_luma = |kind: BPred16x16, l0mv: Mv, l1mv: Mv| -> [i32; 256] {
match kind {
BPred16x16::L0 => build_inter_pred_luma(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
l0mv,
),
BPred16x16::L1 => build_inter_pred_luma(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
l1mv,
),
BPred16x16::Bi => {
let a = build_inter_pred_luma(
ref_l0.recon_y,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
l0mv,
);
let b = build_inter_pred_luma(
ref_l1.recon_y,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
l1mv,
);
let mut out = [0i32; 256];
if let Some(wp) = weighted {
let logwd = wp.log2_wd;
let off_avg = (wp.offset_l0 + wp.offset_l1 + 1) >> 1;
let round = 1i32 << logwd;
for i in 0..256 {
let v =
(a[i] * wp.weight_l0 + b[i] * wp.weight_l1 + round) >> (logwd + 1);
out[i] = (v + off_avg).clamp(0, 255);
}
} else {
for i in 0..256 {
out[i] = (a[i] + b[i] + 1) >> 1;
}
}
out
}
}
};
let build_chroma = |kind: BPred16x16, l0mv: Mv, l1mv: Mv| -> (Vec<i32>, Vec<i32>) {
match kind {
BPred16x16::L0 => (
chroma_plane_pred(ref_l0.recon_u, l0mv),
chroma_plane_pred(ref_l0.recon_v, l0mv),
),
BPred16x16::L1 => (
chroma_plane_pred(ref_l1.recon_u, l1mv),
chroma_plane_pred(ref_l1.recon_v, l1mv),
),
BPred16x16::Bi => {
let au = chroma_plane_pred(ref_l0.recon_u, l0mv);
let bu = chroma_plane_pred(ref_l1.recon_u, l1mv);
let av = chroma_plane_pred(ref_l0.recon_v, l0mv);
let bv = chroma_plane_pred(ref_l1.recon_v, l1mv);
(
au.iter()
.zip(&bu)
.map(|(&x, &y)| (x + y + 1) >> 1)
.collect(),
av.iter()
.zip(&bv)
.map(|(&x, &y)| (x + y + 1) >> 1)
.collect(),
)
}
}
};
let pred_y_l0 = build_luma(BPred16x16::L0, mv_l0, mv_l1);
let pred_y_l1 = build_luma(BPred16x16::L1, mv_l0, mv_l1);
let pred_y_bi = build_luma(BPred16x16::Bi, mv_l0, mv_l1);
let sad_l0 = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &pred_y_l0);
let sad_l1 = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &pred_y_l1);
let sad_bi = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &pred_y_bi);
let direct = if self.cfg.direct_temporal_mv_pred {
b_temporal_direct_derive(
ref_l1.partition_mvs,
width_mbs,
mb_x,
mb_y,
curr_poc,
ref_l0.pic_order_cnt,
ref_l1.pic_order_cnt,
)
} else {
b_spatial_direct_derive(
mv_grid_l0,
mv_grid_l1,
ref_l1.partition_mvs,
width_mbs,
mb_x,
mb_y,
)
};
let direct_pred_kind: Option<BPred16x16> =
match (direct.ref_idx_l0 >= 0, direct.ref_idx_l1 >= 0) {
(true, true) => Some(BPred16x16::Bi),
(true, false) => Some(BPred16x16::L0),
(false, true) => Some(BPred16x16::L1),
(false, false) => None,
};
let (direct_competitive, direct_pred_y, direct_sad) =
match (direct.is_uniform(), direct_pred_kind) {
(true, Some(dpk)) => {
let dy = build_luma(dpk, direct.mv_l0(), direct.mv_l1());
let dsad = sad_16x16_pred(frame.y, src_stride, mb_x, mb_y, &dy);
(true, dy, dsad)
}
_ => (false, [0i32; 256], u32::MAX),
};
let direct_bias_sad: u32 = {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 16
};
let is_direct = direct_competitive
&& direct_sad <= sad_bi.saturating_add(direct_bias_sad)
&& direct_sad <= sad_l0.saturating_add(direct_bias_sad)
&& direct_sad <= sad_l1.saturating_add(direct_bias_sad);
let (pred, pred_y, eff_mv_l0, eff_mv_l1) = if is_direct {
(
direct_pred_kind.expect("Direct chosen → direct_pred_kind set"),
direct_pred_y,
direct.mv_l0(),
direct.mv_l1(),
)
} else if sad_bi <= sad_l0 && sad_bi <= sad_l1 {
(BPred16x16::Bi, pred_y_bi, mv_l0, mv_l1)
} else if sad_l0 <= sad_l1 {
(BPred16x16::L0, pred_y_l0, mv_l0, Mv::ZERO)
} else {
(BPred16x16::L1, pred_y_l1, Mv::ZERO, mv_l1)
};
let mut me_l0_cells: [Mv; 4] = [Mv::ZERO; 4];
let mut me_l1_cells: [Mv; 4] = [Mv::ZERO; 4];
for cell in 0..4usize {
let sx = cell % 2;
let sy = cell / 2;
let r0 = search_quarter_pel_8x8(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l0.recon_y,
ref_l0.width as usize,
ref_l0.width,
ref_l0.height,
mb_x,
mb_y,
sx,
sy,
4,
4,
);
let r1 = search_quarter_pel_8x8(
frame.y,
src_stride,
self.cfg.width,
self.cfg.height,
ref_l1.recon_y,
ref_l1.width as usize,
ref_l1.width,
ref_l1.height,
mb_x,
mb_y,
sx,
sy,
4,
4,
);
me_l0_cells[cell] = Mv::new(r0.mv_x, r0.mv_y);
me_l1_cells[cell] = Mv::new(r1.mv_x, r1.mv_y);
}
let cand_top_l0 = [mv_l0, me_l0_cells[0], me_l0_cells[1]];
let cand_top_l1 = [mv_l1, me_l1_cells[0], me_l1_cells[1]];
let cand_bot_l0 = [mv_l0, me_l0_cells[2], me_l0_cells[3]];
let cand_bot_l1 = [mv_l1, me_l1_cells[2], me_l1_cells[3]];
let (top_mode, top_mv_l0, top_mv_l1, sad_top) =
best_partition_mode_b(&cand_top_l0, &cand_top_l1, |pred, mv0, mv1| {
partition_sad_b_16x8(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 0, mv0, mv1,
)
});
let (bot_mode, bot_mv_l0, bot_mv_l1, sad_bot) =
best_partition_mode_b(&cand_bot_l0, &cand_bot_l1, |pred, mv0, mv1| {
partition_sad_b_16x8(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 8, mv0, mv1,
)
});
let sad_16x8_total = sad_top.saturating_add(sad_bot);
let cand_left_l0 = [mv_l0, me_l0_cells[0], me_l0_cells[2]];
let cand_left_l1 = [mv_l1, me_l1_cells[0], me_l1_cells[2]];
let cand_right_l0 = [mv_l0, me_l0_cells[1], me_l0_cells[3]];
let cand_right_l1 = [mv_l1, me_l1_cells[1], me_l1_cells[3]];
let (left_mode, left_mv_l0, left_mv_l1, sad_left) =
best_partition_mode_b(&cand_left_l0, &cand_left_l1, |pred, mv0, mv1| {
partition_sad_b_8x16(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 0, mv0, mv1,
)
});
let (right_mode, right_mv_l0, right_mv_l1, sad_right) =
best_partition_mode_b(&cand_right_l0, &cand_right_l1, |pred, mv0, mv1| {
partition_sad_b_8x16(
pred, frame.y, src_stride, ref_l0, ref_l1, mb_x, mb_y, 8, mv0, mv1,
)
});
let sad_8x16_total = sad_left.saturating_add(sad_right);
let best_16x16_sad = if is_direct {
direct_sad
} else {
sad_bi.min(sad_l0).min(sad_l1)
};
let part_bias_sad: u32 = {
let lambda = (qp_y / 6).clamp(0, 8) as u32;
(lambda + 1) * 8
};
let part_threshold = best_16x16_sad.saturating_sub(part_bias_sad);
let (part_winner, part_winner_sad) = if sad_16x8_total <= sad_8x16_total {
(PartitionShape::P16x8, sad_16x8_total)
} else {
(PartitionShape::P8x16, sad_8x16_total)
};
let mut partition_choice: Option<PartitionChoiceB> = None;
if !is_direct && part_winner_sad < part_threshold {
partition_choice = Some(match part_winner {
PartitionShape::P16x8 => PartitionChoiceB {
shape: PartitionShape::P16x8,
mode_a: top_mode,
mode_b: bot_mode,
mv_l0_a: top_mv_l0,
mv_l1_a: top_mv_l1,
mv_l0_b: bot_mv_l0,
mv_l1_b: bot_mv_l1,
},
PartitionShape::P8x16 => PartitionChoiceB {
shape: PartitionShape::P8x16,
mode_a: left_mode,
mode_b: right_mode,
mv_l0_a: left_mv_l0,
mv_l1_a: left_mv_l1,
mv_l0_b: right_mv_l0,
mv_l1_b: right_mv_l1,
},
});
}
let (ct_w, ct_h) = chroma_tile_dims(chroma_array_type);
let (pred_y, pred_u, pred_v) = if let Some(pc) = partition_choice.as_ref() {
let mut py = [0i32; 256];
let mut pu = vec![0i32; ct_w * ct_h];
let mut pv = vec![0i32; ct_w * ct_h];
match pc.shape {
PartitionShape::P16x8 => {
let py0 = build_b_partition_pred_luma_16x8(
pc.mode_a, ref_l0, ref_l1, mb_x, mb_y, 0, pc.mv_l0_a, pc.mv_l1_a,
);
let py1 = build_b_partition_pred_luma_16x8(
pc.mode_b, ref_l0, ref_l1, mb_x, mb_y, 8, pc.mv_l0_b, pc.mv_l1_b,
);
for j in 0..8usize {
for i in 0..16usize {
py[j * 16 + i] = py0[j * 16 + i];
py[(8 + j) * 16 + i] = py1[j * 16 + i];
}
}
}
PartitionShape::P8x16 => {
let py0 = build_b_partition_pred_luma_8x16(
pc.mode_a, ref_l0, ref_l1, mb_x, mb_y, 0, pc.mv_l0_a, pc.mv_l1_a,
);
let py1 = build_b_partition_pred_luma_8x16(
pc.mode_b, ref_l0, ref_l1, mb_x, mb_y, 8, pc.mv_l0_b, pc.mv_l1_b,
);
for j in 0..16usize {
for i in 0..8usize {
py[j * 16 + i] = py0[j * 8 + i];
py[j * 16 + 8 + i] = py1[j * 8 + i];
}
}
}
}
for (half, (mode, hmv0, hmv1)) in [
(pc.mode_a, pc.mv_l0_a, pc.mv_l1_a),
(pc.mode_b, pc.mv_l0_b, pc.mv_l1_b),
]
.into_iter()
.enumerate()
{
let (cx0, cy0, rw, rh) = match pc.shape {
PartitionShape::P16x8 => {
(mb_x * ct_w, mb_y * ct_h + half * (ct_h / 2), ct_w, ct_h / 2)
}
PartitionShape::P8x16 => {
(mb_x * ct_w + half * (ct_w / 2), mb_y * ct_h, ct_w / 2, ct_h)
}
};
let ru = build_b_partition_chroma_rect_fmt(
mode,
ref_l0.recon_u,
ref_l1.recon_u,
chroma_w,
chroma_h,
cx0,
cy0,
rw,
rh,
hmv0,
hmv1,
chroma_array_type,
);
let rv = build_b_partition_chroma_rect_fmt(
mode,
ref_l0.recon_v,
ref_l1.recon_v,
chroma_w,
chroma_h,
cx0,
cy0,
rw,
rh,
hmv0,
hmv1,
chroma_array_type,
);
let ox = cx0 - mb_x * ct_w;
let oy = cy0 - mb_y * ct_h;
for j in 0..rh {
for i in 0..rw {
pu[(oy + j) * ct_w + ox + i] = ru[j * rw + i];
pv[(oy + j) * ct_w + ox + i] = rv[j * rw + i];
}
}
}
(py, pu, pv)
} else {
let (pu, pv) = build_chroma(pred, eff_mv_l0, eff_mv_l1);
(pred_y, pu, pv)
};
let mvp_l0 = mvp_for_16x16(mv_grid_l0, mb_x, mb_y, 0);
let mvp_l1 = mvp_for_16x16(mv_grid_l1, mb_x, mb_y, 0);
let (mvd_l0_x, mvd_l0_y) = if !is_direct && partition_choice.is_none() && pred.uses_l0() {
(eff_mv_l0.x - mvp_l0.x, eff_mv_l0.y - mvp_l0.y)
} else {
(0, 0)
};
let (mvd_l1_x, mvd_l1_y) = if !is_direct && partition_choice.is_none() && pred.uses_l1() {
(eff_mv_l1.x - mvp_l1.x, eff_mv_l1.y - mvp_l1.y)
} else {
(0, 0)
};
let (mvd_l0_parts, mvd_l1_parts) = if let Some(pc) = partition_choice.as_ref() {
partition_mvds_b(pc, mv_grid_l0, mv_grid_l1, mb_x, mb_y)
} else {
([(0, 0); 2], [(0, 0); 2])
};
let wq = self.cfg.scaling_matrix.inter_weights();
let inter_luma = forward_inter_luma(frame.y, src_stride, mb_x, mb_y, &pred_y, qp_y, &wq.w4);
let mut luma_4x4_levels_scan = inter_luma.levels_scan;
let mut blk_has_nz = inter_luma.blk_has_nz;
let mut recon_block_residual = inter_luma.recon_residual;
let mut cbp_luma: u8 = 0;
for blk8 in 0..4usize {
if (0..4).any(|sub| blk_has_nz[blk8 * 4 + sub]) {
cbp_luma |= 1u8 << blk8;
}
}
let mut transform_size_8x8 = false;
if self.cfg.transform_8x8 {
let inter_luma8 = forward_inter_luma_8x8(
frame.y,
src_stride,
mb_x,
mb_y,
&pred_y,
qp_y,
&wq.w8,
sw.field_scan(),
);
let mut cbp_luma8: u8 = 0;
for blk8 in 0..4usize {
if inter_luma8.blk_has_nz[blk8 * 4] {
cbp_luma8 |= 1u8 << blk8;
}
}
let lambda = lambda_ssd(qp_y, false);
let j4 = inter_luma_transform_cost(
frame.y,
src_stride,
mb_x,
mb_y,
&pred_y,
&inter_luma,
cbp_luma,
lambda,
);
let j8 = inter_luma_transform_cost(
frame.y,
src_stride,
mb_x,
mb_y,
&pred_y,
&inter_luma8,
cbp_luma8,
lambda,
);
if j8 < j4 {
transform_size_8x8 = true;
luma_4x4_levels_scan = inter_luma8.levels_scan;
blk_has_nz = inter_luma8.blk_has_nz;
recon_block_residual = inter_luma8.recon_residual;
cbp_luma = cbp_luma8;
}
}
let mut u_dc8 = [0i32; 8];
let mut v_dc8 = [0i32; 8];
let mut u_ac8 = [[0i32; 16]; 8];
let mut v_ac8 = [[0i32; 16]; 8];
let mut u_recon_res_422 = [0i32; 128];
let mut v_recon_res_422 = [0i32; 128];
let mut cb444 = InterLumaForward::default();
let mut cr444 = InterLumaForward::default();
let mut cbp_chroma: u8 = 0;
match chroma_array_type {
0 => {}
3 => {
let pu: &[i32; 256] = pred_u[..256].try_into().expect("pred tile");
let pv: &[i32; 256] = pred_v[..256].try_into().expect("pred tile");
if transform_size_8x8 {
cb444 = forward_inter_luma_8x8(
frame.u,
chroma_width,
mb_x,
mb_y,
pu,
qp_c,
&wq.w8,
sw.field_scan(),
);
cr444 = forward_inter_luma_8x8(
frame.v,
chroma_width,
mb_x,
mb_y,
pv,
qp_c,
&wq.w8,
sw.field_scan(),
);
} else {
cb444 = forward_inter_luma(frame.u, chroma_width, mb_x, mb_y, pu, qp_c, &wq.w4);
cr444 = forward_inter_luma(frame.v, chroma_width, mb_x, mb_y, pv, qp_c, &wq.w4);
}
for blk8 in 0..4usize {
let cb_nz = (0..4).any(|sub| cb444.blk_has_nz[blk8 * 4 + sub]);
let cr_nz = (0..4).any(|sub| cr444.blk_has_nz[blk8 * 4 + sub]);
if cb_nz || cr_nz {
cbp_luma |= 1u8 << blk8;
}
}
if cbp_luma == 0 {
transform_size_8x8 = false;
}
}
_ => {
let pu: &[i32; 128] = pred_u[..128].try_into().expect("pred tile");
let pv: &[i32; 128] = pred_v[..128].try_into().expect("pred tile");
let (dcu, acu, ru) =
encode_chroma_residual_422(frame.u, chroma_width, mb_x, mb_y, pu, qp_c, &wq.w4);
let (dcv, acv, rv) =
encode_chroma_residual_422(frame.v, chroma_width, mb_x, mb_y, pv, qp_c, &wq.w4);
u_dc8 = dcu;
v_dc8 = dcv;
u_ac8 = acu;
v_ac8 = acv;
u_recon_res_422 = ru;
v_recon_res_422 = rv;
let any_dc_nz = u_dc8.iter().any(|&v| v != 0) || v_dc8.iter().any(|&v| v != 0);
let any_ac_nz = u_ac8.iter().any(|b| b.iter().any(|&v| v != 0))
|| v_ac8.iter().any(|b| b.iter().any(|&v| v != 0));
cbp_chroma = if any_ac_nz {
2
} else if any_dc_nz {
1
} else {
0
};
}
}
if cbp_luma == 0 && self.cfg.transform_8x8 {
transform_size_8x8 = false;
}
let b_t8x8_flag = if self.cfg.transform_8x8 {
Some(transform_size_8x8)
} else {
None
};
let width = self.cfg.width as usize;
for blk8 in 0..4usize {
let send = (cbp_luma >> blk8) & 1 == 1;
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = LUMA_4X4_BLK[blk];
for j in 0..4usize {
for i in 0..4usize {
let p = pred_y[(by * 4 + j) * 16 + bx * 4 + i];
let r = if send {
recon_block_residual[blk][j * 4 + i]
} else {
0
};
let py = mb_y * 16 + by * 4 + j;
let px = mb_x * 16 + bx * 4 + i;
recon_y[py * width + px] = (p + r).clamp(0, 255) as u8;
}
}
}
}
match chroma_array_type {
0 => {}
3 => {
for (pred_c, fwd, recon) in [
(&pred_u, &cb444, &mut *recon_u),
(&pred_v, &cr444, &mut *recon_v),
] {
for blk8 in 0..4usize {
let send = (cbp_luma >> blk8) & 1 == 1;
for sub in 0..4usize {
let blk = blk8 * 4 + sub;
let (bx, by) = LUMA_4X4_BLK[blk];
for j in 0..4usize {
for i in 0..4usize {
let p = pred_c[(by * 4 + j) * 16 + bx * 4 + i];
let r = if send {
fwd.recon_residual[blk][j * 4 + i]
} else {
0
};
let cy = mb_y * 16 + by * 4 + j;
let cx = mb_x * 16 + bx * 4 + i;
recon[cy * chroma_width + cx] = (p + r).clamp(0, 255) as u8;
}
}
}
}
}
}
_ => {
for j in 0..16usize {
for i in 0..8usize {
let cy = mb_y * 16 + j;
let cx = mb_x * 8 + i;
let pu = pred_u[j * 8 + i];
let pv = pred_v[j * 8 + i];
let res_u = if cbp_chroma == 0 {
0
} else {
u_recon_res_422[j * 8 + i]
};
let res_v = if cbp_chroma == 0 {
0
} else {
v_recon_res_422[j * 8 + i]
};
recon_u[cy * chroma_width + cx] = (pu + res_u).clamp(0, 255) as u8;
recon_v[cy * chroma_width + cx] = (pv + res_v).clamp(0, 255) as u8;
}
}
}
}
let mb_is_skip = is_direct && cbp_luma == 0 && cbp_chroma == 0;
{
let cur = &mut nc_grid.mbs[mb_addr];
cur.is_available = true;
cur.is_intra = false;
cur.is_skip = mb_is_skip;
cur.is_i_pcm = false;
cur.luma_total_coeff = [0u8; 16];
cur.cb_total_coeff = [0u8; 8];
cur.cr_total_coeff = [0u8; 8];
cur.cb_luma_total_coeff = [0u8; 16];
cur.cr_luma_total_coeff = [0u8; 16];
}
let mut luma_4x4_nc = [0i32; 16];
let mut own_totals = [0u8; 16];
for blk in 0..16usize {
let blk8 = blk / 4;
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
let nc = derive_nc_luma(
nc_grid,
mb_addr as u32,
blk as u8,
LumaNcKind::Ac,
true,
false,
);
luma_4x4_nc[blk] = nc;
own_totals[blk] = if (cbp_luma >> blk8) & 1 == 1 {
luma_4x4_levels_scan[blk]
.iter()
.filter(|&&v| v != 0)
.count() as u8
} else {
0
};
}
nc_grid.mbs[mb_addr].luma_total_coeff = own_totals;
if mb_is_skip {
*pending_skip += 1;
} else {
sw.ue(*pending_skip);
*pending_skip = 0;
enum BChromaEmit {
Mono,
C422 { nc_cb: [i32; 8], nc_cr: [i32; 8] },
C444 { cb_nc: [i32; 16], cr_nc: [i32; 16] },
}
let emit = if chroma_array_type == 0 {
BChromaEmit::Mono
} else if chroma_array_type == 3 {
BChromaEmit::C444 {
cb_nc: derive_plane_nc_444_inter(
nc_grid,
mb_addr as u32,
false,
cbp_luma,
&cb444.levels_scan,
),
cr_nc: derive_plane_nc_444_inter(
nc_grid,
mb_addr as u32,
true,
cbp_luma,
&cr444.levels_scan,
),
}
} else {
let (nc_cb, nc_cr) = derive_chroma_ac_nc_and_commit_totals(
nc_grid,
mb_addr as u32,
false,
cbp_chroma,
&u_ac8,
&v_ac8,
2,
);
BChromaEmit::C422 { nc_cb, nc_cr }
};
let chroma_kind = match &emit {
BChromaEmit::Mono => crate::encoder::macroblock::ChromaWriteKind::Mono,
BChromaEmit::C444 { cb_nc, cr_nc } => {
crate::encoder::macroblock::ChromaWriteKind::Yuv444Inter {
cb_levels: &cb444.levels_scan,
cr_levels: &cr444.levels_scan,
cb_nc,
cr_nc,
cbp_luma,
}
}
BChromaEmit::C422 { nc_cb, nc_cr } => {
crate::encoder::macroblock::ChromaWriteKind::Yuv422 {
chroma_dc_cb: &u_dc8,
chroma_dc_cr: &v_dc8,
chroma_ac_cb: &u_ac8,
chroma_ac_cr: &v_ac8,
cb_ac_nc: nc_cb,
cr_ac_nc: nc_cr,
}
}
};
if is_direct {
let dcfg = crate::encoder::macroblock::BDirect16x16McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
cbp_luma,
cbp_chroma,
mb_qp_delta: 0,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: [0; 4],
chroma_dc_cr: [0; 4],
chroma_ac_cb: [[0; 16]; 4],
chroma_ac_cr: [[0; 16]; 4],
chroma_ac_nc_cb: [0; 8],
chroma_ac_nc_cr: [0; 8],
};
crate::encoder::macroblock::write_b_direct_16x16_mb_chroma(
sw,
&dcfg,
chroma_array_type,
chroma_kind,
)
.expect("write B_Direct_16x16 mb (4:2:2/4:4:4)");
} else if let Some(pc) = partition_choice.as_ref() {
match pc.shape {
PartitionShape::P16x8 => {
let cfg = crate::encoder::macroblock::B16x8McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
top: pc.mode_a,
bottom: pc.mode_b,
mvd_l0: mvd_l0_parts,
mvd_l1: mvd_l1_parts,
cbp_luma,
cbp_chroma,
mb_qp_delta: 0,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: [0; 4],
chroma_dc_cr: [0; 4],
chroma_ac_cb: [[0; 16]; 4],
chroma_ac_cr: [[0; 16]; 4],
chroma_ac_nc_cb: [0; 8],
chroma_ac_nc_cr: [0; 8],
};
crate::encoder::macroblock::write_b_16x8_mb_chroma(
sw,
&cfg,
0,
0,
chroma_array_type,
chroma_kind,
)
.expect("write B_*_16x8 mb (4:2:2/4:4:4)");
}
PartitionShape::P8x16 => {
let cfg = crate::encoder::macroblock::B8x16McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
left: pc.mode_a,
right: pc.mode_b,
mvd_l0: mvd_l0_parts,
mvd_l1: mvd_l1_parts,
cbp_luma,
cbp_chroma,
mb_qp_delta: 0,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: [0; 4],
chroma_dc_cr: [0; 4],
chroma_ac_cb: [[0; 16]; 4],
chroma_ac_cr: [[0; 16]; 4],
chroma_ac_nc_cb: [0; 8],
chroma_ac_nc_cr: [0; 8],
};
crate::encoder::macroblock::write_b_8x16_mb_chroma(
sw,
&cfg,
0,
0,
chroma_array_type,
chroma_kind,
)
.expect("write B_*_8x16 mb (4:2:2/4:4:4)");
}
}
} else {
let mb_cfg = B16x16McbConfig {
transform_size_8x8_flag: b_t8x8_flag,
pred,
mvd_l0_x,
mvd_l0_y,
mvd_l1_x,
mvd_l1_y,
cbp_luma,
cbp_chroma,
mb_qp_delta: 0,
luma_4x4_levels: luma_4x4_levels_scan,
luma_4x4_nc,
chroma_dc_cb: [0; 4],
chroma_dc_cr: [0; 4],
chroma_ac_cb: [[0; 16]; 4],
chroma_ac_cr: [[0; 16]; 4],
chroma_ac_nc_cb: [0; 8],
chroma_ac_nc_cr: [0; 8],
};
crate::encoder::macroblock::write_b_16x16_mb_chroma(
sw,
&mb_cfg,
0,
0,
chroma_array_type,
chroma_kind,
)
.expect("write B_*_16x16 mb (4:2:2/4:4:4)");
}
}
let (l0_mvs_8x8, l0_refs_8x8, l0_used) =
grid_state_b(partition_choice.as_ref(), true, pred, eff_mv_l0);
let (l1_mvs_8x8, l1_refs_8x8, l1_used) =
grid_state_b(partition_choice.as_ref(), false, pred, eff_mv_l1);
*mv_grid_l0.slot_mut(mb_x, mb_y) = if l0_used {
MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: l0_refs_8x8,
mv_l0_8x8: l0_mvs_8x8,
}
} else {
MvGridSlot {
available: true,
is_intra: true,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
}
};
*mv_grid_l1.slot_mut(mb_x, mb_y) = if l1_used {
MvGridSlot {
available: true,
is_intra: false,
ref_idx_l0_8x8: l1_refs_8x8,
mv_l0_8x8: l1_mvs_8x8,
}
} else {
MvGridSlot {
available: true,
is_intra: true,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
}
};
let s = intra_grid.slot_mut(mb_x, mb_y);
s.available = true;
s.is_i_nxn = false;
let luma_nonzero_4x4 = luma_nz_mask_from_blocks(&blk_has_nz);
let chroma_nonzero_4x4: u16 = match chroma_array_type {
0 | 3 => 0,
_ => {
let mut m = 0u16;
for blk in 0..8usize {
if cbp_chroma == 2 && u_ac8[blk].iter().any(|&v| v != 0) {
m |= 1 << blk;
}
if cbp_chroma == 2 && v_ac8[blk].iter().any(|&v| v != 0) {
m |= 1 << (8 + blk);
}
}
m
}
};
let (mv_l0_arr, ref_idx_l0, ref_poc_l0) =
deblock_mv_arr_b(partition_choice.as_ref(), true, pred, eff_mv_l0, 0);
let (mv_l1_arr, ref_idx_l1, ref_poc_l1) =
deblock_mv_arr_b(partition_choice.as_ref(), false, pred, eff_mv_l1, 1000);
MbDeblockInfo {
is_intra: false,
qp_y,
luma_nonzero_4x4,
chroma_nonzero_4x4,
transform_size_8x8_flag: transform_size_8x8,
mv_l0: mv_l0_arr,
ref_idx_l0,
ref_poc_l0,
mv_l1: mv_l1_arr,
ref_idx_l1,
ref_poc_l1,
is_intra_4x4: false,
}
}
#[allow(clippy::too_many_arguments)]
fn encode_b_mb_with_intra_fallback(
&self,
frame: &YuvFrame<'_>,
ref_l0: &EncodedFrameRef<'_>,
ref_l1: &EncodedFrameRef<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
src_stride: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid_l0: &mut MvGrid,
mv_grid_l1: &mut MvGrid,
pending_skip: &mut u32,
curr_poc: i32,
weighted: Option<WeightedBipredLuma>,
qp_tracker: &mut MbQpTracker,
) -> MbDeblockInfo {
if !self.cfg.intra_in_inter || self.cfg.chroma_format_idc != 1 {
return self.encode_b_mb(
frame,
ref_l0,
ref_l1,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
src_stride,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
pending_skip,
curr_poc,
weighted,
qp_tracker,
);
}
let width = self.cfg.width as usize;
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let snap_pre = InterMbStateSnapshot::capture_b(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
mb_addr,
width_mbs,
*pending_skip,
);
let pre_bits = sw.bits_emitted();
let qp_tracker_pre = *qp_tracker;
let inter_dbl = self.encode_b_mb(
frame,
ref_l0,
ref_l1,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
src_stride,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
pending_skip,
curr_poc,
weighted,
qp_tracker,
);
let inter_bits = sw.bits_emitted() - pre_bits;
let inter_d = ssd_16x16(frame.y, width, mb_x * 16, mb_y * 16, &{
let mut buf = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
buf[j * 16 + i] = recon_y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
}
}
buf
});
let lambda = lambda_ssd(qp_y, false);
let j_inter = cost_combined(inter_d, inter_bits as u64, lambda);
let snap_inter = InterMbStateSnapshot::capture_b(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
mb_addr,
width_mbs,
*pending_skip,
);
let qp_tracker_inter = *qp_tracker;
snap_pre.restore_b(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
mb_addr,
pending_skip,
);
*qp_tracker = qp_tracker_pre;
let intra_dbl = self.encode_b_mb_intra16x16(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
pending_skip,
qp_tracker,
);
let intra_bits = sw.bits_emitted() - pre_bits;
let intra_d = ssd_16x16(frame.y, width, mb_x * 16, mb_y * 16, &{
let mut buf = [0i32; 256];
for j in 0..16usize {
for i in 0..16usize {
buf[j * 16 + i] = recon_y[(mb_y * 16 + j) * width + (mb_x * 16 + i)] as i32;
}
}
buf
});
let j_intra = cost_combined(intra_d, intra_bits as u64, lambda);
if j_inter <= j_intra {
snap_pre.restore_b(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
mb_addr,
pending_skip,
);
snap_inter.restore_b(
recon_y,
recon_u,
recon_v,
width,
chroma_width,
mb_x,
mb_y,
sw,
nc_grid,
intra_grid,
mv_grid_l0,
mv_grid_l1,
mb_addr,
pending_skip,
);
*qp_tracker = qp_tracker_inter;
inter_dbl
} else {
intra_dbl
}
}
#[allow(clippy::too_many_arguments)]
fn encode_b_mb_intra16x16(
&self,
frame: &YuvFrame<'_>,
mb_x: usize,
mb_y: usize,
qp_y: i32,
qp_c: i32,
chroma_width: usize,
chroma_height: usize,
recon_y: &mut [u8],
recon_u: &mut [u8],
recon_v: &mut [u8],
sw: &mut BitWriter,
nc_grid: &mut CavlcNcGrid,
intra_grid: &mut IntraGrid,
mv_grid_l0: &mut MvGrid,
mv_grid_l1: &mut MvGrid,
pending_skip: &mut u32,
qp_tracker: &mut MbQpTracker,
) -> MbDeblockInfo {
sw.ue(*pending_skip);
*pending_skip = 0;
qp_tracker.write_pending_mbaff_flag(sw);
let mb_qp_delta = qp_y - qp_tracker.cur;
qp_tracker.cur = qp_y;
self.encode_intra16x16_in_inter_slice(
frame,
mb_x,
mb_y,
qp_y,
qp_c,
chroma_width,
chroma_height,
recon_y,
recon_u,
recon_v,
sw,
nc_grid,
intra_grid,
23,
mb_qp_delta,
None,
);
let intra_slot = MvGridSlot {
available: true,
is_intra: true,
ref_idx_l0_8x8: [-1; 4],
mv_l0_8x8: [Mv::ZERO; 4],
};
*mv_grid_l0.slot_mut(mb_x, mb_y) = intra_slot;
*mv_grid_l1.slot_mut(mb_x, mb_y) = intra_slot;
let width_mbs = (self.cfg.width / 16) as usize;
let mb_addr = mb_y * width_mbs + mb_x;
let any_nz = nc_grid.mbs[mb_addr]
.luma_total_coeff
.iter()
.any(|&v| v != 0);
let luma_nonzero_4x4: u16 = if any_nz {
let mut nz = [false; 16];
for (blkz, slot) in nz.iter_mut().enumerate() {
*slot = nc_grid.mbs[mb_addr].luma_total_coeff[blkz] != 0;
}
luma_nz_mask_from_blocks(&nz)
} else {
0
};
MbDeblockInfo {
is_intra: true,
qp_y,
luma_nonzero_4x4,
chroma_nonzero_4x4: 0,
..Default::default()
}
}
}
#[cfg(test)]
mod weighted_bipred_tests {
use super::*;
#[test]
fn weight_selector_finds_alpha_07_on_textured_fade() {
let w = 16usize;
let h = 16usize;
let mut l0 = vec![0u8; w * h];
let mut l1 = vec![0u8; w * h];
let mut src = vec![0u8; w * h];
for j in 0..h {
for i in 0..w {
let tex_a = (((i * 13 + j * 7) % 16) as i32) * 4; let tex_b = ((i * 13 + j * 7) % 16) as i32; let a = 30 + tex_a;
let b = 180 + tex_b;
let v = (0.7 * a as f64 + 0.3 * b as f64).round() as i32;
l0[j * w + i] = a.clamp(0, 255) as u8;
l1[j * w + i] = b.clamp(0, 255) as u8;
src[j * w + i] = v.clamp(0, 255) as u8;
}
}
let table = compute_explicit_bipred_weights(&src, &l0, &l1, w, h, w, w);
assert_eq!(table.luma_log2_weight_denom, 5);
assert_eq!(
table.luma_weight_l0 + table.luma_weight_l1,
64,
"weights must sum to 2^(logwd+1)"
);
assert!(
(40..=48).contains(&table.luma_weight_l0),
"expected w0 ≈ 45 (α=0.7), got {}",
table.luma_weight_l0,
);
}
#[test]
fn weight_selector_returns_default_on_midpoint_fade() {
let w = 16usize;
let h = 16usize;
let l0 = vec![60u8; w * h];
let l1 = vec![160u8; w * h];
let src = vec![110u8; w * h];
let table = compute_explicit_bipred_weights(&src, &l0, &l1, w, h, w, w);
assert_eq!(table.luma_weight_l0, 32);
assert_eq!(table.luma_weight_l1, 32);
assert_eq!(table.luma_offset_l0, 0);
assert_eq!(table.luma_offset_l1, 0);
}
}
#[cfg(test)]
mod level_derivation_tests {
use super::*;
#[test]
fn min_level_idc_picks_smallest_covering_level() {
assert_eq!(min_level_idc_for_picture_size(11, 9), 10);
assert_eq!(min_level_idc_for_picture_size(22, 18), 11);
assert_eq!(min_level_idc_for_picture_size(44, 36), 22);
assert_eq!(min_level_idc_for_picture_size(45, 36), 22);
assert_eq!(min_level_idc_for_picture_size(80, 45), 31);
assert_eq!(min_level_idc_for_picture_size(120, 68), 40);
assert_eq!(min_level_idc_for_picture_size(240, 135), 51);
assert_eq!(min_level_idc_for_picture_size(480, 270), 60);
assert_eq!(min_level_idc_for_picture_size(1024, 1024), 62);
}
#[test]
fn encoder_config_default_level_matches_picture_size() {
let cfg = EncoderConfig::new(64, 64);
assert_eq!(cfg.level_idc, 10);
let cfg = EncoderConfig::new(1280, 720);
assert_eq!(cfg.level_idc, 31);
let cfg = EncoderConfig::new(1920, 1080);
assert_eq!(cfg.level_idc, 40);
let cfg = EncoderConfig::new(3840, 2160);
assert_eq!(cfg.level_idc, 51);
}
#[test]
fn non_aligned_dimensions_round_up_to_whole_mbs() {
assert_eq!(EncoderConfig::new(1920, 1088).level_idc, 40);
assert_eq!(EncoderConfig::new(1920, 1081).level_idc, 40);
}
}