#[cfg(test)]
mod tests;
use std::num::{NonZeroU8, NonZeroUsize};
use anyhow::Result;
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;
use smallvec::SmallVec;
use crate::{
average::average2,
pad::pad_reference_frame,
params::{ReduceFilter, Subpel, SubpelMethod},
reduce::{
ReduceFn, reduce_average, reduce_bilinear, reduce_cubic, reduce_quadratic, reduce_triangle,
},
refine::{
RefineFn, refine_diagonal_bilinear, refine_horizontal_bicubic, refine_horizontal_bilinear,
refine_horizontal_wiener, refine_vertical_bicubic, refine_vertical_bilinear,
refine_vertical_wiener,
},
util::{Pixel, vs_bitblt},
};
#[derive(Debug, Clone)]
pub struct MVPlane {
pub subpel_window_offsets: SmallVec<[usize; 16]>,
pub width: NonZeroUsize,
pub height: NonZeroUsize,
pub padded_width: NonZeroUsize,
pub padded_height: NonZeroUsize,
pub stride: NonZeroUsize,
pub hpad: usize,
pub vpad: usize,
pub offset_padding: usize,
#[allow(dead_code)]
pub hpad_pel: usize,
#[allow(dead_code)]
pub vpad_pel: usize,
pub bits_per_sample: NonZeroU8,
pub pel: Subpel,
pub is_padded: bool,
pub is_refined: bool,
pub is_filled: bool,
}
impl MVPlane {
#[inline]
pub fn new(
width: NonZeroUsize,
height: NonZeroUsize,
pel: Subpel,
hpad: usize,
vpad: usize,
bits_per_sample: NonZeroU8,
plane_offset: usize,
pitch: NonZeroUsize,
) -> Result<Self> {
let pel_val = u8::from(pel) as usize;
let padded_width = width.saturating_add(2 * hpad);
let padded_height = height.saturating_add(2 * vpad);
let offset_padding = pitch.get() * vpad + hpad;
let windows = pel_val * pel_val;
let mut offsets = SmallVec::with_capacity(windows);
for i in 0..windows {
let offset = i * pitch.get() * padded_height.get();
offsets.push(plane_offset + offset);
}
Ok(Self {
width,
height,
padded_width,
padded_height,
hpad,
vpad,
hpad_pel: hpad * pel_val,
vpad_pel: vpad * pel_val,
subpel_window_offsets: offsets,
offset_padding,
stride: pitch,
bits_per_sample,
pel,
is_padded: false,
is_refined: false,
is_filled: false,
})
}
#[inline]
pub fn fill_plane<T: Pixel>(&mut self, src: &[T], src_pitch: NonZeroUsize, dest: &mut [T]) {
if self.is_filled {
return;
}
let offset = *semisafe_get(&self.subpel_window_offsets, 0) + self.offset_padding;
vs_bitblt(
semisafe_get_mut(dest, offset..),
self.stride,
src,
src_pitch,
self.width,
self.height,
);
self.is_filled = true;
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(skip_all, name = "mv_plane::refine_ext", fields(self.pel, self.is_refined))
)]
#[inline]
pub fn refine_ext<T: Pixel>(
&mut self,
src_2x: &[T],
src_2x_pitch: NonZeroUsize,
is_ext_padded: bool,
dest: &mut [T],
) {
if !self.is_refined {
match self.pel {
Subpel::Full => {
}
Subpel::Half => {
self.refine_ext_pel2(src_2x, src_2x_pitch, is_ext_padded, dest);
}
Subpel::Quarter => {
self.refine_ext_pel4(src_2x, src_2x_pitch, is_ext_padded, dest);
}
}
}
self.is_refined = true;
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(skip_all, name = "mv_plane::reduce_to", fields(filter, reduced_plane.is_filled))
)]
#[inline]
pub fn reduce_to<T: Pixel>(
&self,
reduced_plane: &mut MVPlane,
filter: ReduceFilter,
dest: &mut [T],
src: &[T],
dest_pitch: NonZeroUsize,
src_pitch: NonZeroUsize,
dest_width: NonZeroUsize,
dest_height: NonZeroUsize,
) {
if reduced_plane.is_filled {
return;
}
let dest = semisafe_get_mut(
dest,
*semisafe_get(&reduced_plane.subpel_window_offsets, 0) + reduced_plane.offset_padding..,
);
let src = semisafe_get(
src,
*semisafe_get(&self.subpel_window_offsets, 0) + self.offset_padding..,
);
let reduce: ReduceFn<T> = match filter {
ReduceFilter::Average => reduce_average,
ReduceFilter::Triangle => reduce_triangle,
ReduceFilter::Bilinear => reduce_bilinear,
ReduceFilter::Quadratic => reduce_quadratic,
ReduceFilter::Cubic => reduce_cubic,
};
reduce(dest, src, dest_pitch, src_pitch, dest_width, dest_height);
reduced_plane.is_filled = true;
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(skip_all, name = "mv_plane::pad", fields(self.is_padded))
)]
#[inline]
pub fn pad<T: Pixel>(&mut self, src: &mut [T]) {
if !self.is_padded {
pad_reference_frame(
*semisafe_get(&self.subpel_window_offsets, 0),
self.stride,
self.hpad,
self.vpad,
self.width,
self.height,
src,
);
self.is_padded = true;
}
}
fn refine_with_split<T: Pixel>(
plane: &mut [T],
src_offset: usize,
dest_offset: usize,
refine_fn: RefineFn<T>,
pitch: NonZeroUsize,
padded_width: NonZeroUsize,
padded_height: NonZeroUsize,
bits_per_sample: NonZeroU8,
) {
debug_assert!(
bits_per_sample.get() as usize > (size_of::<T>() - 1) * 8
&& (bits_per_sample.get() as usize <= size_of::<T>() * 8)
);
if src_offset <= dest_offset {
let (left, right) = plane.split_at_mut(dest_offset);
refine_fn(
right,
semisafe_get(left, src_offset..),
pitch,
padded_width,
padded_height,
bits_per_sample,
);
} else {
let (left, right) = plane.split_at_mut(src_offset);
refine_fn(
semisafe_get_mut(left, dest_offset..),
right,
pitch,
padded_width,
padded_height,
bits_per_sample,
);
}
}
fn average2_with_split<T: Pixel>(
plane: &mut [T],
src1_offset: usize,
src2_offset: usize,
dest_offset: usize,
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
) {
let mut offsets = [
(src1_offset, 1u8), (src2_offset, 2u8), (dest_offset, 0u8), ];
offsets.sort_by_key(|&(offset, _)| offset);
let _first_offset = offsets[0].0;
let second_offset = offsets[1].0;
let third_offset = offsets[2].0;
let first_type = offsets[0].1;
let second_type = offsets[1].1;
let third_type = offsets[2].1;
let (first_part, rest) = plane.split_at_mut(second_offset);
let (second_part, third_part) = rest.split_at_mut(third_offset - second_offset);
let (src1_slice, src2_slice, dest_slice) = match (first_type, second_type, third_type) {
(1, 2, 0) => (
semisafe_get(first_part, src1_offset..),
semisafe_get(second_part, ..),
semisafe_get_mut(third_part, ..),
),
(1, 0, 2) => (
semisafe_get(first_part, src1_offset..),
semisafe_get(third_part, src2_offset - third_offset..),
semisafe_get_mut(second_part, dest_offset - second_offset..),
),
(2, 1, 0) => (
semisafe_get(second_part, src1_offset - second_offset..),
semisafe_get(first_part, src2_offset..),
semisafe_get_mut(third_part, ..),
),
(2, 0, 1) => (
semisafe_get(third_part, src1_offset - third_offset..),
semisafe_get(first_part, src2_offset..),
semisafe_get_mut(second_part, dest_offset - second_offset..),
),
(0, 1, 2) => (
semisafe_get(second_part, src1_offset - second_offset..),
semisafe_get(third_part, src2_offset - third_offset..),
semisafe_get_mut(first_part, dest_offset..),
),
(0, 2, 1) => (
semisafe_get(third_part, src1_offset - third_offset..),
semisafe_get(second_part, src2_offset - second_offset..),
semisafe_get_mut(first_part, dest_offset..),
),
_ => unreachable!(),
};
average2(dest_slice, src1_slice, src2_slice, pitch, width, height);
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(skip_all, name = "mv_plane::refine", fields(method))
)]
#[inline]
pub fn refine<T: Pixel>(&mut self, method: SubpelMethod, plane: &mut [T]) {
if self.is_refined {
return;
}
if self.pel == Subpel::Full {
self.is_refined = true;
return;
}
let refine: [RefineFn<T>; 3] = match method {
SubpelMethod::Bilinear => [
refine_horizontal_bilinear,
refine_vertical_bilinear,
refine_diagonal_bilinear,
],
SubpelMethod::Bicubic => [
refine_horizontal_bicubic,
refine_vertical_bicubic,
refine_horizontal_bicubic,
],
SubpelMethod::Wiener => [
refine_horizontal_wiener,
refine_vertical_wiener,
refine_horizontal_wiener,
],
};
let mut src_offsets = [0; 3];
let mut dest_offsets = [0; 3];
match self.pel {
Subpel::Full => unreachable!(),
Subpel::Half => {
dest_offsets[0] = *semisafe_get(&self.subpel_window_offsets, 1);
dest_offsets[1] = *semisafe_get(&self.subpel_window_offsets, 2);
dest_offsets[2] = *semisafe_get(&self.subpel_window_offsets, 3);
src_offsets[0] = *semisafe_get(&self.subpel_window_offsets, 0);
src_offsets[1] = *semisafe_get(&self.subpel_window_offsets, 0);
if method == SubpelMethod::Bilinear {
src_offsets[2] = *semisafe_get(&self.subpel_window_offsets, 0);
} else {
src_offsets[2] = *semisafe_get(&self.subpel_window_offsets, 2);
}
}
Subpel::Quarter => {
dest_offsets[0] = *semisafe_get(&self.subpel_window_offsets, 2);
dest_offsets[1] = *semisafe_get(&self.subpel_window_offsets, 8);
dest_offsets[2] = *semisafe_get(&self.subpel_window_offsets, 10);
src_offsets[0] = *semisafe_get(&self.subpel_window_offsets, 0);
src_offsets[1] = *semisafe_get(&self.subpel_window_offsets, 0);
if method == SubpelMethod::Bilinear {
src_offsets[2] = *semisafe_get(&self.subpel_window_offsets, 0);
} else {
src_offsets[2] = *semisafe_get(&self.subpel_window_offsets, 8);
}
}
}
for i in 0..3 {
Self::refine_with_split(
plane,
*semisafe_get(&src_offsets, i),
*semisafe_get(&dest_offsets, i),
*semisafe_get(&refine, i),
self.stride,
self.padded_width,
self.padded_height,
self.bits_per_sample,
);
}
if self.pel == Subpel::Quarter {
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 0),
*semisafe_get(&self.subpel_window_offsets, 2),
*semisafe_get(&self.subpel_window_offsets, 1),
self.stride,
self.padded_width,
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 8),
*semisafe_get(&self.subpel_window_offsets, 10),
*semisafe_get(&self.subpel_window_offsets, 9),
self.stride,
self.padded_width,
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 0),
*semisafe_get(&self.subpel_window_offsets, 8),
*semisafe_get(&self.subpel_window_offsets, 4),
self.stride,
self.padded_width,
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 2),
*semisafe_get(&self.subpel_window_offsets, 10),
*semisafe_get(&self.subpel_window_offsets, 6),
self.stride,
self.padded_width,
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 4),
*semisafe_get(&self.subpel_window_offsets, 6),
*semisafe_get(&self.subpel_window_offsets, 5),
self.stride,
self.padded_width,
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 0) + 1,
*semisafe_get(&self.subpel_window_offsets, 2),
*semisafe_get(&self.subpel_window_offsets, 3),
self.stride,
unsafe { NonZeroUsize::new_unchecked(self.padded_width.get() - 1) },
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 8) + 1,
*semisafe_get(&self.subpel_window_offsets, 10),
*semisafe_get(&self.subpel_window_offsets, 11),
self.stride,
unsafe { NonZeroUsize::new_unchecked(self.padded_width.get() - 1) },
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 0) + self.stride.get(),
*semisafe_get(&self.subpel_window_offsets, 8),
*semisafe_get(&self.subpel_window_offsets, 12),
self.stride,
self.padded_width,
unsafe { NonZeroUsize::new_unchecked(self.padded_height.get() - 1) },
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 2) + self.stride.get(),
*semisafe_get(&self.subpel_window_offsets, 10),
*semisafe_get(&self.subpel_window_offsets, 14),
self.stride,
self.padded_width,
unsafe { NonZeroUsize::new_unchecked(self.padded_height.get() - 1) },
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 12),
*semisafe_get(&self.subpel_window_offsets, 14),
*semisafe_get(&self.subpel_window_offsets, 13),
self.stride,
self.padded_width,
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 4) + 1,
*semisafe_get(&self.subpel_window_offsets, 6),
*semisafe_get(&self.subpel_window_offsets, 7),
self.stride,
unsafe { NonZeroUsize::new_unchecked(self.padded_width.get() - 1) },
self.padded_height,
);
Self::average2_with_split(
plane,
*semisafe_get(&self.subpel_window_offsets, 12) + 1,
*semisafe_get(&self.subpel_window_offsets, 14),
*semisafe_get(&self.subpel_window_offsets, 15),
self.stride,
unsafe { NonZeroUsize::new_unchecked(self.padded_width.get() - 1) },
self.padded_height,
);
}
self.is_refined = true;
}
#[must_use]
#[inline]
pub fn get_absolute_pix_offset(&self, x: i32, y: i32) -> usize {
match self.pel {
Subpel::Full => self.get_absolute_pix_offset_pel1(x, y),
Subpel::Half => self.get_absolute_pix_offset_pel2(x, y),
Subpel::Quarter => self.get_absolute_pix_offset_pel4(x, y),
}
}
#[must_use]
#[inline]
pub fn get_absolute_pix_offset_pel<const PEL: usize>(&self, x: i32, y: i32) -> usize {
match PEL {
1 => self.get_absolute_pix_offset_pel1(x, y),
2 => self.get_absolute_pix_offset_pel2(x, y),
4 => self.get_absolute_pix_offset_pel4(x, y),
_ => unreachable!(),
}
}
#[must_use]
#[inline]
pub fn get_pel_pix_offset(&self, x: i32, y: i32) -> usize {
(*semisafe_get(&self.subpel_window_offsets, 0) as i32 + x + y * self.stride.get() as i32)
as usize
}
#[must_use]
#[inline]
pub fn get_absolute_pix_offset_pel1(&self, x: i32, y: i32) -> usize {
self.get_pel_pix_offset(x, y)
}
#[must_use]
#[inline]
pub fn get_absolute_pix_offset_pel2(&self, mut x: i32, mut y: i32) -> usize {
let idx = ((x & 1) | ((y & 1) << 1)) as usize;
x >>= 1;
y >>= 1;
(*semisafe_get(&self.subpel_window_offsets, idx) as i32 + x + y * self.stride.get() as i32)
as usize
}
#[must_use]
#[inline]
pub fn get_absolute_pix_offset_pel4(&self, mut x: i32, mut y: i32) -> usize {
let idx = ((x & 3) | ((y & 3) << 2)) as usize;
x >>= 2;
y >>= 2;
(*semisafe_get(&self.subpel_window_offsets, idx) as i32 + x + y * self.stride.get() as i32)
as usize
}
#[must_use]
#[inline]
pub fn get_pix_offset(&self, x: i32, y: i32) -> usize {
self.get_absolute_pix_offset(x + self.hpad_pel as i32, y + self.vpad_pel as i32)
}
pub(crate) fn update(&mut self, src_offset: usize, pitch: NonZeroUsize) {
self.stride = pitch;
self.offset_padding = self.stride.get() * self.vpad + self.hpad;
for i in 0..(self.pel as usize).pow(2) {
*semisafe_get_mut(&mut self.subpel_window_offsets, i) =
src_offset + i * self.stride.get() * self.padded_height.get();
}
self.reset_state();
}
const fn reset_state(&mut self) {
self.is_refined = false;
self.is_filled = false;
self.is_padded = false;
}
}
#[must_use]
#[inline]
pub fn plane_height_luma(
src_height: NonZeroUsize,
level: usize,
y_ratio_uv: NonZeroU8,
vpad: usize,
) -> NonZeroUsize {
let mut height = src_height.get();
let y_ratio_uv_val = y_ratio_uv.get() as usize;
for _i in 1..=level {
height = if vpad >= y_ratio_uv_val {
(height / y_ratio_uv_val + 1) / 2 * y_ratio_uv_val
} else {
((height / y_ratio_uv_val) / 2) * y_ratio_uv_val
};
}
debug_assert!(
height > 0,
"Calculated height must be non-zero. src_height: {}, level: {}, y_ratio_uv: {}, vpad: {}",
src_height,
level,
y_ratio_uv,
vpad
);
unsafe { NonZeroUsize::new_unchecked(height) }
}
#[must_use]
#[inline]
pub fn plane_width_luma(
src_width: NonZeroUsize,
level: usize,
x_ratio_uv: NonZeroU8,
hpad: usize,
) -> NonZeroUsize {
let mut width = src_width.get();
let x_ratio_uv_val = x_ratio_uv.get() as usize;
for _i in 1..=level {
width = if hpad >= x_ratio_uv_val {
(width / x_ratio_uv_val + 1) / 2 * x_ratio_uv_val
} else {
((width / x_ratio_uv_val) / 2) * x_ratio_uv_val
};
}
debug_assert!(
width > 0,
"Calculated width must be non-zero. src_width: {}, level: {}, x_ratio_uv: {}, hpad: {}",
src_width,
level,
x_ratio_uv,
hpad
);
unsafe { NonZeroUsize::new_unchecked(width) }
}
#[must_use]
#[inline]
pub fn plane_super_offset(
chroma: bool,
src_height: NonZeroUsize,
level: usize,
pel: Subpel,
vpad: usize,
plane_pitch: NonZeroUsize,
y_ratio_uv: NonZeroU8,
) -> usize {
let mut height;
let mut offset;
if level == 0 {
offset = 0;
} else {
let pel = u8::from(pel) as usize;
let plane_pitch_val = plane_pitch.get();
let src_height_val = src_height.get();
let y_ratio_uv_val = y_ratio_uv.get() as usize;
offset = pel * pel * plane_pitch_val * (src_height_val + vpad * 2);
for i in 1..level {
height = if chroma {
plane_height_luma(
src_height.saturating_mul(y_ratio_uv.into()),
i,
y_ratio_uv,
vpad * y_ratio_uv_val,
)
.get()
/ y_ratio_uv_val
} else {
plane_height_luma(src_height, i, y_ratio_uv, vpad).get()
};
offset += plane_pitch_val * (height + vpad * 2);
}
}
offset
}