#![forbid(unsafe_code)]
#![allow(dead_code)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_sign_loss)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::too_many_arguments)]
use crate::error::{CodecError, CodecResult};
use crate::frame::VideoFrame;
use rayon::prelude::*;
use std::sync::Arc;
pub const VP9_SB_SIZE: u32 = 64;
pub const VP9_MAX_TILE_COLS: u32 = 64;
pub const VP9_MAX_TILE_ROWS: u32 = 4;
#[derive(Clone, Debug)]
pub struct Vp9TileConfig {
pub tile_cols_log2: u32,
pub tile_rows_log2: u32,
pub threads: usize,
pub enable_deblock_sync: bool,
pub base_qindex: u8,
}
impl Default for Vp9TileConfig {
fn default() -> Self {
Self {
tile_cols_log2: 0,
tile_rows_log2: 0,
threads: 0,
enable_deblock_sync: true,
base_qindex: 128,
}
}
}
impl Vp9TileConfig {
pub fn new(tile_cols_log2: u32, tile_rows_log2: u32, threads: usize) -> CodecResult<Self> {
if tile_cols_log2 > 6 {
return Err(CodecError::InvalidParameter(
"tile_cols_log2 must be 0-6".to_string(),
));
}
if tile_rows_log2 > 2 {
return Err(CodecError::InvalidParameter(
"tile_rows_log2 must be 0-2 for VP9 compliance".to_string(),
));
}
Ok(Self {
tile_cols_log2,
tile_rows_log2,
threads,
enable_deblock_sync: true,
base_qindex: 128,
})
}
#[must_use]
pub fn auto(width: u32, height: u32, threads: usize) -> Self {
let max_sb_cols = width.div_ceil(VP9_SB_SIZE);
let max_sb_rows = height.div_ceil(VP9_SB_SIZE);
let mut cols_log2 = 0u32;
while cols_log2 < 6 && (1u32 << (cols_log2 + 1)) <= max_sb_cols {
cols_log2 += 1;
}
if threads > 1 {
let thread_log2 = (threads as f32).log2().floor() as u32;
cols_log2 = cols_log2.min(thread_log2);
}
let mut rows_log2 = 0u32;
while rows_log2 < 2 && (1u32 << (rows_log2 + 1)) <= max_sb_rows {
rows_log2 += 1;
}
Self {
tile_cols_log2: cols_log2,
tile_rows_log2: rows_log2,
threads,
enable_deblock_sync: true,
base_qindex: 128,
}
}
#[must_use]
pub fn tile_cols(&self) -> u32 {
1u32 << self.tile_cols_log2
}
#[must_use]
pub fn tile_rows(&self) -> u32 {
1u32 << self.tile_rows_log2
}
#[must_use]
pub fn tile_count(&self) -> u32 {
self.tile_cols() * self.tile_rows()
}
#[must_use]
pub fn thread_count(&self) -> usize {
if self.threads == 0 {
rayon::current_num_threads()
} else {
self.threads
}
}
pub fn validate(&self) -> CodecResult<()> {
if self.tile_cols_log2 > 6 {
return Err(CodecError::InvalidParameter(
"tile_cols_log2 must be 0-6".to_string(),
));
}
if self.tile_rows_log2 > 2 {
return Err(CodecError::InvalidParameter(
"tile_rows_log2 must be 0-2".to_string(),
));
}
Ok(())
}
}
#[derive(Clone, Debug)]
pub struct Vp9TileRegion {
pub col: u32,
pub row: u32,
pub x: u32,
pub y: u32,
pub width: u32,
pub height: u32,
pub index: u32,
pub sb_col_start: u32,
pub sb_col_end: u32,
pub sb_row_start: u32,
pub sb_row_end: u32,
}
impl Vp9TileRegion {
#[must_use]
#[allow(clippy::too_many_arguments)]
pub fn new(
col: u32,
row: u32,
x: u32,
y: u32,
width: u32,
height: u32,
tile_cols: u32,
) -> Self {
let sb_col_start = x / VP9_SB_SIZE;
let sb_col_end = (x + width).div_ceil(VP9_SB_SIZE);
let sb_row_start = y / VP9_SB_SIZE;
let sb_row_end = (y + height).div_ceil(VP9_SB_SIZE);
Self {
col,
row,
x,
y,
width,
height,
index: row * tile_cols + col,
sb_col_start,
sb_col_end,
sb_row_start,
sb_row_end,
}
}
#[must_use]
pub fn sb_cols(&self) -> u32 {
self.sb_col_end - self.sb_col_start
}
#[must_use]
pub fn sb_rows(&self) -> u32 {
self.sb_row_end - self.sb_row_start
}
#[must_use]
pub fn is_left_edge(&self) -> bool {
self.col == 0
}
#[must_use]
pub fn is_top_edge(&self) -> bool {
self.row == 0
}
}
#[derive(Debug)]
pub struct Vp9TileFrameSplitter {
config: Vp9TileConfig,
frame_width: u32,
frame_height: u32,
regions: Vec<Vp9TileRegion>,
}
impl Vp9TileFrameSplitter {
pub fn new(config: Vp9TileConfig, frame_width: u32, frame_height: u32) -> CodecResult<Self> {
config.validate()?;
let mut splitter = Self {
config,
frame_width,
frame_height,
regions: Vec::new(),
};
splitter.compute_regions();
Ok(splitter)
}
fn compute_regions(&mut self) {
self.regions.clear();
let tile_cols = self.config.tile_cols();
let tile_rows = self.config.tile_rows();
let sb_cols = self.frame_width.div_ceil(VP9_SB_SIZE);
let sb_rows = self.frame_height.div_ceil(VP9_SB_SIZE);
for row in 0..tile_rows {
for col in 0..tile_cols {
let sb_col_start = (col * sb_cols) / tile_cols;
let sb_col_end = ((col + 1) * sb_cols) / tile_cols;
let sb_row_start = (row * sb_rows) / tile_rows;
let sb_row_end = ((row + 1) * sb_rows) / tile_rows;
let x = sb_col_start * VP9_SB_SIZE;
let y = sb_row_start * VP9_SB_SIZE;
let width = (sb_col_end * VP9_SB_SIZE).min(self.frame_width) - x;
let height = (sb_row_end * VP9_SB_SIZE).min(self.frame_height) - y;
self.regions
.push(Vp9TileRegion::new(col, row, x, y, width, height, tile_cols));
}
}
}
#[must_use]
pub fn regions(&self) -> &[Vp9TileRegion] {
&self.regions
}
#[must_use]
pub fn tile_count(&self) -> usize {
self.regions.len()
}
#[must_use]
pub fn region(&self, index: usize) -> Option<&Vp9TileRegion> {
self.regions.get(index)
}
}
#[derive(Clone, Debug)]
pub struct Vp9EncodedTile {
pub region: Vp9TileRegion,
pub data: Vec<u8>,
}
impl Vp9EncodedTile {
#[must_use]
pub fn index(&self) -> u32 {
self.region.index
}
#[must_use]
pub fn encoded_size(&self) -> usize {
self.data.len()
}
}
#[derive(Clone, Debug)]
pub struct Vp9TileEncoder {
region: Vp9TileRegion,
base_qindex: u8,
is_keyframe: bool,
enable_deblock: bool,
}
impl Vp9TileEncoder {
#[must_use]
pub fn new(
region: Vp9TileRegion,
base_qindex: u8,
is_keyframe: bool,
enable_deblock: bool,
) -> Self {
Self {
region,
base_qindex,
is_keyframe,
enable_deblock,
}
}
pub fn encode(&self, frame: &VideoFrame) -> CodecResult<Vp9EncodedTile> {
if self.region.x + self.region.width > frame.width
|| self.region.y + self.region.height > frame.height
{
return Err(CodecError::InvalidParameter(format!(
"VP9 tile region [{},{})+[{},{}] exceeds frame {}x{}",
self.region.x,
self.region.y,
self.region.width,
self.region.height,
frame.width,
frame.height,
)));
}
let tile_data = self.extract_luma_data(frame)?;
let encoded = self.encode_tile_bitstream(&tile_data);
Ok(Vp9EncodedTile {
region: self.region.clone(),
data: encoded,
})
}
fn extract_luma_data(&self, frame: &VideoFrame) -> CodecResult<Vec<u8>> {
let mut pixels = Vec::with_capacity((self.region.width * self.region.height) as usize);
if let Some(y_plane) = frame.planes.first() {
for row in self.region.y..(self.region.y + self.region.height) {
let row_start = (row as usize * y_plane.stride) + self.region.x as usize;
let row_end = row_start + self.region.width as usize;
if row_end <= y_plane.data.len() {
pixels.extend_from_slice(&y_plane.data[row_start..row_end]);
} else {
let available = y_plane.data.len().saturating_sub(row_start);
pixels.extend_from_slice(&y_plane.data[row_start..row_start + available]);
pixels.resize(pixels.len() + (self.region.width as usize - available), 0);
}
}
}
Ok(pixels)
}
fn encode_tile_bitstream(&self, _luma: &[u8]) -> Vec<u8> {
let mut bits = Vec::new();
let flags = (u8::from(self.is_keyframe) << 7)
| (u8::from(self.enable_deblock) << 6)
| (self.region.col as u8 & 0x03);
bits.push(flags);
bits.push(self.base_qindex);
bits.extend_from_slice(&self.region.width.to_le_bytes());
bits.extend_from_slice(&self.region.height.to_le_bytes());
bits.extend_from_slice(&self.region.sb_col_start.to_le_bytes());
bits.extend_from_slice(&self.region.sb_row_start.to_le_bytes());
let sb_count = (self.region.sb_cols() * self.region.sb_rows()) as usize;
bits.resize(bits.len() + sb_count.max(1), 0x00);
bits
}
#[must_use]
pub fn region(&self) -> &Vp9TileRegion {
&self.region
}
}
#[derive(Debug)]
pub struct Vp9DeblockSync {
tile_cols: u32,
sb_rows: u32,
completed_sb_rows: Vec<std::sync::atomic::AtomicU32>,
}
impl Vp9DeblockSync {
#[must_use]
pub fn new(tile_cols: u32, sb_rows: u32) -> Self {
let completed_sb_rows = (0..tile_cols)
.map(|_| std::sync::atomic::AtomicU32::new(0))
.collect();
Self {
tile_cols,
sb_rows,
completed_sb_rows,
}
}
pub fn mark_complete(&self, tile_col: u32, sb_row_count: u32) {
if tile_col < self.tile_cols {
self.completed_sb_rows[tile_col as usize]
.store(sb_row_count, std::sync::atomic::Ordering::Release);
}
}
#[must_use]
pub fn completed_rows(&self, tile_col: u32) -> u32 {
if tile_col < self.tile_cols {
self.completed_sb_rows[tile_col as usize].load(std::sync::atomic::Ordering::Acquire)
} else {
self.sb_rows
}
}
#[must_use]
pub fn can_deblock(&self, tile_col: u32, sb_row: u32) -> bool {
if tile_col == 0 {
return true; }
self.completed_rows(tile_col - 1) > sb_row
}
#[must_use]
pub fn sb_rows(&self) -> u32 {
self.sb_rows
}
}
#[derive(Debug)]
pub struct Vp9FrameTileEncoder {
config: Arc<Vp9TileConfig>,
splitter: Vp9TileFrameSplitter,
frame_width: u32,
frame_height: u32,
}
impl Vp9FrameTileEncoder {
pub fn new(config: Vp9TileConfig, frame_width: u32, frame_height: u32) -> CodecResult<Self> {
let splitter = Vp9TileFrameSplitter::new(config.clone(), frame_width, frame_height)?;
Ok(Self {
config: Arc::new(config),
splitter,
frame_width,
frame_height,
})
}
pub fn encode_frame(
&self,
frame: &VideoFrame,
is_keyframe: bool,
) -> CodecResult<Vec<Vp9EncodedTile>> {
if frame.width != self.frame_width || frame.height != self.frame_height {
return Err(CodecError::InvalidParameter(format!(
"Frame {}x{} does not match encoder {}x{}",
frame.width, frame.height, self.frame_width, self.frame_height,
)));
}
let regions = self.splitter.regions();
let base_qindex = self.config.base_qindex;
let enable_deblock = self.config.enable_deblock_sync;
let results: Vec<CodecResult<Vp9EncodedTile>> = regions
.par_iter()
.map(|region| {
let encoder =
Vp9TileEncoder::new(region.clone(), base_qindex, is_keyframe, enable_deblock);
encoder.encode(frame)
})
.collect();
let mut tiles = Vec::with_capacity(results.len());
for result in results {
tiles.push(result?);
}
tiles.sort_by_key(Vp9EncodedTile::index);
if enable_deblock {
let sb_rows = self.frame_height.div_ceil(VP9_SB_SIZE);
let sync = Vp9DeblockSync::new(self.config.tile_cols(), sb_rows);
for col in 0..self.config.tile_cols() {
sync.mark_complete(col, sb_rows);
}
}
Ok(tiles)
}
pub fn assemble_frame(&self, tiles: &[Vp9EncodedTile]) -> CodecResult<Vec<u8>> {
if tiles.is_empty() {
return Err(CodecError::InvalidParameter(
"Cannot assemble empty tile list".to_string(),
));
}
let mut out = Vec::new();
let header = 0x80u8
| ((self.config.tile_cols_log2 as u8 & 0x07) << 4)
| ((self.config.tile_rows_log2 as u8 & 0x07) << 1);
out.push(header);
for (i, tile) in tiles.iter().enumerate() {
let is_last = i == tiles.len() - 1;
if !is_last {
let size = tile.data.len() as u32;
out.extend_from_slice(&size.to_le_bytes());
}
out.extend_from_slice(&tile.data);
}
Ok(out)
}
#[must_use]
pub fn config(&self) -> &Vp9TileConfig {
&self.config
}
#[must_use]
pub fn tile_count(&self) -> usize {
self.splitter.tile_count()
}
#[must_use]
pub fn regions(&self) -> &[Vp9TileRegion] {
self.splitter.regions()
}
}
#[cfg(test)]
mod tests {
use super::*;
use oximedia_core::PixelFormat;
fn make_frame(w: u32, h: u32) -> VideoFrame {
let mut f = VideoFrame::new(PixelFormat::Yuv420p, w, h);
f.allocate();
f
}
#[test]
fn test_tile_config_default() {
let cfg = Vp9TileConfig::default();
assert_eq!(cfg.tile_cols(), 1);
assert_eq!(cfg.tile_rows(), 1);
assert_eq!(cfg.tile_count(), 1);
}
#[test]
fn test_tile_config_new_valid() {
let cfg = Vp9TileConfig::new(2, 1, 4).expect("should succeed");
assert_eq!(cfg.tile_cols(), 4);
assert_eq!(cfg.tile_rows(), 2);
assert_eq!(cfg.tile_count(), 8);
}
#[test]
fn test_tile_config_new_invalid_cols() {
assert!(Vp9TileConfig::new(7, 0, 1).is_err());
}
#[test]
fn test_tile_config_new_invalid_rows() {
assert!(Vp9TileConfig::new(0, 3, 1).is_err());
}
#[test]
fn test_tile_config_auto() {
let cfg = Vp9TileConfig::auto(1920, 1080, 4);
assert!(cfg.tile_count() >= 1);
assert!(cfg.tile_count() <= VP9_MAX_TILE_COLS * (1u32 << 2));
}
#[test]
fn test_splitter_single_tile() {
let cfg = Vp9TileConfig::default();
let s = Vp9TileFrameSplitter::new(cfg, 1920, 1080).expect("should succeed");
assert_eq!(s.tile_count(), 1);
let r = s.region(0).expect("should succeed");
assert_eq!(r.x, 0);
assert_eq!(r.y, 0);
assert_eq!(r.width, 1920);
assert_eq!(r.height, 1080);
}
#[test]
fn test_splitter_2x1_tiles() {
let cfg = Vp9TileConfig::new(1, 0, 2).expect("should succeed"); let s = Vp9TileFrameSplitter::new(cfg, 1920, 1080).expect("should succeed");
assert_eq!(s.tile_count(), 2);
let r0 = s.region(0).expect("should succeed");
let r1 = s.region(1).expect("should succeed");
assert_eq!(r0.x, 0);
assert!(r1.x > 0);
assert_eq!(r0.y, 0);
assert_eq!(r1.y, 0);
assert_eq!(r0.width + r1.width, 1920);
}
#[test]
fn test_splitter_2x2_tiles() {
let cfg = Vp9TileConfig::new(1, 1, 4).expect("should succeed"); let s = Vp9TileFrameSplitter::new(cfg, 1920, 1088).expect("should succeed");
assert_eq!(s.tile_count(), 4);
let top_height = s.region(0).expect("should succeed").height;
let bot_height = s.region(2).expect("should succeed").height;
assert_eq!(top_height + bot_height, 1088);
}
#[test]
fn test_deblock_sync_left_edge() {
let sync = Vp9DeblockSync::new(4, 16);
assert!(sync.can_deblock(0, 0)); }
#[test]
fn test_deblock_sync_dependency() {
let sync = Vp9DeblockSync::new(4, 16);
assert!(!sync.can_deblock(1, 0)); sync.mark_complete(0, 3);
assert!(sync.can_deblock(1, 0)); assert!(sync.can_deblock(1, 2)); assert!(!sync.can_deblock(1, 3)); }
#[test]
fn test_single_tile_encode() {
let region = Vp9TileRegion::new(0, 0, 0, 0, 320, 240, 1);
let enc = Vp9TileEncoder::new(region, 128, true, true);
let frame = make_frame(1920, 1080);
let result = enc.encode(&frame);
assert!(result.is_ok());
let tile = result.expect("should succeed");
assert!(tile.encoded_size() > 0);
}
#[test]
fn test_single_tile_out_of_bounds() {
let region = Vp9TileRegion::new(0, 0, 1900, 1070, 200, 200, 1);
let enc = Vp9TileEncoder::new(region, 128, false, false);
let frame = make_frame(1920, 1080);
assert!(enc.encode(&frame).is_err());
}
#[test]
fn test_frame_encoder_single_tile() {
let cfg = Vp9TileConfig::default();
let enc = Vp9FrameTileEncoder::new(cfg, 1920, 1080).expect("should succeed");
let frame = make_frame(1920, 1080);
let tiles = enc.encode_frame(&frame, true).expect("should succeed");
assert_eq!(tiles.len(), 1);
}
#[test]
fn test_frame_encoder_4x2_tiles() {
let cfg = Vp9TileConfig::new(2, 1, 8).expect("should succeed"); let enc = Vp9FrameTileEncoder::new(cfg, 1920, 1088).expect("should succeed");
let frame = make_frame(1920, 1088);
let tiles = enc.encode_frame(&frame, false).expect("should succeed");
assert_eq!(tiles.len(), 8);
for (i, tile) in tiles.iter().enumerate() {
assert_eq!(tile.index() as usize, i);
}
}
#[test]
fn test_frame_encoder_wrong_dimensions() {
let cfg = Vp9TileConfig::default();
let enc = Vp9FrameTileEncoder::new(cfg, 1920, 1080).expect("should succeed");
let frame = make_frame(1280, 720);
assert!(enc.encode_frame(&frame, true).is_err());
}
#[test]
fn test_assemble_frame() {
let cfg = Vp9TileConfig::new(1, 0, 2).expect("should succeed"); let enc = Vp9FrameTileEncoder::new(cfg, 1920, 1080).expect("should succeed");
let frame = make_frame(1920, 1080);
let tiles = enc.encode_frame(&frame, true).expect("should succeed");
let assembled = enc.assemble_frame(&tiles).expect("should succeed");
assert!(!assembled.is_empty());
assert_eq!(assembled[0] & 0x80, 0x80);
}
#[test]
fn test_assemble_empty_error() {
let cfg = Vp9TileConfig::default();
let enc = Vp9FrameTileEncoder::new(cfg, 1920, 1080).expect("should succeed");
assert!(enc.assemble_frame(&[]).is_err());
}
#[test]
fn test_tile_region_sb_alignment() {
let region = Vp9TileRegion::new(1, 0, 64, 0, 128, 64, 2);
assert_eq!(region.sb_col_start, 1);
assert_eq!(region.sb_col_end, 3);
assert_eq!(region.sb_cols(), 2);
assert_eq!(region.sb_rows(), 1);
}
}