use crate::dpb::RefEntry;
use crate::inter::{MotionInfo, Mv, PredFlags};
#[inline]
fn clip3(lo: i32, hi: i32, v: i32) -> i32 {
v.clamp(lo, hi)
}
pub(crate) fn scale_mv(mv: Mv, col_dist: i32, cur_dist: i32) -> Mv {
let td = clip3(-128, 127, col_dist);
let tb = clip3(-128, 127, cur_dist);
if td == 0 {
return mv;
}
let tx = (16384 + (td.abs() >> 1)) / td;
let dsf = clip3(-4096, 4095, (tb * tx + 32) >> 6);
let sc = |c: i16| -> i16 {
let p = dsf * c as i32;
let v = p.signum() * ((p.abs() + 127) >> 8);
clip3(-32768, 32767, v) as i16
};
Mv::new(sc(mv.x), sc(mv.y))
}
pub(crate) trait Neighbors {
fn motion_at(&self, x: isize, y: isize) -> Option<MotionInfo>;
fn available(&self, x: isize, y: isize) -> bool;
fn temporal(&self, x: usize, y: usize, list: usize, ref_poc: i32, cur_poc: i32) -> Option<Mv>;
fn cur_poc(&self) -> i32;
fn ctb_log2(&self) -> u32;
}
#[derive(Clone, Copy, Default)]
pub(crate) struct MergeCand {
pub(crate) pred: PredFlags,
pub(crate) mv: [Mv; 2],
pub(crate) ref_idx: [i8; 2],
}
impl MergeCand {
fn from_motion(m: &MotionInfo) -> Self {
MergeCand {
pred: m.pred,
mv: m.mv,
ref_idx: m.ref_idx,
}
}
fn same_motion(&self, o: &MergeCand) -> bool {
self.pred == o.pred && self.mv == o.mv && self.ref_idx == o.ref_idx
}
}
pub(crate) struct PuGeom {
pub(crate) x: usize,
pub(crate) y: usize,
pub(crate) w: usize,
pub(crate) h: usize,
pub(crate) cu_x: usize,
pub(crate) cu_y: usize,
pub(crate) is_b: bool,
pub(crate) part_idx: usize,
pub(crate) cu_w: usize,
pub(crate) cu_h: usize,
pub(crate) curr_pic_ref: bool,
pub(crate) par_mrg_level: u32,
}
#[inline]
fn prediction_block_available<N: Neighbors>(nb: &N, pu: &PuGeom, px: isize, py: isize) -> bool {
if px < 0 || py < 0 {
return false;
}
let (x, y) = (px as usize, py as usize);
let same_cb = x >= pu.cu_x
&& y >= pu.cu_y
&& x < pu.cu_x.saturating_add(pu.cu_w)
&& y < pu.cu_y.saturating_add(pu.cu_h);
if !same_cb {
return nb.available(px, py);
}
!(pu.w.saturating_mul(2) == pu.cu_w
&& pu.h.saturating_mul(2) == pu.cu_h
&& pu.part_idx == 1
&& pu.cu_y.saturating_add(pu.h) <= y
&& pu.cu_x.saturating_add(pu.w) > x)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn derive_merge<N: Neighbors>(
nb: &N,
pu: &PuGeom,
merge_idx: usize,
max_cand: usize,
temporal_enabled: bool,
list0: &[RefEntry],
list1: &[RefEntry],
) -> MergeCand {
let mut cands: Vec<MergeCand> = Vec::with_capacity(max_cand + 2);
let (x, y, w, h) = (pu.x as isize, pu.y as isize, pu.w as isize, pu.h as isize);
let a1 = spatial(nb, x - 1, y + h - 1, pu, SpatialPos::A1);
if let Some(c) = a1 {
push_unique(&mut cands, c, max_cand);
}
let b1 = spatial(nb, x + w - 1, y - 1, pu, SpatialPos::B1);
if let Some(c) = b1
&& a1.is_none_or(|a| !a.same_motion(&c))
{
push_unique(&mut cands, c, max_cand);
}
if let Some(c) = spatial(nb, x + w, y - 1, pu, SpatialPos::B0)
&& b1.is_none_or(|b| !b.same_motion(&c))
{
push_unique(&mut cands, c, max_cand);
}
if let Some(c) = spatial(nb, x - 1, y + h, pu, SpatialPos::A0)
&& a1.is_none_or(|a| !a.same_motion(&c))
{
push_unique(&mut cands, c, max_cand);
}
if cands.len() < 4
&& let Some(c) = spatial(nb, x - 1, y - 1, pu, SpatialPos::B2)
{
let dup = a1.is_some_and(|a| a.same_motion(&c)) || b1.is_some_and(|b| b.same_motion(&c));
if !dup {
push_unique(&mut cands, c, max_cand);
}
}
if temporal_enabled
&& cands.len() < max_cand
&& let Some(tc) = temporal_merge(nb, pu, list0, list1)
{
cands.push(tc);
}
if pu.is_b {
derive_combined(&mut cands, max_cand, list0, list1);
}
let n0 = list0.len().saturating_sub(usize::from(pu.curr_pic_ref));
derive_zero(&mut cands, max_cand, pu.is_b, n0, list1.len());
cands
.get(merge_idx)
.copied()
.unwrap_or_else(|| zero_cand(pu.is_b))
}
#[derive(Clone, Copy)]
enum SpatialPos {
A0,
A1,
B0,
B1,
B2,
}
fn spatial<N: Neighbors>(
nb: &N,
px: isize,
py: isize,
pu: &PuGeom,
pos: SpatialPos,
) -> Option<MergeCand> {
if pu.part_idx == 1 {
let vertical_split = pu.w < pu.cu_w && pu.h == pu.cu_h;
let horizontal_split = pu.h < pu.cu_h && pu.w == pu.cu_w;
match pos {
SpatialPos::A1 if vertical_split => return None,
SpatialPos::B1 if horizontal_split => return None,
_ => {}
}
}
let lvl = pu.par_mrg_level;
if px >= 0
&& py >= 0
&& (pu.x as isize >> lvl) == (px >> lvl)
&& (pu.y as isize >> lvl) == (py >> lvl)
{
return None;
}
if !prediction_block_available(nb, pu, px, py) {
return None;
}
let m = nb.motion_at(px, py)?;
if m.is_intra {
return None;
}
Some(MergeCand::from_motion(&m))
}
fn temporal_merge<N: Neighbors>(
nb: &N,
pu: &PuGeom,
list0: &[RefEntry],
list1: &[RefEntry],
) -> Option<MergeCand> {
let cur_poc = nb.cur_poc();
let ctb_log2 = nb.ctb_log2();
let br_y = pu.y + pu.h;
let br: Option<(usize, usize)> = if (br_y >> ctb_log2) == (pu.y >> ctb_log2) {
Some((pu.x + pu.w, br_y))
} else {
None
};
let ctr = (pu.x + pu.w / 2, pu.y + pu.h / 2);
let ref0_poc = list0.first().map(|r| r.poc).unwrap_or(cur_poc);
let mv0 = br
.and_then(|b| nb.temporal(b.0, b.1, 0, ref0_poc, cur_poc))
.or_else(|| nb.temporal(ctr.0, ctr.1, 0, ref0_poc, cur_poc));
let mv1 = if pu.is_b {
let ref1_poc = list1.first().map(|r| r.poc).unwrap_or(cur_poc);
br.and_then(|b| nb.temporal(b.0, b.1, 1, ref1_poc, cur_poc))
.or_else(|| nb.temporal(ctr.0, ctr.1, 1, ref1_poc, cur_poc))
} else {
None
};
if mv0.is_none() && mv1.is_none() {
return None;
}
let mut cand = MergeCand {
pred: PredFlags::default(),
mv: [Mv::default(); 2],
ref_idx: [-1; 2],
};
if let Some(mv0) = mv0 {
cand.pred.l0 = true;
cand.mv[0] = mv0;
cand.ref_idx[0] = 0;
}
if let Some(mv1) = mv1 {
cand.pred.l1 = true;
cand.mv[1] = mv1;
cand.ref_idx[1] = 0;
}
Some(cand)
}
fn derive_combined(
cands: &mut Vec<MergeCand>,
max_cand: usize,
list0: &[RefEntry],
list1: &[RefEntry],
) {
if cands.len() < 2 || cands.len() >= max_cand {
return;
}
static L0: [usize; 12] = [0, 1, 0, 2, 1, 2, 0, 3, 1, 3, 2, 3];
static L1: [usize; 12] = [1, 0, 2, 0, 2, 1, 3, 0, 3, 1, 3, 2];
let n = cands.len();
let mut k = 0;
while cands.len() < max_cand && k < n * (n - 1) {
let i0 = L0[k % 12];
let i1 = L1[k % 12];
k += 1;
if i0 >= n || i1 >= n {
continue;
}
let a = cands[i0];
let b = cands[i1];
if a.pred.l0 && b.pred.l1 {
let poc0 = list0.get(a.ref_idx[0].max(0) as usize).map(|r| r.poc);
let poc1 = list1.get(b.ref_idx[1].max(0) as usize).map(|r| r.poc);
if poc0 != poc1 || a.mv[0] != b.mv[1] {
cands.push(MergeCand {
pred: PredFlags { l0: true, l1: true },
mv: [a.mv[0], b.mv[1]],
ref_idx: [a.ref_idx[0], b.ref_idx[1]],
});
}
}
}
}
fn derive_zero(cands: &mut Vec<MergeCand>, max_cand: usize, is_b: bool, n0: usize, n1: usize) {
let num_ref = if is_b { n0.min(n1) } else { n0 };
if num_ref == 0 {
return;
}
let mut zero_idx = 0i8;
while cands.len() < max_cand {
let r = if (zero_idx as usize) < num_ref {
zero_idx
} else {
0
};
cands.push(MergeCand {
pred: PredFlags { l0: true, l1: is_b },
mv: [Mv::default(), Mv::default()],
ref_idx: [r, if is_b { r } else { -1 }],
});
zero_idx += 1;
if zero_idx > 32 {
break;
}
}
}
fn zero_cand(is_b: bool) -> MergeCand {
MergeCand {
pred: PredFlags { l0: true, l1: is_b },
mv: [Mv::default(), Mv::default()],
ref_idx: [0, if is_b { 0 } else { -1 }],
}
}
fn push_unique(cands: &mut Vec<MergeCand>, c: MergeCand, max_cand: usize) {
if cands.len() < max_cand {
cands.push(c);
}
}
pub(crate) fn derive_amvp<N: Neighbors>(
nb: &N,
pu: &PuGeom,
list: usize,
ref_poc: i32,
temporal_enabled: bool,
col_list0: &[RefEntry],
col_list1: &[RefEntry],
) -> [Mv; 2] {
let (x, y, w, h) = (pu.x as isize, pu.y as isize, pu.w as isize, pu.h as isize);
let cur_poc = nb.cur_poc();
let mut preds: Vec<Mv> = Vec::with_capacity(2);
let a_pos = [(x - 1, y + h), (x - 1, y + h - 1)];
let b_pos = [(x + w, y - 1), (x + w - 1, y - 1), (x - 1, y - 1)];
let a_inter = |px: isize, py: isize| -> bool {
prediction_block_available(nb, pu, px, py)
&& nb.motion_at(px, py).is_some_and(|m| !m.is_intra)
};
let is_scaled_flag = a_inter(a_pos[0].0, a_pos[0].1) || a_inter(a_pos[1].0, a_pos[1].1);
let target_lt = col_list0
.iter()
.chain(col_list1.iter())
.find(|r| r.poc == ref_poc)
.map(|r| r.long_term)
.unwrap_or(false);
let mut a = amvp_spatial(
nb, pu, &a_pos, list, ref_poc, cur_poc, true, true, target_lt,
);
let mut b = amvp_spatial(
nb, pu, &b_pos, list, ref_poc, cur_poc, false, true, target_lt,
);
if !is_scaled_flag {
if a.is_none() {
a = b;
}
b = amvp_spatial(
nb, pu, &b_pos, list, ref_poc, cur_poc, true, false, target_lt,
);
}
if let Some(mv) = a {
preds.push(mv);
}
if let Some(mv) = b
&& a != Some(mv)
{
preds.push(mv);
}
if preds.len() < 2 && temporal_enabled {
let ctb_log2 = nb.ctb_log2();
let br_y = pu.y + pu.h;
let br: Option<(usize, usize)> = if (br_y >> ctb_log2) == (pu.y >> ctb_log2) {
Some((pu.x + pu.w, br_y))
} else {
None
};
let ctr = (pu.x + pu.w / 2, pu.y + pu.h / 2);
let _ = (col_list0, col_list1);
if let Some(mv) = br
.and_then(|b| nb.temporal(b.0, b.1, list, ref_poc, cur_poc))
.or_else(|| nb.temporal(ctr.0, ctr.1, list, ref_poc, cur_poc))
{
preds.push(mv);
}
}
while preds.len() < 2 {
preds.push(Mv::default());
}
[preds[0], preds[1]]
}
#[allow(clippy::too_many_arguments)]
fn amvp_spatial<N: Neighbors>(
nb: &N,
pu: &PuGeom,
positions: &[(isize, isize)],
list: usize,
ref_poc: i32,
cur_poc: i32,
allow_scaled: bool,
do_same_ref: bool,
target_lt: bool,
) -> Option<Mv> {
if do_same_ref {
for &(px, py) in positions {
if !prediction_block_available(nb, pu, px, py) {
continue;
}
if let Some(m) = nb.motion_at(px, py) {
if m.is_intra {
continue;
}
for l in [list, 1 - list] {
if m.pred_uses(l) && m.ref_poc[l] == ref_poc {
return Some(m.mv[l]);
}
}
}
}
}
if !allow_scaled {
return None;
}
for &(px, py) in positions {
if !prediction_block_available(nb, pu, px, py) {
continue;
}
if let Some(m) = nb.motion_at(px, py) {
if m.is_intra {
continue;
}
for l in [list, 1 - list] {
if m.pred_uses(l) && m.ref_lt[l] == target_lt {
if target_lt {
return Some(m.mv[l]);
}
let col_dist = cur_poc - m.ref_poc[l];
let cur_dist = cur_poc - ref_poc;
return Some(scale_mv(m.mv[l], col_dist, cur_dist));
}
}
}
}
None
}
impl MotionInfo {
#[inline]
fn pred_uses(&self, l: usize) -> bool {
(l == 0 && self.pred.l0) || (l == 1 && self.pred.l1)
}
}
#[cfg(test)]
mod tests {
use super::*;
struct L1OnlyTemporal;
impl Neighbors for L1OnlyTemporal {
fn motion_at(&self, _x: isize, _y: isize) -> Option<MotionInfo> {
None
}
fn available(&self, _x: isize, _y: isize) -> bool {
false
}
fn temporal(
&self,
_x: usize,
_y: usize,
list: usize,
_ref_poc: i32,
_cur_poc: i32,
) -> Option<Mv> {
(list == 1).then_some(Mv::new(12, -4))
}
fn cur_poc(&self) -> i32 {
8
}
fn ctb_log2(&self) -> u32 {
6
}
}
#[test]
fn scale_identity_when_equal_dist() {
let mv = Mv::new(10, -6);
let s = scale_mv(mv, 4, 4);
assert_eq!(s, mv);
}
#[test]
fn scale_halves_on_double_dist() {
let mv = Mv::new(16, 0);
let s = scale_mv(mv, 8, 4);
assert!((s.x - 8).abs() <= 1);
}
#[test]
fn current_picture_amvp_uses_normative_zero_padding() {
let pu = PuGeom {
x: 0,
y: 0,
w: 16,
h: 16,
cu_x: 0,
cu_y: 0,
is_b: false,
part_idx: 0,
cu_w: 16,
cu_h: 16,
par_mrg_level: 2,
curr_pic_ref: false,
};
let current = [RefEntry {
_dpb_index: usize::MAX,
poc: 8,
long_term: true,
}];
assert_eq!(
derive_amvp(&L1OnlyTemporal, &pu, 0, 8, false, ¤t, &[]),
[Mv::default(), Mv::default()]
);
}
#[test]
fn temporal_merge_keeps_l1_only_candidate() {
let pu = PuGeom {
x: 0,
y: 0,
w: 16,
h: 16,
cu_x: 0,
cu_y: 0,
is_b: true,
part_idx: 0,
cu_w: 16,
cu_h: 16,
par_mrg_level: 2,
curr_pic_ref: false,
};
let refs0 = [RefEntry {
_dpb_index: 0,
poc: 4,
long_term: false,
}];
let refs1 = [RefEntry {
_dpb_index: 1,
poc: 12,
long_term: false,
}];
let cand = temporal_merge(&L1OnlyTemporal, &pu, &refs0, &refs1).unwrap();
assert_eq!(
cand.pred,
PredFlags {
l0: false,
l1: true
}
);
assert_eq!(cand.mv, [Mv::default(), Mv::new(12, -4)]);
assert_eq!(cand.ref_idx, [-1, 0]);
}
#[test]
fn default_parallel_merge_level_filters_same_region() {
struct OneNeighbor;
impl Neighbors for OneNeighbor {
fn motion_at(&self, _x: isize, _y: isize) -> Option<MotionInfo> {
Some(MotionInfo {
pred: PredFlags {
l0: true,
l1: false,
},
ref_idx: [0, -1],
..Default::default()
})
}
fn available(&self, _x: isize, _y: isize) -> bool {
true
}
fn temporal(
&self,
_x: usize,
_y: usize,
_list: usize,
_ref_poc: i32,
_cur_poc: i32,
) -> Option<Mv> {
None
}
fn cur_poc(&self) -> i32 {
0
}
fn ctb_log2(&self) -> u32 {
6
}
}
let pu = PuGeom {
x: 10,
y: 10,
w: 4,
h: 4,
cu_x: 8,
cu_y: 8,
is_b: false,
part_idx: 0,
cu_w: 4,
cu_h: 4,
par_mrg_level: 2,
curr_pic_ref: false,
};
assert!(spatial(&OneNeighbor, 9, 10, &pu, SpatialPos::A1).is_none());
}
#[test]
fn previous_pu_in_same_cu_does_not_require_decoded_map() {
struct SameCuNeighbor;
impl Neighbors for SameCuNeighbor {
fn motion_at(&self, _x: isize, _y: isize) -> Option<MotionInfo> {
Some(MotionInfo {
pred: PredFlags {
l0: true,
l1: false,
},
ref_idx: [0, -1],
..Default::default()
})
}
fn available(&self, _x: isize, _y: isize) -> bool {
false
}
fn temporal(
&self,
_x: usize,
_y: usize,
_list: usize,
_ref_poc: i32,
_cur_poc: i32,
) -> Option<Mv> {
None
}
fn cur_poc(&self) -> i32 {
0
}
fn ctb_log2(&self) -> u32 {
6
}
}
let pu = PuGeom {
x: 0,
y: 8,
w: 8,
h: 8,
cu_x: 0,
cu_y: 0,
is_b: false,
part_idx: 2,
cu_w: 16,
cu_h: 16,
par_mrg_level: 0,
curr_pic_ref: false,
};
assert!(prediction_block_available(&SameCuNeighbor, &pu, 7, 7));
assert!(spatial(&SameCuNeighbor, 7, 7, &pu, SpatialPos::B1).is_some());
}
#[test]
fn second_nxn_pu_rejects_unavailable_same_cu_region() {
struct AlwaysAvailable;
impl Neighbors for AlwaysAvailable {
fn motion_at(&self, _x: isize, _y: isize) -> Option<MotionInfo> {
Some(MotionInfo::default())
}
fn available(&self, _x: isize, _y: isize) -> bool {
true
}
fn temporal(
&self,
_x: usize,
_y: usize,
_list: usize,
_ref_poc: i32,
_cur_poc: i32,
) -> Option<Mv> {
None
}
fn cur_poc(&self) -> i32 {
0
}
fn ctb_log2(&self) -> u32 {
6
}
}
let pu = PuGeom {
x: 8,
y: 0,
w: 8,
h: 8,
cu_x: 0,
cu_y: 0,
is_b: false,
part_idx: 1,
cu_w: 16,
cu_h: 16,
par_mrg_level: 0,
curr_pic_ref: false,
};
assert!(!prediction_block_available(&AlwaysAvailable, &pu, 7, 8));
assert!(prediction_block_available(&AlwaysAvailable, &pu, 8, 8));
}
#[test]
fn zero_merge_candidates_include_current_picture_reference_slot() {
let pu = PuGeom {
x: 0,
y: 0,
w: 16,
h: 16,
cu_x: 0,
cu_y: 0,
is_b: false,
part_idx: 0,
cu_w: 16,
cu_h: 16,
par_mrg_level: 2,
curr_pic_ref: false,
};
let refs = [
RefEntry {
_dpb_index: 0,
poc: 4,
long_term: false,
},
RefEntry {
_dpb_index: usize::MAX,
poc: 8,
long_term: true,
},
];
let cand = derive_merge(&L1OnlyTemporal, &pu, 1, 2, false, &refs, &[]);
assert_eq!(cand.ref_idx, [1, -1]);
assert_eq!(cand.mv, [Mv::default(), Mv::default()]);
}
}