#![allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
reason = "every converted sample is clamped to 0..=255 before the u8 cast; \
this reproduces the reference encoder's byte handling exactly"
)]
use crate::lossy::prelude::*;
const YUV_FIX: u32 = 16;
const YUV_HALF: i32 = 1 << (YUV_FIX - 1);
const UV_ROUND: i32 = YUV_HALF << 2;
const UV_OFFSET: i32 = 128 << (YUV_FIX + 2);
pub(crate) struct SourceYuv {
pub(crate) y: Vec<u8>,
pub(crate) u: Vec<u8>,
pub(crate) v: Vec<u8>,
pub(crate) mb_w: usize,
pub(crate) mb_h: usize,
}
impl SourceYuv {
pub(crate) const fn y_stride(&self) -> usize {
self.mb_w * 16
}
pub(crate) const fn uv_stride(&self) -> usize {
self.mb_w * 8
}
}
fn rgb_to_y(r: i32, g: i32, b: i32) -> u8 {
let luma = 16839 * r + 33059 * g + 6420 * b;
((luma + YUV_HALF + (16 << YUV_FIX)) >> YUV_FIX).clamp(0, 255) as u8
}
fn clip_uv_sum(acc: i32) -> u8 {
((acc + UV_ROUND + UV_OFFSET) >> (YUV_FIX + 2)).clamp(0, 255) as u8
}
fn rgb_to_u(sr: i32, sg: i32, sb: i32) -> u8 {
clip_uv_sum(-9719 * sr - 19081 * sg + 28800 * sb)
}
fn rgb_to_v(sr: i32, sg: i32, sb: i32) -> u8 {
clip_uv_sum(28800 * sr - 24116 * sg - 4684 * sb)
}
#[must_use]
#[allow(
clippy::many_single_char_names,
reason = "r/g/b are the conventional color-channel names and x/y the pixel \
coordinates; longer names would obscure the conversion"
)]
pub(crate) fn from_rgba(rgba: &[u8], width: usize, height: usize) -> SourceYuv {
let mb_w = width.div_ceil(16);
let mb_h = height.div_ceil(16);
let y_stride = mb_w * 16;
let uv_stride = mb_w * 8;
let (y_h, uv_h) = (mb_h * 16, mb_h * 8);
let rgb = |x: usize, y: usize| -> (i32, i32, i32) {
let cx = x.min(width - 1);
let cy = y.min(height - 1);
let i = (cy * width + cx) * 4;
(
i32::from(rgba[i]),
i32::from(rgba[i + 1]),
i32::from(rgba[i + 2]),
)
};
let mut y = vec![0u8; y_stride * y_h];
for py in 0..y_h {
for px in 0..y_stride {
let (r, g, b) = rgb(px, py);
y[py * y_stride + px] = rgb_to_y(r, g, b);
}
}
let mut u = vec![0u8; uv_stride * uv_h];
let mut v = vec![0u8; uv_stride * uv_h];
for cy in 0..uv_h {
for cx in 0..uv_stride {
let mut sr = 0;
let mut sg = 0;
let mut sb = 0;
for dy in 0..2 {
for dx in 0..2 {
let (r, g, b) = rgb(2 * cx + dx, 2 * cy + dy);
sr += r;
sg += g;
sb += b;
}
}
u[cy * uv_stride + cx] = rgb_to_u(sr, sg, sb);
v[cy * uv_stride + cx] = rgb_to_v(sr, sg, sb);
}
}
SourceYuv {
y,
u,
v,
mb_w,
mb_h,
}
}
#[cfg(test)]
mod tests {
use super::from_rgba;
fn solid(width: usize, height: usize, rgb: [u8; 3]) -> Vec<u8> {
let mut buf = Vec::with_capacity(width * height * 4);
for _ in 0..width * height {
buf.extend_from_slice(&[rgb[0], rgb[1], rgb[2], 0xff]);
}
buf
}
#[test]
fn neutral_gray_maps_to_studio_swing() {
let src = from_rgba(&solid(16, 16, [128, 128, 128]), 16, 16);
assert!(src.y.iter().all(|&p| p == 126), "gray luma is 126");
assert!(src.u.iter().all(|&p| p == 128), "gray U centered");
assert!(src.v.iter().all(|&p| p == 128), "gray V centered");
}
#[test]
fn black_and_white_hit_the_studio_pedestal_and_ceiling() {
let black = from_rgba(&solid(16, 16, [0, 0, 0]), 16, 16);
assert!(black.y.iter().all(|&p| p == 16), "black luma is 16");
let white = from_rgba(&solid(16, 16, [255, 255, 255]), 16, 16);
assert!(white.y.iter().all(|&p| p == 235), "white luma is 235");
assert!(black.u.iter().all(|&p| p == 128) && white.u.iter().all(|&p| p == 128));
}
#[test]
fn pure_primaries_push_chroma_off_center() {
let blue = from_rgba(&solid(16, 16, [0, 0, 255]), 16, 16);
assert!(blue.u[0] > 200, "blue U high, got {}", blue.u[0]);
let red = from_rgba(&solid(16, 16, [255, 0, 0]), 16, 16);
assert!(red.v[0] > 200, "red V high, got {}", red.v[0]);
}
#[test]
fn chroma_rounding_bias_lands_on_the_nearest_level() {
let src = from_rgba(&solid(16, 16, [0, 0, 2]), 16, 16);
assert!(
src.u.iter().all(|&p| p == 129),
"dark-blue U rounds to 129, got {}",
src.u[0]
);
}
#[test]
fn planes_are_padded_to_whole_macroblocks() {
let src = from_rgba(&solid(17, 13, [64, 200, 32]), 17, 13);
assert_eq!((src.mb_w, src.mb_h), (2, 1));
assert_eq!(src.y.len(), 32 * 16);
assert_eq!(src.u.len(), 16 * 8);
assert_eq!(src.y_stride(), 32);
assert_eq!(src.uv_stride(), 16);
let first = src.y[0];
assert!(src.y.iter().all(|&p| p == first), "edge-extended uniformly");
}
}