use crate::cdf::{
get_tx_size, inter_tx_type_set, tx_size_sqr_index, FindMvStackResult, PartitionWalker,
QuantizerParams, TileCdfContext, TileGeometry, BLOCK_4X4, BLOCK_64X64, BLOCK_8X8, DCT_DCT,
EIGHTTAP, GM_TYPE_IDENTITY, MAX_TX_SIZE_RECT, MAX_VARTX_DEPTH, MODE_GLOBALMV,
MODE_GLOBAL_GLOBALMV, MODE_NEARESTMV, MODE_NEAREST_NEARESTMV, MODE_NEARMV, MODE_NEAR_NEARMV,
MODE_NEWMV, MODE_NEW_NEWMV, MOTION_MODE_OBMC, MOTION_MODE_SIMPLE, MOTION_MODE_WARPED_CAUSAL,
NUM_4X4_BLOCKS_HIGH, NUM_4X4_BLOCKS_WIDE, PARTITION_HORZ, PARTITION_VERT, SPLIT_TX_SIZE,
SWITCHABLE, TX_4X4, TX_HEIGHT, TX_SIZE_SQR_UP, TX_WIDTH,
};
use crate::encoder::block_mode_info::assign_mv_pred_mv;
use crate::encoder::ivf::{IvfWriter, FOURCC_AV01};
use crate::encoder::key_frame::{
encode_leaf_sq, lambda_for, leaf_rate, region_distortion, region_distortion_wh, residual_tx,
restore_region, save_region, save_region_wh, tu_bd_stamp, BlockDecodedMirror, ReconState,
RegionSnapshot,
};
use crate::encoder::obu::{build_temporal_unit, ObuFrame};
use crate::encoder::partition_tree::{
write_partition_tree_syntax, PartitionSyntaxWriter, SyntaxBlock, SyntaxFrameParams,
SyntaxInterBlock, SyntaxInterFrameParams, SyntaxNode, VarTxSyntaxTree,
};
use crate::encoder::rate_twin::{score256, RateModel, RateTwin, TuCtx, TuFork};
use crate::encoder::symbol_writer::SymbolWriter;
use crate::encoder::tile_group_obu::{write_tile_group_obu, TileGroupObu, TilePayload};
use crate::encoder::yuv_frame::YuvFrame;
use crate::encoder::yuv_frame::{
build_intra_only_yuv420_8bit_fh_with_q, sb_grid_origins, Yuv420Frame,
};
use crate::frame_header::{
FrameHeader, FrameType, InterFrameRefs, ALL_FRAMES_PUB, PRIMARY_REF_NONE,
};
use crate::inter_pred::{
reconstruct_inter_leaf_at, InterModeInfoGrid, PlaneReconContext, RefFrameStoreEntry,
};
use crate::obu::ObuType;
use crate::sequence_header::SequenceHeader;
use crate::uncompressed_header_tail::{
FilmGrainParams, InterpolationFilter, TxMode, REFS_PER_FRAME, TOTAL_REFS_PER_FRAME,
};
use crate::uncompressed_header_tail::{GOLDEN_FRAME, LAST_FRAME};
use crate::Error;
#[derive(Debug, Clone)]
pub struct GopFrameRecon {
pub y: Vec<u8>,
pub u: Vec<u8>,
pub v: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct GopFrameReconYuv {
pub y: Vec<u16>,
pub u: Vec<u16>,
pub v: Vec<u16>,
}
#[derive(Debug, Clone)]
pub struct EncodedGopYuv {
pub ivf_bytes: Vec<u8>,
pub temporal_units: Vec<Vec<u8>>,
pub recon: Vec<GopFrameReconYuv>,
pub seq: SequenceHeader,
}
#[doc(hidden)]
#[derive(Debug, Clone)]
pub struct TunedGopYuv {
pub gop: EncodedGopYuv,
pub seg_temporal_updates: Vec<bool>,
pub p_segment_maps: Vec<Vec<i32>>,
pub hp_mv_elections: Vec<bool>,
pub delta_q_elections: Vec<bool>,
pub qm_elections: Vec<bool>,
pub cdef_elections: Vec<bool>,
pub lr_elections: Vec<bool>,
pub ctx_donor_elections: Vec<Option<u32>>,
pub film_grain_elected: bool,
pub superres_inter_elections: Vec<Option<u32>>,
}
#[derive(Debug, Clone)]
pub struct EncodedGop {
pub ivf_bytes: Vec<u8>,
pub temporal_units: Vec<Vec<u8>>,
pub recon: Vec<GopFrameRecon>,
pub seq: SequenceHeader,
}
#[doc(hidden)]
#[derive(Debug, Clone)]
pub struct TunedGop {
pub gop: EncodedGop,
pub seg_temporal_updates: Vec<bool>,
pub p_segment_maps: Vec<Vec<i32>>,
pub hp_mv_elections: Vec<bool>,
pub delta_q_elections: Vec<bool>,
pub qm_elections: Vec<bool>,
pub cdef_elections: Vec<bool>,
pub lr_elections: Vec<bool>,
pub ctx_donor_elections: Vec<Option<u32>>,
pub film_grain_elected: bool,
pub superres_inter_elections: Vec<Option<u32>>,
}
pub const GOP_MAX_FRAMES: usize = 64;
pub(crate) fn narrow_gop_8bit(g: EncodedGopYuv) -> EncodedGop {
let narrow = |p: Vec<u16>| p.into_iter().map(|s| s as u8).collect::<Vec<u8>>();
EncodedGop {
ivf_bytes: g.ivf_bytes,
temporal_units: g.temporal_units,
recon: g
.recon
.into_iter()
.map(|r| GopFrameRecon {
y: narrow(r.y),
u: narrow(r.u),
v: narrow(r.v),
})
.collect(),
seq: g.seq,
}
}
#[derive(Debug, Clone)]
pub(crate) struct SavedMotionField {
mf_mvs: Vec<i16>,
mf_ref_frames: Vec<i8>,
saved_order_hints: [i32; 8],
mi_rows: u32,
mi_cols: u32,
frame_is_intra: bool,
}
impl SavedMotionField {
pub(crate) fn intra(mi_rows: u32, mi_cols: u32) -> Self {
SavedMotionField {
mf_mvs: Vec::new(),
mf_ref_frames: Vec::new(),
saved_order_hints: [0; 8],
mi_rows,
mi_cols,
frame_is_intra: true,
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct RefSlotCarry {
pub(crate) cdfs: Box<TileCdfContext>,
pub(crate) segment_ids: Vec<i32>,
pub(crate) mi_rows: u32,
pub(crate) mi_cols: u32,
pub(crate) gm_params: [[i32; 6]; crate::uncompressed_header_tail::TOTAL_REFS_PER_FRAME],
pub(crate) donor_cdfs: Vec<Box<TileCdfContext>>,
pub(crate) ctx_update_span: Option<(u32, u32)>,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct InterFrameAux {
pub(crate) seg_temporal: bool,
pub(crate) primary_ref: u8,
pub(crate) hp_mv: bool,
pub(crate) delta_q: bool,
pub(crate) qm: bool,
pub(crate) cdef: bool,
pub(crate) lr: bool,
pub(crate) donor_elected: Option<u32>,
}
const SEARCH_RANGE: i32 = 16;
pub(crate) const DELTA_Q_RES: u8 = 3;
pub(crate) fn qm_level_for_q(base_q_idx: u8) -> u8 {
match base_q_idx {
0..=80 => 10,
81..=160 => 8,
_ => 6,
}
}
pub(crate) fn qm_ladder_for_q(base_q_idx: u8) -> Vec<u8> {
let keyed = qm_level_for_q(base_q_idx);
let lighter = (keyed + 2).min(12);
if lighter == keyed {
vec![keyed]
} else {
vec![keyed, lighter]
}
}
pub(crate) const QM_LADDER_REFINES_ONLY: bool = true;
pub(crate) fn qm_arm_allowed(base_q_idx: u8, width: usize, height: usize) -> bool {
(88..=176).contains(&base_q_idx) && width * height >= 96 * 80
}
pub(crate) fn qm_probe(input: &YuvFrame) -> bool {
let (w, h) = (input.width as usize, input.height as usize);
let mut sum = 0u64;
let mut count = 0u64;
for r in 0..h {
for c in 1..w.saturating_sub(1) {
let m = i64::from(input.y[r * w + c]);
sum += (2 * m - i64::from(input.y[r * w + c - 1]) - i64::from(input.y[r * w + c + 1]))
.unsigned_abs();
count += 1;
}
}
for r in 1..h.saturating_sub(1) {
for c in 0..w {
let m = i64::from(input.y[r * w + c]);
sum +=
(2 * m - i64::from(input.y[(r - 1) * w + c]) - i64::from(input.y[(r + 1) * w + c]))
.unsigned_abs();
count += 1;
}
}
if count == 0 {
return false;
}
let mean16 = (sum * 16) / count / (1u64 << (input.bit_depth - 8));
mean16 >= 8
}
pub(crate) fn delta_q_plan_units(input: &YuvFrame, mi_rows: u32, mi_cols: u32) -> Option<Vec<i32>> {
let w = input.width as usize;
let h = input.height as usize;
let sbs = sb_grid_origins(mi_rows, mi_cols);
let mut acts: Vec<f64> = Vec::with_capacity(sbs.len());
for &(sb_r, sb_c) in &sbs {
let y0 = (sb_r as usize) * 4;
let x0 = (sb_c as usize) * 4;
let y1 = (y0 + 64).min(h);
let x1 = (x0 + 64).min(w);
let n = ((y1 - y0) * (x1 - x0)) as u64;
let mut sum = 0u64;
let mut sum2 = 0u128;
for r in y0..y1 {
for c in x0..x1 {
let v = u64::from(input.y[r * w + c]);
sum += v;
sum2 += u128::from(v * v);
}
}
let mean = sum as f64 / n as f64;
let var = (sum2 as f64 / n as f64) - mean * mean;
acts.push(var.max(0.0));
}
let log_mean = acts.iter().map(|&a| a.max(1.0).log2()).sum::<f64>() / acts.len() as f64;
let units: Vec<i32> = acts
.iter()
.map(|&a| (((a.max(1.0).log2() - log_mean) * 0.5).round() as i32).clamp(-2, 2))
.collect();
units.iter().any(|&u| u != 0).then_some(units)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CompoundSel {
ctype: u8,
wedge_index: u8,
wedge_sign: u8,
mask_type: u8,
}
impl CompoundSel {
const AVERAGE: CompoundSel = CompoundSel {
ctype: crate::inter_pred::COMPOUND_AVERAGE,
wedge_index: 0,
wedge_sign: 0,
mask_type: 0,
};
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MotionSel {
Simple,
Obmc,
Warp([i32; 6]),
}
impl MotionSel {
fn ordinal(self) -> u8 {
match self {
MotionSel::Simple => MOTION_MODE_SIMPLE,
MotionSel::Obmc => MOTION_MODE_OBMC,
MotionSel::Warp(_) => MOTION_MODE_WARPED_CAUSAL,
}
}
}
#[derive(Debug, Clone, Copy)]
struct DriverGridCell {
mi_size: usize,
is_inter: u8,
ref_frames: [i8; 2],
mvs: [i16; 4],
interp_filters: [u8; 2],
motion_mode: u8,
y_mode: u8,
compound_type: u8,
wedge_index: u8,
wedge_sign: u8,
mask_type: u8,
interintra_mode: u8,
wedge_interintra: u8,
interintra_wedge_index: u8,
local_warp: [i32; 6],
local_warp_valid: u8,
}
struct DriverGridSnapshot {
mi_row: u32,
mi_col: u32,
r1: u32,
c1: u32,
cells: Vec<DriverGridCell>,
}
#[derive(Clone, Copy)]
struct InterIntraTrial<'a> {
mode: u8,
wedge: Option<u8>,
neigh: &'a ReconState,
}
pub fn encode_gop_yuv420(frames: &[Yuv420Frame]) -> Result<EncodedGop, Error> {
encode_gop_yuv420_with_q(frames, 0)
}
pub fn encode_gop_yuv420_with_q(
frames: &[Yuv420Frame],
base_q_idx: u8,
) -> Result<EncodedGop, Error> {
encode_gop_yuv420_with_q_seg(frames, base_q_idx, &[])
}
pub fn encode_gop_yuv420_with_q_seg(
frames: &[Yuv420Frame],
base_q_idx: u8,
alt_q: &[i16],
) -> Result<EncodedGop, Error> {
encode_gop_yuv420_with_q_seg_rate_model(frames, base_q_idx, alt_q, RateModel::Twin)
}
#[doc(hidden)]
pub fn encode_gop_yuv420_with_q_seg_rate_model(
frames: &[Yuv420Frame],
base_q_idx: u8,
alt_q: &[i16],
model: RateModel,
) -> Result<EncodedGop, Error> {
encode_gop_yuv420_with_q_seg_rate_model_gm(frames, base_q_idx, alt_q, model, true)
}
#[doc(hidden)]
pub fn encode_gop_yuv420_with_q_seg_rate_model_gm(
frames: &[Yuv420Frame],
base_q_idx: u8,
alt_q: &[i16],
model: RateModel,
global_motion: bool,
) -> Result<EncodedGop, Error> {
encode_gop_yuv420_with_q_seg_tuned(
frames,
base_q_idx,
alt_q,
GopTuning {
model,
global_motion,
..GopTuning::default()
},
)
.map(|t| t.gop)
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy)]
pub struct GopTuning {
pub model: RateModel,
pub global_motion: bool,
pub primary_ref: bool,
pub temporal_seg: bool,
pub high_precision_mv: bool,
pub delta_q: bool,
pub qm: bool,
pub cdef: bool,
pub cdef_units: bool,
pub lr: bool,
pub tiles: (u32, u32),
pub tile_groups: u32,
pub ctx_update_elect: bool,
pub superres: bool,
pub film_grain: bool,
pub s_frame_period: u32,
pub error_resilient_period: u32,
pub tile_spans: bool,
pub short_ref_signaling: bool,
}
impl Default for GopTuning {
fn default() -> Self {
GopTuning {
model: RateModel::Twin,
global_motion: true,
primary_ref: true,
temporal_seg: true,
high_precision_mv: true,
delta_q: true,
qm: true,
cdef: true,
cdef_units: true,
lr: true,
tiles: (0, 0),
tile_groups: 1,
ctx_update_elect: true,
superres: true,
film_grain: true,
s_frame_period: 0,
error_resilient_period: 0,
tile_spans: false,
short_ref_signaling: false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LosslessRegion {
pub x: u32,
pub y: u32,
pub width: u32,
pub height: u32,
}
fn build_exact_mask(
regions: &[LosslessRegion],
width: u32,
height: u32,
mi_rows: u32,
mi_cols: u32,
) -> Vec<bool> {
let mut mask = vec![false; (mi_rows as usize) * (mi_cols as usize)];
for reg in regions {
let x0 = reg.x.min(width);
let y0 = reg.y.min(height);
let x1 = reg.x.saturating_add(reg.width).min(width);
let y1 = reg.y.saturating_add(reg.height).min(height);
if x0 >= x1 || y0 >= y1 {
continue;
}
let (c0, r0) = (x0 / 4, y0 / 4);
let (c1, r1) = (x1.div_ceil(4).min(mi_cols), y1.div_ceil(4).min(mi_rows));
for r in r0..r1 {
for c in c0..c1 {
mask[(r * mi_cols + c) as usize] = true;
}
}
}
let mut out = mask.clone();
for r in 0..mi_rows {
for c in 0..mi_cols {
if mask[(r * mi_cols + c) as usize] {
let (gr, gc) = (r & !1, c & !1);
for rr in gr..(gr + 2).min(mi_rows) {
for cc in gc..(gc + 2).min(mi_cols) {
out[(rr * mi_cols + cc) as usize] = true;
}
}
}
}
}
out
}
pub fn encode_gop_yuv420_with_q_lossless_regions(
frames: &[Yuv420Frame],
base_q_idx: u8,
regions: &[LosslessRegion],
auto_detect: bool,
) -> Result<EncodedGop, Error> {
if base_q_idx == 0 {
return encode_gop_yuv420_with_q(frames, 0);
}
let alt_q = [0i16, -i16::from(base_q_idx)];
encode_gop_yuv420_with_q_seg_lossless_tuned(
frames,
base_q_idx,
&alt_q,
regions,
auto_detect,
GopTuning::default(),
)
.map(|t| t.gop)
}
#[doc(hidden)]
pub fn encode_gop_yuv420_with_q_seg_tuned(
frames: &[Yuv420Frame],
base_q_idx: u8,
alt_q: &[i16],
tuning: GopTuning,
) -> Result<TunedGop, Error> {
encode_gop_yuv420_with_q_seg_lossless_tuned(frames, base_q_idx, alt_q, &[], false, tuning)
}
#[doc(hidden)]
pub fn encode_gop_yuv420_with_q_seg_lossless_tuned(
frames: &[Yuv420Frame],
base_q_idx: u8,
alt_q: &[i16],
regions: &[LosslessRegion],
auto_detect: bool,
tuning: GopTuning,
) -> Result<TunedGop, Error> {
encode_gop_yuv420_with_q_seg_extras_tuned(
frames,
base_q_idx,
alt_q,
regions,
auto_detect,
None,
tuning,
)
}
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn encode_gop_yuv420_with_q_seg_extras_tuned(
frames: &[Yuv420Frame],
base_q_idx: u8,
alt_q: &[i16],
regions: &[LosslessRegion],
auto_detect: bool,
extras: Option<&SegExtras>,
tuning: GopTuning,
) -> Result<TunedGop, Error> {
let wide: Vec<YuvFrame> = frames.iter().map(YuvFrame::from_yuv420_8bit).collect();
let t = encode_gop_yuv_seg_extras_tuned(
&wide,
base_q_idx,
alt_q,
regions,
auto_detect,
extras,
tuning,
)?;
Ok(TunedGop {
gop: narrow_gop_8bit(t.gop),
seg_temporal_updates: t.seg_temporal_updates,
p_segment_maps: t.p_segment_maps,
hp_mv_elections: t.hp_mv_elections,
delta_q_elections: t.delta_q_elections,
qm_elections: t.qm_elections,
cdef_elections: t.cdef_elections,
lr_elections: t.lr_elections,
ctx_donor_elections: t.ctx_donor_elections,
film_grain_elected: t.film_grain_elected,
superres_inter_elections: t.superres_inter_elections,
})
}
pub fn encode_gop_yuv(frames: &[YuvFrame]) -> Result<EncodedGopYuv, Error> {
encode_gop_yuv_with_q(frames, 0)
}
pub fn encode_gop_yuv_with_q(frames: &[YuvFrame], base_q_idx: u8) -> Result<EncodedGopYuv, Error> {
encode_gop_yuv_seg_extras_tuned(
frames,
base_q_idx,
&[],
&[],
false,
None,
GopTuning::default(),
)
.map(|t| t.gop)
}
pub fn encode_gop_yuv_with_q_tile_layout(
frames: &[YuvFrame],
base_q_idx: u8,
widths_sb: &[u32],
heights_sb: &[u32],
) -> Result<EncodedGopYuv, Error> {
encode_gop_yuv_seg_extras_tuned_layout(
frames,
base_q_idx,
&[],
&[],
false,
None,
GopTuning::default(),
Some((widths_sb, heights_sb)),
)
.map(|t| t.gop)
}
pub fn encode_gop_yuv420_with_q_tile_layout(
frames: &[Yuv420Frame],
base_q_idx: u8,
widths_sb: &[u32],
heights_sb: &[u32],
) -> Result<EncodedGop, Error> {
let wide: Vec<YuvFrame> = frames.iter().map(YuvFrame::from_yuv420_8bit).collect();
encode_gop_yuv_with_q_tile_layout(&wide, base_q_idx, widths_sb, heights_sb).map(narrow_gop_8bit)
}
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn encode_gop_yuv_seg_extras_tuned(
frames: &[YuvFrame],
base_q_idx: u8,
alt_q: &[i16],
regions: &[LosslessRegion],
auto_detect: bool,
extras: Option<&SegExtras>,
tuning: GopTuning,
) -> Result<TunedGopYuv, Error> {
encode_gop_yuv_seg_extras_tuned_layout(
frames,
base_q_idx,
alt_q,
regions,
auto_detect,
extras,
tuning,
None,
)
}
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn encode_gop_yuv_seg_extras_tuned_layout(
frames: &[YuvFrame],
base_q_idx: u8,
alt_q: &[i16],
regions: &[LosslessRegion],
auto_detect: bool,
extras: Option<&SegExtras>,
tuning: GopTuning,
explicit_tiles: Option<(&[u32], &[u32])>,
) -> Result<TunedGopYuv, Error> {
let plain = encode_gop_yuv_core_fg(
frames,
base_q_idx,
alt_q,
regions,
auto_detect,
extras,
tuning,
explicit_tiles,
None,
)?;
if !(tuning.film_grain
&& base_q_idx > 0
&& regions.is_empty()
&& !auto_detect
&& tuning.model == RateModel::Twin
&& frames.len() >= 2)
{
return Ok(plain);
}
let elected = elect_film_grain(
frames,
base_q_idx,
plain.gop.seq.color_config.matrix_coefficients,
plain.gop.ivf_bytes.len(),
&plain.gop.recon,
|denoised, fg| {
encode_gop_yuv_core_fg(
denoised,
base_q_idx,
alt_q,
regions,
auto_detect,
extras,
tuning,
explicit_tiles,
Some(fg),
)
},
|c: &TunedGopYuv| (c.gop.ivf_bytes.len(), c.gop.recon.as_slice()),
)?;
match elected {
Some((cand, grained)) => {
let mut out = cand;
out.gop.recon = grained;
out.film_grain_elected = true;
Ok(out)
}
None => Ok(plain),
}
}
pub(crate) fn elect_film_grain<T, F, V>(
frames: &[YuvFrame],
base_q_idx: u8,
mc: u8,
plain_bytes: usize,
plain_recon: &[GopFrameReconYuv],
mut encode: F,
view: V,
) -> Result<Option<(T, Vec<GopFrameReconYuv>)>, Error>
where
F: FnMut(&[YuvFrame], &FilmGrainParams) -> Result<T, Error>,
V: Fn(&T) -> (usize, &[GopFrameReconYuv]),
{
if frames.len() < 2 {
return Ok(None);
}
use crate::encoder::film_grain_elect as fge;
let dn01: Vec<YuvFrame> = frames[..2].iter().map(fge::denoise_frame).collect();
if !fge::film_grain_probe(&frames[..2], &dn01) {
return Ok(None);
}
let mut denoised = dn01;
denoised.extend(frames[2..].iter().map(fge::denoise_frame));
let est = fge::noise_estimate(&frames[0], &denoised[0]);
let bit_depth = frames[0].bit_depth;
let (width, height) = (frames[0].width, frames[0].height);
let format = frames[0].format;
let (ssx, ssy) = format.subsampling();
let num_planes = format.num_planes();
let mut chroma_cb =
num_planes > 1 && fge::chroma_noise_gate(&frames[..2], &denoised[..2], 1, est.cb_sigma16);
let mut chroma_cr =
num_planes > 1 && fge::chroma_noise_gate(&frames[..2], &denoised[..2], 2, est.cr_sigma16);
if ssx == 1 && ssy == 1 && !(chroma_cb && chroma_cr) {
chroma_cb = false;
chroma_cr = false;
}
if std::env::var_os("OXIDEAV_AV1_FG_DEBUG").is_some() {
eprintln!(
"fg-elect gates: cb {} cr {} (sigma16 y {} cb {} cr {})",
chroma_cb, chroma_cr, est.sigma16, est.cb_sigma16, est.cr_sigma16
);
}
let fit1 = fge::fit_ar_lag(&frames[0], &denoised[0], 1);
let fit2 = fge::fit_ar_lag(&frames[0], &denoised[0], 2);
let fit3 = {
let (w, h) = (width as usize, height as usize);
(fge::lag_h_rho(&frames[0].y, &denoised[0].y, w, h, 3).abs() >= 0.2)
.then(|| fge::fit_ar_lag(&frames[0], &denoised[0], 3))
.flatten()
};
let mut cand_params: Vec<FilmGrainParams> = vec![fge::build_grain_params(
&est, bit_depth, ssx, ssy, mc, None, false, false, false,
)];
let csfl_ok = chroma_cb && chroma_cr;
let chroma_only: Option<fge::ArFit> = ((chroma_cb || chroma_cr) && fit1.is_none())
.then(|| fge::fit_chroma_corr(&frames[0], &denoised[0]));
for fit in [
fit1.as_ref(),
fit2.as_ref(),
fit3.as_ref(),
chroma_only.as_ref(),
]
.into_iter()
.flatten()
{
cand_params.push(fge::build_grain_params(
&est,
bit_depth,
ssx,
ssy,
mc,
Some(fit),
chroma_cb,
chroma_cr,
false,
));
if csfl_ok {
cand_params.push(fge::build_grain_params(
&est,
bit_depth,
ssx,
ssy,
mc,
Some(fit),
false,
false,
true,
));
}
}
let sse_pair = |a: &GopFrameReconYuv, f: &YuvFrame| -> u64 {
let mut d = 0u64;
for (x, y) in
a.y.iter()
.zip(&f.y)
.chain(a.u.iter().zip(&f.u))
.chain(a.v.iter().zip(&f.v))
{
let diff = i64::from(*x) - i64::from(*y);
d += (diff * diff) as u64;
}
d
};
let d_plain: u64 = plain_recon
.iter()
.zip(frames)
.map(|(rc, f)| sse_pair(rc, f))
.sum();
let lambda = lambda_for(&QuantizerParams::neutral(base_q_idx, bit_depth));
let score_plain = score256(d_plain, lambda, (plain_bytes as u64) * 8 * 256);
let shift = u32::from(bit_depth - 8);
let mut best: Option<(T, Vec<GopFrameReconYuv>, u64)> = None;
for fg_base in &cand_params {
let cand = encode(&denoised, fg_base)?;
let (cand_bytes, cand_recon) = view(&cand);
let grained: Vec<GopFrameReconYuv> = cand_recon
.iter()
.enumerate()
.map(|(k, rc)| {
let mut fg = fg_base.clone();
fg.grain_seed = fge::grain_seed_for(k as u32);
let (gy, gu, gv) = fge::apply_grain_to_recon(
&fg, &rc.y, &rc.u, &rc.v, width, height, bit_depth, ssx, ssy, num_planes, mc,
);
GopFrameReconYuv {
y: gy,
u: gu,
v: gv,
}
})
.collect();
let mut d_grain: u64 = cand_recon
.iter()
.zip(&denoised)
.map(|(rc, f)| sse_pair(rc, f))
.sum();
for k in 0..frames.len() {
type PlaneRefs<'p> = (&'p [u16], &'p [u16], &'p [u16], &'p [u16], usize);
let planes: [PlaneRefs<'_>; 3] = [
(
&frames[k].y,
&denoised[k].y,
&grained[k].y,
&cand_recon[k].y,
frames[k].y.len(),
),
(
&frames[k].u,
&denoised[k].u,
&grained[k].u,
&cand_recon[k].u,
frames[k].u.len(),
),
(
&frames[k].v,
&denoised[k].v,
&grained[k].v,
&cand_recon[k].v,
frames[k].v.len(),
),
];
for (src, den, syn, pre, samples) in planes {
if samples == 0 {
continue;
}
let sig_src = fge::plane_sigma16(src, den, bit_depth);
let sig_syn = fge::plane_sigma16(syn, pre, bit_depth);
let diff = sig_src.abs_diff(sig_syn);
d_grain += (samples as u64) * diff * diff * (1u64 << (2 * shift)) / 256;
}
let w = width as usize;
let h = height as usize;
let sig_src = fge::plane_sigma16(&frames[k].y, &denoised[k].y, bit_depth);
let mut gap2 = 0.0f64;
for lag in 1..=3usize {
let rho_src = fge::lag_h_rho(&frames[k].y, &denoised[k].y, w, h, lag);
let rho_syn = fge::lag_h_rho(&grained[k].y, &cand_recon[k].y, w, h, lag);
let gap = rho_src - rho_syn;
gap2 += gap * gap;
}
let term = (frames[k].y.len() as f64) * gap2 * ((sig_src * sig_src) as f64) / 256.0
* f64::from(1u32 << (2 * shift));
d_grain += term.round() as u64;
}
let score = score256(d_grain, lambda, (cand_bytes as u64) * 8 * 256);
if std::env::var_os("OXIDEAV_AV1_FG_DEBUG").is_some() {
eprintln!(
"fg-elect cand: lag {} cb {} cr {} csfl {} bytes {} d_grain {} score {} | plain bytes {} d {} score {}",
fg_base.ar_coeff_lag,
fg_base.num_cb_points,
fg_base.num_cr_points,
fg_base.chroma_scaling_from_luma,
cand_bytes,
d_grain,
score,
plain_bytes,
d_plain,
score_plain
);
}
if cand_bytes < plain_bytes && best.as_ref().map(|(_, _, s)| score < *s).unwrap_or(true) {
best = Some((cand, grained, score));
}
}
match best {
Some((cand, grained, score_grain)) if score_grain < score_plain => {
Ok(Some((cand, grained)))
}
_ => Ok(None),
}
}
#[allow(clippy::too_many_arguments)]
fn encode_gop_yuv_core_fg(
frames: &[YuvFrame],
base_q_idx: u8,
alt_q: &[i16],
regions: &[LosslessRegion],
auto_detect: bool,
extras: Option<&SegExtras>,
tuning: GopTuning,
explicit_tiles: Option<(&[u32], &[u32])>,
film_grain: Option<&FilmGrainParams>,
) -> Result<TunedGopYuv, Error> {
if let Some(x) = extras {
let ll = seg_lossless_array(base_q_idx, alt_q);
if alt_q.is_empty()
|| x.ref_frame[0].is_some()
|| x.skip[0]
|| x.globalmv[0]
|| x.ref_frame.iter().flatten().any(|rf| !(1..=7).contains(rf))
|| (0..crate::uncompressed_header_tail::MAX_SEGMENTS).any(|s| x.skip[s] && ll[s])
{
return Err(Error::PartitionWalkOutOfRange);
}
}
let model = tuning.model;
let global_motion = tuning.global_motion;
if frames.is_empty() || frames.len() > GOP_MAX_FRAMES {
return Err(Error::PartitionWalkOutOfRange);
}
if !alt_q.is_empty() {
if alt_q.len() == 1
|| alt_q.len() > crate::uncompressed_header_tail::MAX_SEGMENTS
|| alt_q[0] != 0
|| base_q_idx == 0
{
return Err(Error::PartitionWalkOutOfRange);
}
for &d in alt_q {
if !(-255..=255).contains(&d) || i32::from(base_q_idx) + i32::from(d) > 255 {
return Err(Error::PartitionWalkOutOfRange);
}
}
}
let (width, height) = (frames[0].width, frames[0].height);
let (bit_depth, format) = (frames[0].bit_depth, frames[0].format);
if frames.iter().any(|f| {
f.width != width || f.height != height || f.bit_depth != bit_depth || f.format != format
}) {
return Err(Error::PartitionWalkOutOfRange);
}
let key_mask: Option<Vec<bool>> = (!regions.is_empty()).then(|| {
build_exact_mask(regions, width, height, height / 4, width / 4)
});
let donor_armed = tuning.ctx_update_elect
&& tuning.primary_ref
&& (tuning.tiles != (0, 0) || explicit_tiles.is_some());
let key_fg: Option<FilmGrainParams> = film_grain.map(|base| {
let mut f = base.clone();
f.grain_seed = crate::encoder::film_grain_elect::grain_seed_for(0);
f
});
let sr_inter_armed = tuning.superres
&& base_q_idx > 0
&& alt_q.is_empty()
&& regions.is_empty()
&& !auto_detect
&& model == RateModel::Twin
&& crate::encoder::superres_elect::superres_arm_allowed(
base_q_idx,
width as usize,
height as usize,
)
&& frames
.iter()
.skip(1)
.any(crate::encoder::superres_elect::superres_probe);
let (key, key_carry) = crate::encoder::key_frame::encode_key_frame_yuv_full(
&frames[0],
base_q_idx,
model,
if key_mask.is_some() { alt_q } else { &[] },
key_mask.as_deref(),
tuning.cdef,
tuning.cdef_units,
tuning.lr,
&crate::encoder::key_frame::KeyExtras {
tiles: tuning.tiles,
tile_groups: tuning.tile_groups,
tile_spans: tuning.tile_spans,
explicit_tiles,
delta_q: tuning.delta_q,
qm: tuning.qm,
superres_elect: tuning.superres,
film_grain: key_fg.as_ref(),
collect_donor_cdfs: donor_armed,
superres_gate: sr_inter_armed,
..crate::encoder::key_frame::KeyExtras::default()
},
)?;
let seq = key.seq.clone();
let key_sr_denom: Option<u32> = key
.fh
.frame_size
.as_ref()
.and_then(|fs| fs.use_superres.then_some(fs.superres_denom));
let sr_lambda =
crate::encoder::key_frame::lambda_for(&QuantizerParams::neutral(base_q_idx, bit_depth));
let mut temporal_units = vec![key.temporal_unit_bytes.clone()];
let mut recon = vec![GopFrameReconYuv {
y: key.recon_y,
u: key.recon_u,
v: key.recon_v,
}];
let key_mi = {
let fh0 = build_intra_only_yuv420_8bit_fh_with_q(&seq, width, height, base_q_idx);
let fs0 = fh0.frame_size.expect("KEY builder always sizes");
(fs0.mi_rows, fs0.mi_cols)
};
if (!regions.is_empty() || auto_detect)
&& !seg_lossless_array(base_q_idx, alt_q)[..alt_q.len()]
.iter()
.any(|&l| l)
{
return Err(Error::PartitionWalkOutOfRange);
}
let exact_mask: Option<Vec<bool>> =
(!regions.is_empty()).then(|| build_exact_mask(regions, width, height, key_mi.0, key_mi.1));
let mut mf_store: [SavedMotionField; 8] =
core::array::from_fn(|_| SavedMotionField::intra(key_mi.0, key_mi.1));
let key_carry = std::rc::Rc::new(key_carry);
let mut carry_store: [std::rc::Rc<RefSlotCarry>; 8] =
core::array::from_fn(|_| key_carry.clone());
let mut carry_tu: [usize; 8] = [0; 8];
let mut seg_temporal_updates: Vec<bool> = Vec::new();
let mut p_segment_maps: Vec<Vec<i32>> = Vec::new();
let mut hp_mv_elections: Vec<bool> = Vec::new();
let mut delta_q_elections: Vec<bool> = Vec::new();
let mut qm_elections: Vec<bool> = Vec::new();
let mut cdef_elections: Vec<bool> = Vec::new();
let mut lr_elections: Vec<bool> = Vec::new();
let mut ctx_donor_elections: Vec<Option<u32>> = Vec::new();
let mut superres_inter_elections: Vec<Option<u32>> = Vec::new();
let mut floor: u32 = 0;
for (k, input) in frames[1..].iter().enumerate() {
let p_index = (k + 1) as u32;
let is_s = tuning.s_frame_period > 0 && p_index % tuning.s_frame_period == 0;
let is_er = !is_s
&& tuning.error_resilient_period > 0
&& p_index % tuning.error_resilient_period == 0;
let prev = recon.last().expect("at least the KEY recon");
let prevprev = &recon[p_index.saturating_sub(2).max(floor) as usize];
let primary = if tuning.primary_ref && !is_s && !is_er {
Some(carry_store[(p_index & 1) as usize].clone())
} else {
None
};
let p_fg: Option<FilmGrainParams> = film_grain.map(|base| {
let mut f = base.clone();
f.grain_seed = crate::encoder::film_grain_elect::grain_seed_for(p_index);
f
});
let encode_arm = |arm_input: &YuvFrame, superres: Option<(u32, u32)>| {
encode_p_frame_yuv(
arm_input,
prev,
prevprev,
&seq,
base_q_idx,
p_index,
alt_q,
&mf_store,
model,
global_motion,
primary.as_deref(),
tuning.temporal_seg,
exact_mask.as_deref(),
auto_detect,
extras,
tuning.high_precision_mv,
tuning.delta_q,
tuning.qm,
tuning.cdef,
tuning.cdef_units,
tuning.lr,
tuning.tiles,
tuning.tile_groups,
tuning.tile_spans,
explicit_tiles,
donor_armed,
p_fg.as_ref(),
floor,
&recon[floor as usize],
is_s,
is_er,
tuning.short_ref_signaling,
superres,
superres.is_some().then_some(input),
)
};
let mut best = encode_arm(input, None)?;
let mut sr_elected: Option<u32> = None;
if sr_inter_armed
&& !is_s
&& !is_er
&& crate::encoder::superres_elect::superres_probe(input)
{
let sr_score = |tu: &[u8], rc: &GopFrameReconYuv| -> u64 {
let mut d = 0u64;
for (a, b) in
rc.y.iter()
.zip(&input.y)
.chain(rc.u.iter().zip(&input.u))
.chain(rc.v.iter().zip(&input.v))
{
let diff = i64::from(*a) - i64::from(*b);
d += (diff * diff) as u64;
}
score256(d, sr_lambda, (tu.len() as u64) * 8 * 256)
};
let mut best_score = sr_score(&best.0, &best.1);
for denom in crate::encoder::superres_elect::inter_candidate_denoms(
width,
height,
tuning.tiles,
explicit_tiles,
key_sr_denom,
) {
let wd = crate::encoder::superres_elect::superres_coded_width(width, denom);
let down = crate::encoder::superres_elect::downscale_width(input, wd);
let cand = encode_arm(&down, Some((width, denom)))?;
let sc = sr_score(&cand.0, &cand.1);
if sc < best_score {
best_score = sc;
best = cand;
sr_elected = Some(denom);
}
}
}
superres_inter_elections.push(sr_elected);
let (tu, rc, saved_mf, carry, aux) = best;
temporal_units.push(tu);
recon.push(rc);
if let Some(t) = aux.donor_elected {
let src_slot = (p_index & 1) as usize;
let elected_from = carry_store[src_slot].clone();
let span = elected_from
.ctx_update_span
.expect("donor election only fires on multi-tile primaries");
patch_ctx_update_in_tu(&mut temporal_units[carry_tu[src_slot]], span, t);
let fixed = std::rc::Rc::new(RefSlotCarry {
cdfs: elected_from.donor_cdfs[t as usize].clone(),
segment_ids: elected_from.segment_ids.clone(),
mi_rows: elected_from.mi_rows,
mi_cols: elected_from.mi_cols,
gm_params: elected_from.gm_params,
donor_cdfs: Vec::new(),
ctx_update_span: None,
});
for slot in carry_store.iter_mut() {
if std::rc::Rc::ptr_eq(slot, &elected_from) {
*slot = fixed.clone();
}
}
}
seg_temporal_updates.push(aux.seg_temporal);
hp_mv_elections.push(aux.hp_mv);
delta_q_elections.push(aux.delta_q);
qm_elections.push(aux.qm);
cdef_elections.push(aux.cdef);
lr_elections.push(aux.lr);
ctx_donor_elections.push(aux.donor_elected);
if !alt_q.is_empty() {
p_segment_maps.push(carry.segment_ids.clone());
}
if is_s {
floor = p_index;
let carry_rc = std::rc::Rc::new(carry);
for slot in 0..8usize {
mf_store[slot] = saved_mf.clone();
carry_store[slot] = carry_rc.clone();
carry_tu[slot] = p_index as usize;
}
} else {
mf_store[((p_index - 1) & 1) as usize] = saved_mf;
carry_store[((p_index - 1) & 1) as usize] = std::rc::Rc::new(carry);
carry_tu[((p_index - 1) & 1) as usize] = p_index as usize;
}
}
let mut ivf_bytes: Vec<u8> = Vec::new();
{
let cursor = std::io::Cursor::new(&mut ivf_bytes);
let mut iw = IvfWriter::new(cursor, FOURCC_AV01, width as u16, height as u16, 25, 1)
.map_err(|_| Error::PartitionWalkOutOfRange)?;
for (idx, tu) in temporal_units.iter().enumerate() {
iw.write_frame(tu, idx as u64)
.map_err(|_| Error::PartitionWalkOutOfRange)?;
}
iw.patch_frame_count()
.map_err(|_| Error::PartitionWalkOutOfRange)?;
}
Ok(TunedGopYuv {
gop: EncodedGopYuv {
ivf_bytes,
temporal_units,
recon,
seq,
},
seg_temporal_updates,
p_segment_maps,
hp_mv_elections,
delta_q_elections,
qm_elections,
cdef_elections,
lr_elections,
ctx_donor_elections,
film_grain_elected: false,
superres_inter_elections,
})
}
#[cfg(test)]
fn gop_ref_hints(p_index: u32) -> (u32, u32) {
gop_ref_hints_floor(p_index, 0)
}
fn gop_ref_hints_floor(p_index: u32, floor: u32) -> (u32, u32) {
(p_index - 1, p_index.saturating_sub(2).max(floor))
}
#[cfg(test)]
fn gop_order_hints(p_index: u32, order_hint_bits: u8) -> crate::cdf::FrameInterOrderHints {
let (last_hint, golden_hint) = gop_ref_hints(p_index);
let mut by_ref = [0i32; crate::uncompressed_header_tail::ALTREF_FRAME + 1];
for (r, hint) in by_ref.iter_mut().enumerate().skip(1) {
*hint = if r == 4 {
golden_hint as i32
} else {
last_hint as i32
};
}
crate::cdf::FrameInterOrderHints {
order_hint_bits: u32::from(order_hint_bits),
current_order_hint: p_index as i32,
order_hints_by_ref: by_ref,
}
}
pub(crate) struct InterFrameConfig<'a> {
pub order_hint: u32,
pub show_frame: bool,
pub refresh_frame_flags: u8,
pub ref_frame_idx: [u8; 7],
pub short_ref_signaling: bool,
pub slot_hints: [u32; 8],
pub single_refs: Vec<i8>,
pub compound_pairs: Vec<[i8; 2]>,
pub refs: Vec<&'a GopFrameReconYuv>,
pub slot_to_plane: [usize; 8],
pub primary_ref_frame: u8,
pub primary_carry: Option<&'a RefSlotCarry>,
pub allow_temporal_seg: bool,
pub alt_primaries: Vec<(u8, Option<&'a RefSlotCarry>)>,
pub exact_mask: Option<&'a [bool]>,
pub auto_lossless: bool,
pub seg_extras: Option<&'a SegExtras>,
pub high_precision_mv: bool,
pub delta_q: bool,
pub qm: bool,
pub cdef: bool,
pub cdef_units: bool,
pub lr: bool,
pub freeze_cdfs: bool,
pub tiles: (u32, u32),
pub tile_groups: u32,
pub collect_donor_cdfs: bool,
pub elect_donor: bool,
pub explicit_tiles: Option<(&'a [u32], &'a [u32])>,
pub film_grain: Option<&'a FilmGrainParams>,
pub s_frame: bool,
pub error_resilient: bool,
pub tile_spans: bool,
pub superres: Option<(u32, u32)>,
pub superres_source: Option<&'a YuvFrame>,
pub ref_dims: Vec<(u32, u32)>,
}
fn build_inter_frame_fh(
seq: &SequenceHeader,
width: u32,
height: u32,
base_q_idx: u8,
cfg: &InterFrameConfig<'_>,
alt_q: &[i16],
) -> FrameHeader {
let mut fh = build_intra_only_yuv420_8bit_fh_with_q(seq, width, height, base_q_idx);
fh.frame_type = FrameType::Inter;
fh.frame_is_intra = false;
fh.frame_size_override_flag =
width != seq.max_frame_width_minus_1 + 1 || height != seq.max_frame_height_minus_1 + 1;
fh.show_frame = cfg.show_frame;
fh.showable_frame = true;
fh.error_resilient_mode = false;
fh.force_integer_mv = false;
fh.disable_cdf_update = cfg.freeze_cdfs;
fh.disable_frame_end_update_cdf = cfg.freeze_cdfs;
fh.primary_ref_frame = cfg.primary_ref_frame;
fh.refresh_frame_flags = cfg.refresh_frame_flags;
fh.order_hint = cfg.order_hint & ((1 << u32::from(seq.order_hint_bits)) - 1);
fh.ref_order_hints = Some(cfg.slot_hints);
if let Some(fg) = cfg.film_grain {
fh.film_grain_params = Some(fg.clone());
}
fh.tx_mode = Some(if base_q_idx == 0 {
TxMode::Only4x4
} else {
TxMode::TxModeSelect
});
let reference_select = !cfg.compound_pairs.is_empty();
fh.reference_select = Some(reference_select);
let (skip_allowed, _) = crate::encoder::frame_obu::skip_mode_params_twin(
&cfg.slot_hints,
&cfg.ref_frame_idx,
fh.order_hint,
seq.order_hint_bits,
);
fh.skip_mode_present = Some(skip_allowed && reference_select && seq.enable_order_hint);
fh.allow_warped_motion = Some(seq.enable_warped_motion);
fh.inter_refs = Some(InterFrameRefs {
frame_refs_short_signaling: cfg.short_ref_signaling,
last_frame_idx: cfg.short_ref_signaling.then_some(cfg.ref_frame_idx[0]),
gold_frame_idx: cfg
.short_ref_signaling
.then_some(cfg.ref_frame_idx[GOLDEN_FRAME - LAST_FRAME]),
ref_frame_idx: cfg.ref_frame_idx,
allow_high_precision_mv: cfg.high_precision_mv,
interpolation_filter: InterpolationFilter::Switchable,
is_motion_mode_switchable: true,
use_ref_frame_mvs: seq.enable_order_hint
&& seq.enable_ref_frame_mvs
&& !cfg.error_resilient,
});
if cfg.error_resilient {
fh.error_resilient_mode = true;
fh.primary_ref_frame = PRIMARY_REF_NONE;
fh.allow_warped_motion = Some(false);
}
if cfg.s_frame {
fh.frame_type = FrameType::Switch;
fh.error_resilient_mode = true;
fh.frame_size_override_flag = true;
fh.refresh_frame_flags = ALL_FRAMES_PUB;
fh.primary_ref_frame = PRIMARY_REF_NONE;
fh.allow_warped_motion = Some(false);
}
if !alt_q.is_empty() {
fh.segmentation_params = Some(match cfg.seg_extras {
Some(x) => segmentation_params_for_plan(alt_q, x),
None => segmentation_params_for(alt_q),
});
}
fh
}
pub(crate) fn segmentation_params_for(
alt_q: &[i16],
) -> crate::uncompressed_header_tail::SegmentationParams {
segmentation_params_for_plan(alt_q, &SegExtras::default())
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SegExtras {
pub ref_frame: [Option<i8>; crate::uncompressed_header_tail::MAX_SEGMENTS],
pub skip: [bool; crate::uncompressed_header_tail::MAX_SEGMENTS],
pub globalmv: [bool; crate::uncompressed_header_tail::MAX_SEGMENTS],
}
impl SegExtras {
pub(crate) fn any(&self) -> bool {
self.ref_frame.iter().any(|r| r.is_some())
|| self.skip.iter().any(|&b| b)
|| self.globalmv.iter().any(|&b| b)
}
pub(crate) fn top_segment(&self) -> usize {
(0..crate::uncompressed_header_tail::MAX_SEGMENTS)
.rev()
.find(|&s| self.ref_frame[s].is_some() || self.skip[s] || self.globalmv[s])
.unwrap_or(0)
}
}
pub(crate) fn segmentation_params_for_plan(
alt_q: &[i16],
extras: &SegExtras,
) -> crate::uncompressed_header_tail::SegmentationParams {
use crate::uncompressed_header_tail::{
SegmentationParams, SEG_LVL_ALT_Q, SEG_LVL_GLOBALMV, SEG_LVL_REF_FRAME, SEG_LVL_SKIP,
};
let mut sp = SegmentationParams::disabled();
sp.enabled = true;
sp.update_map = true;
sp.temporal_update = false;
sp.update_data = true;
for (seg, &d) in alt_q.iter().enumerate() {
sp.segment_feature_active[seg][SEG_LVL_ALT_Q] = true;
sp.segment_feature_data[seg][SEG_LVL_ALT_Q] = d;
}
for (seg, rf) in extras.ref_frame.iter().enumerate() {
if let Some(data) = rf {
sp.segment_feature_active[seg][SEG_LVL_REF_FRAME] = true;
sp.segment_feature_data[seg][SEG_LVL_REF_FRAME] = i16::from(*data);
}
}
for (seg, &on) in extras.skip.iter().enumerate() {
sp.segment_feature_active[seg][SEG_LVL_SKIP] = on;
}
for (seg, &on) in extras.globalmv.iter().enumerate() {
sp.segment_feature_active[seg][SEG_LVL_GLOBALMV] = on;
}
sp.seg_id_pre_skip = extras.any();
sp.last_active_seg_id = alt_q.len().saturating_sub(1).max(if extras.any() {
extras.top_segment()
} else {
0
}) as u8;
sp
}
pub(crate) fn seg_lossless_array(
base_q_idx: u8,
alt_q: &[i16],
) -> [bool; crate::uncompressed_header_tail::MAX_SEGMENTS] {
core::array::from_fn(|s| {
let q = match alt_q.get(s) {
Some(&d) => (i32::from(base_q_idx) + i32::from(d)).clamp(0, 255),
None => i32::from(base_q_idx),
};
q == 0
})
}
#[must_use]
pub(crate) fn delta_q_lossless_up_cap(base_q_idx: u8, alt_q: &[i16]) -> i32 {
let ll = seg_lossless_array(base_q_idx, alt_q);
let step = 1i32 << DELTA_Q_RES;
alt_q
.iter()
.enumerate()
.filter(|&(sid, _)| ll[sid])
.map(|(_, &d)| ((-i32::from(d)) - i32::from(base_q_idx)) / step)
.min()
.unwrap_or(2)
.clamp(0, 2)
}
#[cfg(test)]
fn p_frame_config<'a>(
prev: &'a GopFrameReconYuv,
prevprev: &'a GopFrameReconYuv,
p_index: u32,
) -> InterFrameConfig<'a> {
p_frame_config_primary(prev, prevprev, p_index, None, 0, false)
}
fn p_frame_config_primary<'a>(
prev: &'a GopFrameReconYuv,
prevprev: &'a GopFrameReconYuv,
p_index: u32,
primary_carry: Option<&'a RefSlotCarry>,
floor: u32,
s_frame: bool,
) -> InterFrameConfig<'a> {
let last_slot = (p_index & 1) as usize;
let golden_slot = ((p_index - 1) & 1) as usize;
let (last_hint, golden_hint) = gop_ref_hints_floor(p_index, floor);
let mut slot_hints = [floor; 8];
slot_hints[last_slot] = last_hint;
slot_hints[golden_slot] = golden_hint;
let mut ref_frame_idx = [last_slot as u8; REFS_PER_FRAME];
ref_frame_idx[3] = golden_slot as u8;
let mut slot_to_plane = [0usize; 8];
slot_to_plane[golden_slot] = 1;
InterFrameConfig {
order_hint: p_index,
show_frame: true,
refresh_frame_flags: if s_frame {
ALL_FRAMES_PUB
} else {
1 << ((p_index - 1) & 1)
},
ref_frame_idx,
short_ref_signaling: false,
slot_hints,
single_refs: vec![1, 4],
compound_pairs: vec![[1, 4]],
refs: vec![prev, prevprev],
slot_to_plane,
primary_ref_frame: if primary_carry.is_some() && !s_frame {
0 } else {
PRIMARY_REF_NONE
},
primary_carry: if s_frame { None } else { primary_carry },
allow_temporal_seg: true,
alt_primaries: Vec::new(),
exact_mask: None,
auto_lossless: false,
seg_extras: None,
film_grain: None,
high_precision_mv: true,
delta_q: true,
qm: true,
cdef: true,
cdef_units: true,
lr: true,
freeze_cdfs: false,
tiles: (0, 0),
tile_groups: 1,
collect_donor_cdfs: false,
elect_donor: false,
explicit_tiles: None,
s_frame,
error_resilient: s_frame,
tile_spans: false,
superres: None,
superres_source: None,
ref_dims: Vec::new(),
}
}
#[cfg(test)]
fn build_p_frame_yuv420_8bit_fh_with_q(
seq: &SequenceHeader,
width: u32,
height: u32,
base_q_idx: u8,
p_index: u32,
alt_q: &[i16],
) -> FrameHeader {
let dummy = GopFrameReconYuv {
y: Vec::new(),
u: Vec::new(),
v: Vec::new(),
};
let cfg = p_frame_config(&dummy, &dummy, p_index);
build_inter_frame_fh(seq, width, height, base_q_idx, &cfg, alt_q)
}
pub(crate) fn patch_ctx_update_in_tu(tu: &mut [u8], span: (u32, u32), id: u32) {
patch_ctx_update_in_tu_frame(tu, 0, span, id);
}
pub(crate) fn patch_ctx_update_in_tu_frame(
tu: &mut [u8],
frame_ord: usize,
span: (u32, u32),
id: u32,
) {
let payload_off = {
let base = tu.as_ptr() as usize;
let mut found: Option<usize> = None;
let mut seen = 0usize;
for desc in crate::obu::ObuIter::new(tu) {
let desc = desc.expect("own temporal unit walks");
if matches!(
desc.obu_type,
crate::obu::ObuType::Frame | crate::obu::ObuType::FrameHeader
) {
if seen == frame_ord {
found = Some(desc.payload.as_ptr() as usize - base);
break;
}
seen += 1;
}
}
found.expect("temporal unit carries the addressed frame OBU")
};
patch_bits_at(tu, payload_off, span, id);
}
pub(crate) fn patch_bits_at(buf: &mut [u8], payload_off: usize, span: (u32, u32), id: u32) {
let (start, n) = span;
for i in 0..n {
let bitpos = payload_off * 8 + (start + i) as usize;
let byte = bitpos / 8;
let mask = 1u8 << (7 - (bitpos % 8));
if (id >> (n - 1 - i)) & 1 == 1 {
buf[byte] |= mask;
} else {
buf[byte] &= !mask;
}
}
}
#[allow(clippy::too_many_arguments)]
fn encode_p_frame_yuv(
input: &YuvFrame,
prev: &GopFrameReconYuv,
prevprev: &GopFrameReconYuv,
seq: &SequenceHeader,
base_q_idx: u8,
p_index: u32,
alt_q: &[i16],
mf_store: &[SavedMotionField; 8],
model: RateModel,
global_motion: bool,
primary_carry: Option<&RefSlotCarry>,
allow_temporal_seg: bool,
exact_mask: Option<&[bool]>,
auto_lossless: bool,
seg_extras: Option<&SegExtras>,
high_precision_mv: bool,
delta_q: bool,
qm: bool,
cdef: bool,
cdef_units: bool,
lr: bool,
tiles: (u32, u32),
tile_groups: u32,
tile_spans: bool,
explicit_tiles: Option<(&[u32], &[u32])>,
donor_armed: bool,
film_grain: Option<&FilmGrainParams>,
floor: u32,
floor_recon: &GopFrameReconYuv,
s_frame: bool,
error_resilient: bool,
short_ref_signaling: bool,
superres: Option<(u32, u32)>,
superres_source: Option<&YuvFrame>,
) -> Result<
(
Vec<u8>,
GopFrameReconYuv,
SavedMotionField,
RefSlotCarry,
InterFrameAux,
),
Error,
> {
let error_resilient = error_resilient || s_frame;
let primary_carry = if error_resilient { None } else { primary_carry };
let mut cfg = p_frame_config_primary(prev, prevprev, p_index, primary_carry, floor, s_frame);
cfg.error_resilient = error_resilient;
if short_ref_signaling && !error_resilient && seq.enable_order_hint {
cfg.ref_frame_idx = crate::frame_header::set_frame_refs(
cfg.ref_frame_idx[0],
cfg.ref_frame_idx[GOLDEN_FRAME - LAST_FRAME],
cfg.order_hint,
seq.order_hint_bits,
seq.enable_order_hint,
&cfg.slot_hints,
);
cfg.short_ref_signaling = true;
cfg.refs.push(floor_recon);
let floor_plane = cfg.refs.len() - 1;
let last_slot = cfg.ref_frame_idx[0] as usize;
let gold_slot = cfg.ref_frame_idx[GOLDEN_FRAME - LAST_FRAME] as usize;
for (slot, plane) in cfg.slot_to_plane.iter_mut().enumerate() {
if slot != last_slot && slot != gold_slot {
*plane = floor_plane;
}
}
}
cfg.film_grain = film_grain;
cfg.allow_temporal_seg = allow_temporal_seg;
cfg.exact_mask = exact_mask;
cfg.auto_lossless = auto_lossless;
cfg.seg_extras = seg_extras;
cfg.high_precision_mv = high_precision_mv;
cfg.delta_q = delta_q;
cfg.qm = qm;
cfg.cdef = cdef;
cfg.cdef_units = cdef_units;
cfg.lr = lr;
cfg.tiles = tiles;
cfg.tile_groups = tile_groups;
cfg.tile_spans = tile_spans;
cfg.explicit_tiles = explicit_tiles;
cfg.collect_donor_cdfs = donor_armed;
cfg.elect_donor = donor_armed && !error_resilient;
cfg.superres = superres;
cfg.superres_source = superres_source;
let (obus, recon, saved, carry, aux) = encode_inter_frame_generic_gm(
input,
seq,
base_q_idx,
&cfg,
alt_q,
mf_store,
model,
global_motion,
)?;
Ok((build_temporal_unit(None, &obus), recon, saved, carry, aux))
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn encode_inter_frame_generic(
input: &YuvFrame,
seq: &SequenceHeader,
base_q_idx: u8,
cfg: &InterFrameConfig<'_>,
alt_q: &[i16],
mf_store: &[SavedMotionField; 8],
model: RateModel,
) -> Result<
(
Vec<ObuFrame>,
GopFrameReconYuv,
SavedMotionField,
RefSlotCarry,
InterFrameAux,
),
Error,
> {
encode_inter_frame_generic_gm(input, seq, base_q_idx, cfg, alt_q, mf_store, model, true)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn encode_inter_frame_generic_gm(
input: &YuvFrame,
seq: &SequenceHeader,
base_q_idx: u8,
cfg: &InterFrameConfig<'_>,
alt_q: &[i16],
mf_store: &[SavedMotionField; 8],
model: RateModel,
global_motion: bool,
) -> Result<
(
Vec<ObuFrame>,
GopFrameReconYuv,
SavedMotionField,
RefSlotCarry,
InterFrameAux,
),
Error,
> {
input.validate()?;
if cfg.primary_ref_frame != PRIMARY_REF_NONE
&& (usize::from(cfg.primary_ref_frame) >= REFS_PER_FRAME || cfg.primary_carry.is_none())
{
return Err(Error::PartitionWalkOutOfRange);
}
if cfg.s_frame && (cfg.refresh_frame_flags != ALL_FRAMES_PUB || !cfg.error_resilient) {
return Err(Error::PartitionWalkOutOfRange);
}
if cfg.error_resilient
&& (cfg.primary_ref_frame != PRIMARY_REF_NONE
|| cfg.primary_carry.is_some()
|| !cfg.alt_primaries.is_empty())
{
return Err(Error::PartitionWalkOutOfRange);
}
let bit_depth = input.bit_depth;
let (ssx, ssy) = input.format.subsampling();
let mono = input.format == crate::encoder::yuv_frame::ChromaFormat::Monochrome;
let num_planes = input.format.num_planes();
let width = input.width as usize;
let height = input.height as usize;
let chroma_w = input.chroma_width() as usize;
let chroma_h = input.chroma_height() as usize;
let ref_w: usize = match cfg.superres {
Some((up_w, denom)) => {
if !(crate::frame_header::SUPERRES_DENOM_MIN..=16).contains(&denom)
|| crate::encoder::superres_elect::superres_coded_width(up_w, denom) != input.width
|| up_w <= input.width
|| !seq.enable_superres
|| up_w > seq.max_frame_width_minus_1 + 1
|| 2 * input.width < up_w
|| cfg.exact_mask.is_some()
|| cfg.auto_lossless
{
return Err(Error::PartitionWalkOutOfRange);
}
match cfg.superres_source {
Some(src)
if src.width == up_w
&& src.height == input.height
&& src.bit_depth == bit_depth
&& src.format == input.format => {}
_ => return Err(Error::PartitionWalkOutOfRange),
}
up_w as usize
}
None => width,
};
let ref_dims: Vec<(u32, u32)> = if cfg.ref_dims.is_empty() {
vec![(ref_w as u32, height as u32); cfg.refs.len()]
} else {
cfg.ref_dims.clone()
};
if ref_dims.len() != cfg.refs.len() {
return Err(Error::PartitionWalkOutOfRange);
}
for (reference, &(rw, rh)) in cfg.refs.iter().zip(&ref_dims) {
let (rw, rh) = (rw as usize, rh as usize);
let (rcw, rch) = if mono {
(0, 0)
} else {
(
(rw + usize::from(ssx)) >> ssx,
(rh + usize::from(ssy)) >> ssy,
)
};
if rw == 0
|| rh == 0
|| reference.y.len() != rw * rh
|| reference.u.len() != rcw * rch
|| reference.v.len() != rcw * rch
|| 2 * width < rw
|| 2 * height < rh
|| width > 16 * rw
|| height > 16 * rh
{
return Err(Error::PartitionWalkOutOfRange);
}
}
let mut scaled_per_ref = [false; 7];
for (r, flag) in scaled_per_ref.iter_mut().enumerate() {
let (rw, rh) = ref_dims[cfg.slot_to_plane[cfg.ref_frame_idx[r] as usize]];
*flag = crate::encoder::superres_elect::is_scaled(rw, rh, width as u32, height as u32);
}
let any_scaled = scaled_per_ref.iter().any(|&f| f);
let mut fh = build_inter_frame_fh(seq, input.width, input.height, base_q_idx, cfg, alt_q);
fh.allow_screen_content_tools = fh.allow_screen_content_tools
&& bit_depth == 8
&& input.format == crate::encoder::yuv_frame::ChromaFormat::Yuv420;
if let Some((up_w, denom)) = cfg.superres {
let fs = fh
.frame_size
.as_mut()
.ok_or(Error::PartitionWalkOutOfRange)?;
fs.upscaled_width = up_w;
fs.render_width = up_w;
fs.use_superres = true;
fs.superres_denom = denom;
fs.coded_denom = (denom - crate::frame_header::SUPERRES_DENOM_MIN) as u8;
if !cfg.s_frame {
fh.frame_size_override_flag = up_w != seq.max_frame_width_minus_1 + 1
|| input.height != seq.max_frame_height_minus_1 + 1;
}
}
let fs = fh
.frame_size
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let remapped_widths: Option<Vec<u32>> = match (cfg.superres, cfg.explicit_tiles) {
(Some((up_w, denom)), Some((ws, hs))) => Some(
crate::encoder::superres_elect::denom_explicit_ok(up_w, input.height, denom, ws, hs)
.ok_or(Error::PartitionWalkOutOfRange)?,
),
_ => None,
};
let explicit_tiles: Option<(&[u32], &[u32])> =
match (remapped_widths.as_deref(), cfg.explicit_tiles) {
(Some(r), Some((_, hs))) => Some((r, hs)),
_ => cfg.explicit_tiles,
};
let mut ti = match explicit_tiles {
Some((widths_sb, heights_sb)) => crate::tile_info::TileInfo::explicit_layout(
mi_cols,
mi_rows,
seq.use_128x128_superblock,
widths_sb,
heights_sb,
),
None => crate::tile_info::TileInfo::uniform_layout(
mi_cols,
mi_rows,
seq.use_128x128_superblock,
cfg.tiles.0,
cfg.tiles.1,
),
}
.ok_or(Error::PartitionWalkOutOfRange)?;
fh.tile_info = Some(ti.clone());
let num_tiles = ti.tile_cols * ti.tile_rows;
let tile_plan: Vec<(TileGeometry, Vec<(u32, u32)>)> = (0..num_tiles)
.map(|t| {
let g = ti.tile_geometry(t).expect("tile_num < NumTiles");
let mut origins: Vec<(u32, u32)> = Vec::new();
let mut r = g.mi_row_start;
while r < g.mi_row_end {
let mut c = g.mi_col_start;
while c < g.mi_col_end {
origins.push((r, c));
c += 16;
}
r += 16;
}
(g, origins)
})
.collect();
let flat_sb: Vec<(u32, u32)> = tile_plan
.iter()
.flat_map(|(_, o)| o.iter().copied())
.collect();
let (layout_cols_log2, layout_rows_log2, ctx_update_id) = (
ti.tile_cols_log2,
ti.tile_rows_log2,
ti.context_update_tile_id,
);
let assemble = move |payloads: &[Vec<u8>]| -> Result<Vec<u8>, Error> {
let tsb = crate::encoder::key_frame::min_tile_size_bytes(payloads);
write_tile_group_obu(&TileGroupObu::whole_frame(
num_tiles,
layout_cols_log2,
layout_rows_log2,
tsb,
payloads.iter().cloned().map(TilePayload::new).collect(),
))
};
let lossless = base_q_idx == 0;
let mut qp = QuantizerParams::neutral(base_q_idx, bit_depth);
if !alt_q.is_empty() {
qp.segmentation_enabled = true;
for (seg, &d) in alt_q.iter().enumerate() {
qp.seg_alt_q_active[seg] = true;
qp.seg_alt_q_data[seg] = d;
}
}
let mut ip = SyntaxInterFrameParams::single_ref_baseline(
mi_rows, mi_cols, false,
);
ip.allow_high_precision_mv = cfg.high_precision_mv;
ip.is_scaled_per_ref = scaled_per_ref;
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = !cfg.compound_pairs.is_empty();
ip.enable_masked_compound = seq.enable_masked_compound;
ip.enable_interintra_compound = seq.enable_interintra_compound;
ip.enable_jnt_comp = seq.enable_jnt_comp;
ip.is_motion_mode_switchable = true;
ip.allow_warped_motion = seq.enable_warped_motion && !fh.error_resilient_mode;
ip.order_hints = crate::cdf::FrameInterOrderHints {
order_hint_bits: u32::from(seq.order_hint_bits),
current_order_hint: fh.order_hint as i32,
order_hints_by_ref: {
let mut by_ref = [0i32; crate::uncompressed_header_tail::ALTREF_FRAME + 1];
for (i, hint) in by_ref.iter_mut().enumerate().skip(1) {
*hint = cfg.slot_hints[cfg.ref_frame_idx[i - 1] as usize] as i32;
}
by_ref
},
};
for r in 1..=7usize {
ip.ref_frame_sign_bias[r] = i32::from(
crate::inter_pred::get_relative_dist(
ip.order_hints.order_hints_by_ref[r],
fh.order_hint as i32,
u32::from(seq.order_hint_bits),
) > 0,
);
}
{
let (allowed, pair) = crate::encoder::frame_obu::skip_mode_params_twin(
&cfg.slot_hints,
&cfg.ref_frame_idx,
fh.order_hint,
seq.order_hint_bits,
);
ip.skip_mode_present = fh.skip_mode_present == Some(true) && allowed;
ip.skip_mode_frame = pair;
}
if !alt_q.is_empty() {
ip.segmentation_update_map = true;
ip.segmentation_temporal_update = false;
let cells = (mi_rows as usize) * (mi_cols as usize);
let prev_ids = match cfg.primary_carry {
Some(c)
if cfg.primary_ref_frame != PRIMARY_REF_NONE
&& c.mi_rows == mi_rows
&& c.mi_cols == mi_cols
&& c.segment_ids.len() == cells =>
{
c.segment_ids.clone()
}
_ => vec![0i32; cells],
};
ip.prev_segment_ids = Some(std::sync::Arc::new(prev_ids));
}
ip.use_ref_frame_mvs =
seq.enable_order_hint && seq.enable_ref_frame_mvs && !fh.error_resilient_mode;
if ip.use_ref_frame_mvs {
use crate::inter_pred::{motion_field_estimation_core, MotionFieldSlot};
let ref_frame_idx = cfg.ref_frame_idx;
let mut order_hints = [0i32; 8];
order_hints.copy_from_slice(&ip.order_hints.order_hints_by_ref);
let mut slots: [Option<MotionFieldSlot<'_>>; 8] = [None; 8];
for (i, slot) in mf_store.iter().enumerate() {
slots[i] = Some(MotionFieldSlot {
mf_mvs: &slot.mf_mvs,
mf_ref_frames: &slot.mf_ref_frames,
saved_order_hints: slot.saved_order_hints,
mi_rows: slot.mi_rows,
mi_cols: slot.mi_cols,
frame_is_intra: slot.frame_is_intra,
});
}
ip.motion_field_mvs = motion_field_estimation_core(
&slots,
&ref_frame_idx,
&order_hints,
fh.order_hint as i32,
u32::from(seq.order_hint_bits),
mi_rows,
mi_cols,
);
}
if global_motion && !any_scaled {
let hp = fh
.inter_refs
.as_ref()
.map(|ir| ir.allow_high_precision_mv)
.unwrap_or(false);
let mut per_plane: Vec<Option<(crate::uncompressed_header_tail::WarpModelType, [i32; 6])>> =
vec![None; cfg.refs.len()];
for &rf in &cfg.single_refs {
let plane = cfg.slot_to_plane[cfg.ref_frame_idx[(rf - 1) as usize] as usize];
let (t, p) = match per_plane[plane] {
Some(tp) => tp,
None => {
let samples = crate::encoder::global_motion::collect_motion_samples(
&input.y,
&cfg.refs[plane].y,
width,
height,
);
let tp = crate::encoder::global_motion::estimate_global_motion(&samples, hp);
per_plane[plane] = Some(tp);
tp
}
};
ip.gm_type[rf as usize] = i32::from(t.as_u8());
ip.gm_params[rf as usize] = p;
}
let mut gmp = crate::uncompressed_header_tail::GlobalMotionParams::identity();
gmp.short_circuited = false;
for r in 1..=7usize {
gmp.gm_type[r] = match ip.gm_type[r] {
1 => crate::uncompressed_header_tail::WarpModelType::Translation,
2 => crate::uncompressed_header_tail::WarpModelType::RotZoom,
3 => crate::uncompressed_header_tail::WarpModelType::Affine,
_ => crate::uncompressed_header_tail::WarpModelType::Identity,
};
gmp.gm_params[r] = ip.gm_params[r];
}
fh.global_motion_params = Some(gmp);
}
let x = cfg.seg_extras.copied().unwrap_or_default();
let seg_id_pre_skip = x.any();
let last_active_seg_id =
alt_q
.len()
.saturating_sub(1)
.max(if x.any() { x.top_segment() } else { 0 }) as u8;
let mut params = SyntaxFrameParams {
subsampling_x: ssx,
subsampling_y: ssy,
num_planes,
seg_id_pre_skip,
segmentation_enabled: !alt_q.is_empty(),
seg_ref_frame: x.ref_frame,
seg_skip: x.skip,
seg_globalmv: x.globalmv,
last_active_seg_id,
lossless_array: seg_lossless_array(base_q_idx, alt_q),
coded_lossless: lossless,
enable_cdef: seq.enable_cdef,
allow_intrabc: false,
cdef_bits: 0,
read_deltas: false,
use_128x128_superblock: seq.use_128x128_superblock,
delta_q_res: 0,
delta_lf_present: false,
delta_lf_multi: false,
mono_chrome: mono,
delta_lf_res: 0,
allow_screen_content_tools: fh.allow_screen_content_tools,
enable_filter_intra: seq.enable_filter_intra,
bit_depth,
tx_mode_select: !lossless,
quant: qp,
reduced_tx_set: fh.reduced_tx_set.unwrap_or(false),
inter: Some(ip.clone()),
};
let frame_start_cdfs: Box<TileCdfContext> = match cfg.primary_carry {
Some(carry) if cfg.primary_ref_frame != PRIMARY_REF_NONE => {
let mut loaded = carry.cdfs.clone();
loaded.zero_counts();
loaded
}
_ => {
let mut c = TileCdfContext::new_from_defaults();
c.init_coeff_cdfs(base_q_idx);
Box::new(c)
}
};
let ll_up_cap = delta_q_lossless_up_cap(base_q_idx, alt_q);
let delta_plan: Option<Vec<i32>> = if cfg.delta_q && !lossless {
delta_q_plan_units(input, mi_rows, mi_cols).and_then(|mut units| {
if ll_up_cap < 2 {
for u in units.iter_mut() {
*u = (*u).min(ll_up_cap);
}
}
units.iter().any(|&u| u != 0).then_some(units)
})
} else {
None
};
struct FrameArm {
payloads: Vec<Vec<u8>>,
tile: Vec<u8>,
all_cdfs: Vec<Box<TileCdfContext>>,
state: PartitionSyntaxWriter,
trees: Vec<SyntaxNode>,
recon: Box<ReconState>,
ictx: Box<PSearchCtx>,
}
let fh_sct = fh.allow_screen_content_tools;
let fh_disable_cdf_update = fh.disable_cdf_update;
let run_arm =
|arm_params: &SyntaxFrameParams, plan: Option<&[i32]>| -> Result<FrameArm, Error> {
let mut recon = Box::new(ReconState {
y: vec![0u16; width * height],
u: vec![0u16; chroma_w * chroma_h],
v: vec![0u16; chroma_w * chroma_h],
width,
height,
chroma_w,
chroma_h,
bit_depth,
subsampling_x: ssx,
subsampling_y: ssy,
num_planes,
mi_rows,
mi_cols,
lossless,
allow_screen_content_tools: fh_sct,
allow_intrabc: false,
qp: arm_params.quant,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(ssx, ssy, num_planes),
dv_hash: Default::default(),
});
let mut ictx = Box::new(PSearchCtx::with_refs(
&cfg.refs,
cfg.slot_to_plane,
cfg.ref_frame_idx,
cfg.single_refs.clone(),
cfg.compound_pairs.clone(),
mi_rows,
mi_cols,
width,
height,
&ref_dims,
ip.clone(),
base_q_idx,
alt_q,
bit_depth,
ssx,
ssy,
num_planes,
)?);
if let Some(x) = cfg.seg_extras {
if alt_q.is_empty()
|| x.ref_frame
.iter()
.flatten()
.any(|rf| !cfg.single_refs.contains(rf))
{
return Err(Error::PartitionWalkOutOfRange);
}
ictx.extras = Some(*x);
ictx.seg_id_pre_skip = x.any();
}
if cfg.exact_mask.is_some() || cfg.auto_lossless {
if ictx.lossless_seg.is_none() {
return Err(Error::PartitionWalkOutOfRange);
}
if let Some(mask) = cfg.exact_mask {
if mask.len() != (mi_rows as usize) * (mi_cols as usize) {
return Err(Error::PartitionWalkOutOfRange);
}
ictx.exact_mask = Some(mask.to_vec());
}
ictx.auto_lossless = cfg.auto_lossless;
}
let mut state = PartitionSyntaxWriter::new(
mi_rows,
mi_cols,
TileGeometry {
mi_row_start: 0,
mi_row_end: mi_rows,
mi_col_start: 0,
mi_col_end: mi_cols,
},
)
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut trees: Vec<SyntaxNode> = Vec::new();
let mut payloads: Vec<Vec<u8>> = Vec::with_capacity(num_tiles as usize);
let mut all_cdfs: Vec<Box<TileCdfContext>> = Vec::with_capacity(num_tiles as usize);
let sb_cols_frame = mi_cols.div_ceil(16);
for (geo, origins) in tile_plan.iter() {
if !state.begin_tile(*geo) {
return Err(Error::PartitionWalkOutOfRange);
}
ictx.begin_tile(*geo);
recon.bd.set_tile(*geo);
let mut writer = SymbolWriter::new(fh_disable_cdf_update);
let mut cdfs: Box<TileCdfContext> = frame_start_cdfs.clone();
recon.qp.current_q_index = base_q_idx;
for &(sb_r, sb_c) in origins {
recon.bd.clear_for_sb(sb_r, sb_c, mi_rows, mi_cols);
let prev_q = recon.qp.current_q_index;
if let Some(p) = plan {
let k = ((sb_r / 16) * sb_cols_frame + (sb_c / 16)) as usize;
let step = 1i32 << arm_params.delta_q_res;
let want = (i32::from(base_q_idx) + p[k] * step).clamp(1, 255);
let units = (want - i32::from(prev_q)) / step;
recon.qp.current_q_index = (i32::from(prev_q) + units * step) as u8;
ictx.delta_q_units = units;
}
state.arm_read_deltas();
let mut twin = RateTwin::snapshot(&cdfs, &state, &writer);
twin.arm_read_deltas();
let (tree, _cost) = build_p_search_tree(
sb_r,
sb_c,
BLOCK_64X64,
input,
&mut recon,
&mut ictx,
&mut twin,
arm_params,
model,
)?;
write_partition_tree_syntax(
&mut writer,
&mut cdfs,
&mut state,
&tree,
sb_r,
sb_c,
BLOCK_64X64,
arm_params,
)?;
debug_assert!(
twin.matches(&cdfs, &writer),
"rate twin desynced from the writer after superblock ({sb_r},{sb_c})"
);
if plan.is_some() {
if let SyntaxNode::Leaf(b) = &tree {
if b.skip == 1 {
recon.qp.current_q_index = prev_q;
}
}
}
trees.push(tree);
}
payloads.push(writer.finish());
all_cdfs.push(cdfs);
}
let body = assemble(&payloads)?;
Ok(FrameArm {
payloads,
tile: body,
all_cdfs,
state,
trees,
recon,
ictx,
})
};
let mut arm = run_arm(¶ms, None)?;
let qm_arm_levels: Vec<u8> = if cfg.qm
&& !lossless
&& cfg.exact_mask.is_none()
&& !cfg.auto_lossless
&& qm_arm_allowed(base_q_idx, width, height)
&& qm_probe(input)
{
qm_ladder_for_q(base_q_idx)
} else {
Vec::new()
};
let mut qm_elected = false;
if !qm_arm_levels.is_empty() {
let seg_ll = seg_lossless_array(base_q_idx, alt_q);
let lambda = lambda_for(&qp);
let prev_gm_score: Option<[[i32; 6]; TOTAL_REFS_PER_FRAME]> =
if fh.primary_ref_frame != PRIMARY_REF_NONE {
cfg.primary_carry.map(|c| c.gm_params)
} else {
None
};
let header_len = |fh_arm: &FrameHeader| -> usize {
let mut bw = crate::encoder::bitwriter::BitWriter::new();
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
&mut bw,
fh_arm,
seq,
prev_gm_score.as_ref(),
);
bw.byte_align();
bw.finish().len()
};
let plain_sse = |rc: &ReconState| -> u64 {
let mut d = 0u64;
for (a, b) in rc.y.iter().zip(&input.y) {
let diff = i64::from(*a) - i64::from(*b);
d += (diff * diff) as u64;
}
if num_planes > 1 {
for (a, b) in rc.u.iter().zip(&input.u).chain(rc.v.iter().zip(&input.v)) {
let diff = i64::from(*a) - i64::from(*b);
d += (diff * diff) as u64;
}
}
d
};
let mut best_score = score256(
plain_sse(&arm.recon),
lambda,
((header_len(&fh) + arm.tile.len()) as u64) * 8 * 256,
);
for &lvl in &qm_arm_levels {
let mut qp_m = qp;
qp_m.using_qmatrix = true;
for plane in 0..3 {
for (seg, &ll) in seg_ll.iter().enumerate() {
qp_m.seg_qm_level[plane][seg] = if ll { 15 } else { lvl };
}
}
let mut params_m = params.clone();
params_m.quant = qp_m;
let cand = run_arm(¶ms_m, None)?;
let mut fh_m = fh.clone();
if let Some(q) = fh_m.quantization_params.as_mut() {
q.using_qmatrix = true;
q.qm_y = lvl;
q.qm_u = lvl;
q.qm_v = lvl;
}
let score_qm = score256(
plain_sse(&cand.recon),
lambda,
((header_len(&fh_m) + cand.tile.len()) as u64) * 8 * 256,
);
if std::env::var_os("OXIDEAV_AV1_QM_DEBUG").is_some() {
eprintln!(
"qm-frame oh={} lvl={} qm: tile {} B sse {} score {} | best so far {}",
cfg.order_hint,
lvl,
cand.tile.len(),
plain_sse(&cand.recon),
score_qm,
best_score,
);
}
if score_qm < best_score {
best_score = score_qm;
qm_elected = true;
arm = cand;
fh = fh_m;
params = params_m;
qp = qp_m;
}
if QM_LADDER_REFINES_ONLY && !qm_elected {
break;
}
}
}
let mut delta_q_elected = false;
if let Some(plan) = &delta_plan {
let mut qp_d = qp;
qp_d.delta_q_present = true;
qp_d.current_q_index = base_q_idx;
let mut params_d = params.clone();
params_d.read_deltas = true;
params_d.delta_q_res = DELTA_Q_RES;
params_d.quant = qp_d;
let cand = run_arm(¶ms_d, Some(plan))?;
let mut fh_d = fh.clone();
fh_d.delta_q_params = Some(crate::uncompressed_header_tail::DeltaQParams {
delta_q_present: true,
delta_q_res: DELTA_Q_RES,
});
let lambda = lambda_for(&qp);
let prev_gm_score: Option<[[i32; 6]; TOTAL_REFS_PER_FRAME]> =
if fh.primary_ref_frame != PRIMARY_REF_NONE {
cfg.primary_carry.map(|c| c.gm_params)
} else {
None
};
let header_len = |fh_arm: &FrameHeader| -> usize {
let mut bw = crate::encoder::bitwriter::BitWriter::new();
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
&mut bw,
fh_arm,
seq,
prev_gm_score.as_ref(),
);
bw.byte_align();
bw.finish().len()
};
let weighted_sse = |rc: &ReconState| -> u64 {
let mut d = 0u64;
for (sb_index, &(sb_r, sb_c)) in sb_grid_origins(mi_rows, mi_cols).iter().enumerate() {
let units = plan[sb_index];
let y0 = (sb_r as usize) * 4;
let x0 = (sb_c as usize) * 4;
let y1 = (y0 + 64).min(height);
let x1 = (x0 + 64).min(width);
let mut sse = 0u64;
for r in y0..y1 {
for c in x0..x1 {
let diff =
i64::from(rc.y[r * width + c]) - i64::from(input.y[r * width + c]);
sse += (diff * diff) as u64;
}
}
if num_planes > 1 {
let (cy0, cx0) = (y0 >> ssy, x0 >> ssx);
let (cy1, cx1) = (
(y1 + usize::from(ssy)) >> ssy,
(x1 + usize::from(ssx)) >> ssx,
);
for r in cy0..cy1.min(chroma_h) {
for c in cx0..cx1.min(chroma_w) {
let du = i64::from(rc.u[r * chroma_w + c])
- i64::from(input.u[r * chroma_w + c]);
let dv = i64::from(rc.v[r * chroma_w + c])
- i64::from(input.v[r * chroma_w + c]);
sse += (du * du + dv * dv) as u64;
}
}
}
d += if units >= 0 {
sse >> (2 * units)
} else {
sse << (-2 * units)
};
}
d
};
let score_base = score256(
weighted_sse(&arm.recon),
lambda,
((header_len(&fh) + arm.tile.len()) as u64) * 8 * 256,
);
let score_delta = score256(
weighted_sse(&cand.recon),
lambda,
((header_len(&fh_d) + cand.tile.len()) as u64) * 8 * 256,
);
if std::env::var_os("OXIDEAV_AV1_DQ_DEBUG").is_some() {
eprintln!(
"dq-frame oh={} plan={:?} base: tile {} B Dw {} score {} | delta: tile {} B Dw {} score {}",
cfg.order_hint,
plan,
arm.tile.len(),
weighted_sse(&arm.recon),
score_base,
cand.tile.len(),
weighted_sse(&cand.recon),
score_delta,
);
}
if score_delta < score_base {
delta_q_elected = true;
arm = cand;
fh = fh_d;
params = params_d;
}
}
let FrameArm {
payloads: mut tile_payloads,
tile: mut tile_bytes,
mut all_cdfs,
mut state,
mut trees,
mut recon,
ictx,
} = arm;
let disable_cdf_update = fh.disable_cdf_update;
type EmitOut = (
Vec<Vec<u8>>,
Vec<Box<TileCdfContext>>,
PartitionSyntaxWriter,
);
let emit_tiles = |trees: &[SyntaxNode],
emit_params: &SyntaxFrameParams,
start_cdfs: &dyn Fn() -> Box<TileCdfContext>,
lr_plan: Option<&crate::encoder::lr_elect::LrPlan>|
-> Result<EmitOut, Error> {
let mut re_state = PartitionSyntaxWriter::new(
mi_rows,
mi_cols,
TileGeometry {
mi_row_start: 0,
mi_row_end: mi_rows,
mi_col_start: 0,
mi_col_end: mi_cols,
},
)
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut payloads: Vec<Vec<u8>> = Vec::with_capacity(num_tiles as usize);
let mut all: Vec<Box<TileCdfContext>> = Vec::with_capacity(num_tiles as usize);
let mut tree_idx = 0usize;
for (geo, origins) in tile_plan.iter() {
if !re_state.begin_tile(*geo) {
return Err(Error::PartitionWalkOutOfRange);
}
let mut re_writer = SymbolWriter::new(disable_cdf_update);
let mut re_cdfs: Box<TileCdfContext> = start_cdfs();
let mut lr_write_state = crate::encoder::loop_restoration_write::LrWriteState::new();
for &(sb_r, sb_c) in origins {
re_state.arm_read_deltas();
if let Some(plan) = lr_plan {
crate::encoder::loop_restoration_write::write_lr(
&mut re_writer,
&mut re_cdfs,
&mut lr_write_state,
sb_r,
sb_c,
BLOCK_64X64,
&plan.write_params,
&plan.units,
)?;
}
write_partition_tree_syntax(
&mut re_writer,
&mut re_cdfs,
&mut re_state,
&trees[tree_idx],
sb_r,
sb_c,
BLOCK_64X64,
emit_params,
)?;
tree_idx += 1;
}
payloads.push(re_writer.finish());
all.push(re_cdfs);
}
Ok((payloads, all, re_state))
};
let mut hp_elected = cfg.high_precision_mv;
let gm_all_identity = params
.inter
.as_ref()
.is_some_and(|i| (1..=7).all(|r| i.gm_type[r] == 0));
if cfg.high_precision_mv && gm_all_identity {
let mut alt_ip = params.inter.clone().ok_or(Error::PartitionWalkOutOfRange)?;
alt_ip.allow_high_precision_mv = false;
let mut alt_params = params.clone();
alt_params.inter = Some(alt_ip);
let replay = emit_tiles(&trees, &alt_params, &|| frame_start_cdfs.clone(), None)
.and_then(|(p, c, s)| Ok((assemble(&p)?, p, c, s)));
if let Ok((alt_body, alt_payloads, alt_cdfs, alt_state)) = replay {
if alt_body.len() <= tile_bytes.len() {
hp_elected = false;
tile_payloads = alt_payloads;
tile_bytes = alt_body;
all_cdfs = alt_cdfs;
state = alt_state;
params.inter = alt_params.inter;
if let Some(ir) = fh.inter_refs.as_mut() {
ir.allow_high_precision_mv = false;
}
}
}
}
let main_ip = params
.inter
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut seg_temporal_elected = false;
let temporal_codable =
!alt_q.is_empty() && cfg.primary_ref_frame != PRIMARY_REF_NONE && cfg.allow_temporal_seg;
if temporal_codable {
let mut alt_ip = main_ip.clone();
alt_ip.segmentation_temporal_update = true;
let mut alt_params = params.clone();
alt_params.inter = Some(alt_ip);
let (alt_payloads, alt_cdfs, alt_state) =
emit_tiles(&trees, &alt_params, &|| frame_start_cdfs.clone(), None)?;
let alt_body = assemble(&alt_payloads)?;
if alt_body.len() < tile_bytes.len() {
seg_temporal_elected = true;
tile_payloads = alt_payloads;
tile_bytes = alt_body;
all_cdfs = alt_cdfs;
state = alt_state;
}
}
if let Some(sp) = fh.segmentation_params.as_mut() {
sp.temporal_update = seg_temporal_elected;
}
let mut prev_gm_for_header: Option<[[i32; 6]; TOTAL_REFS_PER_FRAME]> =
if fh.primary_ref_frame != PRIMARY_REF_NONE {
cfg.primary_carry.map(|c| c.gm_params)
} else {
None
};
if !cfg.alt_primaries.is_empty() {
if temporal_codable {
return Err(Error::PartitionWalkOutOfRange);
}
let header_len = |ord: u8, prev: Option<&[[i32; 6]; TOTAL_REFS_PER_FRAME]>| -> usize {
let mut fh_c = fh.clone();
fh_c.primary_ref_frame = ord;
let mut bw = crate::encoder::bitwriter::BitWriter::new();
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
&mut bw, &fh_c, seq, prev,
);
bw.byte_align();
bw.finish().len()
};
let mut best_total =
header_len(fh.primary_ref_frame, prev_gm_for_header.as_ref()) + tile_bytes.len();
#[allow(clippy::type_complexity)]
let mut best: Option<(
u8,
Vec<Vec<u8>>,
Vec<u8>,
Vec<Box<TileCdfContext>>,
PartitionSyntaxWriter,
Option<[[i32; 6]; TOTAL_REFS_PER_FRAME]>,
)> = None;
for (ord, carry_opt) in &cfg.alt_primaries {
if *ord != PRIMARY_REF_NONE
&& (usize::from(*ord) >= REFS_PER_FRAME || carry_opt.is_none())
{
return Err(Error::PartitionWalkOutOfRange);
}
let cand_start = || -> Box<TileCdfContext> {
match carry_opt {
Some(carry) if *ord != PRIMARY_REF_NONE => {
let mut loaded = carry.cdfs.clone();
loaded.zero_counts();
loaded
}
_ => {
let mut c = Box::new(TileCdfContext::new_from_defaults());
c.init_coeff_cdfs(base_q_idx);
c
}
}
};
let (cand_payloads, cand_cdfs, cand_state) =
emit_tiles(&trees, ¶ms, &cand_start, None)?;
let cand_body = assemble(&cand_payloads)?;
let cand_prev = if *ord != PRIMARY_REF_NONE {
carry_opt.map(|c| c.gm_params)
} else {
None
};
let total = header_len(*ord, cand_prev.as_ref()) + cand_body.len();
if total < best_total {
best_total = total;
best = Some((
*ord,
cand_payloads,
cand_body,
cand_cdfs,
cand_state,
cand_prev,
));
}
}
if let Some((ord, e_payloads, e_body, e_cdfs, e_state, e_prev)) = best {
fh.primary_ref_frame = ord;
tile_payloads = e_payloads;
tile_bytes = e_body;
all_cdfs = e_cdfs;
state = e_state;
prev_gm_for_header = e_prev;
}
}
let donor_src: Option<&RefSlotCarry> = if fh.primary_ref_frame == PRIMARY_REF_NONE {
None
} else if fh.primary_ref_frame == cfg.primary_ref_frame {
cfg.primary_carry
} else {
cfg.alt_primaries
.iter()
.find(|(o, _)| *o == fh.primary_ref_frame)
.and_then(|(_, c)| *c)
};
let mut donor_elected: Option<u32> = None;
if cfg.elect_donor {
if let Some(pc) = donor_src {
if pc.donor_cdfs.len() > 1 {
let mut d_params = params.clone();
if let Some(ip) = d_params.inter.as_mut() {
ip.segmentation_temporal_update = seg_temporal_elected;
}
for (tid, cand) in pc.donor_cdfs.iter().enumerate().skip(1) {
let cand_start = || -> Box<TileCdfContext> {
let mut loaded = cand.clone();
loaded.zero_counts();
loaded
};
let (cand_payloads, cand_all, cand_state) =
emit_tiles(&trees, &d_params, &cand_start, None)?;
let cand_body = assemble(&cand_payloads)?;
if cand_body.len() < tile_bytes.len() {
donor_elected = Some(tid as u32);
tile_payloads = cand_payloads;
tile_bytes = cand_body;
all_cdfs = cand_all;
state = cand_state;
}
}
}
}
}
let start_cdfs_for = |ord: u8| -> Box<TileCdfContext> {
if let (Some(t), Some(pc)) = (donor_elected, donor_src) {
let mut loaded = pc.donor_cdfs[t as usize].clone();
loaded.zero_counts();
return loaded;
}
if ord == cfg.primary_ref_frame {
frame_start_cdfs.clone()
} else {
match cfg
.alt_primaries
.iter()
.find(|(o, _)| *o == ord)
.and_then(|(_, c)| *c)
{
Some(carry) => {
let mut loaded = carry.cdfs.clone();
loaded.zero_counts();
loaded
}
None => {
let mut c = Box::new(TileCdfContext::new_from_defaults());
c.init_coeff_cdfs(base_q_idx);
c
}
}
}
};
let lr_armed = cfg.lr
&& !lossless
&& seq.enable_restoration
&& cfg.exact_mask.is_none()
&& !cfg.auto_lossless;
let pre_cdef: Option<(Vec<u16>, Vec<u16>, Vec<u16>)> =
lr_armed.then(|| (recon.y.clone(), recon.u.clone(), recon.v.clone()));
let mut committed_cdef_bits: u32 = 0;
let mut cdef_elected = false;
if cfg.cdef && !lossless && seq.enable_cdef && cfg.exact_mask.is_none() && !cfg.auto_lossless {
if let Some(election) =
crate::encoder::cdef_elect::elect_cdef(&crate::encoder::cdef_elect::CdefElectInput {
mirror: state.mirror(),
input,
recon_y: &recon.y,
recon_u: &recon.u,
recon_v: &recon.v,
width,
height,
chroma_w,
chroma_h,
bit_depth,
subsampling_x: ssx,
subsampling_y: ssy,
num_planes,
lambda: crate::encoder::key_frame::lambda_for(&recon.qp),
max_bits: if cfg.cdef_units { 3 } else { 0 },
fast: None,
threads: 1,
})
{
let lambda = crate::encoder::key_frame::lambda_for(&recon.qp);
let cdef_header_len = |cdef_params: &crate::uncompressed_header_tail::CdefParams| {
let mut fh_c = fh.clone();
fh_c.cdef_params = Some(*cdef_params);
let mut bw = crate::encoder::bitwriter::BitWriter::new();
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
&mut bw,
&fh_c,
seq,
prev_gm_for_header.as_ref(),
);
bw.byte_align();
bw.finish().len()
};
let adopted: Option<crate::encoder::cdef_elect::CdefPlan>;
if election.best.params.cdef_bits > 0 {
let plan = &election.best;
let elected_ord = fh.primary_ref_frame;
let mut re_params = params.clone();
if let Some(ip) = re_params.inter.as_mut() {
ip.segmentation_temporal_update = seg_temporal_elected;
}
re_params.cdef_bits = u32::from(plan.params.cdef_bits);
let sb_cols_frame = mi_cols.div_ceil(16);
for (tree, &(sb_r, sb_c)) in trees.iter_mut().zip(flat_sb.iter()) {
let k = ((sb_r / 16) * sb_cols_frame + (sb_c / 16)) as usize;
tree.stamp_cdef_idx(plan.unit_idx[k].max(0));
}
let (re_payloads, re_cdfs, re_state) =
emit_tiles(&trees, &re_params, &|| start_cdfs_for(elected_ord), None)?;
let re_body = assemble(&re_payloads)?;
let unit_score = plan.d * 256
+ lambda * 8 * 256 * ((cdef_header_len(&plan.params) + re_body.len()) as u64);
let (fb_d, fb_hdr) = match &election.frame_level {
Some(p) => (p.d, cdef_header_len(&p.params)),
None => (
election.base_d,
cdef_header_len(
fh.cdef_params
.as_ref()
.expect("lossy inter header carries cdef params"),
),
),
};
let fb_score = fb_d * 256 + lambda * 8 * 256 * ((fb_hdr + tile_bytes.len()) as u64);
if unit_score < fb_score {
tile_payloads = re_payloads;
tile_bytes = re_body;
all_cdfs = re_cdfs;
state = re_state;
committed_cdef_bits = u32::from(plan.params.cdef_bits);
debug_assert!(
{
let grid = state.mirror().cdef_idx();
let sb_cols_dbg = mi_cols.div_ceil(16);
plan.unit_idx.iter().enumerate().all(|(k, &idx)| {
let (ur, uc) = (k as u32 / sb_cols_dbg, k as u32 % sb_cols_dbg);
grid[((ur * 16) * mi_cols + uc * 16) as usize] == idx
})
},
"re-emitted cdef_idx grid diverged from the elected plan"
);
adopted = Some(election.best);
} else {
for tree in trees.iter_mut() {
tree.stamp_cdef_idx(0);
}
adopted = election.frame_level;
}
} else {
adopted = Some(election.best);
}
if let Some(plan) = adopted {
crate::encoder::cdef_elect::apply_cdef_plan(
state.mirror(),
&plan,
input,
&mut recon.y,
&mut recon.u,
&mut recon.v,
width,
height,
chroma_w,
chroma_h,
bit_depth,
ssx,
ssy,
num_planes,
1,
);
fh.cdef_params = Some(plan.params);
cdef_elected = true;
}
}
}
let mut pre_cdef = pre_cdef;
let (lr_w, lr_chroma_w, lr_src): (usize, usize, &YuvFrame) =
if let Some((up_w, _)) = cfg.superres {
let np = usize::from(num_planes);
let upscale = |y: &[u16], u: &[u16], v: &[u16]| {
crate::encoder::superres_elect::upscale_recon(
[y, u, v],
width as u32,
height as u32,
up_w,
mi_cols,
bit_depth,
ssx,
ssy,
np,
)
};
let (uy, uu, uv) = upscale(&recon.y, &recon.u, &recon.v)?;
recon.y = uy;
recon.u = uu;
recon.v = uv;
if let Some(pc) = pre_cdef.as_mut() {
let (py, pu, pv) = upscale(&pc.0, &pc.1, &pc.2)?;
*pc = (py, pu, pv);
}
let src = cfg.superres_source.ok_or(Error::PartitionWalkOutOfRange)?;
(
up_w as usize,
((up_w + u32::from(ssx)) >> ssx) as usize,
src,
)
} else {
(width, chroma_w, input)
};
let mut lr_elected = false;
if let Some((pcy, pcu, pcv)) = pre_cdef.as_ref() {
let price_cdfs = start_cdfs_for(fh.primary_ref_frame);
let lambda = crate::encoder::key_frame::lambda_for(&recon.qp);
let lr_header_len = |lrp: &crate::uncompressed_header_tail::LrParams| {
let mut fh_c = fh.clone();
fh_c.lr_params = Some(*lrp);
let mut bw = crate::encoder::bitwriter::BitWriter::new();
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
&mut bw,
&fh_c,
seq,
prev_gm_for_header.as_ref(),
);
bw.byte_align();
bw.finish().len()
};
let off_hdr = lr_header_len(
fh.lr_params
.as_ref()
.expect("lossy inter header carries lr_params"),
);
type LrCand = (
crate::encoder::lr_elect::LrPlan,
Vec<Vec<u8>>,
Vec<u8>,
Vec<Box<TileCdfContext>>,
PartitionSyntaxWriter,
u64,
);
let mut best_lr: Option<LrCand> = None;
let mut off_score: Option<u64> = None;
let min_shift: u8 = if seq.use_128x128_superblock { 1 } else { 0 };
for unit_shift in min_shift..=2u8 {
let Some(plan) =
crate::encoder::lr_elect::elect_lr(&crate::encoder::lr_elect::LrElectInput {
input: lr_src,
curr_y: crate::encoder::lr_elect::LrCurr::Full(pcy),
curr_u: crate::encoder::lr_elect::LrCurr::Full(pcu),
curr_v: crate::encoder::lr_elect::LrCurr::Full(pcv),
cdef_y: &recon.y,
cdef_u: &recon.u,
cdef_v: &recon.v,
width: lr_w,
height,
chroma_w: lr_chroma_w,
chroma_h,
bit_depth,
subsampling_x: ssx,
subsampling_y: ssy,
num_planes,
mi_rows,
mi_cols,
lambda,
price_cdfs: &price_cdfs,
disable_cdf_update: fh.disable_cdf_update,
frame_width: lr_w,
frame_height: height,
sgr_step: 1,
unit_shift,
use_superres: cfg.superres.is_some(),
superres_denom: cfg
.superres
.map_or(crate::frame_header::SUPERRES_NUM, |(_, d)| d),
threads: 1,
wiener_rounds: 2,
})
else {
continue;
};
let mut re_params = params.clone();
if let Some(ip) = re_params.inter.as_mut() {
ip.segmentation_temporal_update = seg_temporal_elected;
}
re_params.cdef_bits = committed_cdef_bits;
let elected_ord = fh.primary_ref_frame;
let re_out = emit_tiles(
&trees,
&re_params,
&|| start_cdfs_for(elected_ord),
Some(&plan),
);
let Ok((re_payloads, re_cdfs, re_state)) = re_out else {
continue;
};
let re_body = assemble(&re_payloads)?;
let on_score = plan.d * 256
+ lambda * 8 * 256 * ((lr_header_len(&plan.header) + re_body.len()) as u64);
let off = *off_score.get_or_insert(
plan.d_pre * 256 + lambda * 8 * 256 * ((off_hdr + tile_bytes.len()) as u64),
);
if on_score < off && best_lr.as_ref().map_or(true, |b| on_score < b.5) {
best_lr = Some((plan, re_payloads, re_body, re_cdfs, re_state, on_score));
}
}
if let Some((plan, re_payloads, re_body, re_cdfs, re_state, _)) = best_lr {
tile_payloads = re_payloads;
tile_bytes = re_body;
all_cdfs = re_cdfs;
state = re_state;
let applied_d = crate::encoder::lr_elect::apply_lr_plan(
&plan,
lr_src,
crate::encoder::lr_elect::LrCurr::Full(pcy),
crate::encoder::lr_elect::LrCurr::Full(pcu),
crate::encoder::lr_elect::LrCurr::Full(pcv),
&mut recon.y,
&mut recon.u,
&mut recon.v,
lr_w,
height,
lr_chroma_w,
chroma_h,
bit_depth,
ssx,
ssy,
num_planes,
mi_rows,
mi_cols,
lr_w,
height,
);
debug_assert_eq!(
applied_d, plan.d,
"applied ยง7.17 SSD diverged from the elected plan"
);
fh.lr_params = Some(plan.header);
lr_elected = true;
}
}
ti.tile_size_bytes = crate::encoder::key_frame::min_tile_size_bytes(&tile_payloads) as u8;
fh.tile_info = Some(ti.clone());
let tile_group_body = tile_bytes;
let ctx_update_span: Option<(u32, u32)> = if ti.tile_cols_log2 + ti.tile_rows_log2 > 0 {
let write_with = |ctx: u32| -> Vec<u8> {
let mut fh_c = fh.clone();
if let Some(t) = fh_c.tile_info.as_mut() {
t.context_update_tile_id = ctx;
}
let mut bw = crate::encoder::bitwriter::BitWriter::new();
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
&mut bw,
&fh_c,
seq,
prev_gm_for_header.as_ref(),
);
bw.byte_align();
bw.finish()
};
let a = write_with(0);
let b = write_with(1);
debug_assert_eq!(a.len(), b.len(), "f(n) field cannot move the header length");
let n = ti.tile_cols_log2 + ti.tile_rows_log2;
let mut lsb: Option<u32> = None;
for (i, (&x, &y)) in a.iter().zip(b.iter()).enumerate() {
let d = x ^ y;
if d != 0 {
debug_assert!(
d.count_ones() == 1 && lsb.is_none(),
"exactly one bit distinguishes ctx ids 0 and 1"
);
lsb = Some((i as u32) * 8 + (7 - d.trailing_zeros()));
}
}
let lsb = lsb.ok_or(Error::PartitionWalkOutOfRange)?;
Some((lsb + 1 - n, n))
} else {
None
};
let write_fh_bits = |bw: &mut crate::encoder::bitwriter::BitWriter| {
crate::encoder::frame_obu::encode_uncompressed_header_with_prev_gm(
bw,
&fh,
seq,
prev_gm_for_header.as_ref(),
);
};
let effective_groups = cfg.tile_groups.clamp(1, num_tiles);
let spans_split = cfg.tile_spans && effective_groups == 1 && num_tiles > 1;
let frame_obus: Vec<ObuFrame> = if effective_groups > 1 || spans_split {
let fh_payload = {
let mut bw = crate::encoder::bitwriter::BitWriter::new();
write_fh_bits(&mut bw);
bw.trailing_bits_to_alignment();
bw.finish()
};
let tg_bodies = if spans_split {
vec![crate::encoder::tile_group_obu::single_group_with_spans(
&tile_group_body,
num_tiles,
ti.tile_cols_log2,
ti.tile_rows_log2,
u32::from(ti.tile_size_bytes),
)?]
} else {
crate::encoder::tile_group_obu::split_whole_frame_body(
&tile_group_body,
num_tiles,
ti.tile_cols_log2,
ti.tile_rows_log2,
u32::from(ti.tile_size_bytes),
effective_groups,
)?
};
let mut obus = vec![ObuFrame::new(ObuType::FrameHeader, fh_payload)];
obus.extend(
tg_bodies
.into_iter()
.map(|b| ObuFrame::new(ObuType::TileGroup, b)),
);
obus
} else {
let frame_body = {
let mut bw = crate::encoder::bitwriter::BitWriter::new();
write_fh_bits(&mut bw);
bw.byte_align();
let mut body = bw.finish();
body.extend_from_slice(&tile_group_body);
body
};
vec![ObuFrame::new(ObuType::Frame, frame_body)]
};
const REFMVS_LIMIT: i16 = (1 << 12) - 1;
let cells = (mi_rows as usize) * (mi_cols as usize);
let mut mf_ref_frames: Vec<i8> = vec![-1; cells];
let mut mf_mvs: Vec<i16> = vec![0; cells * 2];
{
let raw_refs = ictx.mirror.ref_frames();
let raw_mvs = ictx.mirror.mvs();
let hint_bits = u32::from(seq.order_hint_bits);
let by_ref = ictx.ip.order_hints.order_hints_by_ref;
for cell in 0..cells {
for list in 0..2usize {
let r = raw_refs[cell * 2 + list];
if r > 0 {
let dist = crate::inter_pred::get_relative_dist(
by_ref[r as usize],
fh.order_hint as i32,
hint_bits,
);
if dist < 0 {
let mv_row = raw_mvs[(cell * 2 + list) * 2];
let mv_col = raw_mvs[(cell * 2 + list) * 2 + 1];
if mv_row.abs() <= REFMVS_LIMIT && mv_col.abs() <= REFMVS_LIMIT {
mf_ref_frames[cell] = r;
mf_mvs[cell * 2] = mv_row;
mf_mvs[cell * 2 + 1] = mv_col;
}
}
}
}
}
}
let saved_mf = SavedMotionField {
mf_mvs,
mf_ref_frames,
saved_order_hints: ictx.ip.order_hints.order_hints_by_ref,
mi_rows,
mi_cols,
frame_is_intra: false,
};
let (cdfs, donor_cdfs) = if cfg.collect_donor_cdfs && all_cdfs.len() > 1 {
(all_cdfs[ctx_update_id as usize].clone(), all_cdfs)
} else {
(all_cdfs.swap_remove(ctx_update_id as usize), Vec::new())
};
let carry = RefSlotCarry {
cdfs,
segment_ids: state.mirror().segment_ids().to_vec(),
mi_rows,
mi_cols,
gm_params: fh
.global_motion_params
.as_ref()
.map(|g| g.gm_params)
.unwrap_or_else(crate::uncompressed_header_tail::prev_gm_params_default),
donor_cdfs,
ctx_update_span,
};
Ok((
frame_obus,
GopFrameReconYuv {
y: recon.y,
u: recon.u,
v: recon.v,
},
saved_mf,
carry,
InterFrameAux {
seg_temporal: seg_temporal_elected,
primary_ref: fh.primary_ref_frame,
hp_mv: hp_elected,
delta_q: delta_q_elected,
qm: qm_elected,
cdef: cdef_elected,
lr: lr_elected,
donor_elected,
},
))
}
struct PSearchCtx {
mi_rows: u32,
mi_cols: u32,
luma_w: u32,
luma_h: u32,
ref_dims: Vec<(u32, u32)>,
search_planes: Vec<Option<Vec<u16>>>,
bit_depth: u8,
subsampling_x: u8,
subsampling_y: u8,
num_planes: u8,
ref_planes: Vec<[Vec<u16>; 3]>,
slot_to_plane: [usize; 8],
ref_frame_idx: [u8; 7],
single_refs: Vec<i8>,
compound_pairs: Vec<[i8; 2]>,
mi_sizes: Vec<usize>,
is_inters: Vec<u8>,
ref_frames: Vec<i8>,
mvs: Vec<i16>,
interp_filters: Vec<u8>,
motion_modes: Vec<u8>,
y_modes: Vec<u8>,
avail_us: Vec<u8>,
avail_ls: Vec<u8>,
compound_types: Vec<u8>,
wedge_indices: Vec<u8>,
wedge_signs: Vec<u8>,
mask_types: Vec<u8>,
interintra_modes: Vec<u8>,
wedge_interintras: Vec<u8>,
interintra_wedge_indices: Vec<u8>,
local_warp: Vec<i32>,
local_warp_valid: Vec<u8>,
scratch: [Vec<u16>; 3],
gm_types: [u8; 8],
gm_flat: [i32; 48],
ref_hints: [i32; 8],
hint_bits: u32,
current_hint: i32,
no_scaled: [bool; 8],
seg_alt_q: Vec<i16>,
base_q_idx: u8,
exact_mask: Option<Vec<bool>>,
auto_lossless: bool,
lossless_seg: Option<u8>,
extras: Option<SegExtras>,
seg_id_pre_skip: bool,
mirror: PartitionWalker,
ip: SyntaxInterFrameParams,
delta_q_units: i32,
}
impl PSearchCtx {
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
fn new(
prev: &GopFrameReconYuv,
prevprev: &GopFrameReconYuv,
mi_rows: u32,
mi_cols: u32,
width: usize,
height: usize,
ip: SyntaxInterFrameParams,
p_index: u32,
base_q_idx: u8,
seg_alt_q: &[i16],
) -> Result<Self, Error> {
let last_slot = (p_index & 1) as usize;
let golden_slot = ((p_index - 1) & 1) as usize;
let mut slot_to_plane = [0usize; 8];
slot_to_plane[golden_slot] = 1;
let mut ref_frame_idx = [last_slot as u8; 7];
ref_frame_idx[3] = golden_slot as u8;
Self::with_refs(
&[prev, prevprev],
slot_to_plane,
ref_frame_idx,
vec![1, 4],
vec![[1, 4]],
mi_rows,
mi_cols,
width,
height,
&[(width as u32, height as u32); 2],
ip,
base_q_idx,
seg_alt_q,
8,
1,
1,
3,
)
}
#[allow(clippy::too_many_arguments)]
fn with_refs(
refs: &[&GopFrameReconYuv],
slot_to_plane: [usize; 8],
ref_frame_idx: [u8; 7],
single_refs: Vec<i8>,
compound_pairs: Vec<[i8; 2]>,
mi_rows: u32,
mi_cols: u32,
width: usize,
height: usize,
ref_dims: &[(u32, u32)],
ip: SyntaxInterFrameParams,
base_q_idx: u8,
seg_alt_q: &[i16],
bit_depth: u8,
subsampling_x: u8,
subsampling_y: u8,
num_planes: u8,
) -> Result<Self, Error> {
if refs.is_empty()
|| slot_to_plane.iter().any(|&p| p >= refs.len())
|| single_refs.iter().any(|&r| !(1..=7).contains(&r))
|| compound_pairs.iter().any(|pr| {
!single_refs.contains(&pr[0]) || !single_refs.contains(&pr[1]) || pr[0] >= pr[1]
})
{
return Err(Error::PartitionWalkOutOfRange);
}
let cells = (mi_rows as usize) * (mi_cols as usize);
let mut gm_flat = [0i32; 48];
let mut gm_types_u8 = [GM_TYPE_IDENTITY as u8; 8];
for r in 0..8 {
gm_flat[r * 6..r * 6 + 6].copy_from_slice(&ip.gm_params[r]);
gm_types_u8[r] = ip.gm_type[r] as u8;
}
let mut ref_frames = vec![0i8; cells * 2];
for cell in 0..cells {
ref_frames[cell * 2] = 0; ref_frames[cell * 2 + 1] = -1; }
let mirror = PartitionWalker::new(
mi_rows,
mi_cols,
TileGeometry {
mi_row_start: 0,
mi_row_end: mi_rows,
mi_col_start: 0,
mi_col_end: mi_cols,
},
)
.ok_or(Error::PartitionWalkOutOfRange)?;
if ref_dims.len() != refs.len() {
return Err(Error::PartitionWalkOutOfRange);
}
let scaled_of = |plane: usize| -> bool {
let (rw, rh) = ref_dims[plane];
crate::encoder::superres_elect::is_scaled(rw, rh, width as u32, height as u32)
};
let search_planes: Vec<Option<Vec<u16>>> = refs
.iter()
.enumerate()
.map(|(plane, r)| {
scaled_of(plane).then(|| {
let (rw, rh) = ref_dims[plane];
crate::encoder::superres_elect::resample_plane_bilinear(
&r.y,
rw as usize,
rh as usize,
width,
height,
)
})
})
.collect();
let mut no_scaled = [false; 8];
for (slot, flag) in no_scaled.iter_mut().enumerate() {
*flag = scaled_of(slot_to_plane[slot]);
}
Ok(PSearchCtx {
mi_rows,
mi_cols,
luma_w: width as u32,
luma_h: height as u32,
ref_dims: ref_dims.to_vec(),
search_planes,
bit_depth,
subsampling_x,
subsampling_y,
num_planes,
ref_planes: refs
.iter()
.map(|r| [r.y.clone(), r.u.clone(), r.v.clone()])
.collect(),
slot_to_plane,
ref_frame_idx,
single_refs,
compound_pairs,
mi_sizes: vec![BLOCK_4X4; cells],
is_inters: vec![0u8; cells],
ref_frames,
mvs: vec![0i16; cells * 4],
interp_filters: vec![EIGHTTAP; cells * 2],
motion_modes: vec![MOTION_MODE_SIMPLE; cells],
y_modes: vec![0u8; cells],
avail_us: {
let mut v = vec![1u8; cells];
v[..mi_cols as usize].fill(0);
v
},
avail_ls: {
let mut v = vec![1u8; cells];
for r in 0..mi_rows as usize {
v[r * mi_cols as usize] = 0;
}
v
},
compound_types: vec![crate::cdf::COMPOUND_AVERAGE; cells],
wedge_indices: vec![0u8; cells],
wedge_signs: vec![0u8; cells],
mask_types: vec![0u8; cells],
interintra_modes: vec![0u8; cells],
wedge_interintras: vec![0u8; cells],
interintra_wedge_indices: vec![0u8; cells],
local_warp: vec![0i32; cells * 6],
local_warp_valid: vec![0u8; cells],
scratch: {
let (cw, ch) = if num_planes > 1 {
(width >> subsampling_x, height >> subsampling_y)
} else {
(0, 0)
};
[
vec![0u16; width * height],
vec![0u16; cw * ch],
vec![0u16; cw * ch],
]
},
gm_types: gm_types_u8,
gm_flat,
ref_hints: {
let mut h = [0i32; 8];
h.copy_from_slice(&ip.order_hints.order_hints_by_ref);
h
},
hint_bits: ip.order_hints.order_hint_bits,
current_hint: ip.order_hints.current_order_hint,
seg_alt_q: seg_alt_q.to_vec(),
base_q_idx,
exact_mask: None,
auto_lossless: false,
extras: None,
seg_id_pre_skip: false,
lossless_seg: {
let ll = seg_lossless_array(base_q_idx, seg_alt_q);
(0..seg_alt_q.len()).find(|&s| ll[s]).map(|s| s as u8)
},
no_scaled,
mirror,
ip,
delta_q_units: 0,
})
}
fn segment_for_block(&self, input: &YuvFrame, mi_row: u32, mi_col: u32, b_size: usize) -> u8 {
if self.seg_alt_q.is_empty() {
return 0;
}
if let Some(ll) = self.lossless_seg {
if self.demands_exact(mi_row, mi_col, b_size) {
return ll;
}
if self.exact_mask.is_some() || self.auto_lossless {
let cap = u64::from(ll.saturating_sub(1));
let act = self.block_activity(input, mi_row, mi_col, b_size);
return (act / 6).min(cap) as u8;
}
}
let act = self.block_activity(input, mi_row, mi_col, b_size);
let mut top = (self.seg_alt_q.len() - 1) as u64;
if let Some(x) = self.extras.as_ref() {
let first_feature = (0..self.seg_alt_q.len())
.find(|&s| x.ref_frame[s].is_some() || x.skip[s] || x.globalmv[s])
.unwrap_or(self.seg_alt_q.len());
top = top.min(first_feature.saturating_sub(1) as u64);
}
(act / 6).min(top) as u8
}
fn block_activity(&self, input: &YuvFrame, mi_row: u32, mi_col: u32, b_size: usize) -> u64 {
let (bw, bh) = (
NUM_4X4_BLOCKS_WIDE[b_size] * 4,
NUM_4X4_BLOCKS_HIGH[b_size] * 4,
);
let (row0, col0) = ((mi_row as usize) * 4, (mi_col as usize) * 4);
let w = input.width as usize;
let mut sum = 0u64;
for i in 0..bh {
for j in 0..bw {
sum += u64::from(input.y[(row0 + i) * w + (col0 + j)]);
}
}
let n = (bw * bh) as u64;
let mean = sum / n;
let mut mad = 0u64;
for i in 0..bh {
for j in 0..bw {
let v = u64::from(input.y[(row0 + i) * w + (col0 + j)]);
mad += v.abs_diff(mean);
}
}
mad / n
}
fn seg_qp(&self, frame_qp: &QuantizerParams, segment_id: u8) -> QuantizerParams {
if self.seg_alt_q.is_empty() || segment_id == 0 {
return *frame_qp;
}
let base = if frame_qp.delta_q_present {
i32::from(frame_qp.current_q_index)
} else {
i32::from(self.base_q_idx)
};
let q = (base + i32::from(self.seg_alt_q[segment_id as usize])).clamp(0, 255) as u8;
let mut qp = QuantizerParams::neutral(q, frame_qp.bit_depth);
if frame_qp.using_qmatrix {
qp.using_qmatrix = true;
let row = frame_qp.seg_qm_level_for(segment_id);
for (plane, &lvl) in row.iter().enumerate() {
qp.seg_qm_level[plane] = [lvl; crate::uncompressed_header_tail::MAX_SEGMENTS];
}
}
qp
}
fn seg_lossless(&self, segment_id: u8) -> bool {
let s = segment_id as usize;
let q = if self.seg_alt_q.is_empty() || s >= self.seg_alt_q.len() {
i32::from(self.base_q_idx)
} else {
(i32::from(self.base_q_idx) + i32::from(self.seg_alt_q[s])).clamp(0, 255)
};
q == 0
}
fn synthetic_leaf(&self, input: &YuvFrame, mi_row: u32, mi_col: u32, b_size: usize) -> bool {
let (bw, bh) = (
NUM_4X4_BLOCKS_WIDE[b_size] * 4,
NUM_4X4_BLOCKS_HIGH[b_size] * 4,
);
let (row0, col0) = ((mi_row as usize) * 4, (mi_col as usize) * 4);
let w = input.width as usize;
let mut seen = [false; 256];
let mut distinct = 0usize;
for i in 0..bh {
for j in 0..bw {
let v = input.y[(row0 + i) * w + (col0 + j)] as usize;
if !seen[v] {
seen[v] = true;
distinct += 1;
if distinct > 16 {
return false;
}
}
}
}
true
}
fn demands_exact(&self, mi_row: u32, mi_col: u32, b_size: usize) -> bool {
let Some(mask) = self.exact_mask.as_ref() else {
return false;
};
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size] as u32;
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size] as u32;
let r1 = (mi_row + bh4).min(self.mi_rows);
let c1 = (mi_col + bw4).min(self.mi_cols);
(mi_row..r1).any(|r| (mi_col..c1).any(|c| mask[(r * self.mi_cols + c) as usize]))
}
fn begin_tile(&mut self, geometry: TileGeometry) {
self.mirror.begin_tile(geometry);
for r in 0..self.mi_rows {
for c in 0..self.mi_cols {
let i = (r * self.mi_cols + c) as usize;
self.avail_us[i] = u8::from(geometry.is_inside(r as i32 - 1, c as i32));
self.avail_ls[i] = u8::from(geometry.is_inside(r as i32, c as i32 - 1));
}
}
}
fn segment_pred(&self, mi_row: u32, mi_col: u32) -> u8 {
let at = |r: i64, c: i64| -> i32 {
if r < 0 || c < 0 || r >= i64::from(self.mi_rows) || c >= i64::from(self.mi_cols) {
return -1;
}
self.mirror.segment_ids()[(r as u32 * self.mi_cols + c as u32) as usize]
};
let cell = (mi_row * self.mi_cols + mi_col) as usize;
let (avail_u, avail_l) = (self.avail_us[cell] != 0, self.avail_ls[cell] != 0);
let (r, c) = (i64::from(mi_row), i64::from(mi_col));
let prev_ul = if avail_u && avail_l {
at(r - 1, c - 1)
} else {
-1
};
let prev_u = if avail_u { at(r - 1, c) } else { -1 };
let prev_l = if avail_l { at(r, c - 1) } else { -1 };
#[allow(clippy::if_same_then_else)]
let pred: i32 = if prev_u == -1 {
if prev_l == -1 {
0
} else {
prev_l
}
} else if prev_l == -1 {
prev_u
} else if prev_ul == prev_u {
prev_u
} else {
prev_l
};
pred as u8
}
fn search_luma(&self, plane: usize) -> &[u16] {
self.search_planes[plane]
.as_deref()
.unwrap_or(&self.ref_planes[plane][0])
}
fn plane_of_ref(&self, rf: i8) -> usize {
self.slot_to_plane[self.ref_frame_idx[(rf - 1) as usize] as usize]
}
fn find_stack(
&self,
mi_row: u32,
mi_col: u32,
b_size: usize,
ref_frame: [i8; 2],
) -> Result<FindMvStackResult, Error> {
self.mirror.find_mv_stack(
mi_row,
mi_col,
b_size,
[i32::from(ref_frame[0]), i32::from(ref_frame[1])],
ref_frame[1] > 0,
self.ip.use_ref_frame_mvs,
self.ip.gm_type,
self.ip.gm_params,
self.ip.ref_frame_sign_bias,
self.ip.allow_high_precision_mv,
self.ip.force_integer_mv,
&self.ip.motion_field_mvs,
)
}
fn snapshot_grids_rect(&self, mi_row: u32, mi_col: u32, n4: u32) -> DriverGridSnapshot {
let r1 = (mi_row + n4).min(self.mi_rows);
let c1 = (mi_col + n4).min(self.mi_cols);
let mut cells = Vec::with_capacity(((r1 - mi_row) * (c1 - mi_col)) as usize);
for rr in mi_row..r1 {
for cc in mi_col..c1 {
let cell = (rr * self.mi_cols + cc) as usize;
cells.push(DriverGridCell {
mi_size: self.mi_sizes[cell],
is_inter: self.is_inters[cell],
ref_frames: [self.ref_frames[cell * 2], self.ref_frames[cell * 2 + 1]],
mvs: [
self.mvs[cell * 4],
self.mvs[cell * 4 + 1],
self.mvs[cell * 4 + 2],
self.mvs[cell * 4 + 3],
],
interp_filters: [
self.interp_filters[cell * 2],
self.interp_filters[cell * 2 + 1],
],
motion_mode: self.motion_modes[cell],
y_mode: self.y_modes[cell],
compound_type: self.compound_types[cell],
wedge_index: self.wedge_indices[cell],
wedge_sign: self.wedge_signs[cell],
mask_type: self.mask_types[cell],
interintra_mode: self.interintra_modes[cell],
wedge_interintra: self.wedge_interintras[cell],
interintra_wedge_index: self.interintra_wedge_indices[cell],
local_warp: {
let mut w = [0i32; 6];
w.copy_from_slice(&self.local_warp[cell * 6..cell * 6 + 6]);
w
},
local_warp_valid: self.local_warp_valid[cell],
});
}
}
DriverGridSnapshot {
mi_row,
mi_col,
r1,
c1,
cells,
}
}
fn restore_grids_rect(&mut self, snap: &DriverGridSnapshot) {
let mut it = snap.cells.iter();
for rr in snap.mi_row..snap.r1 {
for cc in snap.mi_col..snap.c1 {
let cell = (rr * self.mi_cols + cc) as usize;
let s = it.next().expect("snapshot covers the rect");
self.mi_sizes[cell] = s.mi_size;
self.is_inters[cell] = s.is_inter;
self.ref_frames[cell * 2] = s.ref_frames[0];
self.ref_frames[cell * 2 + 1] = s.ref_frames[1];
self.mvs[cell * 4..cell * 4 + 4].copy_from_slice(&s.mvs);
self.interp_filters[cell * 2] = s.interp_filters[0];
self.interp_filters[cell * 2 + 1] = s.interp_filters[1];
self.motion_modes[cell] = s.motion_mode;
self.y_modes[cell] = s.y_mode;
self.compound_types[cell] = s.compound_type;
self.wedge_indices[cell] = s.wedge_index;
self.wedge_signs[cell] = s.wedge_sign;
self.mask_types[cell] = s.mask_type;
self.interintra_modes[cell] = s.interintra_mode;
self.wedge_interintras[cell] = s.wedge_interintra;
self.interintra_wedge_indices[cell] = s.interintra_wedge_index;
self.local_warp[cell * 6..cell * 6 + 6].copy_from_slice(&s.local_warp);
self.local_warp_valid[cell] = s.local_warp_valid;
}
}
}
fn warp_fit(
&self,
mi_row: u32,
mi_col: u32,
b_size: usize,
ref0: i8,
mv: [[i32; 2]; 2],
) -> Option<[i32; 6]> {
let mut cands: Vec<crate::inter_pred::WarpSampleCand> = Vec::new();
let ns = self.mirror.find_warp_samples_list(
mi_row,
mi_col,
b_size,
i32::from(ref0),
mv,
&mut cands,
);
if ns == 0 {
return None;
}
let lw = crate::inter_pred::warp_estimation(
&cands,
mi_row as i32,
mi_col as i32,
NUM_4X4_BLOCKS_WIDE[b_size] as i32,
NUM_4X4_BLOCKS_HIGH[b_size] as i32,
[mv[0][0] as i16, mv[0][1] as i16],
);
if !lw.local_valid {
return None;
}
match crate::inter_pred::setup_shear(lw.local_warp_params) {
Some(s) if s.warp_valid => Some(lw.local_warp_params),
_ => None,
}
}
fn warp_num_samples(
&self,
mi_row: u32,
mi_col: u32,
b_size: usize,
ref0: i8,
mv: [[i32; 2]; 2],
) -> u32 {
let mut cands: Vec<crate::inter_pred::WarpSampleCand> = Vec::new();
self.mirror
.find_warp_samples_list(mi_row, mi_col, b_size, i32::from(ref0), mv, &mut cands)
}
fn stamp_leaf(
&mut self,
block: &SyntaxBlock,
mi_row: u32,
mi_col: u32,
b_size: usize,
lossless: bool,
) {
use crate::cdf::EncoderBlockSyntaxStamp;
let tx_pre = if lossless {
TX_4X4 as u8
} else {
MAX_TX_SIZE_RECT[b_size] as u8
};
let stamp = match block.inter.as_ref() {
Some(ib) => EncoderBlockSyntaxStamp {
mi_row,
mi_col,
sub_size: b_size,
skip: block.skip,
skip_mode: ib.skip_mode,
segment_id: block.segment_id,
is_inter: 1,
y_mode: ib.y_mode,
ref_frame: if ib.interintra_mode.is_some() {
[
ib.ref_frame[0],
crate::uncompressed_header_tail::INTRA_FRAME as i8,
]
} else {
ib.ref_frame
},
mv: ib.mv[0],
mv2: ib.mv[1],
interp_filter: ib.interp_filter,
motion_mode: ib.motion_mode,
palette_size_y: 0,
palette_colors_y: &[],
palette_size_uv: 0,
palette_colors_u: &[],
palette_colors_v: &[],
cdef: None,
tx_size: tx_pre,
comp_group_idx: u8::from(
ib.ref_frame[1] > 0
&& ib.skip_mode == 0
&& matches!(
ib.compound_type,
crate::inter_pred::COMPOUND_WEDGE | crate::inter_pred::COMPOUND_DIFFWTD
),
),
compound_idx: u8::from(
!(ib.ref_frame[1] > 0
&& ib.skip_mode == 0
&& ib.compound_type == crate::inter_pred::COMPOUND_DISTANCE),
),
compound_type: if ib.interintra_mode.is_some()
|| (ib.ref_frame[1] > 0 && ib.skip_mode == 0)
{
ib.compound_type
} else {
crate::inter_pred::COMPOUND_AVERAGE
},
wedge_index: ib.wedge_index,
wedge_sign: ib.wedge_sign,
mask_type: ib.mask_type,
interintra_mode: ib.interintra_mode.unwrap_or(0),
wedge_interintra: u8::from(
ib.interintra_mode.is_some() && ib.wedge_interintra == 1,
),
interintra_wedge_index: ib.interintra_wedge_index,
},
None => EncoderBlockSyntaxStamp {
mi_row,
mi_col,
sub_size: b_size,
skip: block.skip,
skip_mode: 0,
segment_id: block.segment_id,
is_inter: 0,
y_mode: block.y_mode,
ref_frame: [0, -1],
mv: [0, 0],
mv2: [0, 0],
interp_filter: [0, 0],
motion_mode: MOTION_MODE_SIMPLE,
palette_size_y: block.palette.size_y,
palette_colors_y: &block.palette.colors_y,
palette_size_uv: block.palette.size_uv,
palette_colors_u: &block.palette.colors_u,
palette_colors_v: &block.palette.colors_v,
cdef: None,
tx_size: tx_pre,
comp_group_idx: 0,
compound_idx: 1,
compound_type: crate::inter_pred::COMPOUND_AVERAGE,
wedge_index: 0,
wedge_sign: 0,
mask_type: 0,
interintra_mode: 0,
wedge_interintra: 0,
interintra_wedge_index: 0,
},
};
self.mirror.stamp_encoder_block_syntax(&stamp);
}
fn stamp_intra_leaf_grids(&mut self, mi_row: u32, mi_col: u32, b_size: usize, y_mode: u8) {
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size] as u32;
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size] as u32;
for dr in 0..bh4 {
let rr = mi_row + dr;
if rr >= self.mi_rows {
break;
}
for dc in 0..bw4 {
let cc = mi_col + dc;
if cc >= self.mi_cols {
break;
}
let cell = (rr * self.mi_cols + cc) as usize;
self.mi_sizes[cell] = b_size;
self.is_inters[cell] = 0;
self.ref_frames[cell * 2] = 0;
self.ref_frames[cell * 2 + 1] = -1;
self.mvs[cell * 4] = 0;
self.mvs[cell * 4 + 1] = 0;
self.mvs[cell * 4 + 2] = 0;
self.mvs[cell * 4 + 3] = 0;
self.interp_filters[cell * 2] = EIGHTTAP;
self.interp_filters[cell * 2 + 1] = EIGHTTAP;
self.motion_modes[cell] = MOTION_MODE_SIMPLE;
self.y_modes[cell] = y_mode;
self.compound_types[cell] = crate::inter_pred::COMPOUND_AVERAGE;
self.wedge_indices[cell] = 0;
self.wedge_signs[cell] = 0;
self.mask_types[cell] = 0;
self.interintra_modes[cell] = 0;
self.wedge_interintras[cell] = 0;
self.interintra_wedge_indices[cell] = 0;
}
}
}
#[allow(clippy::too_many_arguments)]
fn predict_leaf(
&mut self,
mi_row: u32,
mi_col: u32,
b_size: usize,
ref_frame: [i8; 2],
y_mode: u8,
mv: [[i32; 2]; 2],
filter: u8,
comp: CompoundSel,
ii: Option<InterIntraTrial<'_>>,
motion: MotionSel,
) -> Result<(), Error> {
if ii.is_some() && ref_frame[1] != -1 {
return Err(Error::PartitionWalkOutOfRange);
}
if motion != MotionSel::Simple && (ref_frame[1] != -1 || ii.is_some()) {
return Err(Error::PartitionWalkOutOfRange);
}
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size] as u32;
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size] as u32;
let stamp_rf1: i8 = if ii.is_some() {
crate::uncompressed_header_tail::INTRA_FRAME as i8
} else {
ref_frame[1]
};
let stamp_ctype: u8 = match &ii {
Some(t) if t.wedge.is_some() => crate::inter_pred::COMPOUND_WEDGE,
Some(_) => crate::inter_pred::COMPOUND_INTRA,
None => comp.ctype,
};
for dr in 0..bh4 {
let rr = mi_row + dr;
if rr >= self.mi_rows {
break;
}
for dc in 0..bw4 {
let cc = mi_col + dc;
if cc >= self.mi_cols {
break;
}
let cell = (rr * self.mi_cols + cc) as usize;
self.mi_sizes[cell] = b_size;
self.is_inters[cell] = 1;
self.ref_frames[cell * 2] = ref_frame[0];
self.ref_frames[cell * 2 + 1] = stamp_rf1;
self.mvs[cell * 4] = mv[0][0] as i16;
self.mvs[cell * 4 + 1] = mv[0][1] as i16;
self.mvs[cell * 4 + 2] = mv[1][0] as i16;
self.mvs[cell * 4 + 3] = mv[1][1] as i16;
self.interp_filters[cell * 2] = filter;
self.interp_filters[cell * 2 + 1] = filter;
self.motion_modes[cell] = motion.ordinal();
self.y_modes[cell] = y_mode;
self.compound_types[cell] = stamp_ctype;
self.wedge_indices[cell] = comp.wedge_index;
self.wedge_signs[cell] = comp.wedge_sign;
self.mask_types[cell] = comp.mask_type;
self.interintra_modes[cell] = ii.as_ref().map_or(0, |t| t.mode);
self.wedge_interintras[cell] =
u8::from(ii.as_ref().is_some_and(|t| t.wedge.is_some()));
self.interintra_wedge_indices[cell] =
ii.as_ref().and_then(|t| t.wedge).unwrap_or(0);
}
}
{
let origin = (mi_row * self.mi_cols + mi_col) as usize;
if let MotionSel::Warp(p) = motion {
self.local_warp[origin * 6..origin * 6 + 6].copy_from_slice(&p);
self.local_warp_valid[origin] = 1;
} else {
self.local_warp[origin * 6..origin * 6 + 6].fill(0);
self.local_warp_valid[origin] = 0;
}
}
if let Some(t) = &ii {
let mode = match t.mode {
crate::cdf::II_V_PRED => crate::cdf::V_PRED,
crate::cdf::II_H_PRED => crate::cdf::H_PRED,
crate::cdf::II_DC_PRED => crate::cdf::DC_PRED,
_ => crate::cdf::SMOOTH_PRED,
};
let cell = (mi_row * self.mi_cols + mi_col) as usize;
let avail_u = self.avail_us[cell] != 0;
let avail_l = self.avail_ls[cell] != 0;
for plane in 0..usize::from(self.num_planes) {
let (sub_x, sub_y): (u8, u8) = if plane > 0 {
(self.subsampling_x, self.subsampling_y)
} else {
(0, 0)
};
let plane_sz =
crate::cdf::get_plane_residual_size(b_size, plane as u8, sub_x, sub_y)
.ok_or(Error::PartitionWalkOutOfRange)?;
let w = NUM_4X4_BLOCKS_WIDE[plane_sz] * 4;
let h = NUM_4X4_BLOCKS_HIGH[plane_sz] * 4;
let tx_sz = crate::cdf::find_tx_size(w, h).ok_or(Error::PartitionWalkOutOfRange)?;
let base_x = ((mi_col >> sub_x) * 4) as usize;
let base_y = ((mi_row >> sub_y) * 4) as usize;
let (have_ar, have_bl) =
super::key_frame::tu_corner_avail(&t.neigh.bd, plane, base_x, base_y, w, h);
let (buf, pw, ph) = t.neigh.plane(plane);
let read = |yy: u32, xx: u32| -> u16 { buf[(yy as usize) * pw + (xx as usize)] };
let out = &mut self.scratch[plane];
PartitionWalker::predict_intra_into_u16_plane(
&read,
out,
pw,
(pw - 1) as u32,
(ph - 1) as u32,
base_x as u32,
base_y as u32,
tx_sz,
mode,
0,
self.bit_depth,
avail_u,
avail_l,
have_ar,
have_bl,
false,
0,
None,
);
}
}
let (luma_w, luma_h) = (self.luma_w, self.luma_h);
let (mi_rows, mi_cols) = (self.mi_rows, self.mi_cols);
let slot_to_plane = self.slot_to_plane;
let ref_frame_idx = self.ref_frame_idx;
let bit_depth = self.bit_depth;
let (ssx, ssy) = (self.subsampling_x, self.subsampling_y);
let num_planes = usize::from(self.num_planes);
let PSearchCtx {
ref_planes,
compound_types,
wedge_indices,
wedge_signs,
mask_types,
interintra_modes,
wedge_interintras,
interintra_wedge_indices,
mi_sizes,
is_inters,
ref_frames,
mvs,
interp_filters,
motion_modes,
y_modes,
avail_us,
avail_ls,
local_warp,
local_warp_valid,
scratch,
gm_types,
gm_flat,
ref_hints,
hint_bits,
current_hint,
no_scaled,
ref_dims,
..
} = &mut *self;
let store_y = make_store(ref_planes, ref_dims, &slot_to_plane, 0, 0);
let store_u = make_store(ref_planes, ref_dims, &slot_to_plane, 1, ssx);
let store_v = make_store(ref_planes, ref_dims, &slot_to_plane, 2, ssx);
let grid = InterModeInfoGrid {
mi_sizes,
is_inters,
ref_frames,
mvs,
interp_filters,
compound_types,
wedge_indices,
wedge_signs,
mask_types,
interintra_modes,
wedge_interintras,
interintra_wedge_indices,
order_hint_bits: *hint_bits,
current_order_hint: *current_hint,
order_hints_by_ref: ref_hints,
mi_rows,
mi_cols,
bit_depth,
warp: Some(crate::GridWarpContext {
motion_modes,
local_warp_params: local_warp,
local_warp_valid,
y_modes,
gm_types,
gm_params: gm_flat,
force_integer_mv: false,
is_scaled: no_scaled,
}),
obmc: Some(crate::GridObmcContext {
motion_modes,
avail_u: avail_us,
avail_l: avail_ls,
}),
};
let [s0, s1, s2] = scratch;
let mut planes = [
PlaneReconContext {
plane: 0,
subsampling_x: 0,
subsampling_y: 0,
frame_store: &store_y,
frame_width: luma_w,
frame_height: luma_h,
curr: s0,
curr_stride: luma_w as usize,
},
PlaneReconContext {
plane: 1,
subsampling_x: ssx,
subsampling_y: ssy,
frame_store: &store_u,
frame_width: luma_w,
frame_height: luma_h,
curr: s1,
curr_stride: (luma_w as usize) >> ssx,
},
PlaneReconContext {
plane: 2,
subsampling_x: ssx,
subsampling_y: ssy,
frame_store: &store_v,
frame_width: luma_w,
frame_height: luma_h,
curr: s2,
curr_stride: (luma_w as usize) >> ssx,
},
];
reconstruct_inter_leaf_at(
&grid,
&ref_frame_idx,
&mut planes[..num_planes],
mi_row as usize,
mi_col as usize,
)
}
}
fn make_store<'a>(
ref_planes: &'a [[Vec<u16>; 3]],
ref_dims: &[(u32, u32)],
slot_to_plane: &[usize; 8],
plane: usize,
ssx: u8,
) -> [RefFrameStoreEntry<'a>; 8] {
core::array::from_fn(|slot| {
let idx = slot_to_plane[slot];
let (rw, rh) = ref_dims[idx];
RefFrameStoreEntry {
plane: &ref_planes[idx][plane],
stride: ((rw + u32::from(ssx)) >> ssx) as usize,
upscaled_width: rw,
width: rw,
height: rh,
}
})
}
#[allow(clippy::too_many_arguments)]
fn motion_search_luma(
input: &YuvFrame,
ref_y: &[u16],
width: usize,
height: usize,
row0: usize,
col0: usize,
bw: usize,
bh: usize,
) -> [i32; 2] {
let n = bw.max(bh) as u64;
let cost_at = |dy: i32, dx: i32| -> u64 {
let mut ssd = 0u64;
for i in 0..bh {
let sy = ((row0 + i) as i32 + dy).clamp(0, height as i32 - 1) as usize;
for j in 0..bw {
let sx = ((col0 + j) as i32 + dx).clamp(0, width as i32 - 1) as usize;
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ref_y[sy * width + sx]);
ssd += (d * d) as u64;
}
}
ssd + (dy.unsigned_abs() as u64 + dx.unsigned_abs() as u64) * n
};
let mut best = (cost_at(0, 0), [0i32, 0]);
let mut dy = -SEARCH_RANGE;
while dy <= SEARCH_RANGE {
let mut dx = -SEARCH_RANGE;
while dx <= SEARCH_RANGE {
if !(dy == 0 && dx == 0) {
let c = cost_at(dy, dx);
if c < best.0 {
best = (c, [dy, dx]);
}
}
dx += 4;
}
dy += 4;
}
let center = best.1;
for dy in center[0] - 3..=center[0] + 3 {
for dx in center[1] - 3..=center[1] + 3 {
if dy == center[0] && dx == center[1] {
continue;
}
let c = cost_at(dy, dx);
if c < best.0 {
best = (c, [dy, dx]);
}
}
}
[best.1[0] * 8, best.1[1] * 8]
}
#[allow(clippy::too_many_arguments)]
fn refine_mv_subpel(
input: &YuvFrame,
ictx: &mut PSearchCtx,
mi_r: u32,
mi_c: u32,
b_size: usize,
ref_frame: i8,
row0: usize,
col0: usize,
bw: usize,
bh: usize,
width: usize,
mv_int: [i32; 2],
) -> Result<[i32; 2], Error> {
let score = |ictx: &mut PSearchCtx, mv: [i32; 2]| -> Result<u64, Error> {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
[ref_frame, -1],
MODE_NEWMV,
[mv, [0, 0]],
EIGHTTAP,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)?;
let mut ssd = 0u64;
for i in 0..bh {
for j in 0..bw {
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ictx.scratch[0][(row0 + i) * width + col0 + j]);
ssd += (d * d) as u64;
}
}
Ok(ssd + ((mv[0].unsigned_abs() + mv[1].unsigned_abs()) as u64) * (bw.max(bh) as u64) / 8)
};
let mut best = (score(ictx, mv_int)?, mv_int);
let steps: &[i32] = if ictx.ip.allow_high_precision_mv {
&[4, 2, 1]
} else {
&[4, 2]
};
for &step in steps {
let center = best.1;
for dy in [-step, 0, step] {
for dx in [-step, 0, step] {
if dy == 0 && dx == 0 {
continue;
}
let cand = [center[0] + dy, center[1] + dx];
let c = score(ictx, cand)?;
if c < best.0 {
best = (c, cand);
}
}
}
}
Ok(best.1)
}
#[allow(clippy::too_many_arguments)]
fn encode_inter_leaf(
mi_r: u32,
mi_c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
) -> Result<SyntaxBlock, Error> {
let leaf = encode_inter_leaf_base(mi_r, mi_c, b_size, input, recon, ictx, pricing)?;
let mut leaf = if ictx.extras.is_some()
&& crate::cdf::block_width(b_size) >= 8
&& crate::cdf::block_height(b_size) >= 8
{
seg_extras_trials(mi_r, mi_c, b_size, input, recon, ictx, pricing, leaf)?
} else {
leaf
};
stamp_delta_q(&mut leaf, b_size, ictx, pricing);
Ok(leaf)
}
fn stamp_delta_q(
leaf: &mut SyntaxBlock,
b_size: usize,
ictx: &PSearchCtx,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
) {
if ictx.delta_q_units == 0 {
return;
}
let pending = pricing.is_some_and(|(twin, _)| twin.deltas_pending());
leaf.reduced_delta_q_index = if pending && !(b_size == BLOCK_64X64 && leaf.skip == 1) {
ictx.delta_q_units
} else {
0
};
}
#[allow(clippy::too_many_arguments)]
fn encode_inter_leaf_base(
mi_r: u32,
mi_c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
) -> Result<SyntaxBlock, Error> {
let leaf = encode_inter_leaf_modes(mi_r, mi_c, b_size, input, recon, ictx, pricing)?;
if !ictx.ip.skip_mode_present
|| crate::cdf::block_width(b_size) < 8
|| crate::cdf::block_height(b_size) < 8
{
return Ok(leaf);
}
let sm_ref: [i8; 2] = ictx.ip.skip_mode_frame;
let stack = ictx.find_stack(mi_r, mi_c, b_size, sm_ref)?;
let sm_mv = [stack.ref_stack_mv[0][0], stack.ref_stack_mv[0][1]];
if sm_mv
.iter()
.any(|m| m[0].unsigned_abs() >= (1 << 14) || m[1].unsigned_abs() >= (1 << 14))
{
return Ok(leaf);
}
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size];
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size];
let lambda = lambda_for(&recon.qp);
let rate_of = |blk: &SyntaxBlock| -> Result<u64, Error> {
match pricing {
Some((twin, params)) => twin.price_block(blk, mi_r, mi_c, b_size, params),
None => Ok(p_leaf_rate(blk) * 256),
}
};
let d_normal = region_distortion_wh(recon, input, mi_r, mi_c, bw4, bh4);
let score_normal = score256(d_normal, lambda, rate_of(&leaf)?);
let after_normal = save_region_wh(recon, mi_r, mi_c, bw4, bh4);
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
sm_ref,
MODE_NEAREST_NEARESTMV,
sm_mv,
EIGHTTAP,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)?;
let (bw, bh) = (bw4 * 4, bh4 * 4);
let (row0, col0) = ((mi_r as usize) * 4, (mi_c as usize) * 4);
let (crow0, ccol0) = recon.chroma_origin(mi_r, mi_c);
let (cbw, cbh) = if recon.num_planes > 1 {
(bw >> recon.subsampling_x, bh >> recon.subsampling_y)
} else {
(0, 0)
};
let width = recon.width;
let cw = recon.chroma_w;
for i in 0..bh {
for j in 0..bw {
recon.y[(row0 + i) * width + (col0 + j)] =
ictx.scratch[0][(row0 + i) * width + (col0 + j)];
}
}
for i in 0..cbh {
for j in 0..cbw {
recon.u[(crow0 + i) * cw + (ccol0 + j)] =
ictx.scratch[1][(crow0 + i) * cw + (ccol0 + j)];
recon.v[(crow0 + i) * cw + (ccol0 + j)] =
ictx.scratch[2][(crow0 + i) * cw + (ccol0 + j)];
}
}
tu_bd_stamp(&mut recon.bd, 0, col0, row0, bw, bh);
if recon.num_planes > 1 {
tu_bd_stamp(&mut recon.bd, 1, ccol0, crow0, cbw, cbh);
tu_bd_stamp(&mut recon.bd, 2, ccol0, crow0, cbw, cbh);
}
let mut sm_leaf = SyntaxBlock::skip_leaf(0, None);
sm_leaf.segment_id = if ictx.seg_id_pre_skip {
0
} else {
ictx.segment_pred(mi_r, mi_c)
};
sm_leaf.inter = Some(SyntaxInterBlock {
ref_frame: sm_ref,
y_mode: MODE_NEAREST_NEARESTMV,
mv: sm_mv,
ref_mv_idx: 0,
interp_filter: [EIGHTTAP; 2],
skip_mode: 1,
compound_type: crate::inter_pred::COMPOUND_AVERAGE,
wedge_index: 0,
wedge_sign: 0,
mask_type: 0,
interintra_mode: None,
wedge_interintra: 0,
interintra_wedge_index: 0,
motion_mode: MOTION_MODE_SIMPLE,
});
let d_sm = region_distortion_wh(recon, input, mi_r, mi_c, bw4, bh4);
let score_sm = score256(d_sm, lambda, rate_of(&sm_leaf)?);
let must_be_exact =
ictx.seg_lossless(sm_leaf.segment_id) || ictx.demands_exact(mi_r, mi_c, b_size);
if score_sm < score_normal && (!must_be_exact || d_sm == 0) {
return Ok(sm_leaf);
}
restore_region(recon, mi_r, mi_c, &after_normal);
if let Some(ib) = leaf.inter.as_ref() {
let ii = ib.interintra_mode.map(|m| InterIntraTrial {
mode: m,
wedge: (ib.wedge_interintra == 1).then_some(ib.interintra_wedge_index),
neigh: &*recon,
});
let sel = if ib.interintra_mode.is_some() {
CompoundSel::AVERAGE
} else {
CompoundSel {
ctype: ib.compound_type,
wedge_index: ib.wedge_index,
wedge_sign: ib.wedge_sign,
mask_type: ib.mask_type,
}
};
let motion = match ib.motion_mode {
MOTION_MODE_OBMC => MotionSel::Obmc,
MOTION_MODE_WARPED_CAUSAL => MotionSel::Warp(
ictx.warp_fit(mi_r, mi_c, b_size, ib.ref_frame[0], ib.mv)
.ok_or(Error::PartitionWalkOutOfRange)?,
),
_ => MotionSel::Simple,
};
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ib.ref_frame,
ib.y_mode,
ib.mv,
ib.interp_filter[0],
sel,
ii,
motion,
)?;
}
Ok(leaf)
}
fn repredict_leaf_arm(
mi_r: u32,
mi_c: u32,
b_size: usize,
recon: &ReconState,
ictx: &mut PSearchCtx,
leaf: &SyntaxBlock,
) -> Result<(), Error> {
let Some(ib) = leaf.inter.as_ref() else {
return Ok(());
};
let ii = ib.interintra_mode.map(|m| InterIntraTrial {
mode: m,
wedge: (ib.wedge_interintra == 1).then_some(ib.interintra_wedge_index),
neigh: recon,
});
let sel = if ib.interintra_mode.is_some() {
CompoundSel::AVERAGE
} else {
CompoundSel {
ctype: ib.compound_type,
wedge_index: ib.wedge_index,
wedge_sign: ib.wedge_sign,
mask_type: ib.mask_type,
}
};
let motion = match ib.motion_mode {
MOTION_MODE_OBMC => MotionSel::Obmc,
MOTION_MODE_WARPED_CAUSAL => MotionSel::Warp(
ictx.warp_fit(mi_r, mi_c, b_size, ib.ref_frame[0], ib.mv)
.ok_or(Error::PartitionWalkOutOfRange)?,
),
_ => MotionSel::Simple,
};
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ib.ref_frame,
ib.y_mode,
ib.mv,
ib.interp_filter[0],
sel,
ii,
motion,
)
}
#[allow(clippy::too_many_arguments)]
fn seg_extras_trials(
mi_r: u32,
mi_c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
mut best: SyntaxBlock,
) -> Result<SyntaxBlock, Error> {
let x = ictx.extras.expect("caller gated on extras");
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size];
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size];
let lambda = lambda_for(&recon.qp);
let rate_of = |blk: &SyntaxBlock| -> Result<u64, Error> {
match pricing {
Some((twin, params)) => twin.price_block(blk, mi_r, mi_c, b_size, params),
None => Ok(p_leaf_rate(blk) * 256),
}
};
let d_now = region_distortion_wh(recon, input, mi_r, mi_c, bw4, bh4);
let mut score_best = score256(d_now, lambda, rate_of(&best)?);
for (s, rf) in x.ref_frame.iter().enumerate() {
let Some(rf) = rf else { continue };
if x.skip[s] || x.globalmv[s] {
continue;
}
let Some(ib) = best.inter.as_ref() else {
break;
};
if ib.skip_mode != 0
|| ib.ref_frame[0] != *rf
|| ib.ref_frame[1] > 0
|| best.segment_id as usize == s
{
continue;
}
if ictx.seg_qp(&recon.qp, s as u8).base_q_idx
!= ictx.seg_qp(&recon.qp, best.segment_id).base_q_idx
{
continue;
}
let mut alt = best.clone();
alt.segment_id = s as u8;
let sc = score256(d_now, lambda, rate_of(&alt)?);
if sc < score_best {
score_best = sc;
best = alt;
}
}
let mut after_best = save_region_wh(recon, mi_r, mi_c, bw4, bh4);
for s in 0..crate::uncompressed_header_tail::MAX_SEGMENTS {
let (is_skip_seg, is_gmv_seg) = (x.skip[s], x.globalmv[s]);
if !is_skip_seg && !is_gmv_seg {
continue;
}
let stack = ictx.find_stack(mi_r, mi_c, b_size, [1, -1])?;
let gmv = [stack.global_mvs[0], [0, 0]];
if gmv[0][0].unsigned_abs() >= (1 << 14) || gmv[0][1].unsigned_abs() >= (1 << 14) {
continue;
}
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
[1, -1],
MODE_GLOBALMV,
gmv,
EIGHTTAP,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)?;
let trial: SyntaxBlock = if is_skip_seg {
let (bw, bh) = (bw4 * 4, bh4 * 4);
let (row0, col0) = ((mi_r as usize) * 4, (mi_c as usize) * 4);
let (crow0, ccol0) = recon.chroma_origin(mi_r, mi_c);
let (cbw, cbh) = if recon.num_planes > 1 {
(bw >> recon.subsampling_x, bh >> recon.subsampling_y)
} else {
(0, 0)
};
let width = recon.width;
let cw = recon.chroma_w;
for i in 0..bh {
for j in 0..bw {
recon.y[(row0 + i) * width + (col0 + j)] =
ictx.scratch[0][(row0 + i) * width + (col0 + j)];
}
}
for i in 0..cbh {
for j in 0..cbw {
recon.u[(crow0 + i) * cw + (ccol0 + j)] =
ictx.scratch[1][(crow0 + i) * cw + (ccol0 + j)];
recon.v[(crow0 + i) * cw + (ccol0 + j)] =
ictx.scratch[2][(crow0 + i) * cw + (ccol0 + j)];
}
}
tu_bd_stamp(&mut recon.bd, 0, col0, row0, bw, bh);
if recon.num_planes > 1 {
tu_bd_stamp(&mut recon.bd, 1, ccol0, crow0, cbw, cbh);
tu_bd_stamp(&mut recon.bd, 2, ccol0, crow0, cbw, cbh);
}
let mut l = SyntaxBlock::skip_leaf(0, None);
l.segment_id = s as u8;
l.inter = Some(SyntaxInterBlock {
ref_frame: [1, -1],
y_mode: MODE_GLOBALMV,
mv: gmv,
ref_mv_idx: 0,
interp_filter: [EIGHTTAP; 2],
skip_mode: 0,
compound_type: crate::inter_pred::COMPOUND_AVERAGE,
wedge_index: 0,
wedge_sign: 0,
mask_type: 0,
interintra_mode: None,
wedge_interintra: 0,
interintra_wedge_index: 0,
motion_mode: MOTION_MODE_SIMPLE,
});
l
} else {
let seg_qp = ictx.seg_qp(&recon.qp, s as u8);
let mut max_depth = 0u32;
let mut t = MAX_TX_SIZE_RECT[b_size];
while max_depth < MAX_VARTX_DEPTH && t != TX_4X4 {
t = SPLIT_TX_SIZE[t];
max_depth += 1;
}
let before_trial = save_region_wh(recon, mi_r, mi_c, bw4, bh4);
let mut inner: Option<(SyntaxBlock, RegionSnapshot, u64)> = None;
for depth in 0..=max_depth {
let cand = encode_inter_leaf_residual(
mi_r,
mi_c,
b_size,
input,
recon,
ictx,
[1, -1],
MODE_GLOBALMV,
gmv,
0,
EIGHTTAP,
CompoundSel::AVERAGE,
None,
MOTION_MODE_SIMPLE,
depth,
s as u8,
&seg_qp,
pricing,
)?;
let d = region_distortion_wh(recon, input, mi_r, mi_c, bw4, bh4);
let sc = score256(d, lambda, rate_of(&cand)?);
if inner.as_ref().map(|(_, _, s0)| sc < *s0).unwrap_or(true) {
inner = Some((cand, save_region_wh(recon, mi_r, mi_c, bw4, bh4), sc));
}
restore_region(recon, mi_r, mi_c, &before_trial);
}
let (cand, after_cand, _) = inner.expect("at least depth 0");
restore_region(recon, mi_r, mi_c, &after_cand);
cand
};
let d = region_distortion_wh(recon, input, mi_r, mi_c, bw4, bh4);
let sc = score256(d, lambda, rate_of(&trial)?);
if sc < score_best {
score_best = sc;
best = trial;
after_best = save_region_wh(recon, mi_r, mi_c, bw4, bh4);
} else {
restore_region(recon, mi_r, mi_c, &after_best);
}
}
restore_region(recon, mi_r, mi_c, &after_best);
repredict_leaf_arm(mi_r, mi_c, b_size, recon, ictx, &best)?;
Ok(best)
}
#[allow(clippy::too_many_arguments)]
fn encode_inter_leaf_modes(
mi_r: u32,
mi_c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
) -> Result<SyntaxBlock, Error> {
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size];
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size];
let (bw, bh) = (bw4 * 4, bh4 * 4);
let row0 = (mi_r as usize) * 4;
let col0 = (mi_c as usize) * 4;
let width = recon.width;
struct ModeCand {
ref_frame: [i8; 2],
y_mode: u8,
mv: [[i32; 2]; 2],
ref_mv_idx: u32,
rate: u64,
bits256: Option<u64>,
}
let diff_bits = |mv: [i32; 2], pred: [i32; 2]| -> u64 {
let b = |d: i32| u64::from(34 - d.unsigned_abs().leading_zeros());
b(mv[0] - pred[0]) + b(mv[1] - pred[1])
};
let nbr_cell = (mi_r * ictx.mi_cols + mi_c) as usize;
let (nbr_avail_u, nbr_avail_l) = (ictx.avail_us[nbr_cell] != 0, ictx.avail_ls[nbr_cell] != 0);
let nbr_ref = |r: u32, c: u32| -> [i32; 2] {
let cell = ((r * ictx.mi_cols + c) as usize) * 2;
let rfg = ictx.mirror.ref_frames();
[i32::from(rfg[cell]), i32::from(rfg[cell + 1])]
};
let above_ref_frame = if nbr_avail_u {
nbr_ref(mi_r - 1, mi_c)
} else {
[0, -1]
};
let left_ref_frame = if nbr_avail_l {
nbr_ref(mi_r, mi_c - 1)
} else {
[0, -1]
};
let (p_reference_select, p_fimv, p_hp) = (
ictx.ip.reference_select,
ictx.ip.force_integer_mv,
ictx.ip.allow_high_precision_mv,
);
let price_mode = |cand_ref: [i8; 2],
y_mode: u8,
mv: [[i32; 2]; 2],
ref_mv_idx: u32,
stack: &crate::cdf::FindMvStackResult|
-> Result<Option<u64>, Error> {
match pricing {
Some((twin, _)) => Ok(Some(twin.price_inter_mode(
[i32::from(cand_ref[0]), i32::from(cand_ref[1])],
y_mode,
mv,
ref_mv_idx,
stack,
b_size,
p_reference_select,
nbr_avail_u,
nbr_avail_l,
above_ref_frame,
left_ref_frame,
p_fimv,
p_hp,
)?)),
None => Ok(None),
}
};
let mut cands: Vec<ModeCand> = Vec::with_capacity(18);
let mut searched_mv: [[i32; 2]; 8] = [[0, 0]; 8];
let single_refs = ictx.single_refs.clone();
let compound_pairs = if bw4.min(bh4) >= 2 && ictx.ip.reference_select {
ictx.compound_pairs.clone()
} else {
Vec::new()
};
for (ref_ord, &rf) in single_refs.iter().enumerate() {
let ref_bias = ref_ord as u64;
let stack = ictx.find_stack(mi_r, mi_c, b_size, [rf, -1])?;
let ref_plane = ictx.plane_of_ref(rf);
let mv_int = motion_search_luma(
input,
ictx.search_luma(ref_plane),
width,
recon.height,
row0,
col0,
bw,
bh,
);
let mv_new = refine_mv_subpel(
input, ictx, mi_r, mi_c, b_size, rf, row0, col0, bw, bh, width, mv_int,
)?;
searched_mv[rf as usize] = mv_new;
let nfound = stack.num_mv_found;
{
let window: u32 = match nfound {
0 | 1 => 0,
2 => 1,
_ => 2,
};
let mut best: Option<(u64, u32, u64, Option<u64>)> = None;
for idx in 0..=window {
let pred = assign_mv_pred_mv(&stack, MODE_NEWMV, 0, idx)?;
let rate = 5 + ref_bias + u64::from(idx) + diff_bits(mv_new, pred);
let exact = price_mode([rf, -1], MODE_NEWMV, [mv_new, [0, 0]], idx, &stack)?;
let key = exact.unwrap_or(rate);
if best.map_or(true, |(k, _, _, _)| key < k) {
best = Some((key, idx, rate, exact));
}
}
let (_, idx, rate, bits256) = best.expect("slot 0 is always reachable");
cands.push(ModeCand {
ref_frame: [rf, -1],
y_mode: MODE_NEWMV,
mv: [mv_new, [0, 0]],
ref_mv_idx: idx,
rate,
bits256,
});
}
let nearest_mv = [stack.ref_stack_mv[0][0], [0, 0]];
cands.push(ModeCand {
ref_frame: [rf, -1],
y_mode: MODE_NEARESTMV,
mv: nearest_mv,
ref_mv_idx: 0,
rate: 3 + ref_bias,
bits256: price_mode([rf, -1], MODE_NEARESTMV, nearest_mv, 0, &stack)?,
});
let near_top: u32 = match nfound {
0..=2 => 1,
3 => 2,
_ => 3,
};
for idx in 1..=near_top {
let near_mv = [stack.ref_stack_mv[idx as usize][0], [0, 0]];
cands.push(ModeCand {
ref_frame: [rf, -1],
y_mode: MODE_NEARMV,
mv: near_mv,
ref_mv_idx: idx,
rate: 4 + ref_bias + u64::from(idx - 1),
bits256: price_mode([rf, -1], MODE_NEARMV, near_mv, idx, &stack)?,
});
}
let global_mv = [stack.global_mvs[0], [0, 0]];
cands.push(ModeCand {
ref_frame: [rf, -1],
y_mode: MODE_GLOBALMV,
mv: global_mv,
ref_mv_idx: 0,
rate: 4 + ref_bias,
bits256: price_mode([rf, -1], MODE_GLOBALMV, global_mv, 0, &stack)?,
});
}
for &crf in &compound_pairs {
let stack = ictx.find_stack(mi_r, mi_c, b_size, crf)?;
let nfound = stack.num_mv_found;
let nn_mv = [stack.ref_stack_mv[0][0], stack.ref_stack_mv[0][1]];
cands.push(ModeCand {
ref_frame: crf,
y_mode: MODE_NEAREST_NEARESTMV,
mv: nn_mv,
ref_mv_idx: 0,
rate: 6,
bits256: price_mode(crf, MODE_NEAREST_NEARESTMV, nn_mv, 0, &stack)?,
});
let near_top: u32 = match nfound {
0..=2 => 1,
3 => 2,
_ => 3,
};
for idx in 1..=near_top {
let nn_mv = [
stack.ref_stack_mv[idx as usize][0],
stack.ref_stack_mv[idx as usize][1],
];
cands.push(ModeCand {
ref_frame: crf,
y_mode: MODE_NEAR_NEARMV,
mv: nn_mv,
ref_mv_idx: idx,
rate: 7 + u64::from(idx - 1),
bits256: price_mode(crf, MODE_NEAR_NEARMV, nn_mv, idx, &stack)?,
});
}
let gg_mv = [stack.global_mvs[0], stack.global_mvs[1]];
cands.push(ModeCand {
ref_frame: crf,
y_mode: MODE_GLOBAL_GLOBALMV,
mv: gg_mv,
ref_mv_idx: 0,
rate: 7,
bits256: price_mode(crf, MODE_GLOBAL_GLOBALMV, gg_mv, 0, &stack)?,
});
{
let window: u32 = match nfound {
0 | 1 => 0,
2 => 1,
_ => 2,
};
let pair_mv = [searched_mv[crf[0] as usize], searched_mv[crf[1] as usize]];
let mut best: Option<(u64, u32, u64, Option<u64>)> = None;
for idx in 0..=window {
let pred0 = assign_mv_pred_mv(&stack, MODE_NEW_NEWMV, 0, idx)?;
let pred1 = assign_mv_pred_mv(&stack, MODE_NEW_NEWMV, 1, idx)?;
let rate = 8
+ u64::from(idx)
+ diff_bits(pair_mv[0], pred0)
+ diff_bits(pair_mv[1], pred1);
let exact = price_mode(crf, MODE_NEW_NEWMV, pair_mv, idx, &stack)?;
let key = exact.unwrap_or(rate);
if best.map_or(true, |(k, _, _, _)| key < k) {
best = Some((key, idx, rate, exact));
}
}
let (_, idx, rate, bits256) = best.expect("slot 0 is always reachable");
cands.push(ModeCand {
ref_frame: crf,
y_mode: MODE_NEW_NEWMV,
mv: pair_mv,
ref_mv_idx: idx,
rate,
bits256,
});
}
}
cands.retain(|c| {
c.mv.iter()
.all(|m| m[0].unsigned_abs() < (1 << 14) && m[1].unsigned_abs() < (1 << 14))
});
let lambda = lambda_for(&recon.qp);
let mut best: Option<(u64, usize)> = None;
let mut best_compound: Option<(u64, usize)> = None;
let mut best_single: Option<(u64, usize)> = None;
for (ci, cand) in cands.iter().enumerate() {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
cand.ref_frame,
cand.y_mode,
cand.mv,
EIGHTTAP,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)?;
let mut ssd = 0u64;
for i in 0..bh {
for j in 0..bw {
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ictx.scratch[0][(row0 + i) * width + col0 + j]);
ssd += (d * d) as u64;
}
}
let score = match cand.bits256 {
Some(b) => score256(ssd, lambda, b),
None => ssd + lambda * cand.rate,
};
if best.map_or(true, |(s, _)| score < s) {
best = Some((score, ci));
}
if cand.ref_frame[1] > 0 && best_compound.map_or(true, |(s, _)| score < s) {
best_compound = Some((score, ci));
}
if cand.ref_frame[1] < 0 && best_single.map_or(true, |(s, _)| score < s) {
best_single = Some((score, ci));
}
}
let (mut best_score, mut best_ci) = best.ok_or(Error::PartitionWalkOutOfRange)?;
let mut comp = CompoundSel::AVERAGE;
if ictx.ip.enable_masked_compound {
if let Some((_, cci)) = best_compound {
let (c_ref, c_mode, c_mv, c_rate, c_bits256) = {
let c = &cands[cci];
(c.ref_frame, c.y_mode, c.mv, c.rate, c.bits256)
};
let mut trials: Vec<(CompoundSel, u64)> = Vec::new();
if crate::cdf::wedge_bits(b_size) > 0 {
for wi in 0..16u8 {
for ws in 0..=1u8 {
trials.push((
CompoundSel {
ctype: crate::inter_pred::COMPOUND_WEDGE,
wedge_index: wi,
wedge_sign: ws,
mask_type: 0,
},
6,
));
}
}
}
for mt in 0..=1u8 {
trials.push((
CompoundSel {
ctype: crate::inter_pred::COMPOUND_DIFFWTD,
wedge_index: 0,
wedge_sign: 0,
mask_type: mt,
},
3,
));
}
for (sel, extra) in trials {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
c_ref,
c_mode,
c_mv,
EIGHTTAP,
sel,
None,
MotionSel::Simple,
)?;
let mut ssd = 0u64;
for i in 0..bh {
for j in 0..bw {
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ictx.scratch[0][(row0 + i) * width + col0 + j]);
ssd += (d * d) as u64;
}
}
let score = match c_bits256 {
Some(b) => score256(ssd, lambda, b + extra * 256),
None => ssd + lambda * (c_rate + extra),
};
if score < best_score {
best_score = score;
best_ci = cci;
comp = sel;
}
}
}
}
if ictx.ip.enable_jnt_comp {
if let Some((_, cci)) = best_compound {
let (c_ref, c_mode, c_mv, c_rate, c_bits256) = {
let c = &cands[cci];
(c.ref_frame, c.y_mode, c.mv, c.rate, c.bits256)
};
let sel = CompoundSel {
ctype: crate::inter_pred::COMPOUND_DISTANCE,
wedge_index: 0,
wedge_sign: 0,
mask_type: 0,
};
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
c_ref,
c_mode,
c_mv,
EIGHTTAP,
sel,
None,
MotionSel::Simple,
)?;
let mut ssd = 0u64;
for i in 0..bh {
for j in 0..bw {
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ictx.scratch[0][(row0 + i) * width + col0 + j]);
ssd += (d * d) as u64;
}
}
let score = match c_bits256 {
Some(b) => score256(ssd, lambda, b),
None => ssd + lambda * c_rate,
};
if score < best_score {
best_score = score;
best_ci = cci;
comp = sel;
}
}
}
let mut ii_sel: Option<(u8, Option<u8>)> = None;
if ictx.ip.enable_interintra_compound
&& (crate::cdf::BLOCK_8X8..=crate::cdf::BLOCK_32X32).contains(&b_size)
{
if let Some((_, sci)) = best_single {
let (s_ref, s_mode, s_mv, s_rate, s_bits256) = {
let c = &cands[sci];
(c.ref_frame, c.y_mode, c.mv, c.rate, c.bits256)
};
let leaf_ssd = |ictx: &PSearchCtx| -> u64 {
let mut ssd = 0u64;
for i in 0..bh {
for j in 0..bw {
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ictx.scratch[0][(row0 + i) * width + col0 + j]);
ssd += (d * d) as u64;
}
}
ssd
};
let mut stage: Option<(u64, u8)> = None;
for m in 0..crate::cdf::INTERINTRA_MODES as u8 {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
s_ref,
s_mode,
s_mv,
EIGHTTAP,
CompoundSel::AVERAGE,
Some(InterIntraTrial {
mode: m,
wedge: None,
neigh: &*recon,
}),
MotionSel::Simple,
)?;
let ssd = leaf_ssd(ictx);
if stage.map_or(true, |(s, _)| ssd < s) {
stage = Some((ssd, m));
}
let score = match s_bits256 {
Some(b) => score256(ssd, lambda, b + 4 * 256),
None => ssd + lambda * (s_rate + 4),
};
if score < best_score {
best_score = score;
best_ci = sci;
comp = CompoundSel::AVERAGE;
ii_sel = Some((m, None));
}
}
if crate::cdf::wedge_bits(b_size) > 0 {
let (_, base_mode) = stage.expect("stage 1 ran");
for wi in 0..16u8 {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
s_ref,
s_mode,
s_mv,
EIGHTTAP,
CompoundSel::AVERAGE,
Some(InterIntraTrial {
mode: base_mode,
wedge: Some(wi),
neigh: &*recon,
}),
MotionSel::Simple,
)?;
let score = match s_bits256 {
Some(b) => score256(leaf_ssd(ictx), lambda, b + 8 * 256),
None => leaf_ssd(ictx) + lambda * (s_rate + 8),
};
if score < best_score {
best_score = score;
best_ci = sci;
comp = CompoundSel::AVERAGE;
ii_sel = Some((base_mode, Some(wi)));
}
}
}
}
}
let _ = best_score;
let ModeCand {
ref_frame,
y_mode,
mv,
ref_mv_idx,
..
} = cands[best_ci];
let leaf_luma_ssd = |ictx: &PSearchCtx| -> u64 {
let mut ssd = 0u64;
for i in 0..bh {
for j in 0..bw {
let d = i64::from(input.y[(row0 + i) * width + col0 + j])
- i64::from(ictx.scratch[0][(row0 + i) * width + col0 + j]);
ssd += (d * d) as u64;
}
}
ssd
};
let (mut filter, simple_ssd) = if ictx.ip.interpolation_filter != SWITCHABLE {
let f = ictx.ip.interpolation_filter;
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ref_frame,
y_mode,
mv,
f,
comp,
None,
MotionSel::Simple,
)?;
(f, leaf_luma_ssd(ictx))
} else if y_mode == MODE_GLOBALMV || y_mode == MODE_GLOBAL_GLOBALMV {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ref_frame,
y_mode,
mv,
EIGHTTAP,
comp,
None,
MotionSel::Simple,
)?;
(EIGHTTAP, leaf_luma_ssd(ictx))
} else {
use crate::inter_pred::{EIGHTTAP_SHARP, EIGHTTAP_SMOOTH};
let mut best_f = (u64::MAX, EIGHTTAP);
for f in [EIGHTTAP, EIGHTTAP_SMOOTH, EIGHTTAP_SHARP] {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ref_frame,
y_mode,
mv,
f,
comp,
None,
MotionSel::Simple,
)?;
let ssd = leaf_luma_ssd(ictx);
if ssd < best_f.0 {
best_f = (ssd, f);
}
}
(best_f.1, best_f.0)
};
let mut motion = MotionSel::Simple;
if ictx.ip.is_motion_mode_switchable
&& ii_sel.is_none()
&& ref_frame[1] == -1
&& bw.min(bh) >= 8
&& !(!ictx.ip.force_integer_mv
&& y_mode == MODE_GLOBALMV
&& ictx.ip.gm_type[ref_frame[0] as usize] > crate::cdf::GM_TYPE_TRANSLATION)
&& ictx.mirror.has_overlappable_candidates(mi_r, mi_c, b_size)
{
let num_samples = ictx.warp_num_samples(mi_r, mi_c, b_size, ref_frame[0], mv);
let (mm_switchable, mm_warp, mm_fimv, mm_gm_type, mm_scaled) = (
ictx.ip.is_motion_mode_switchable,
ictx.ip.allow_warped_motion,
ictx.ip.force_integer_mv,
ictx.ip.gm_type,
ictx.ip.is_scaled_per_ref,
);
let mm_bits = move |mm: u8| -> Result<u64, Error> {
match pricing {
Some((twin, _)) => twin.price_motion_mode(
mm,
b_size,
false,
[i32::from(ref_frame[0]), i32::from(ref_frame[1])],
y_mode,
num_samples,
mm_switchable,
mm_warp,
mm_fimv,
mm_gm_type,
mm_scaled,
true,
),
None => Ok(0),
}
};
let mut best_score = score256(simple_ssd, lambda, mm_bits(MOTION_MODE_SIMPLE)?);
{
use crate::inter_pred::{EIGHTTAP_SHARP, EIGHTTAP_SMOOTH};
let obmc_bits = mm_bits(MOTION_MODE_OBMC)?;
let fixed = [ictx.ip.interpolation_filter];
let obmc_filters: &[u8] = if ictx.ip.interpolation_filter != SWITCHABLE {
&fixed
} else if y_mode == MODE_GLOBALMV {
&[EIGHTTAP]
} else {
&[EIGHTTAP, EIGHTTAP_SMOOTH, EIGHTTAP_SHARP]
};
for &f in obmc_filters {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ref_frame,
y_mode,
mv,
f,
comp,
None,
MotionSel::Obmc,
)?;
let score = score256(leaf_luma_ssd(ictx), lambda, obmc_bits);
if score < best_score {
best_score = score;
motion = MotionSel::Obmc;
filter = f;
}
}
}
if ictx.ip.allow_warped_motion
&& !ictx.ip.force_integer_mv
&& !ictx.ip.is_scaled_per_ref[(ref_frame[0] - 1) as usize]
{
if let Some(params) = ictx.warp_fit(mi_r, mi_c, b_size, ref_frame[0], mv) {
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ref_frame,
y_mode,
mv,
EIGHTTAP,
comp,
None,
MotionSel::Warp(params),
)?;
let score = score256(
leaf_luma_ssd(ictx),
lambda,
mm_bits(MOTION_MODE_WARPED_CAUSAL)?,
);
if score < best_score {
motion = MotionSel::Warp(params);
filter = if ictx.ip.interpolation_filter != SWITCHABLE {
ictx.ip.interpolation_filter
} else {
EIGHTTAP
};
}
}
}
}
ictx.predict_leaf(
mi_r,
mi_c,
b_size,
ref_frame,
y_mode,
mv,
filter,
comp,
ii_sel.map(|(m, w)| InterIntraTrial {
mode: m,
wedge: w,
neigh: &*recon,
}),
motion,
)?;
let segment_id = ictx.segment_for_block(input, mi_r, mi_c, b_size);
let seg_qp = ictx.seg_qp(&recon.qp, segment_id);
if ictx.seg_lossless(segment_id) {
return encode_inter_leaf_residual(
mi_r,
mi_c,
b_size,
input,
recon,
ictx,
ref_frame,
y_mode,
mv,
ref_mv_idx,
filter,
comp,
ii_sel,
motion.ordinal(),
0,
segment_id,
&seg_qp,
pricing,
);
}
let max_tx = MAX_TX_SIZE_RECT[b_size];
let mut max_depth = 0u32;
let mut t = max_tx;
while max_depth < MAX_VARTX_DEPTH && t != TX_4X4 {
t = SPLIT_TX_SIZE[t];
max_depth += 1;
}
let (sr, sc, sw4, sh4) = {
let (mut sr, mut sc, mut sw4, mut sh4) = (mi_r, mi_c, bw4, bh4);
if bh4 == 1 {
sr = mi_r & !1;
sh4 = 2;
}
if bw4 == 1 {
sc = mi_c & !1;
sw4 = 2;
}
(sr, sc, sw4, sh4)
};
let before = save_region_wh(recon, sr, sc, sw4, sh4);
let mut best: Option<(SyntaxBlock, RegionSnapshot, u64)> = None;
for depth in 0..=max_depth {
let leaf = encode_inter_leaf_residual(
mi_r,
mi_c,
b_size,
input,
recon,
ictx,
ref_frame,
y_mode,
mv,
ref_mv_idx,
filter,
comp,
ii_sel,
motion.ordinal(),
depth,
segment_id,
&seg_qp,
pricing,
)?;
let d = region_distortion_wh(recon, input, sr, sc, sw4, sh4);
let rate256 = match pricing {
Some((twin, params)) => twin.price_block(&leaf, mi_r, mi_c, b_size, params)?,
None => (p_leaf_rate(&leaf) + 2 * u64::from(depth)) * 256,
};
let score = score256(d, lambda, rate256);
let improves = match best.as_ref() {
Some((_, _, s)) => score < *s,
None => true,
};
if improves {
best = Some((leaf, save_region_wh(recon, sr, sc, sw4, sh4), score));
}
restore_region(recon, sr, sc, &before);
}
if ictx.auto_lossless && ictx.synthetic_leaf(input, mi_r, mi_c, b_size) {
if let Some(ll_seg) = ictx.lossless_seg {
let ll_qp = ictx.seg_qp(&recon.qp, ll_seg);
let leaf = encode_inter_leaf_residual(
mi_r,
mi_c,
b_size,
input,
recon,
ictx,
ref_frame,
y_mode,
mv,
ref_mv_idx,
filter,
comp,
ii_sel,
motion.ordinal(),
0,
ll_seg,
&ll_qp,
pricing,
)?;
let d = region_distortion_wh(recon, input, sr, sc, sw4, sh4);
let rate256 = match pricing {
Some((twin, params)) => twin.price_block(&leaf, mi_r, mi_c, b_size, params)?,
None => p_leaf_rate(&leaf) * 256,
};
let score = score256(d, lambda, rate256);
let improves = match best.as_ref() {
Some((_, _, s)) => score < *s,
None => true,
};
if improves {
best = Some((leaf, save_region_wh(recon, sr, sc, sw4, sh4), score));
}
restore_region(recon, sr, sc, &before);
}
}
let (leaf, after, _) = best.expect("at least depth 0");
restore_region(recon, sr, sc, &after);
Ok(leaf)
}
fn transform_tree_tu_order(
x: usize,
y: usize,
tx: usize,
depth: u32,
out: &mut Vec<(usize, usize)>,
) {
if depth == 0 {
out.push((x, y));
return;
}
let sub = SPLIT_TX_SIZE[tx];
let (sw, sh) = (TX_WIDTH[sub], TX_HEIGHT[sub]);
for i in 0..TX_HEIGHT[tx] / sh {
for j in 0..TX_WIDTH[tx] / sw {
transform_tree_tu_order(x + j * sw, y + i * sh, sub, depth - 1, out);
}
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn residual_tx_search_luma_inter(
input_plane: &[u16],
recon_plane: &mut [u16],
pw: usize,
row0: usize,
col0: usize,
tx_sz: usize,
pred: &[u16],
qp: &QuantizerParams,
fork: Option<(&mut TuFork, &TuCtx<'_>)>,
) -> Result<(Vec<i32>, u8), Error> {
use crate::cdf::{dequantize_step1, is_tx_type_in_set, tx_size_sqr_index, TX_SIZE_SQR_UP};
use crate::encoder::forward_quantize::forward_quantize;
use crate::encoder::forward_transform_2d::forward_transform_2d;
use crate::encoder::key_frame::repack_compact;
use crate::transform::inverse_transform_2d;
let w = TX_WIDTH[tx_sz];
let h = TX_HEIGHT[tx_sz];
let max_val = (1i64 << qp.bit_depth) - 1;
let mut residual = vec![0i64; w * h];
for i in 0..h {
for j in 0..w {
residual[i * w + j] =
i64::from(input_plane[(row0 + i) * pw + (col0 + j)]) - i64::from(pred[i * w + j]);
}
}
let set = inter_tx_type_set(
tx_size_sqr_index(tx_sz) as u32,
TX_SIZE_SQR_UP[tx_sz] as u32,
false,
);
let lambda = lambda_for(qp);
let fork_xy = fork.as_ref().map(|(_, ctx)| {
(
(col0 as u32 - ctx.base_x) / 4,
(row0 as u32 - ctx.base_y) / 4,
)
});
let mut best: Option<(Vec<i32>, Vec<i64>, u8, u64)> = None;
for t in 0..crate::cdf::TX_TYPES {
let admissible = t == DCT_DCT || (set > 0 && is_tx_type_in_set(true, set, t));
if !admissible {
continue;
}
let coeffs = forward_transform_2d(&residual, tx_sz, t, false);
let qm = qp.seg_qm_level[0][0];
let quant = repack_compact(forward_quantize(&coeffs, tx_sz, 0, 0, t, qm, qp), w, h);
let all_zero = quant.iter().all(|&q| q == 0);
let dequant = dequantize_step1(&quant, tx_sz, 0, 0, t, qm, qp);
let res_back = inverse_transform_2d(&dequant, tx_sz, t, u32::from(qp.bit_depth), false);
let (flip_ud, flip_lr) = crate::encoder::key_frame::step3_flips(t);
let mut d = 0u64;
for i in 0..h {
let yy = if flip_ud { h - 1 - i } else { i };
for j in 0..w {
let xx = if flip_lr { w - 1 - j } else { j };
let rec = (i64::from(pred[yy * w + xx]) + res_back[i * w + j]).clamp(0, max_val);
let diff = i64::from(input_plane[(row0 + yy) * pw + (col0 + xx)]) - rec;
d += (diff * diff) as u64;
}
}
let label = if all_zero { DCT_DCT as u8 } else { t as u8 };
let score = match (&fork, fork_xy) {
(Some((tu_fork, ctx)), Some((fx, fy))) => {
let bits256 = tu_fork.price_luma_tu(ctx, tx_sz, fx, fy, &quant, label)?;
crate::encoder::rate_twin::score256(d, lambda, bits256)
}
_ => {
let mut rate = 0u64;
for &qv in &quant {
if qv != 0 {
rate += 3 + u64::from(32 - qv.unsigned_abs().leading_zeros());
}
}
d + lambda * rate
}
};
let improves = match best.as_ref() {
Some((_, _, _, s)) => score < *s,
None => true,
};
if improves {
best = Some((quant, res_back, label, score));
}
if t == DCT_DCT && all_zero {
break;
}
}
let (quant, res_back, label, _) = best.expect("DCT_DCT is always admissible");
if let (Some((tu_fork, ctx)), Some((fx, fy))) = (fork, fork_xy) {
tu_fork.commit_luma_tu(ctx, tx_sz, fx, fy, &quant, label)?;
}
let (flip_ud, flip_lr) = crate::encoder::key_frame::step3_flips(label as usize);
for i in 0..h {
let yy = if flip_ud { h - 1 - i } else { i };
for j in 0..w {
let xx = if flip_lr { w - 1 - j } else { j };
let p = i64::from(pred[yy * w + xx]) + res_back[i * w + j];
recon_plane[(row0 + yy) * pw + (col0 + xx)] = p.clamp(0, max_val) as u16;
}
}
Ok((quant, label))
}
pub(crate) fn uniform_var_tx_tree(tx: usize, depth: u32) -> VarTxSyntaxTree {
if depth == 0 {
VarTxSyntaxTree::Leaf
} else {
let sub = SPLIT_TX_SIZE[tx];
let count = (TX_HEIGHT[tx] / TX_HEIGHT[sub]) * (TX_WIDTH[tx] / TX_WIDTH[sub]);
VarTxSyntaxTree::Split(
(0..count)
.map(|_| uniform_var_tx_tree(sub, depth - 1))
.collect(),
)
}
}
#[allow(clippy::too_many_arguments)]
fn encode_inter_leaf_residual(
mi_r: u32,
mi_c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &PSearchCtx,
ref_frame: [i8; 2],
y_mode: u8,
mv: [[i32; 2]; 2],
ref_mv_idx: u32,
filter: u8,
comp: CompoundSel,
ii: Option<(u8, Option<u8>)>,
motion_mode: u8,
depth: u32,
segment_id: u8,
seg_qp: &QuantizerParams,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
) -> Result<SyntaxBlock, Error> {
let bw4 = NUM_4X4_BLOCKS_WIDE[b_size];
let bh4 = NUM_4X4_BLOCKS_HIGH[b_size];
let (bw, bh) = (bw4 * 4, bh4 * 4);
let row0 = (mi_r as usize) * 4;
let col0 = (mi_c as usize) * 4;
let width = recon.width;
let lossless = ictx.seg_lossless(segment_id);
let qp = *seg_qp;
let mut tu_fork = match (&pricing, lossless) {
(Some((twin, _)), false) => Some(twin.tu_fork()),
_ => None,
};
let tu_ctx = pricing.map(|(_, params)| TuCtx {
params,
mi_row: mi_r,
mi_col: mi_c,
mi_size: b_size,
base_x: col0 as u32,
base_y: row0 as u32,
is_inter: true,
segment_id,
y_mode: 0,
use_filter_intra: false,
filter_intra_mode: None,
});
let luma_tx = if lossless {
TX_4X4
} else {
let mut t = MAX_TX_SIZE_RECT[b_size];
for _ in 0..depth {
t = SPLIT_TX_SIZE[t];
}
t
};
let (ltw, lth) = (TX_WIDTH[luma_tx], TX_HEIGHT[luma_tx]);
let mut tu_order: Vec<(usize, usize)> = Vec::new();
if lossless {
let mut ty = 0usize;
while ty < bh {
let mut tx = 0usize;
while tx < bw {
tu_order.push((tx, ty));
tx += ltw;
}
ty += lth;
}
} else {
transform_tree_tu_order(0, 0, MAX_TX_SIZE_RECT[b_size], depth, &mut tu_order);
}
let mut residual_quant: Vec<Vec<i32>> = Vec::new();
let mut luma_tx_types: Vec<u8> = Vec::new();
let tu_cols = bw / ltw;
let mut tu_type_grid = vec![DCT_DCT as u8; tu_cols * (bh / lth)];
for &(tx, ty) in &tu_order {
let (tr, tc) = (row0 + ty, col0 + tx);
let mut pred = vec![0u16; ltw * lth];
for i in 0..lth {
for j in 0..ltw {
pred[i * ltw + j] = ictx.scratch[0][(tr + i) * width + (tc + j)];
}
}
let (q, tt) = if lossless {
(
residual_tx(
&input.y,
&mut recon.y,
width,
tr,
tc,
luma_tx,
&pred,
0,
lossless,
DCT_DCT,
&qp,
),
DCT_DCT as u8,
)
} else {
residual_tx_search_luma_inter(
&input.y,
&mut recon.y,
width,
tr,
tc,
luma_tx,
&pred,
&qp,
match (&mut tu_fork, &tu_ctx) {
(Some(f), Some(c)) => Some((f, c)),
_ => None,
},
)?
};
tu_bd_stamp(&mut recon.bd, 0, tc, tr, ltw, lth);
residual_quant.push(q);
luma_tx_types.push(tt);
tu_type_grid[(ty / lth) * tu_cols + tx / ltw] = tt;
}
let has_chroma = recon.num_planes > 1
&& (bh4 != 1 || recon.subsampling_y == 0 || (mi_r & 1) == 1)
&& (bw4 != 1 || recon.subsampling_x == 0 || (mi_c & 1) == 1);
let (crow0, ccol0) = recon.chroma_origin(mi_r, mi_c);
let (cbw, cbh) = match crate::cdf::get_plane_residual_size(
b_size,
1,
recon.subsampling_x,
recon.subsampling_y,
) {
Some(psz) => (NUM_4X4_BLOCKS_WIDE[psz] * 4, NUM_4X4_BLOCKS_HIGH[psz] * 4),
None => (0, 0),
};
let chroma_tx = if lossless {
TX_4X4
} else {
get_tx_size(1, luma_tx, b_size, recon.subsampling_x, recon.subsampling_y).unwrap_or(TX_4X4)
};
let (ctw, cth) = (TX_WIDTH[chroma_tx], TX_HEIGHT[chroma_tx]);
let chroma_set = inter_tx_type_set(
tx_size_sqr_index(chroma_tx) as u32,
TX_SIZE_SQR_UP[chroma_tx] as u32,
false,
);
let cw = recon.chroma_w;
let (lift_x_sh, lift_y_sh) = (
2 + u32::from(recon.subsampling_x),
2 + u32::from(recon.subsampling_y),
);
for plane in 1..usize::from(recon.num_planes) {
if !has_chroma {
break;
}
let mut ty = 0usize;
while ty < cbh {
let mut tx = 0usize;
while tx < cbw {
let (tr, tc) = (crow0 + ty, ccol0 + tx);
let chroma_tt = if lossless || TX_SIZE_SQR_UP[chroma_tx] > crate::cdf::TX_32X32 {
DCT_DCT
} else {
let lifted_x = ((tc >> 2) << lift_x_sh).saturating_sub(col0);
let lifted_y = ((tr >> 2) << lift_y_sh).saturating_sub(row0);
let luma_tt =
tu_type_grid[(lifted_y / lth) * tu_cols + lifted_x / ltw] as usize;
if crate::cdf::is_tx_type_in_set(true, chroma_set, luma_tt) {
luma_tt
} else {
DCT_DCT
}
};
let mut pred = vec![0u16; ctw * cth];
for i in 0..cth {
for j in 0..ctw {
pred[i * ctw + j] = ictx.scratch[plane][(tr + i) * cw + (tc + j)];
}
}
let plane_buf = if plane == 1 {
&mut recon.u
} else {
&mut recon.v
};
let q = residual_tx(
if plane == 1 { &input.u } else { &input.v },
plane_buf,
cw,
tr,
tc,
chroma_tx,
&pred,
plane as u8,
lossless,
chroma_tt,
&qp,
);
tu_bd_stamp(&mut recon.bd, plane, tc, tr, ctw, cth);
residual_quant.push(q);
tx += ctw;
}
ty += cth;
}
}
let all_zero = residual_quant.iter().all(|tu| tu.iter().all(|&q| q == 0));
let skip = u8::from(all_zero);
if all_zero {
residual_quant.clear();
luma_tx_types.clear();
}
if luma_tx_types.iter().all(|&t| t == DCT_DCT as u8) {
luma_tx_types.clear();
}
let mut block = SyntaxBlock::skip_leaf(0, None);
block.skip = skip;
block.segment_id = if skip == 1 && !ictx.seg_id_pre_skip {
ictx.segment_pred(mi_r, mi_c)
} else {
segment_id
};
block.residual_quant = residual_quant.into();
block.residual_tx_type = luma_tx_types;
if !lossless && skip == 0 && b_size > crate::cdf::BLOCK_4X4 {
block.var_tx_trees = vec![uniform_var_tx_tree(MAX_TX_SIZE_RECT[b_size], depth)];
}
block.inter = Some(SyntaxInterBlock {
ref_frame,
y_mode,
mv,
ref_mv_idx,
interp_filter: [filter; 2],
skip_mode: 0,
compound_type: match ii {
Some((_, Some(_))) => crate::inter_pred::COMPOUND_WEDGE,
Some(_) => crate::inter_pred::COMPOUND_INTRA,
None => comp.ctype,
},
wedge_index: comp.wedge_index,
wedge_sign: comp.wedge_sign,
mask_type: comp.mask_type,
interintra_mode: ii.map(|(m, _)| m),
wedge_interintra: u8::from(matches!(ii, Some((_, Some(_))))),
interintra_wedge_index: ii.and_then(|(_, w)| w).unwrap_or(0),
motion_mode,
});
Ok(block)
}
fn p_leaf_rate(block: &SyntaxBlock) -> u64 {
if block.inter.as_ref().is_some_and(|ib| ib.skip_mode != 0) {
return 2;
}
let mut rate = leaf_rate(block);
match &block.inter {
Some(ib) => {
if ib.ref_frame[0] != 1 {
rate += 1;
}
match ib.compound_type {
crate::inter_pred::COMPOUND_WEDGE => rate += 6,
crate::inter_pred::COMPOUND_DIFFWTD => rate += 3,
_ => {}
}
match ib.y_mode {
MODE_NEWMV => {
let bits = |v: i32| u64::from(34 - v.unsigned_abs().leading_zeros());
rate += 6 + bits(ib.mv[0][0]) + bits(ib.mv[0][1]);
}
MODE_NEARESTMV => rate += 3,
MODE_NEARMV => rate += 4 + u64::from(ib.ref_mv_idx.saturating_sub(1)),
MODE_NEAREST_NEARESTMV => rate += 6,
MODE_NEAR_NEARMV => rate += 7 + u64::from(ib.ref_mv_idx.saturating_sub(1)),
MODE_NEW_NEWMV => {
let bits = |v: i32| u64::from(34 - v.unsigned_abs().leading_zeros());
rate += 8
+ bits(ib.mv[0][0])
+ bits(ib.mv[0][1])
+ bits(ib.mv[1][0])
+ bits(ib.mv[1][1]);
}
_ => rate += 4 + 3 * u64::from(ib.ref_frame[1] > 0),
}
}
None => rate += 16,
}
rate
}
fn p_tree_rate(node: &SyntaxNode) -> u64 {
match node {
SyntaxNode::Leaf(b) => p_leaf_rate(b),
SyntaxNode::Split(children) => 4 + children.iter().map(|c| p_tree_rate(c)).sum::<u64>(),
SyntaxNode::Horz(_) | SyntaxNode::Vert(_) => {
4 + node
.asymmetric_blocks()
.iter()
.map(|b| p_leaf_rate(b))
.sum::<u64>()
}
_ => {
5 + node
.asymmetric_blocks()
.iter()
.map(|b| p_leaf_rate(b))
.sum::<u64>()
}
}
}
#[allow(clippy::too_many_arguments)]
fn encode_intra_candidate(
r: u32,
c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &PSearchCtx,
pricing: Option<(&RateTwin, &SyntaxFrameParams)>,
) -> Result<SyntaxBlock, Error> {
if !ictx.demands_exact(r, c, b_size) {
let mut leaf = encode_leaf_sq(r, c, b_size, input, recon, pricing)?;
if !ictx.seg_alt_q.is_empty() && leaf.skip == 1 && !ictx.seg_id_pre_skip {
leaf.segment_id = ictx.segment_pred(r, c);
if ictx.seg_lossless(leaf.segment_id) {
leaf.tx_size = None;
}
}
stamp_delta_q(&mut leaf, b_size, ictx, pricing);
return Ok(leaf);
}
let ll_seg = ictx
.lossless_seg
.expect("demand mask requires a lossless segment (validated at arm time)");
let saved_qp = recon.qp;
let saved_lossless = recon.lossless;
recon.qp = QuantizerParams::neutral(0, saved_qp.bit_depth);
recon.lossless = true;
let out = crate::encoder::key_frame::encode_leaf_sq_seg(
r,
c,
b_size,
input,
recon,
pricing,
Some(ll_seg),
);
recon.qp = saved_qp;
recon.lossless = saved_lossless;
let mut leaf = out?;
if leaf.skip == 1 && !ictx.seg_id_pre_skip {
leaf.segment_id = ictx.segment_pred(r, c);
}
stamp_delta_q(&mut leaf, b_size, ictx, pricing);
Ok(leaf)
}
#[allow(clippy::too_many_arguments)]
fn build_p_search_tree(
r: u32,
c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
twin: &mut RateTwin,
params: &SyntaxFrameParams,
model: RateModel,
) -> Result<(SyntaxNode, u64), Error> {
if r >= recon.mi_rows || c >= recon.mi_cols {
return Ok((SyntaxNode::dummy_oob(), 0));
}
if b_size == BLOCK_4X4 {
let lambda = lambda_for(&recon.qp);
let (sr, sc) = (r & !1, c & !1);
let before = save_region_wh(recon, sr, sc, 2, 2);
let m_before = ictx.mirror.snapshot_encoder_stamp_rect(r, c, 1);
let inter_leaf = encode_inter_leaf(
r,
c,
b_size,
input,
recon,
ictx,
(model == RateModel::Twin).then_some((&*twin, params)),
)?;
ictx.stamp_leaf(
&inter_leaf,
r,
c,
b_size,
ictx.seg_lossless(inter_leaf.segment_id),
);
let d_inter = region_distortion_wh(recon, input, sr, sc, 2, 2);
let h_rate_inter = p_leaf_rate(&inter_leaf);
let node_inter = SyntaxNode::Leaf(Box::new(inter_leaf));
let mut twin_inter = twin.clone();
let cost_inter = twin_inter.commit_subtree(&node_inter, r, c, b_size, params)?;
let score_inter = match model {
RateModel::Twin => score256(d_inter, lambda, cost_inter),
RateModel::Heuristic => score256(d_inter, lambda, h_rate_inter * 256),
};
let after_inter = save_region_wh(recon, sr, sc, 2, 2);
let m_after_inter = ictx.mirror.snapshot_encoder_stamp_rect(r, c, 1);
restore_region(recon, sr, sc, &before);
ictx.mirror.restore_encoder_stamp_rect(&m_before);
let intra_leaf = encode_intra_candidate(
r,
c,
b_size,
input,
recon,
ictx,
(model == RateModel::Twin).then_some((&*twin, params)),
)?;
ictx.stamp_leaf(
&intra_leaf,
r,
c,
b_size,
ictx.seg_lossless(intra_leaf.segment_id),
);
let d_intra = region_distortion_wh(recon, input, sr, sc, 2, 2);
let h_rate_intra = p_leaf_rate(&intra_leaf);
let node_intra = SyntaxNode::Leaf(Box::new(intra_leaf));
let mut twin_intra = twin.clone();
let cost_intra = twin_intra.commit_subtree(&node_intra, r, c, b_size, params)?;
let score_intra = match model {
RateModel::Twin => score256(d_intra, lambda, cost_intra),
RateModel::Heuristic => score256(d_intra, lambda, h_rate_intra * 256),
};
if score_inter <= score_intra {
restore_region(recon, sr, sc, &after_inter);
ictx.mirror.restore_encoder_stamp_rect(&m_after_inter);
let SyntaxNode::Leaf(inter_leaf) = &node_inter else {
unreachable!("candidate A is a leaf");
};
ictx.stamp_leaf(
inter_leaf,
r,
c,
b_size,
ictx.seg_lossless(inter_leaf.segment_id),
);
*twin = twin_inter;
return Ok((node_inter, cost_inter));
}
let SyntaxNode::Leaf(intra_leaf) = &node_intra else {
unreachable!("candidate B is a leaf");
};
ictx.stamp_intra_leaf_grids(r, c, b_size, intra_leaf.y_mode);
*twin = twin_intra;
return Ok((node_intra, cost_intra));
}
let n4 = NUM_4X4_BLOCKS_WIDE[b_size] as u32;
let half = n4 >> 1;
let sub = crate::cdf::partition_subsize(crate::cdf::PARTITION_SPLIT, b_size)
.ok_or(Error::PartitionWalkOutOfRange)?;
let fully_inside = r + n4 <= recon.mi_rows && c + n4 <= recon.mi_cols;
if !fully_inside {
let mut cost = twin.commit_partition_symbol(crate::cdf::PARTITION_SPLIT, r, c, b_size)?;
let (nw, c0) = build_p_search_tree(r, c, sub, input, recon, ictx, twin, params, model)?;
let (ne, c1) =
build_p_search_tree(r, c + half, sub, input, recon, ictx, twin, params, model)?;
let (sw, c2) =
build_p_search_tree(r + half, c, sub, input, recon, ictx, twin, params, model)?;
let (se, c3) = build_p_search_tree(
r + half,
c + half,
sub,
input,
recon,
ictx,
twin,
params,
model,
)?;
cost += c0 + c1 + c2 + c3;
let children = [Box::new(nw), Box::new(ne), Box::new(sw), Box::new(se)];
return Ok((SyntaxNode::Split(children), cost));
}
let lambda = lambda_for(&recon.qp);
let before = save_region(recon, r, c, n4 as usize);
let m_before = ictx.mirror.snapshot_encoder_stamp_rect(r, c, n4);
let g_before = ictx.snapshot_grids_rect(r, c, n4);
let score_of = |d: u64, cost256: u64, h_rate: u64| -> u64 {
match model {
RateModel::Twin => score256(d, lambda, cost256),
RateModel::Heuristic => score256(d, lambda, h_rate * 256),
}
};
let inter_leaf = encode_inter_leaf(
r,
c,
b_size,
input,
recon,
ictx,
(model == RateModel::Twin).then_some((&*twin, params)),
)?;
ictx.stamp_leaf(
&inter_leaf,
r,
c,
b_size,
ictx.seg_lossless(inter_leaf.segment_id),
);
let d_inter = region_distortion(recon, input, r, c, n4 as usize);
let h_rate_inter = p_leaf_rate(&inter_leaf);
let node_inter = SyntaxNode::Leaf(Box::new(inter_leaf));
let mut twin_inter = twin.clone();
let cost_inter = twin_inter.commit_subtree(&node_inter, r, c, b_size, params)?;
let score_inter = score_of(d_inter, cost_inter, h_rate_inter);
let after_inter = save_region(recon, r, c, n4 as usize);
let m_after_inter = ictx.mirror.snapshot_encoder_stamp_rect(r, c, n4);
let g_after_inter = ictx.snapshot_grids_rect(r, c, n4);
restore_region(recon, r, c, &before);
ictx.mirror.restore_encoder_stamp_rect(&m_before);
ictx.restore_grids_rect(&g_before);
let intra_leaf = encode_intra_candidate(
r,
c,
b_size,
input,
recon,
ictx,
(model == RateModel::Twin).then_some((&*twin, params)),
)?;
ictx.stamp_leaf(
&intra_leaf,
r,
c,
b_size,
ictx.seg_lossless(intra_leaf.segment_id),
);
ictx.stamp_intra_leaf_grids(r, c, b_size, intra_leaf.y_mode);
let d_intra = region_distortion(recon, input, r, c, n4 as usize);
let h_rate_intra = p_leaf_rate(&intra_leaf);
let node_intra = SyntaxNode::Leaf(Box::new(intra_leaf));
let mut twin_intra = twin.clone();
let cost_intra = twin_intra.commit_subtree(&node_intra, r, c, b_size, params)?;
let score_intra = score_of(d_intra, cost_intra, h_rate_intra);
let after_intra = save_region(recon, r, c, n4 as usize);
let m_after_intra = ictx.mirror.snapshot_encoder_stamp_rect(r, c, n4);
let g_after_intra = ictx.snapshot_grids_rect(r, c, n4);
restore_region(recon, r, c, &before);
ictx.mirror.restore_encoder_stamp_rect(&m_before);
ictx.restore_grids_rect(&g_before);
let mut best_node: (
SyntaxNode,
RegionSnapshot,
_,
DriverGridSnapshot,
u64,
RateTwin,
u64,
) = if score_inter <= score_intra {
(
node_inter,
after_inter,
m_after_inter,
g_after_inter,
score_inter,
twin_inter,
cost_inter,
)
} else {
(
node_intra,
after_intra,
m_after_intra,
g_after_intra,
score_intra,
twin_intra,
cost_intra,
)
};
let quarter = half >> 1;
let split_sz = sub;
for part in [
PARTITION_HORZ,
PARTITION_VERT,
crate::cdf::PARTITION_HORZ_A,
crate::cdf::PARTITION_HORZ_B,
crate::cdf::PARTITION_VERT_A,
crate::cdf::PARTITION_VERT_B,
crate::cdf::PARTITION_HORZ_4,
crate::cdf::PARTITION_VERT_4,
] {
let psub = match crate::cdf::partition_subsize(part, b_size) {
Some(sz) => sz,
None => continue,
};
if b_size == BLOCK_8X8 && !matches!(part, PARTITION_HORZ | PARTITION_VERT) {
continue;
}
if !crate::encoder::key_frame::chroma_partition_ok(recon, psub) {
continue;
}
let cells: Vec<(u32, u32, usize)> = match part {
PARTITION_HORZ => vec![(r, c, psub), (r + half, c, psub)],
PARTITION_VERT => vec![(r, c, psub), (r, c + half, psub)],
crate::cdf::PARTITION_HORZ_A => vec![
(r, c, split_sz),
(r, c + half, split_sz),
(r + half, c, psub),
],
crate::cdf::PARTITION_HORZ_B => vec![
(r, c, psub),
(r + half, c, split_sz),
(r + half, c + half, split_sz),
],
crate::cdf::PARTITION_VERT_A => vec![
(r, c, split_sz),
(r + half, c, split_sz),
(r, c + half, psub),
],
crate::cdf::PARTITION_VERT_B => vec![
(r, c, psub),
(r, c + half, split_sz),
(r + half, c + half, split_sz),
],
crate::cdf::PARTITION_HORZ_4 => (0..4u32).map(|i| (r + quarter * i, c, psub)).collect(),
_ => (0..4u32).map(|i| (r, c + quarter * i, psub)).collect(),
};
let mut twin_shape = twin.clone();
let mut cost_shape = twin_shape.commit_partition_symbol(part, r, c, b_size)?;
let mut blocks: Vec<SyntaxBlock> = Vec::with_capacity(cells.len());
let mut h_rate = if cells.len() == 2 { 4 } else { 5 };
for &(rr, cc, sz) in &cells {
let blk = encode_inter_leaf(
rr,
cc,
sz,
input,
recon,
ictx,
(model == RateModel::Twin).then_some((&twin_shape, params)),
)?;
ictx.stamp_leaf(&blk, rr, cc, sz, ictx.seg_lossless(blk.segment_id));
h_rate += p_leaf_rate(&blk);
cost_shape += twin_shape.commit_block(&blk, rr, cc, sz, params)?;
blocks.push(blk);
}
let d = region_distortion(recon, input, r, c, n4 as usize);
let bx = |b: SyntaxBlock| Box::new(b);
let node = match part {
PARTITION_HORZ | PARTITION_VERT => {
let mut it = blocks.into_iter();
let pair = [bx(it.next().unwrap()), bx(it.next().unwrap())];
if part == PARTITION_HORZ {
SyntaxNode::Horz(pair)
} else {
SyntaxNode::Vert(pair)
}
}
crate::cdf::PARTITION_HORZ_A
| crate::cdf::PARTITION_HORZ_B
| crate::cdf::PARTITION_VERT_A
| crate::cdf::PARTITION_VERT_B => {
let mut it = blocks.into_iter();
let trio = [
bx(it.next().unwrap()),
bx(it.next().unwrap()),
bx(it.next().unwrap()),
];
match part {
crate::cdf::PARTITION_HORZ_A => SyntaxNode::HorzA(trio),
crate::cdf::PARTITION_HORZ_B => SyntaxNode::HorzB(trio),
crate::cdf::PARTITION_VERT_A => SyntaxNode::VertA(trio),
_ => SyntaxNode::VertB(trio),
}
}
_ => {
let mut it = blocks.into_iter();
let four = [
bx(it.next().unwrap()),
bx(it.next().unwrap()),
bx(it.next().unwrap()),
bx(it.next().unwrap()),
];
if part == crate::cdf::PARTITION_HORZ_4 {
SyntaxNode::Horz4(four)
} else {
SyntaxNode::Vert4(four)
}
}
};
let score = score_of(d, cost_shape, h_rate);
if score < best_node.4 {
best_node = (
node,
save_region(recon, r, c, n4 as usize),
ictx.mirror.snapshot_encoder_stamp_rect(r, c, n4),
ictx.snapshot_grids_rect(r, c, n4),
score,
twin_shape,
cost_shape,
);
}
restore_region(recon, r, c, &before);
ictx.mirror.restore_encoder_stamp_rect(&m_before);
ictx.restore_grids_rect(&g_before);
}
let (
node_leafish,
after_leafish,
m_after_leafish,
g_after_leafish,
score_leafish,
twin_leafish,
cost_leafish,
) = best_node;
let mut twin_split = twin.clone();
let mut cost_split =
twin_split.commit_partition_symbol(crate::cdf::PARTITION_SPLIT, r, c, b_size)?;
let (nw, c0) = build_p_search_tree(
r,
c,
sub,
input,
recon,
ictx,
&mut twin_split,
params,
model,
)?;
let (ne, c1) = build_p_search_tree(
r,
c + half,
sub,
input,
recon,
ictx,
&mut twin_split,
params,
model,
)?;
let (sw, c2) = build_p_search_tree(
r + half,
c,
sub,
input,
recon,
ictx,
&mut twin_split,
params,
model,
)?;
let (se, c3) = build_p_search_tree(
r + half,
c + half,
sub,
input,
recon,
ictx,
&mut twin_split,
params,
model,
)?;
cost_split += c0 + c1 + c2 + c3;
let children = [Box::new(nw), Box::new(ne), Box::new(sw), Box::new(se)];
let d_split = region_distortion(recon, input, r, c, n4 as usize);
let h_rate_split = children.iter().map(|ch| p_tree_rate(ch)).sum::<u64>() + 4;
let score_split = score_of(d_split, cost_split, h_rate_split);
if score_leafish <= score_split {
restore_region(recon, r, c, &after_leafish);
ictx.mirror.restore_encoder_stamp_rect(&m_after_leafish);
ictx.restore_grids_rect(&g_after_leafish);
*twin = twin_leafish;
Ok((node_leafish, cost_leafish))
} else {
*twin = twin_split;
Ok((SyntaxNode::Split(children), cost_split))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::encoder::yuv_frame::ChromaFormat;
fn plane8(p: &[u16]) -> Vec<u8> {
p.iter().map(|&s| s as u8).collect()
}
fn gop_seg_yuv(frames: &[YuvFrame], q: u8, alt_q: &[i16]) -> Result<EncodedGopYuv, Error> {
encode_gop_yuv_seg_extras_tuned(frames, q, alt_q, &[], false, None, GopTuning::default())
.map(|t| t.gop)
}
struct TestTileState {
cdfs: TileCdfContext,
state: PartitionSyntaxWriter,
writer: SymbolWriter,
params: SyntaxFrameParams,
}
fn test_tile_state(recon: &ReconState, ictx: &PSearchCtx) -> TestTileState {
let writer = SymbolWriter::new(false);
let mut cdfs = TileCdfContext::new_from_defaults();
cdfs.init_coeff_cdfs(recon.qp.base_q_idx);
let state = PartitionSyntaxWriter::new(
recon.mi_rows,
recon.mi_cols,
TileGeometry {
mi_row_start: 0,
mi_row_end: recon.mi_rows,
mi_col_start: 0,
mi_col_end: recon.mi_cols,
},
)
.expect("test geometry is valid");
let params = SyntaxFrameParams {
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
seg_id_pre_skip: false,
segmentation_enabled: !ictx.seg_alt_q.is_empty(),
seg_ref_frame: [None; crate::uncompressed_header_tail::MAX_SEGMENTS],
seg_skip: [false; crate::uncompressed_header_tail::MAX_SEGMENTS],
seg_globalmv: [false; crate::uncompressed_header_tail::MAX_SEGMENTS],
last_active_seg_id: ictx.seg_alt_q.len().saturating_sub(1) as u8,
lossless_array: seg_lossless_array(recon.qp.base_q_idx, &ictx.seg_alt_q),
coded_lossless: recon.lossless,
enable_cdef: false,
allow_intrabc: false,
cdef_bits: 0,
read_deltas: false,
use_128x128_superblock: false,
delta_q_res: 0,
delta_lf_present: false,
delta_lf_multi: false,
mono_chrome: false,
delta_lf_res: 0,
allow_screen_content_tools: recon.allow_screen_content_tools,
enable_filter_intra: true,
bit_depth: 8,
tx_mode_select: !recon.lossless,
quant: recon.qp,
reduced_tx_set: false,
inter: Some(ictx.ip.clone()),
};
TestTileState {
cdfs,
state,
writer,
params,
}
}
fn search_and_commit_sb(
ts: &mut TestTileState,
r: u32,
c: u32,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
) -> Result<SyntaxNode, Error> {
ts.state.arm_read_deltas();
let mut twin = RateTwin::snapshot(&ts.cdfs, &ts.state, &ts.writer);
twin.arm_read_deltas();
let (node, _) = build_p_search_tree(
r,
c,
BLOCK_64X64,
input,
recon,
ictx,
&mut twin,
&ts.params,
RateModel::Twin,
)?;
write_partition_tree_syntax(
&mut ts.writer,
&mut ts.cdfs,
&mut ts.state,
&node,
r,
c,
BLOCK_64X64,
&ts.params,
)?;
assert!(
twin.matches(&ts.cdfs, &ts.writer),
"rate twin desynced from the test tile writer after superblock ({r},{c})"
);
Ok(node)
}
fn build_p_search_tree_for_tests(
r: u32,
c: u32,
b_size: usize,
input: &YuvFrame,
recon: &mut ReconState,
ictx: &mut PSearchCtx,
) -> Result<SyntaxNode, Error> {
assert_eq!(b_size, BLOCK_64X64, "tests search whole superblocks");
let mut ts = test_tile_state(recon, ictx);
search_and_commit_sb(&mut ts, r, c, input, recon, ictx)
}
fn moving_gradient(w: u32, h: u32, shift_y: usize, shift_x: usize, seed: u32) -> YuvFrame {
let (wu, hu) = (w as usize, h as usize);
let s = seed as usize;
let mut f = YuvFrame::filled(w, h, 8, ChromaFormat::Yuv420, 0);
for i in 0..hu {
for j in 0..wu {
let (si, sj) = (i + shift_y, j + shift_x);
f.y[i * wu + j] = ((si * 5 + sj * 3 + (si / 16) * (sj / 16) + s) % 256) as u16;
}
}
let (cw, ch) = (wu / 2, hu / 2);
for i in 0..ch {
for j in 0..cw {
let (si, sj) = (i + shift_y / 2, j + shift_x / 2);
f.u[i * cw + j] = ((128 + si * 2 + sj + s) % 256) as u16;
f.v[i * cw + j] = ((64 + si + sj * 2 + s) % 256) as u16;
}
}
f
}
fn count_modes(node: &SyntaxNode, counts: &mut [u32; 4]) {
match node {
SyntaxNode::Leaf(b) => {
if let Some(ib) = b.inter.as_ref() {
match ib.y_mode {
MODE_NEARESTMV => counts[0] += 1,
MODE_NEARMV => counts[1] += 1,
MODE_GLOBALMV => counts[2] += 1,
MODE_NEWMV => counts[3] += 1,
_ => {}
}
}
}
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_modes(ch, counts);
}
}
other => {
let blocks = other.asymmetric_blocks();
for b in blocks.iter() {
if let Some(ib) = b.inter.as_ref() {
match ib.y_mode {
MODE_NEARESTMV => counts[0] += 1,
MODE_NEARMV => counts[1] += 1,
MODE_GLOBALMV => counts[2] += 1,
MODE_NEWMV => counts[3] += 1,
_ => {}
}
}
}
}
}
}
#[test]
fn r412_p_search_selects_stack_predicted_modes_on_uniform_motion() {
let f0 = moving_gradient(64, 64, 0, 0, 7);
let f1 = moving_gradient(64, 64, 3, 5, 7);
let base_q_idx = 60u8;
let key = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f0, base_q_idx).unwrap();
let reference = GopFrameReconYuv {
y: key.recon_y,
u: key.recon_u,
v: key.recon_v,
};
let fh = build_p_frame_yuv420_8bit_fh_with_q(&key.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let qp = QuantizerParams::neutral(base_q_idx, 8);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp,
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
let mut ictx = PSearchCtx::new(
&reference,
&reference,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f1, &mut recon, &mut ictx).unwrap();
let mut counts = [0u32; 4];
count_modes(&tree, &mut counts);
let total: u32 = counts.iter().sum();
assert!(total > 0, "uniform motion must produce inter leaves");
assert!(
counts[0] + counts[1] > 0,
"at least one NEARESTMV / NEARMV leaf must be selected on \
uniformly translating content (got NEAREST={} NEAR={} \
GLOBAL={} NEW={})",
counts[0],
counts[1],
counts[2],
counts[3]
);
assert!(
counts[3] > 0,
"the first block of the shared-motion region still codes NEWMV \
(its stack is empty); got NEW={}",
counts[3]
);
}
#[test]
fn r412_p_search_selects_sharp_filter_on_kernel_matched_content() {
use crate::inter_pred::EIGHTTAP_SHARP;
let f0 = moving_gradient(64, 64, 0, 0, 77);
let base_q_idx = 90u8;
let key = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f0, base_q_idx).unwrap();
let reference = GopFrameReconYuv {
y: key.recon_y,
u: key.recon_u,
v: key.recon_v,
};
let mut f1 = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
{
let ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
let mut probe = PSearchCtx::new(
&reference,
&reference,
16,
16,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
probe
.predict_leaf(
0,
0,
BLOCK_64X64,
[1, -1],
MODE_NEWMV,
[[0, 4], [0, 0]],
EIGHTTAP_SHARP,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)
.unwrap();
for (dst, &src) in f1.y.iter_mut().zip(probe.scratch[0].iter()) {
*dst = src;
}
for (dst, &src) in f1.u.iter_mut().zip(probe.scratch[1].iter()) {
*dst = src;
}
for (dst, &src) in f1.v.iter_mut().zip(probe.scratch[2].iter()) {
*dst = src;
}
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&key.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
let mut ictx = PSearchCtx::new(
&reference,
&reference,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f1, &mut recon, &mut ictx).unwrap();
fn count_filters(node: &SyntaxNode, counts: &mut [u32; 3]) {
match node {
SyntaxNode::Leaf(b) => {
if let Some(ib) = b.inter.as_ref() {
counts[usize::from(ib.interp_filter[0].min(2))] += 1;
}
}
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_filters(ch, counts);
}
}
other => {
let blocks = other.asymmetric_blocks();
for b in blocks.iter() {
if let Some(ib) = b.inter.as_ref() {
counts[usize::from(ib.interp_filter[0].min(2))] += 1;
}
}
}
}
}
let mut counts = [0u32; 3];
count_filters(&tree, &mut counts);
let sharp = counts[usize::from(EIGHTTAP_SHARP)];
assert!(
sharp > 0,
"kernel-matched half-pel content must commit EIGHTTAP_SHARP \
leaves (got EIGHTTAP={} SMOOTH={} SHARP={sharp})",
counts[0],
counts[1]
);
}
#[test]
fn r412_p_search_selects_horz_partition_on_band_motion() {
let band = |k: usize| -> YuvFrame {
let (w, h) = (64usize, 64usize);
let mut f = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let tex = |x: i64, y: i64| -> u16 {
let (xu, yu) = ((x + 512) as usize, (y + 512) as usize);
((xu * 7 + yu * 3 + (xu / 8) * (yu / 8) * 5) % 256) as u16
};
for i in 0..h {
let s: i64 = if i < 16 {
4 * k as i64
} else {
-4 * (k as i64)
};
for j in 0..w {
f.y[i * w + j] = tex(j as i64 - s, i as i64);
}
}
let (cw, ch) = (w / 2, h / 2);
for i in 0..ch {
let s: i64 = if i < 8 { 2 * k as i64 } else { -2 * (k as i64) };
for j in 0..cw {
let x = ((j as i64 - s) + 512) as usize;
f.u[i * cw + j] = ((128 + x * 2 + i) % 256) as u16;
f.v[i * cw + j] = ((64 + x + i * 2) % 256) as u16;
}
}
f
};
let frames: Vec<YuvFrame> = (0..3).map(band).collect();
let base_q_idx = 60u8;
let key =
crate::encoder::key_frame::encode_key_frame_yuv_with_q(&frames[0], base_q_idx).unwrap();
let reference = GopFrameReconYuv {
y: key.recon_y.clone(),
u: key.recon_u.clone(),
v: key.recon_v.clone(),
};
let fh = build_p_frame_yuv420_8bit_fh_with_q(&key.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
let mut ictx = PSearchCtx::new(
&reference,
&reference,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &frames[1], &mut recon, &mut ictx)
.unwrap();
fn count_rect(node: &SyntaxNode, horz: &mut u32, vert: &mut u32) {
match node {
SyntaxNode::Leaf(_) => {}
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_rect(ch, horz, vert);
}
}
SyntaxNode::Horz(_) => *horz += 1,
SyntaxNode::Vert(_) => *vert += 1,
_ => {}
}
}
let (mut horz, mut vert) = (0u32, 0u32);
count_rect(&tree, &mut horz, &mut vert);
assert!(
horz > 0,
"band-split motion must select at least one PARTITION_HORZ node \
(got HORZ={horz} VERT={vert})"
);
let enc = encode_gop_yuv_with_q(&frames, base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), frames.len());
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(f.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r412_p_search_selects_golden_reference_across_flash() {
let calm = moving_gradient(64, 64, 0, 0, 5);
let mut flash = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let mut state = 0x1234_5678u32;
let mut next = move || {
state ^= state << 13;
state ^= state >> 17;
state ^= state << 5;
state
};
for v in flash.y.iter_mut() {
*v = (next() & 0xFF) as u16;
}
let frames = vec![calm.clone(), flash, calm];
let base_q_idx = 90u8;
let enc = encode_gop_yuv_with_q(&frames, base_q_idx).unwrap();
let prev = enc.recon[1].clone();
let prevprev = enc.recon[0].clone();
let fh = build_p_frame_yuv420_8bit_fh_with_q(&enc.seq, 64, 64, base_q_idx, 2, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
let mut ictx = PSearchCtx::new(
&prev,
&prevprev,
mi_rows,
mi_cols,
64,
64,
ip,
2,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &frames[2], &mut recon, &mut ictx)
.unwrap();
fn count_refs(node: &SyntaxNode, golden: &mut u32, last: &mut u32) {
let leafs = |b: &SyntaxBlock, golden: &mut u32, last: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.ref_frame[0] == 4 {
*golden += 1;
} else {
*last += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafs(b, golden, last),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_refs(ch, golden, last);
}
}
other => {
let blocks = other.asymmetric_blocks();
for b in blocks.iter() {
leafs(b, golden, last);
}
}
}
}
let (mut golden, mut last) = (0u32, 0u32);
count_refs(&tree, &mut golden, &mut last);
assert!(
golden > 0,
"frame after a flash must select GOLDEN_FRAME leaves \
(got GOLDEN={golden} LAST={last})"
);
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), frames.len());
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(f.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r412_p_search_selects_compound_average_on_blended_content() {
let f0 = moving_gradient(64, 64, 0, 0, 40);
let f1 = moving_gradient(64, 64, 0, 0, 140);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let mut f2 = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let avg = |a: &[u16], b: &[u16], out: &mut [u16]| {
for ((o, &x), &y) in out.iter_mut().zip(a).zip(b) {
*o = (u32::from(x) + u32::from(y)).div_ceil(2) as u16;
}
};
avg(&pre.recon[0].y, &pre.recon[1].y, &mut f2.y);
avg(&pre.recon[0].u, &pre.recon[1].u, &mut f2.u);
avg(&pre.recon[0].v, &pre.recon[1].v, &mut f2.v);
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 2, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
let mut ictx = PSearchCtx::new(
&pre.recon[1],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
2,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f2, &mut recon, &mut ictx).unwrap();
fn count_compound(node: &SyntaxNode, compound: &mut u32, single: &mut u32) {
let leafc = |b: &SyntaxBlock, compound: &mut u32, single: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.ref_frame[1] > 0 {
*compound += 1;
} else {
*single += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, compound, single),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_compound(ch, compound, single);
}
}
other => {
let blocks = other.asymmetric_blocks();
for b in blocks.iter() {
leafc(b, compound, single);
}
}
}
}
let (mut compound, mut single) = (0u32, 0u32);
count_compound(&tree, &mut compound, &mut single);
assert!(
compound > 0,
"average-blend content must select COMPOUND_AVERAGE leaves \
(got COMPOUND={compound} SINGLE={single})"
);
let enc = encode_gop_yuv_with_q(&[f0, f2], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(f.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r413_p_search_selects_skip_mode_on_static_content() {
let f = moving_gradient(64, 64, 2, 3, 21);
let base_q_idx = 60u8;
let frames = vec![f.clone(), f.clone(), f.clone(), f.clone()];
let pre = encode_gop_yuv_with_q(&frames[..3], base_q_idx).unwrap();
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 3, &[]);
assert_eq!(fh.skip_mode_present, Some(true), "P3 header presence");
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.order_hints = gop_order_hints(3, pre.seq.order_hint_bits);
ip.skip_mode_present = true;
ip.skip_mode_frame = [1, 4];
let mut ictx = PSearchCtx::new(
&pre.recon[2],
&pre.recon[1],
mi_rows,
mi_cols,
64,
64,
ip,
3,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f, &mut recon, &mut ictx).unwrap();
fn count_skip_mode(node: &SyntaxNode, sm: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, sm: &mut u32, other: &mut u32| {
if b.inter.as_ref().is_some_and(|ib| ib.skip_mode != 0) {
*sm += 1;
} else {
*other += 1;
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, sm, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_skip_mode(ch, sm, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, sm, other);
}
}
}
}
let (mut sm, mut other) = (0u32, 0u32);
count_skip_mode(&tree, &mut sm, &mut other);
assert!(
sm > 0,
"static content on a skip-mode frame must select skip_mode leaves \
(got SKIP_MODE={sm} OTHER={other})"
);
let enc = encode_gop_yuv_with_q(&frames, base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 4);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r413_segmentation_alt_q_gop_round_trips() {
let mut frames: Vec<YuvFrame> = Vec::new();
for k in 0..3u32 {
let mut f = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 100);
for i in 0..64usize {
for j in 32..64usize {
f.y[i * 64 + j] = ((i * 17 + j * 31 + (k as usize) * 5) % 256) as u16;
}
}
for v in f.u.iter_mut() {
*v = 90;
}
for v in f.v.iter_mut() {
*v = 160;
}
frames.push(f);
}
let base_q_idx = 60u8;
let alt_q: [i16; 3] = [0, 24, 48];
let enc = gop_seg_yuv(&frames, base_q_idx, &alt_q).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 3);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&enc.seq, 64, 64, base_q_idx, 1, &alt_q);
let sp = fh.segmentation_params.as_ref().unwrap();
assert!(sp.enabled && sp.update_map && !sp.temporal_update && sp.update_data);
assert_eq!(sp.last_active_seg_id, 2);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: {
let mut q = QuantizerParams::neutral(base_q_idx, 8);
q.segmentation_enabled = true;
for (seg, &d) in alt_q.iter().enumerate() {
q.seg_alt_q_active[seg] = true;
q.seg_alt_q_data[seg] = d;
}
q
},
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.order_hints = gop_order_hints(1, enc.seq.order_hint_bits);
ip.segmentation_update_map = true;
let mut ictx = PSearchCtx::new(
&enc.recon[0],
&enc.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&alt_q,
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &frames[1], &mut recon, &mut ictx)
.unwrap();
fn collect_segments(node: &SyntaxNode, seen: &mut [u32; 8]) {
match node {
SyntaxNode::Leaf(b) => seen[b.segment_id as usize] += 1,
SyntaxNode::Split(children) => {
for ch in children.iter() {
collect_segments(ch, seen);
}
}
other => {
let blocks = other.asymmetric_blocks();
for b in blocks.iter() {
seen[b.segment_id as usize] += 1;
}
}
}
}
let mut seen = [0u32; 8];
collect_segments(&tree, &mut seen);
let distinct = seen.iter().filter(|&&n| n > 0).count();
assert!(
distinct >= 2,
"mixed content must code multiple segments (histogram: {seen:?})"
);
assert!(gop_seg_yuv(&frames, 60, &[1, 2]).is_err());
assert!(gop_seg_yuv(&frames, 60, &[0]).is_err());
assert!(gop_seg_yuv(&frames, 60, &[0, -256]).is_err());
assert!(gop_seg_yuv(&frames, 200, &[0, 100]).is_err());
assert!(gop_seg_yuv(&frames, 0, &[0, 10]).is_err());
}
#[test]
fn r413_p_search_selects_t_shape_partitions() {
let gen = |k: usize| -> YuvFrame {
let mut f = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let tex = |i: usize, j: usize| ((i * 7 + j * 13 + (i / 8) * (j / 8)) % 256) as u16;
for i in 0..64usize {
for j in 0..64usize {
let (si, sj) = if i < 32 {
(i, j + 2 * k)
} else if j < 32 {
(i + 2 * k, j)
} else {
(i + 2 * k, j + 2 * k)
};
f.y[i * 64 + j] = tex(si, sj);
}
}
for v in f.u.iter_mut() {
*v = 110;
}
for v in f.v.iter_mut() {
*v = 150;
}
f
};
let frames: Vec<YuvFrame> = (0..3).map(gen).collect();
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&frames[..1], base_q_idx).unwrap();
let key_recon = GopFrameReconYuv {
y: pre.recon[0].y.clone(),
u: pre.recon[0].u.clone(),
v: pre.recon[0].v.clone(),
};
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.order_hints = gop_order_hints(1, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&key_recon,
&key_recon,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &frames[1], &mut recon, &mut ictx)
.unwrap();
fn count_ext(node: &SyntaxNode, ext: &mut u32) {
match node {
SyntaxNode::HorzA(_)
| SyntaxNode::HorzB(_)
| SyntaxNode::VertA(_)
| SyntaxNode::VertB(_)
| SyntaxNode::Horz4(_)
| SyntaxNode::Vert4(_) => *ext += 1,
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_ext(ch, ext);
}
}
_ => {}
}
}
let mut ext = 0u32;
count_ext(&tree, &mut ext);
assert!(
ext > 0,
"tri-motion content must select an EXT-alphabet partition (tree: {tree:?})"
);
let enc = encode_gop_yuv_with_q(&frames, base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 3);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r413_motion_field_estimation_projects_from_p2() {
let frames: Vec<YuvFrame> = (0..4)
.map(|k| moving_gradient(64, 64, k * 3, k * 2, 9))
.collect();
let enc = encode_gop_yuv_with_q(&frames, 60).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 4);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
let fh1 = build_p_frame_yuv420_8bit_fh_with_q(&enc.seq, 64, 64, 60, 1, &[]);
let fs = fh1.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let key_recon = GopFrameReconYuv {
y: enc.recon[0].y.clone(),
u: enc.recon[0].u.clone(),
v: enc.recon[0].v.clone(),
};
let mf_store: [SavedMotionField; 8] =
core::array::from_fn(|_| SavedMotionField::intra(mi_rows, mi_cols));
let (_, _, saved1, _, _) = encode_p_frame_yuv(
&frames[1],
&key_recon,
&key_recon,
&enc.seq,
60,
1,
&[],
&mf_store,
RateModel::Twin,
true,
None,
true,
None,
false,
None,
true,
true,
true,
true,
true,
true,
(0, 0),
1,
false,
None,
false,
None,
0,
&key_recon,
false,
false,
false,
None,
None,
)
.unwrap();
assert!(!saved1.frame_is_intra);
assert!(
saved1.mf_ref_frames.iter().any(|&r| r > 0),
"P1 must store ยง7.19 motion-field candidates on moving content"
);
use crate::inter_pred::{motion_field_estimation_core, MotionFieldSlot};
let mut slots: [Option<MotionFieldSlot<'_>>; 8] = [None; 8];
let intra = SavedMotionField::intra(mi_rows, mi_cols);
for (i, s) in slots.iter_mut().enumerate() {
let src = if i == 0 { &saved1 } else { &intra };
*s = Some(MotionFieldSlot {
mf_mvs: &src.mf_mvs,
mf_ref_frames: &src.mf_ref_frames,
saved_order_hints: src.saved_order_hints,
mi_rows: src.mi_rows,
mi_cols: src.mi_cols,
frame_is_intra: src.frame_is_intra,
});
}
let mut ref_frame_idx = [0u8; 7];
ref_frame_idx[3] = 1;
let hints2 = gop_order_hints(2, enc.seq.order_hint_bits);
let mut order_hints = [0i32; 8];
order_hints.copy_from_slice(&hints2.order_hints_by_ref);
let mfmvs = motion_field_estimation_core(
&slots,
&ref_frame_idx,
&order_hints,
2,
u32::from(enc.seq.order_hint_bits),
mi_rows,
mi_cols,
);
let mut valid = 0u32;
for r in 1..=7usize {
for y8 in 0..mfmvs.h8() {
for x8 in 0..mfmvs.w8() {
if mfmvs.get(r, y8, x8)[0] != crate::cdf::MFMV_INVALID {
valid += 1;
}
}
}
}
assert!(
valid > 0,
"P2 ยง7.9 projection must land valid temporal candidates"
);
}
#[test]
fn r412_predict_leaf_threads_interp_filter() {
use crate::inter_pred::{EIGHTTAP_SHARP, EIGHTTAP_SMOOTH};
let mut refr = GopFrameReconYuv {
y: vec![0u16; 64 * 64],
u: vec![128u16; 32 * 32],
v: vec![128u16; 32 * 32],
};
for i in 0..64 {
for j in 0..64 {
refr.y[i * 64 + j] = if j == 10 { 255 } else { 0 };
}
}
let ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
let mut ictx = PSearchCtx::new(&refr, &refr, 16, 16, 64, 64, ip, 1, 0, &[]).unwrap();
let mut outs = Vec::new();
for f in [EIGHTTAP, EIGHTTAP_SMOOTH, EIGHTTAP_SHARP] {
ictx.predict_leaf(
2,
2,
BLOCK_8X8,
[1, -1],
MODE_NEWMV,
[[0, 4], [0, 0]],
f,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)
.unwrap();
let mut v = Vec::new();
for i in 8..16 {
for j in 8..16 {
v.push(ictx.scratch[0][i * 64 + j]);
}
}
outs.push(v);
}
assert_ne!(
outs[0], outs[1],
"EIGHTTAP vs SMOOTH must differ at half-pel"
);
assert_ne!(
outs[0], outs[2],
"EIGHTTAP vs SHARP must differ at half-pel"
);
}
#[test]
fn r417_predict_leaf_inter_intra_blends_intra_half() {
let refr = GopFrameReconYuv {
y: vec![100u16; 64 * 64],
u: vec![128u16; 32 * 32],
v: vec![128u16; 32 * 32],
};
let mut recon = fresh_recon(64, 64, 60);
recon.y.fill(200);
recon.u.fill(60);
recon.v.fill(60);
let ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
let mut ictx = PSearchCtx::new(&refr, &refr, 16, 16, 64, 64, ip, 1, 60, &[]).unwrap();
ictx.predict_leaf(
2,
2,
BLOCK_8X8,
[1, -1],
MODE_NEWMV,
[[0, 0], [0, 0]],
EIGHTTAP,
CompoundSel::AVERAGE,
Some(InterIntraTrial {
mode: crate::cdf::II_V_PRED,
wedge: None,
neigh: &recon,
}),
MotionSel::Simple,
)
.unwrap();
let mut smooth_luma = Vec::new();
for i in 8..16usize {
for j in 8..16usize {
let v = ictx.scratch[0][i * 64 + j];
assert!(
(100..=200).contains(&v),
"blend out of the [inter, intra] envelope at ({i}, {j}): {v}"
);
smooth_luma.push(v);
}
}
assert!(
smooth_luma.iter().any(|&v| v > 100),
"intra half must contribute to the luma blend"
);
assert_ne!(
&smooth_luma[0..8],
&smooth_luma[56..64],
"the ยง7.11.3.14 smooth II mask varies down the block"
);
let cu = ictx.scratch[1][4 * 32 + 4];
assert!(
(60..128).contains(&cu),
"chroma blend must pull toward the intra half: {cu}"
);
let cell = 2 * 16 + 2;
assert_eq!(
ictx.ref_frames[cell * 2 + 1],
0,
"RefFrame[1] = INTRA_FRAME"
);
assert_eq!(ictx.interintra_modes[cell], crate::cdf::II_V_PRED);
assert_eq!(ictx.wedge_interintras[cell], 0);
assert_eq!(
ictx.compound_types[cell],
crate::inter_pred::COMPOUND_INTRA,
"ยง5.11.29 bit-silent derivation on the non-wedge arm"
);
ictx.predict_leaf(
2,
2,
BLOCK_8X8,
[1, -1],
MODE_NEWMV,
[[0, 0], [0, 0]],
EIGHTTAP,
CompoundSel::AVERAGE,
Some(InterIntraTrial {
mode: crate::cdf::II_V_PRED,
wedge: Some(7),
neigh: &recon,
}),
MotionSel::Simple,
)
.unwrap();
let mut wedge_luma = Vec::new();
for i in 8..16usize {
for j in 8..16usize {
wedge_luma.push(ictx.scratch[0][i * 64 + j]);
}
}
assert_ne!(smooth_luma, wedge_luma, "wedge mask differs from smooth");
assert_eq!(ictx.wedge_interintras[cell], 1);
assert_eq!(ictx.interintra_wedge_indices[cell], 7);
assert_eq!(ictx.compound_types[cell], crate::inter_pred::COMPOUND_WEDGE);
ictx.predict_leaf(
2,
2,
BLOCK_8X8,
[1, -1],
MODE_NEWMV,
[[0, 0], [0, 0]],
EIGHTTAP,
CompoundSel::AVERAGE,
None,
MotionSel::Simple,
)
.unwrap();
assert_eq!(ictx.ref_frames[cell * 2 + 1], -1);
assert_eq!(ictx.wedge_interintras[cell], 0);
assert_eq!(ictx.interintra_modes[cell], 0);
for i in 8..16usize {
for j in 8..16usize {
assert_eq!(
ictx.scratch[0][i * 64 + j],
100,
"plain zero-MV inter pred is the flat reference"
);
}
}
}
#[test]
fn r412_gop_with_stack_modes_round_trips_through_spec_driver() {
let frames: Vec<YuvFrame> = (0..4)
.map(|k| moving_gradient(96, 80, 2 * k, 4 * k, 31))
.collect();
let enc = encode_gop_yuv_with_q(&frames, 80).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), frames.len());
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(f.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
fn b_frame_ctx<'a>(
anchor: &'a GopFrameReconYuv,
future: &'a GopFrameReconYuv,
mi_rows: u32,
mi_cols: u32,
base_q_idx: u8,
skip_mode: bool,
) -> PSearchCtx {
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.order_hints = crate::cdf::FrameInterOrderHints {
order_hint_bits: 7,
current_order_hint: 1,
order_hints_by_ref: [0, 0, 0, 0, 0, 2, 2, 2],
};
for r in 5..=7usize {
ip.ref_frame_sign_bias[r] = 1;
}
ip.skip_mode_present = skip_mode;
ip.skip_mode_frame = [1, 5];
let mut slot_to_plane = [0usize; 8];
slot_to_plane[1] = 1;
let mut ref_frame_idx = [0u8; 7];
ref_frame_idx[4] = 1;
ref_frame_idx[5] = 1;
ref_frame_idx[6] = 1;
PSearchCtx::with_refs(
&[anchor, future],
slot_to_plane,
ref_frame_idx,
vec![1, 5],
vec![[1, 5]],
mi_rows,
mi_cols,
(mi_cols * 4) as usize,
(mi_rows * 4) as usize,
&[(mi_cols * 4, mi_rows * 4); 2],
ip,
base_q_idx,
&[],
8,
1,
1,
3,
)
.unwrap()
}
fn fresh_recon(w: usize, h: usize, q: u8) -> ReconState {
ReconState {
y: vec![0u16; w * h],
u: vec![0u16; (w / 2) * (h / 2)],
v: vec![0u16; (w / 2) * (h / 2)],
width: w,
height: h,
chroma_w: w / 2,
chroma_h: h / 2,
mi_rows: (h / 4) as u32,
mi_cols: (w / 4) as u32,
lossless: q == 0,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(q, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
}
}
#[test]
fn r415_b_search_selects_bidir_compound_on_blended_content() {
let f0 = moving_gradient(64, 64, 0, 0, 40);
let f2 = moving_gradient(64, 64, 0, 0, 140);
let base_q_idx = 60u8;
let k0 = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f0, base_q_idx).unwrap();
let k2 = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f2, base_q_idx).unwrap();
let anchor = GopFrameReconYuv {
y: k0.recon_y,
u: k0.recon_u,
v: k0.recon_v,
};
let future = GopFrameReconYuv {
y: k2.recon_y,
u: k2.recon_u,
v: k2.recon_v,
};
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let avg = |a: &[u16], b: &[u16], out: &mut [u16]| {
for ((o, &x), &y) in out.iter_mut().zip(a).zip(b) {
*o = (u32::from(x) + u32::from(y)).div_ceil(2) as u16;
}
};
avg(&anchor.y, &future.y, &mut target.y);
avg(&anchor.u, &future.u, &mut target.u);
avg(&anchor.v, &future.v, &mut target.v);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ictx = b_frame_ctx(&anchor, &future, 16, 16, base_q_idx, false);
recon.bd.clear_for_sb(0, 0, 16, 16);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
fn count(node: &SyntaxNode, compound: &mut u32, single: &mut u32) {
let leafc = |b: &SyntaxBlock, compound: &mut u32, single: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.ref_frame == [1, 5] {
*compound += 1;
} else {
*single += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, compound, single),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, compound, single);
}
}
other => {
for b in other.asymmetric_blocks().iter() {
leafc(b, compound, single);
}
}
}
}
let (mut compound, mut single) = (0u32, 0u32);
count(&tree, &mut compound, &mut single);
assert!(
compound > 0,
"average-blend content must select {{ LAST, BWDREF }} bidirectional \
compound leaves (got COMPOUND={compound} SINGLE={single})"
);
}
#[test]
fn r415_b_search_selects_backward_reference_on_future_matching_content() {
let f0 = moving_gradient(64, 64, 0, 0, 5);
let f2 = moving_gradient(64, 64, 9, 13, 505);
let base_q_idx = 60u8;
let k0 = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f0, base_q_idx).unwrap();
let k2 = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f2, base_q_idx).unwrap();
let anchor = GopFrameReconYuv {
y: k0.recon_y,
u: k0.recon_u,
v: k0.recon_v,
};
let future = GopFrameReconYuv {
y: k2.recon_y,
u: k2.recon_u,
v: k2.recon_v,
};
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ictx = b_frame_ctx(&anchor, &future, 16, 16, base_q_idx, false);
recon.bd.clear_for_sb(0, 0, 16, 16);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f2, &mut recon, &mut ictx).unwrap();
fn count(node: &SyntaxNode, bwd: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, bwd: &mut u32, other: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.ref_frame == [5, -1] {
*bwd += 1;
} else {
*other += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, bwd, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, bwd, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, bwd, other);
}
}
}
}
let (mut bwd, mut other) = (0u32, 0u32);
count(&tree, &mut bwd, &mut other);
assert!(
bwd > 0,
"future-matching content must select BWDREF single-reference \
leaves (got BWD={bwd} OTHER={other})"
);
}
#[test]
fn r415_search_selects_wedge_compound_on_wedge_blend_content() {
use crate::cdf::BLOCK_32X32;
use crate::inter_pred::COMPOUND_WEDGE;
let f0 = moving_gradient(64, 64, 0, 0, 40);
let f1 = moving_gradient(64, 64, 0, 0, 140);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
{
let ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
let mut probe =
PSearchCtx::new(&pre.recon[1], &pre.recon[0], 16, 16, 64, 64, ip, 2, 0, &[])
.unwrap();
let sel = CompoundSel {
ctype: COMPOUND_WEDGE,
wedge_index: 7,
wedge_sign: 0,
mask_type: 0,
};
for (r, c) in [(0u32, 0u32), (0, 8), (8, 0), (8, 8)] {
probe
.predict_leaf(
r,
c,
BLOCK_32X32,
[1, 4],
MODE_NEW_NEWMV,
[[0, 0], [0, 0]],
EIGHTTAP,
sel,
None,
MotionSel::Simple,
)
.unwrap();
}
for (dst, &src) in target.y.iter_mut().zip(probe.scratch[0].iter()) {
*dst = src;
}
for (dst, &src) in target.u.iter_mut().zip(probe.scratch[1].iter()) {
*dst = src;
}
for (dst, &src) in target.v.iter_mut().zip(probe.scratch[2].iter()) {
*dst = src;
}
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 2, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_masked_compound = true;
ip.order_hints = gop_order_hints(2, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&pre.recon[1],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
2,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
fn count(node: &SyntaxNode, wedge: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, wedge: &mut u32, other: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.compound_type == crate::inter_pred::COMPOUND_WEDGE {
*wedge += 1;
} else {
*other += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, wedge, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, wedge, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, wedge, other);
}
}
}
}
let (mut wedge, mut other) = (0u32, 0u32);
count(&tree, &mut wedge, &mut other);
assert!(
wedge > 0,
"wedge-blend content must commit COMPOUND_WEDGE leaves \
(got WEDGE={wedge} OTHER={other})"
);
let enc = encode_gop_yuv_with_q(&[f0, f1, target], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 3);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_WEDGE_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
format!("{dir}/self-gop-64x64-q60-wedge.ivf"),
&enc.ivf_bytes,
)
.unwrap();
}
}
#[test]
fn r418_p_search_selects_palette_on_pasted_screen_content() {
const COLORS: [u16; 4] = [16, 80, 160, 240];
let f0 = moving_gradient(64, 64, 0, 0, 40);
let mut f1 = moving_gradient(64, 64, 0, 0, 40);
let mut state = 0x0DDB_A115u32;
let mut next = move || {
state ^= state << 13;
state ^= state >> 17;
state ^= state << 5;
state
};
for i in 16..48usize {
for j in 16..48usize {
f1.y[i * 64 + j] = COLORS[(next() & 3) as usize];
}
}
for i in 8..24usize {
for j in 8..24usize {
let par = ((i / 2) + (j / 2)) & 1;
f1.u[i * 32 + j] = if par == 0 { 96 } else { 168 };
f1.v[i * 32 + j] = if ((i / 4) & 1) == par { 108 } else { 152 };
}
}
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.order_hints = gop_order_hints(1, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&pre.recon[0],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f1, &mut recon, &mut ictx).unwrap();
fn count(node: &SyntaxNode, pal: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, pal: &mut u32, other: &mut u32| {
if b.inter.is_none() && b.palette.size_y > 0 {
*pal += 1;
} else {
*other += 1;
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, pal, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, pal, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, pal, other);
}
}
}
}
let (mut pal, mut other) = (0u32, 0u32);
count(&tree, &mut pal, &mut other);
assert!(
pal > 0,
"pasted dither content must commit PaletteSizeY > 0 intra leaves in the \
P-frame search (got PALETTE={pal} OTHER={other})"
);
let decoded = crate::decoder::decode_av1_spec(&pre.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&pre.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&pre.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&pre.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r417_search_selects_wedge_inter_intra_on_wedge_blend_content() {
use crate::cdf::{BLOCK_32X32, II_V_PRED};
let f0 = moving_gradient(64, 64, 0, 0, 40);
let f1 = moving_gradient(64, 64, 0, 0, 220);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
{
let ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
let mut probe =
PSearchCtx::new(&pre.recon[1], &pre.recon[0], 16, 16, 64, 64, ip, 2, 0, &[])
.unwrap();
let mut running = fresh_recon(64, 64, base_q_idx);
for (r, c) in [(0u32, 0u32), (0, 8), (8, 0), (8, 8)] {
probe
.predict_leaf(
r,
c,
BLOCK_32X32,
[1, -1],
MODE_NEARESTMV,
[[0, 0], [0, 0]],
EIGHTTAP,
CompoundSel::AVERAGE,
Some(InterIntraTrial {
mode: II_V_PRED,
wedge: Some(7),
neigh: &running,
}),
MotionSel::Simple,
)
.unwrap();
let (row0, col0) = ((r as usize) * 4, (c as usize) * 4);
for i in 0..32usize {
for j in 0..32usize {
let v = probe.scratch[0][(row0 + i) * 64 + col0 + j];
target.y[(row0 + i) * 64 + col0 + j] = v;
running.y[(row0 + i) * 64 + col0 + j] = v;
}
}
let (cr0, cc0) = (row0 / 2, col0 / 2);
for i in 0..16usize {
for j in 0..16usize {
let u = probe.scratch[1][(cr0 + i) * 32 + cc0 + j];
let v = probe.scratch[2][(cr0 + i) * 32 + cc0 + j];
target.u[(cr0 + i) * 32 + cc0 + j] = u;
target.v[(cr0 + i) * 32 + cc0 + j] = v;
running.u[(cr0 + i) * 32 + cc0 + j] = u;
running.v[(cr0 + i) * 32 + cc0 + j] = v;
}
}
}
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 2, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_interintra_compound = true;
ip.order_hints = gop_order_hints(2, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&pre.recon[1],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
2,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
fn count(node: &SyntaxNode, wedge_ii: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, wedge_ii: &mut u32, other: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.interintra_mode.is_some() && ib.wedge_interintra == 1 {
*wedge_ii += 1;
} else {
*other += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, wedge_ii, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, wedge_ii, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, wedge_ii, other);
}
}
}
}
let (mut wedge_ii, mut other) = (0u32, 0u32);
count(&tree, &mut wedge_ii, &mut other);
assert!(
wedge_ii > 0,
"wedge inter-intra blend content must commit wedge_interintra = 1 leaves \
(got WEDGE_II={wedge_ii} OTHER={other})"
);
let enc = encode_gop_yuv_with_q(&[f0, f1, target], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 3);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
}
#[test]
fn r417_search_selects_sub8_intra_in_mixed_groups() {
let f0 = moving_gradient(64, 64, 0, 0, 71);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(std::slice::from_ref(&f0), base_q_idx).unwrap();
let kr = &pre.recon[0];
let mut f2 = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let at = |p: &[u16], w: usize, h: usize, i: i64, j: i64| -> u16 {
let ii = i.clamp(0, h as i64 - 1) as usize;
let jj = j.clamp(0, w as i64 - 1) as usize;
p[ii * w + jj]
};
for i in 0..64i64 {
for j in 0..64i64 {
let par = ((i / 4) + (j / 4)) & 1;
let (dy, dx) = if par == 0 { (0, 0) } else { (4, 4) };
f2.y[(i * 64 + j) as usize] = at(&kr.y, 64, 64, i + dy, j + dx);
}
}
f2.u.copy_from_slice(&kr.u);
f2.v.copy_from_slice(&kr.v);
for (r0, c0) in [(8usize, 8usize), (8, 40), (40, 8), (40, 40)] {
for i in 0..4 {
for j in 0..4 {
f2.y[(r0 + i) * 64 + c0 + j] = f2.y[(r0 - 1) * 64 + c0 + j];
}
}
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_interintra_compound = true;
ip.order_hints = gop_order_hints(1, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&pre.recon[0],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f2, &mut recon, &mut ictx).unwrap();
fn walk(
node: &SyntaxNode,
r: u32,
c: u32,
b_size: usize,
found: &mut Vec<(u32, u32, bool)>,
) {
match node {
SyntaxNode::Leaf(b) if b_size == BLOCK_4X4 => {
found.push((r, c, b.inter.is_some()));
}
SyntaxNode::Split(children) => {
let sub =
crate::cdf::partition_subsize(crate::cdf::PARTITION_SPLIT, b_size).unwrap();
let half = (NUM_4X4_BLOCKS_WIDE[b_size] as u32) >> 1;
walk(&children[0], r, c, sub, found);
walk(&children[1], r, c + half, sub, found);
walk(&children[2], r + half, c, sub, found);
walk(&children[3], r + half, c + half, sub, found);
}
_ => {}
}
}
let mut sub4: Vec<(u32, u32, bool)> = Vec::new();
walk(&tree, 0, 0, BLOCK_64X64, &mut sub4);
let n_intra4 = sub4.iter().filter(|&&(_, _, inter)| !inter).count();
assert!(
n_intra4 > 0,
"V-replicated cells must commit BLOCK_4X4 intra leaves (got {sub4:?})"
);
let mixed = sub4.iter().any(|&(r, c, inter)| {
!inter
&& sub4
.iter()
.any(|&(r2, c2, inter2)| inter2 && (r2 & !1, c2 & !1) == (r & !1, c & !1))
});
assert!(mixed, "some 2x2 group must mix intra and inter ({sub4:?})");
let enc = encode_gop_yuv_with_q(&[f0, f2], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_SUB8INTRA_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
format!("{dir}/self-gop-64x64-q60-sub8-intra.ivf"),
&enc.ivf_bytes,
)
.unwrap();
}
}
#[test]
fn r417_search_selects_inter_intra_on_blended_content() {
use crate::cdf::{BLOCK_32X32, II_V_PRED};
let f0 = moving_gradient(64, 64, 0, 0, 40);
let f1 = moving_gradient(64, 64, 0, 0, 140);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
{
let ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
let mut probe =
PSearchCtx::new(&pre.recon[1], &pre.recon[0], 16, 16, 64, 64, ip, 2, 0, &[])
.unwrap();
let mut running = fresh_recon(64, 64, base_q_idx);
for (r, c) in [(0u32, 0u32), (0, 8), (8, 0), (8, 8)] {
probe
.predict_leaf(
r,
c,
BLOCK_32X32,
[1, -1],
MODE_NEARESTMV,
[[0, 0], [0, 0]],
EIGHTTAP,
CompoundSel::AVERAGE,
Some(InterIntraTrial {
mode: II_V_PRED,
wedge: None,
neigh: &running,
}),
MotionSel::Simple,
)
.unwrap();
let (row0, col0) = ((r as usize) * 4, (c as usize) * 4);
for i in 0..32usize {
for j in 0..32usize {
let v = probe.scratch[0][(row0 + i) * 64 + col0 + j];
target.y[(row0 + i) * 64 + col0 + j] = v;
running.y[(row0 + i) * 64 + col0 + j] = v;
}
}
let (cr0, cc0) = (row0 / 2, col0 / 2);
for i in 0..16usize {
for j in 0..16usize {
let u = probe.scratch[1][(cr0 + i) * 32 + cc0 + j];
let v = probe.scratch[2][(cr0 + i) * 32 + cc0 + j];
target.u[(cr0 + i) * 32 + cc0 + j] = u;
target.v[(cr0 + i) * 32 + cc0 + j] = v;
running.u[(cr0 + i) * 32 + cc0 + j] = u;
running.v[(cr0 + i) * 32 + cc0 + j] = v;
}
}
}
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 2, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_interintra_compound = true;
ip.order_hints = gop_order_hints(2, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&pre.recon[1],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
2,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
fn count(node: &SyntaxNode, ii: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, ii: &mut u32, other: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.interintra_mode.is_some() {
*ii += 1;
} else {
*other += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, ii, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, ii, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, ii, other);
}
}
}
}
let (mut ii, mut other) = (0u32, 0u32);
count(&tree, &mut ii, &mut other);
assert!(
ii > 0,
"inter-intra blend content must commit ยง5.11.28 leaves \
(got II={ii} OTHER={other})"
);
let enc = encode_gop_yuv_with_q(&[f0, f1, target], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 3);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_INTERINTRA_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
format!("{dir}/self-gop-64x64-q60-interintra.ivf"),
&enc.ivf_bytes,
)
.unwrap();
}
}
#[test]
fn r416_search_selects_distance_compound_on_weighted_blend_content() {
use crate::cdf::BLOCK_32X32;
use crate::inter_pred::COMPOUND_DISTANCE;
let f0 = moving_gradient(64, 64, 0, 0, 40);
let f1 = moving_gradient(64, 64, 0, 0, 140);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
{
let mut ip = SyntaxInterFrameParams::single_ref_baseline(16, 16, false);
ip.order_hints = gop_order_hints(2, pre.seq.order_hint_bits);
let mut probe =
PSearchCtx::new(&pre.recon[1], &pre.recon[0], 16, 16, 64, 64, ip, 2, 0, &[])
.unwrap();
let sel = CompoundSel {
ctype: COMPOUND_DISTANCE,
wedge_index: 0,
wedge_sign: 0,
mask_type: 0,
};
for (r, c) in [(0u32, 0u32), (0, 8), (8, 0), (8, 8)] {
probe
.predict_leaf(
r,
c,
BLOCK_32X32,
[1, 4],
MODE_NEW_NEWMV,
[[0, 0], [0, 0]],
EIGHTTAP,
sel,
None,
MotionSel::Simple,
)
.unwrap();
}
for (dst, &src) in target.y.iter_mut().zip(probe.scratch[0].iter()) {
*dst = src;
}
for (dst, &src) in target.u.iter_mut().zip(probe.scratch[1].iter()) {
*dst = src;
}
for (dst, &src) in target.v.iter_mut().zip(probe.scratch[2].iter()) {
*dst = src;
}
}
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 2, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_masked_compound = true;
ip.enable_jnt_comp = true;
ip.order_hints = gop_order_hints(2, pre.seq.order_hint_bits);
let mut ictx = PSearchCtx::new(
&pre.recon[1],
&pre.recon[0],
mi_rows,
mi_cols,
64,
64,
ip,
2,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
fn count(node: &SyntaxNode, dist: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, dist: &mut u32, other: &mut u32| {
if let Some(ib) = b.inter.as_ref() {
if ib.compound_type == crate::inter_pred::COMPOUND_DISTANCE {
*dist += 1;
} else {
*other += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, dist, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, dist, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, dist, other);
}
}
}
}
let (mut dist, mut other) = (0u32, 0u32);
count(&tree, &mut dist, &mut other);
assert!(
dist > 0,
"distance-blend content must commit COMPOUND_DISTANCE leaves \
(got DISTANCE={dist} OTHER={other})"
);
let enc = encode_gop_yuv_with_q(&[f0, f1, target], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 3);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_JNT_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(format!("{dir}/self-gop-64x64-q60-jnt.ivf"), &enc.ivf_bytes).unwrap();
}
}
fn count_sub8(node: &SyntaxNode, b_size: usize, n_sub8: &mut u32, n_4x4: &mut u32) {
use crate::cdf::{
block_height, block_width, partition_subsize, BLOCK_4X4, PARTITION_HORZ,
PARTITION_HORZ_4, PARTITION_HORZ_A, PARTITION_HORZ_B, PARTITION_SPLIT, PARTITION_VERT,
PARTITION_VERT_4, PARTITION_VERT_A, PARTITION_VERT_B,
};
let leaf_hit = |sz: usize, n_sub8: &mut u32, n_4x4: &mut u32| {
if block_width(sz) < 8 || block_height(sz) < 8 {
*n_sub8 += 1;
}
if sz == BLOCK_4X4 {
*n_4x4 += 1;
}
};
match node {
SyntaxNode::Leaf(_) => leaf_hit(b_size, n_sub8, n_4x4),
SyntaxNode::Split(children) => {
let s = partition_subsize(PARTITION_SPLIT, b_size).unwrap();
for ch in children.iter() {
count_sub8(ch, s, n_sub8, n_4x4);
}
}
SyntaxNode::Horz(_) | SyntaxNode::Vert(_) => {
let part = if matches!(node, SyntaxNode::Horz(_)) {
PARTITION_HORZ
} else {
PARTITION_VERT
};
let s = partition_subsize(part, b_size).unwrap();
for _b in node.asymmetric_blocks() {
leaf_hit(s, n_sub8, n_4x4);
}
}
SyntaxNode::HorzA(_)
| SyntaxNode::HorzB(_)
| SyntaxNode::VertA(_)
| SyntaxNode::VertB(_) => {
let (part, half_first) = match node {
SyntaxNode::HorzA(_) => (PARTITION_HORZ_A, false),
SyntaxNode::HorzB(_) => (PARTITION_HORZ_B, true),
SyntaxNode::VertA(_) => (PARTITION_VERT_A, false),
_ => (PARTITION_VERT_B, true),
};
let half = partition_subsize(part, b_size).unwrap();
let quarter = partition_subsize(PARTITION_SPLIT, b_size).unwrap();
let sizes: [usize; 3] = if half_first {
[half, quarter, quarter]
} else {
[quarter, quarter, half]
};
for (i, _b) in node.asymmetric_blocks().iter().enumerate() {
leaf_hit(sizes[i], n_sub8, n_4x4);
}
}
SyntaxNode::Horz4(_) | SyntaxNode::Vert4(_) => {
let part = if matches!(node, SyntaxNode::Horz4(_)) {
PARTITION_HORZ_4
} else {
PARTITION_VERT_4
};
let s = partition_subsize(part, b_size).unwrap();
for _b in node.asymmetric_blocks() {
leaf_hit(s, n_sub8, n_4x4);
}
}
}
}
#[test]
fn r416_search_selects_sub8_split_on_fine_motion() {
let f0 = moving_gradient(64, 64, 0, 0, 71);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(std::slice::from_ref(&f0), base_q_idx).unwrap();
let kr = &pre.recon[0];
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let at = |p: &[u16], w: usize, h: usize, i: i64, j: i64| -> u16 {
let ii = i.clamp(0, h as i64 - 1) as usize;
let jj = j.clamp(0, w as i64 - 1) as usize;
p[ii * w + jj]
};
for i in 0..64i64 {
for j in 0..64i64 {
let par = ((i / 4) + (j / 4)) & 1;
let (dy, dx) = if par == 0 { (0, 0) } else { (4, 4) };
target.y[(i * 64 + j) as usize] = at(&kr.y, 64, 64, i + dy, j + dx);
}
}
target.u.copy_from_slice(&kr.u);
target.v.copy_from_slice(&kr.v);
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_masked_compound = true;
ip.enable_jnt_comp = true;
ip.order_hints = gop_order_hints(1, pre.seq.order_hint_bits);
let key_recon = GopFrameReconYuv {
y: kr.y.clone(),
u: kr.u.clone(),
v: kr.v.clone(),
};
let mut ictx = PSearchCtx::new(
&key_recon,
&key_recon,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
let (mut n_sub8, mut n_4x4) = (0u32, 0u32);
count_sub8(&tree, BLOCK_64X64, &mut n_sub8, &mut n_4x4);
assert!(
n_4x4 > 0,
"4ร4-granular motion must commit BLOCK_4X4 leaves (sub8={n_sub8} 4x4={n_4x4})"
);
let enc = encode_gop_yuv_with_q(&[f0, target], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_SUB8_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
format!("{dir}/self-gop-64x64-q60-sub8-split.ivf"),
&enc.ivf_bytes,
)
.unwrap();
}
}
#[test]
fn r416_search_selects_sub8_strips_on_band_motion() {
let f0 = moving_gradient(64, 64, 0, 0, 93);
let base_q_idx = 60u8;
let pre = encode_gop_yuv_with_q(std::slice::from_ref(&f0), base_q_idx).unwrap();
let kr = &pre.recon[0];
let mut target = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 0);
let at = |p: &[u16], w: usize, h: usize, i: i64, j: i64| -> u16 {
let ii = i.clamp(0, h as i64 - 1) as usize;
let jj = j.clamp(0, w as i64 - 1) as usize;
p[ii * w + jj]
};
for i in 0..64i64 {
for j in 0..64i64 {
let dx = if (i / 4) & 1 == 0 { 0 } else { 5 };
target.y[(i * 64 + j) as usize] = at(&kr.y, 64, 64, i, j + dx);
}
}
target.u.copy_from_slice(&kr.u);
target.v.copy_from_slice(&kr.v);
let fh = build_p_frame_yuv420_8bit_fh_with_q(&pre.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.reference_select = true;
ip.enable_masked_compound = true;
ip.enable_jnt_comp = true;
ip.order_hints = gop_order_hints(1, pre.seq.order_hint_bits);
let key_recon = GopFrameReconYuv {
y: kr.y.clone(),
u: kr.u.clone(),
v: kr.v.clone(),
};
let mut ictx = PSearchCtx::new(
&key_recon,
&key_recon,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree = build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &target, &mut recon, &mut ictx)
.unwrap();
let (mut n_sub8, mut n_4x4) = (0u32, 0u32);
count_sub8(&tree, BLOCK_64X64, &mut n_sub8, &mut n_4x4);
assert!(
n_sub8 > 0,
"band motion must commit sub-8 strip leaves (sub8={n_sub8} 4x4={n_4x4})"
);
fn shape_census(node: &SyntaxNode, h4: &mut u32, v4: &mut u32, hz: &mut u32, vt: &mut u32) {
match node {
SyntaxNode::Horz4(_) => *h4 += 1,
SyntaxNode::Vert4(_) => *v4 += 1,
SyntaxNode::Horz(_) => *hz += 1,
SyntaxNode::Vert(_) => *vt += 1,
SyntaxNode::Split(ch) => {
for c in ch.iter() {
shape_census(c, h4, v4, hz, vt);
}
}
_ => {}
}
}
let (mut h4, mut v4, mut hz, mut vt) = (0u32, 0u32, 0u32, 0u32);
shape_census(&tree, &mut h4, &mut v4, &mut hz, &mut vt);
let _ = (v4, vt);
assert!(
h4 > 0 && hz > 0,
"band motion must commit both HORZ_4 strip nodes and HORZ-at-8x8 halves \
(horz4={h4} horz={hz} sub8={n_sub8} 4x4={n_4x4})"
);
let enc = encode_gop_yuv_with_q(&[f0, target], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, fr) in decoded.iter().enumerate() {
assert_eq!(fr.planes[0], plane8(&enc.recon[idx].y), "frame {idx} luma");
assert_eq!(fr.planes[1], plane8(&enc.recon[idx].u), "frame {idx} U");
assert_eq!(fr.planes[2], plane8(&enc.recon[idx].v), "frame {idx} V");
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_SUB8_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
format!("{dir}/self-gop-64x64-q60-sub8-strips.ivf"),
&enc.ivf_bytes,
)
.unwrap();
}
}
#[test]
fn r415_b_search_selects_skip_mode_between_converged_references() {
let f = moving_gradient(64, 64, 2, 3, 21);
let base_q_idx = 60u8;
let k = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f, base_q_idx).unwrap();
let anchor = GopFrameReconYuv {
y: k.recon_y,
u: k.recon_u,
v: k.recon_v,
};
let future = anchor.clone();
let mut recon = fresh_recon(64, 64, base_q_idx);
let mut ictx = b_frame_ctx(&anchor, &future, 16, 16, base_q_idx, true);
recon.bd.clear_for_sb(0, 0, 16, 16);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f, &mut recon, &mut ictx).unwrap();
fn count(node: &SyntaxNode, sm: &mut u32, other: &mut u32) {
let leafc = |b: &SyntaxBlock, sm: &mut u32, other: &mut u32| {
if b.inter.as_ref().is_some_and(|ib| ib.skip_mode != 0) {
*sm += 1;
} else {
*other += 1;
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, sm, other),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count(ch, sm, other);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, sm, other);
}
}
}
}
let (mut sm, mut other) = (0u32, 0u32);
count(&tree, &mut sm, &mut other);
assert!(
sm > 0,
"converged references on a skip-mode B frame must select \
skip_mode leaves (got SKIP_MODE={sm} OTHER={other})"
);
}
fn count_motion_modes(node: &SyntaxNode, counts: &mut [u32; 3]) {
let leafc = |b: &SyntaxBlock, counts: &mut [u32; 3]| {
if let Some(ib) = b.inter.as_ref() {
counts[(ib.motion_mode as usize).min(2)] += 1;
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, counts),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_motion_modes(ch, counts);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, counts);
}
}
}
}
fn shear_pair(w: u32, h: u32, dx_top: usize, ramp0: usize, ramp1: usize) -> [YuvFrame; 2] {
let (wu, hu) = (w as usize, h as usize);
let tex = |i: usize, j: usize| -> u16 {
((i * 7 + j * 11 + (i / 4) * (j / 8) + ((i * j) / 13)) % 256) as u16
};
let mut f0 = YuvFrame::filled(w, h, 8, ChromaFormat::Yuv420, 128);
for i in 0..hu {
for j in 0..wu {
f0.y[i * wu + j] = tex(i, j);
}
}
let dx_of = |i: usize| -> usize {
if i < ramp0 {
dx_top
} else if i >= ramp1 {
0
} else {
dx_top * (ramp1 - i) / (ramp1 - ramp0)
}
};
let mut f1 = YuvFrame::filled(w, h, 8, ChromaFormat::Yuv420, 128);
for i in 0..hu {
let dx = dx_of(i);
for j in 0..wu {
f1.y[i * wu + j] = tex(i, j + dx);
}
}
[f0, f1]
}
#[test]
fn r419_p_search_selects_obmc_on_sheared_motion() {
let [f0, f1] = shear_pair(64, 64, 8, 24, 48);
let base_q_idx = 100u8;
let key = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f0, base_q_idx).unwrap();
let reference = GopFrameReconYuv {
y: key.recon_y,
u: key.recon_u,
v: key.recon_v,
};
let fh = build_p_frame_yuv420_8bit_fh_with_q(&key.seq, 64, 64, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 64 * 64],
u: vec![0u16; 32 * 32],
v: vec![0u16; 32 * 32],
width: 64,
height: 64,
chroma_w: 32,
chroma_h: 32,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.is_motion_mode_switchable = true;
let mut ictx = PSearchCtx::new(
&reference,
&reference,
mi_rows,
mi_cols,
64,
64,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, &f1, &mut recon, &mut ictx).unwrap();
let mut counts = [0u32; 3];
count_motion_modes(&tree, &mut counts);
assert!(
counts[MOTION_MODE_OBMC as usize] > 0,
"sheared motion must commit at least one ยง5.11.27 OBMC leaf \
(got SIMPLE={} OBMC={} WARP={})",
counts[0],
counts[1],
counts[2]
);
let enc = encode_gop_yuv_with_q(&[f0, f1], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(
f.planes[0],
plane8(&enc.recon[idx].y),
"display {idx}: luma"
);
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "display {idx}: U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "display {idx}: V");
}
}
fn zoom_pair(w: u32, h: u32, num: i64, den: i64) -> [YuvFrame; 2] {
let (wu, hu) = (w as usize, h as usize);
let tex = |y: i64, x: i64| -> u16 {
let (y, x) = (y as f64 / 8.0, x as f64 / 8.0);
let v = 96.0
+ 60.0 * ((y * 0.11).sin() * (x * 0.13).cos())
+ 40.0 * ((y * 0.05 + x * 0.07).sin());
v.clamp(0.0, 255.0) as u16
};
let mut f0 = YuvFrame::filled(w, h, 8, ChromaFormat::Yuv420, 128);
let mut f1 = YuvFrame::filled(w, h, 8, ChromaFormat::Yuv420, 128);
let (cy, cx) = ((hu / 2) as i64, (wu / 2) as i64);
for i in 0..hu {
for j in 0..wu {
f0.y[i * wu + j] = tex((i as i64) * 8, (j as i64) * 8);
let sy = cy * 8 + ((i as i64) - cy) * 8 * num / den;
let sx = cx * 8 + ((j as i64) - cx) * 8 * num / den;
f1.y[i * wu + j] = tex(sy, sx);
}
}
[f0, f1]
}
#[test]
fn r419_p_search_selects_warped_causal_on_zooming_motion() {
let [f0, f1] = zoom_pair(96, 80, 17, 16);
let base_q_idx = 100u8;
let key = crate::encoder::key_frame::encode_key_frame_yuv_with_q(&f0, base_q_idx).unwrap();
assert!(
key.seq.enable_warped_motion,
"r419: the sequence gate must be open for ยง5.11.27 arm B"
);
let reference = GopFrameReconYuv {
y: key.recon_y,
u: key.recon_u,
v: key.recon_v,
};
let fh = build_p_frame_yuv420_8bit_fh_with_q(&key.seq, 96, 80, base_q_idx, 1, &[]);
assert_eq!(fh.allow_warped_motion, Some(true));
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; 96 * 80],
u: vec![0u16; 48 * 40],
v: vec![0u16; 48 * 40],
width: 96,
height: 80,
chroma_w: 48,
chroma_h: 40,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.is_motion_mode_switchable = true;
ip.allow_warped_motion = true;
let mut ictx = PSearchCtx::new(
&reference,
&reference,
mi_rows,
mi_cols,
96,
80,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
let mut counts = [0u32; 3];
let mut ts = test_tile_state(&recon, &ictx);
for (sb_r, sb_c) in [(0u32, 0u32), (0, 16), (16, 0), (16, 16)] {
if sb_r >= mi_rows || sb_c >= mi_cols {
continue;
}
recon.bd.clear_for_sb(sb_r, sb_c, mi_rows, mi_cols);
let tree =
search_and_commit_sb(&mut ts, sb_r, sb_c, &f1, &mut recon, &mut ictx).unwrap();
count_motion_modes(&tree, &mut counts);
}
assert!(
counts[MOTION_MODE_WARPED_CAUSAL as usize] > 0,
"zooming motion must commit at least one ยง5.11.27 \
WARPED_CAUSAL leaf (got SIMPLE={} OBMC={} WARP={})",
counts[0],
counts[1],
counts[2]
);
let enc = encode_gop_yuv_with_q(&[f0, f1], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(
f.planes[0],
plane8(&enc.recon[idx].y),
"display {idx}: luma"
);
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "display {idx}: U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "display {idx}: V");
}
}
fn noise_frame(w: u32, h: u32, seed: u32) -> YuvFrame {
let mut state = 0x1357_9BDFu32 ^ seed;
let mut next = move || {
state ^= state << 13;
state ^= state >> 17;
state ^= state << 5;
state
};
let mut f = YuvFrame::filled(w, h, 8, ChromaFormat::Yuv420, 0);
for p in f.y.iter_mut() {
*p = (next() & 0xFF) as u16;
}
for p in f.u.iter_mut() {
*p = 96 + (next() & 0x3F) as u16;
}
for p in f.v.iter_mut() {
*p = 64 + (next() & 0x3F) as u16;
}
f
}
fn count_intra_tools(node: &SyntaxNode, fi: &mut u32, cfl: &mut u32) {
let leafc = |b: &SyntaxBlock, fi: &mut u32, cfl: &mut u32| {
if b.inter.is_none() {
if b.use_filter_intra == 1 {
*fi += 1;
}
if b.uv_mode == Some(crate::cdf::UV_CFL_PRED as u8) {
*cfl += 1;
}
}
};
match node {
SyntaxNode::Leaf(b) => leafc(b, fi, cfl),
SyntaxNode::Split(children) => {
for ch in children.iter() {
count_intra_tools(ch, fi, cfl);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafc(b, fi, cfl);
}
}
}
}
fn p_frame_intra_tool_census(f0: &YuvFrame, f1: &YuvFrame, base_q_idx: u8) -> (u32, u32) {
let enc = encode_gop_yuv_with_q(&[f0.clone(), f1.clone()], base_q_idx).unwrap();
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), 2);
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(
f.planes[0],
plane8(&enc.recon[idx].y),
"display {idx}: luma"
);
assert_eq!(f.planes[1], plane8(&enc.recon[idx].u), "display {idx}: U");
assert_eq!(f.planes[2], plane8(&enc.recon[idx].v), "display {idx}: V");
}
let reference = enc.recon[0].clone();
let (w, h) = (f0.width, f0.height);
let key = crate::encoder::key_frame::encode_key_frame_yuv_with_q(f0, base_q_idx).unwrap();
let fh = build_p_frame_yuv420_8bit_fh_with_q(&key.seq, w, h, base_q_idx, 1, &[]);
let fs = fh.frame_size.as_ref().unwrap();
let (mi_rows, mi_cols) = (fs.mi_rows, fs.mi_cols);
let mut recon = ReconState {
y: vec![0u16; (w * h) as usize],
u: vec![0u16; ((w / 2) * (h / 2)) as usize],
v: vec![0u16; ((w / 2) * (h / 2)) as usize],
width: w as usize,
height: h as usize,
chroma_w: (w / 2) as usize,
chroma_h: (h / 2) as usize,
mi_rows,
mi_cols,
lossless: false,
allow_screen_content_tools: true,
allow_intrabc: false,
qp: QuantizerParams::neutral(base_q_idx, 8),
bit_depth: 8,
subsampling_x: 1,
subsampling_y: 1,
num_planes: 3,
search: crate::encoder::key_frame::SearchLimits::default(),
bd: BlockDecodedMirror::new(1, 1, 3),
dv_hash: Default::default(),
};
let mut ip = SyntaxInterFrameParams::single_ref_baseline(mi_rows, mi_cols, false);
ip.interpolation_filter = SWITCHABLE;
ip.is_motion_mode_switchable = true;
ip.allow_warped_motion = true;
let mut ictx = PSearchCtx::new(
&reference,
&reference,
mi_rows,
mi_cols,
w as usize,
h as usize,
ip,
1,
base_q_idx,
&[],
)
.unwrap();
recon.bd.clear_for_sb(0, 0, mi_rows, mi_cols);
let tree =
build_p_search_tree_for_tests(0, 0, BLOCK_64X64, f1, &mut recon, &mut ictx).unwrap();
if std::env::var("OXIDEAV_AV1_DEBUG_CENSUS").is_ok() {
fn dump(node: &SyntaxNode, depth: usize) {
let leafd = |b: &SyntaxBlock, depth: usize| {
eprintln!(
"{:indent$}leaf inter={} y={} uv={:?} fi={} cfl=({:?},{:?})",
"",
b.inter.is_some(),
b.y_mode,
b.uv_mode,
b.use_filter_intra,
b.cfl_alpha_u,
b.cfl_alpha_v,
indent = depth * 2
);
};
match node {
SyntaxNode::Leaf(b) => leafd(b, depth),
SyntaxNode::Split(children) => {
for ch in children.iter() {
dump(ch, depth + 1);
}
}
rest => {
for b in rest.asymmetric_blocks().iter() {
leafd(b, depth);
}
}
}
}
dump(&tree, 0);
}
let (mut fi, mut cfl) = (0u32, 0u32);
count_intra_tools(&tree, &mut fi, &mut cfl);
(fi, cfl)
}
#[test]
fn r419_p_search_selects_filter_intra_inside_inter_frame() {
let base_q_idx = 60u8;
let f0 = noise_frame(64, 64, 40);
let enc1 = encode_gop_yuv_with_q(&[f0.clone(), f0.clone()], base_q_idx).unwrap();
let rec1 = &enc1.recon[1];
let (above_ext, left_ext, _, _) = super::super::key_frame::build_tu_neighbours(
&rec1.y, 64, 64, 32, 32, 0, 0, 32, 32, false, false, 8,
);
let target = super::super::key_frame::predict_filter_intra_from_neighbours(
2, 32, 32, &above_ext, &left_ext, 8,
)
.expect("32x32 filter-intra prediction");
let mut f1 = f0.clone();
for i in 0..32usize {
for j in 0..32usize {
f1.y[(32 + i) * 64 + 32 + j] = target[i * 32 + j];
}
}
let (fi, _) = p_frame_intra_tool_census(&f0, &f1, base_q_idx);
assert!(
fi > 0,
"the FI-constructed fresh region must commit at least one \
ยง5.11.24 use_filter_intra = 1 intra leaf inside the \
P-frame (got {fi})"
);
}
#[test]
fn r419_p_search_selects_cfl_inside_inter_frame() {
let base_q_idx = 60u8;
let mut f0 = noise_frame(64, 64, 41);
f0.u.fill(128);
f0.v.fill(128);
let mut f1 = f0.clone();
let tex = |i: usize, j: usize| -> u16 {
let v = 128.0
+ 70.0 * ((i as f64) * 0.35).sin() * ((j as f64) * 0.29).cos()
+ 30.0 * (((i + j) as f64) * 0.11).sin();
v.clamp(0.0, 255.0) as u16
};
for i in 0..32usize {
for j in 0..32usize {
f1.y[(32 + i) * 64 + 32 + j] = tex(i, j);
}
}
let mut acc = 0i64;
for i in 0..16usize {
for j in 0..16usize {
let s = i64::from(tex(2 * i, 2 * j))
+ i64::from(tex(2 * i + 1, 2 * j))
+ i64::from(tex(2 * i, 2 * j + 1))
+ i64::from(tex(2 * i + 1, 2 * j + 1));
acc += s;
}
}
let mean4 = acc / 256; for i in 0..16usize {
for j in 0..16usize {
let s = i64::from(tex(2 * i, 2 * j))
+ i64::from(tex(2 * i + 1, 2 * j))
+ i64::from(tex(2 * i, 2 * j + 1))
+ i64::from(tex(2 * i + 1, 2 * j + 1));
let ac = (s - mean4) / 16; let cell = (16 + i) * 32 + 16 + j;
f1.u[cell] = (128 + ac).clamp(0, 255) as u16;
f1.v[cell] = (128 - ac).clamp(0, 255) as u16;
}
}
let (_, cfl) = p_frame_intra_tool_census(&f0, &f1, base_q_idx);
assert!(
cfl > 0,
"luma-tracking chroma on a fresh region must commit at \
least one ยง5.11.22 UV_CFL_PRED intra leaf inside the \
P-frame (got {cfl})"
);
}
#[test]
fn r419_motion_mode_and_intra_tool_fixture_streams() {
let dump = |name: &str, enc: &EncodedGopYuv| {
let decoded = crate::decoder::decode_av1_spec(&enc.ivf_bytes).unwrap();
assert_eq!(decoded.len(), enc.recon.len(), "{name}");
for (idx, f) in decoded.iter().enumerate() {
assert_eq!(
f.planes[0],
plane8(&enc.recon[idx].y),
"{name} display {idx}: luma"
);
assert_eq!(
f.planes[1],
plane8(&enc.recon[idx].u),
"{name} display {idx}: U"
);
assert_eq!(
f.planes[2],
plane8(&enc.recon[idx].v),
"{name} display {idx}: V"
);
}
if let Ok(dir) = std::env::var("OXIDEAV_AV1_R419_FIXDIR") {
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(format!("{dir}/{name}.ivf"), &enc.ivf_bytes).unwrap();
let mut yuv = Vec::new();
for rc in &enc.recon {
yuv.extend_from_slice(&plane8(&rc.y));
yuv.extend_from_slice(&plane8(&rc.u));
yuv.extend_from_slice(&plane8(&rc.v));
}
std::fs::write(format!("{dir}/{name}.yuv"), &yuv).unwrap();
}
};
{
let tex = |i: usize, j: usize| -> u16 {
((i * 7 + j * 11 + (i / 4) * (j / 8) + ((i * j) / 13)) % 256) as u16
};
let frames: Vec<YuvFrame> = (0..3)
.map(|k| {
let mut f = YuvFrame::filled(64, 64, 8, ChromaFormat::Yuv420, 128);
for i in 0..64usize {
let (ramp0, ramp1, top) = (24usize, 48usize, 4 * k);
let dx = if i < ramp0 {
top
} else if i >= ramp1 {
0
} else {
top * (ramp1 - i) / (ramp1 - ramp0)
};
for j in 0..64usize {
f.y[i * 64 + j] = tex(i, j + dx);
}
}
f
})
.collect();
let enc = encode_gop_yuv_with_q(&frames, 100).unwrap();
dump("self-gop-64x64-q100-shear-obmc", &enc);
}
{
let [f0, f1] = zoom_pair(96, 80, 17, 16);
let enc = encode_gop_yuv_with_q(&[f0, f1], 100).unwrap();
dump("self-gop-96x80-q100-zoom-warp", &enc);
}
{
let base_q_idx = 60u8;
let mut f0 = noise_frame(64, 64, 41);
f0.u.fill(128);
f0.v.fill(128);
let enc1 = encode_gop_yuv_with_q(&[f0.clone(), f0.clone()], base_q_idx).unwrap();
let rec1 = &enc1.recon[1];
let (above_ext, left_ext, _, _) = super::super::key_frame::build_tu_neighbours(
&rec1.y, 64, 64, 32, 32, 0, 0, 32, 32, false, false, 8,
);
let target = super::super::key_frame::predict_filter_intra_from_neighbours(
2, 32, 32, &above_ext, &left_ext, 8,
)
.expect("32x32 filter-intra prediction");
let mut f1 = f0.clone();
for i in 0..32usize {
for j in 0..32usize {
f1.y[(32 + i) * 64 + 32 + j] = target[i * 32 + j];
}
}
let tex = |i: usize, j: usize| -> u16 {
let v = 128.0
+ 70.0 * ((i as f64) * 0.35).sin() * ((j as f64) * 0.29).cos()
+ 30.0 * (((i + j) as f64) * 0.11).sin();
v.clamp(0.0, 255.0) as u16
};
let mut f2 = f1.clone();
for i in 0..32usize {
for j in 0..32usize {
f2.y[(32 + i) * 64 + 32 + j] = tex(i, j);
}
}
let mut acc = 0i64;
for i in 0..16usize {
for j in 0..16usize {
acc += i64::from(tex(2 * i, 2 * j))
+ i64::from(tex(2 * i + 1, 2 * j))
+ i64::from(tex(2 * i, 2 * j + 1))
+ i64::from(tex(2 * i + 1, 2 * j + 1));
}
}
let mean4 = acc / 256;
for i in 0..16usize {
for j in 0..16usize {
let s = i64::from(tex(2 * i, 2 * j))
+ i64::from(tex(2 * i + 1, 2 * j))
+ i64::from(tex(2 * i, 2 * j + 1))
+ i64::from(tex(2 * i + 1, 2 * j + 1));
let ac = (s - mean4) / 16;
let cell = (16 + i) * 32 + 16 + j;
f2.u[cell] = (128 + ac).clamp(0, 255) as u16;
f2.v[cell] = (128 - ac).clamp(0, 255) as u16;
}
}
let enc = encode_gop_yuv_with_q(&[f0, f1, f2], base_q_idx).unwrap();
dump("self-gop-64x64-q60-fi-cfl", &enc);
}
}
}
#[cfg(test)]
mod delta_q_lossless_cap_tests {
use super::delta_q_lossless_up_cap;
#[test]
fn cap_follows_the_lossless_headroom() {
assert_eq!(delta_q_lossless_up_cap(60, &[0, -60]), 0);
assert_eq!(delta_q_lossless_up_cap(60, &[0, -68]), 1);
assert_eq!(delta_q_lossless_up_cap(60, &[0, -76]), 2);
assert_eq!(delta_q_lossless_up_cap(100, &[0, -255]), 2);
assert_eq!(delta_q_lossless_up_cap(72, &[0, 24, -72]), 0);
assert_eq!(delta_q_lossless_up_cap(80, &[0, -40]), 2);
assert_eq!(delta_q_lossless_up_cap(80, &[]), 2);
}
}