#![cfg_attr(not(feature = "std"), no_std)]
#[allow(unused_imports)]
use rusty_h264_common::once::OnceLock;
extern crate alloc;
#[cfg(test)]
extern crate std;
#[cfg(all(not(feature = "std"), not(feature = "libm")))]
compile_error!(
"rusty_h264-encoder without `std` needs the `libm` feature for its floating-point math"
);
#[cfg(not(feature = "std"))]
#[allow(unused_macros)]
macro_rules! eprintln {
($($t:tt)*) => {{
let _ = ::core::format_args!($($t)*);
}};
}
#[cfg(not(feature = "std"))]
#[allow(unused_macros)]
macro_rules! println {
($($t:tt)*) => {{
let _ = ::core::format_args!($($t)*);
}};
}
#[cfg(not(feature = "std"))]
macro_rules! thread_local {
() => {};
($(#[$m:meta])* $vis:vis static $name:ident: $ty:ty = const { $init:expr }; $($rest:tt)*) => {
$(#[$m])* #[allow(non_camel_case_types)] $vis struct $name;
impl $name {
#[allow(dead_code)]
pub fn with<R>(&self, f: impl FnOnce(&$ty) -> R) -> R {
let v: $ty = $init;
f(&v)
}
}
thread_local!($($rest)*);
};
($(#[$m:meta])* $vis:vis static $name:ident: $ty:ty = $init:expr; $($rest:tt)*) => {
$(#[$m])* #[allow(non_camel_case_types)] $vis struct $name;
impl $name {
#[allow(dead_code)]
pub fn with<R>(&self, f: impl FnOnce(&$ty) -> R) -> R {
let v: $ty = $init;
f(&v)
}
}
thread_local!($($rest)*);
};
}
#[allow(unused_imports)]
use alloc::string::{String, ToString};
#[allow(unused_imports)]
use alloc::vec;
#[allow(unused_imports)]
use alloc::vec::Vec;
#[allow(unused_imports)]
use rusty_h264_common::fmath::{F32Ext as _, F64Ext as _};
pub mod bitacct;
mod cabac;
mod config;
mod fastmath;
mod lookahead;
pub mod mb16;
mod mbtree;
#[cfg(feature = "prometheus-telemetry")]
pub mod telemetry;
#[cfg(feature = "prometheus-telemetry")]
pub mod prometheus_telemetry {
pub use crate::telemetry::{enable, p_zero_q8, take, CabacBin, SliceTap};
#[allow(unused_imports)]
use alloc::{
boxed::Box,
format,
string::{String, ToString},
vec,
vec::Vec,
};
#[allow(unused_imports)]
use rusty_h264_common::once::OnceLock;
}
pub fn mbtree_satd_calls() -> u64 {
mbtree::SATD_CALLS.load(core::sync::atomic::Ordering::Relaxed)
}
pub fn mbtree_satd_reset() {
mbtree::SATD_CALLS.store(0, core::sync::atomic::Ordering::Relaxed)
}
pub use mbtree::gopstats;
pub fn gate_census() -> Vec<(u64, u64)> {
signals::census::snapshot().to_vec()
}
pub fn gate_census_by_t8() -> [Vec<(u64, u64)>; 2] {
let s = signals::census::snapshot_by_t8();
[s[0].to_vec(), s[1].to_vec()]
}
fn estimate_luma_weights(
cfg: &EncoderConfig,
frame: &YuvPlanes<'_>,
refs: &[RefFrame],
) -> Vec<(i32, i32)> {
let cw = cfg.mb_width() * 16;
let (w, h) = (frame.width.min(cw), frame.height);
let mut cur: Vec<u8> = Vec::new();
let mut sc = 0u64;
let mut y = 0;
while y < h {
let row = &frame.y[y * frame.width..][..w];
let mut x = 0;
while x < w {
let p = row[x];
sc += p as u64;
cur.push(p);
x += 4;
}
y += 4;
}
let n = cur.len() as u64;
let mut rbuf: Vec<u8> = Vec::with_capacity(cur.len());
refs.iter()
.map(|r| {
rbuf.clear();
let mut sr = 0u64;
let mut y = 0;
while y < h {
let row = &r.y[y * cw..][..w]; let mut x = 0;
while x < w {
let p = row[x];
sr += p as u64;
rbuf.push(p);
x += 4;
}
y += 4;
}
if n == 0 || sr == 0 {
return (64, 0);
}
let (mc, mr) = (sc as f64 / n as f64, sr as f64 / n as f64);
let lw = (rusty_h264_common::fmath::round(mc * 64.0 / mr) as i32).clamp(1, 127);
let lo = (rusty_h264_common::fmath::round(mc - (lw as f64) * mr / 64.0) as i32)
.clamp(-128, 127);
if (lw, lo) == (64, 0) {
return (64, 0);
}
let (mut sad_u, mut sad_w) = (0u64, 0u64);
for (&c, &rr) in cur.iter().zip(rbuf.iter()) {
let (c, rr) = (c as i32, rr as i32);
let rw = (((rr * lw + 32) >> 6) + lo).clamp(0, 255);
sad_u += c.abs_diff(rr) as u64;
sad_w += c.abs_diff(rw) as u64;
}
if sad_w * 100 < sad_u * 99 {
(lw, lo)
} else {
(64, 0)
}
})
.collect()
}
pub fn bframes_gate_signals(
cfg: &EncoderConfig,
frames: &[YuvFrame],
) -> (f64, f64, f64, f64, bool, bool) {
let (w, h) = (cfg.width, cfg.height);
let bi = gop_bi_residual(frames, w, h, 1);
if frames.len() < 2 || w % 16 != 0 || h % 16 != 0 {
return (bi, f64::INFINITY, 0.0, 0.0, false, false);
}
let sig = signals::FrameSignals::new(&frames[1].y, w, w / 16, h / 16, Some(&frames[0].y));
let (mg, dc) = sig.mgain_dc();
(
bi,
sig.gmc_residual(),
mg,
dc,
sig.is_screen(),
sig.grain_signature(),
)
}
pub fn scene_cut_ratios(cfg: &EncoderConfig, frames: &[YuvFrame]) -> Vec<f64> {
lookahead::all_pair_ratios(cfg, frames)
}
pub fn gate_work() -> Vec<u64> {
signals::census::work_snapshot().to_vec()
}
pub fn gate_work_names() -> &'static [&'static str] {
&signals::census::WORK_NAMES
}
pub fn gate_census_names() -> &'static [&'static str] {
&signals::census::NAMES
}
pub fn gate_census_reset() {
signals::census::reset()
}
#[cfg(not(feature = "std"))]
pub fn gate_census_dump_csv() {}
#[cfg(feature = "std")]
pub fn gate_census_dump_csv() {
use core::fmt::Write as _;
use std::io::Write as _;
let Ok(path) = std::env::var("RFF_CENSUS_CSV") else {
return;
};
let snap = signals::census::snapshot();
let mut s = String::from("gate,fired,seen\n");
for (i, name) in signals::census::NAMES.iter().enumerate() {
let _ = writeln!(s, "{name},{},{}", snap[i].0, snap[i].1);
}
if let Ok(mut f) = std::fs::File::create(&path) {
let _ = f.write_all(s.as_bytes());
}
}
mod mvd_cost_tab;
mod params;
mod rc;
mod signals;
mod slice;
#[cfg(feature = "std")]
pub use crate::mb16::EXT_MV;
pub use crate::mb16::{ME_PROBE, MVCMP, MVCMP_FRAME};
#[doc(hidden)]
pub mod cabac_enc_test {
pub use crate::cabac::CabacEncoder;
pub use crate::mb16::b_part_mb_type;
pub use crate::mb16::cb_cbp;
pub use crate::mb16::cb_mb_qp_delta;
pub use crate::mb16::cb_mb_type_b;
pub use crate::mb16::cb_ref_idx;
#[allow(unused_imports)]
use alloc::{
boxed::Box,
format,
string::{String, ToString},
vec,
vec::Vec,
};
#[allow(unused_imports)]
use rusty_h264_common::once::OnceLock;
}
pub use config::{EncoderConfig, LookaheadMode, MemoryEstimate, Preset};
pub use params::{Pps, Sps};
pub use rc::RateControl;
use rusty_h264_common::{
BitWriter, ChromaFormat, NalUnit, NalUnitType, Profile, YuvFrame, YuvPlanes,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EncodeError {
Unsupported(&'static str),
FrameMismatch,
BufferTooSmall {
needed: usize,
},
}
impl core::fmt::Display for EncodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
EncodeError::Unsupported(s) => write!(f, "unsupported: {s}"),
EncodeError::FrameMismatch => write!(f, "frame dimensions do not match encoder config"),
EncodeError::BufferTooSmall { needed } => {
write!(
f,
"output buffer too small: the access unit needs {needed} bytes"
)
}
}
}
}
impl core::error::Error for EncodeError {}
enum Sink<'a> {
Vec(&'a mut Vec<u8>),
Slice {
buf: &'a mut [u8],
len: usize,
needed: usize,
},
}
impl Sink<'_> {
fn nal(&mut self, ref_idc: u8, nal_type: NalUnitType, rbsp: &[u8]) {
let header = rusty_h264_common::nal::nal_header_byte(ref_idc, nal_type);
match self {
Sink::Vec(v) => {
v.extend_from_slice(&[0, 0, 0, 1, header]);
rusty_h264_common::nal::emulation_prevent_into(rbsp, v);
}
Sink::Slice { buf, len, needed } => {
let mut put = |b: u8| {
if *len < buf.len() {
buf[*len] = b;
*len += 1;
}
*needed += 1;
};
for b in [0, 0, 0, 1, header] {
put(b);
}
rusty_h264_common::nal::emulation_prevent_with(rbsp, put);
}
}
}
}
fn copy_out(au: &[u8], out: &mut [u8]) -> Result<usize, EncodeError> {
if out.len() < au.len() {
return Err(EncodeError::BufferTooSmall { needed: au.len() });
}
out[..au.len()].copy_from_slice(au);
Ok(au.len())
}
#[derive(Debug)]
pub struct Encoder {
cfg: EncoderConfig,
#[allow(dead_code)]
sps: Sps,
#[allow(dead_code)]
pps: Pps,
frame_index: u32,
next_frame_num: u32,
gop_index: u32,
refs: Vec<RefFrame>,
rc: Option<RateControl>,
pending_qpo: Option<Vec<i32>>,
pending_aq_probe: Option<YuvFrame>,
la_queue: Vec<YuvFrame>,
since_idr: u32,
force_idr: bool,
keyframe_requested: bool,
last_src: Option<YuvFrame>,
last_prep: Option<mbtree::PairPrep>,
cut_hist: [f64; 2],
gather: Option<YuvFrame>,
bw: BitWriter,
sps_nal: NalUnit,
pps_nal: NalUnit,
}
#[cfg(feature = "std")]
fn panicking() -> bool {
std::thread::panicking()
}
#[cfg(not(feature = "std"))]
fn panicking() -> bool {
false
}
impl Drop for Encoder {
fn drop(&mut self) {
debug_assert!(
self.la_queue.is_empty() || panicking(),
"Encoder dropped with {} frame(s) still in the lookahead queue — call flush()",
self.la_queue.len()
);
}
}
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub(crate) struct RefFrame {
pub y: rusty_h264_common::aligned::AlignedBytes,
pub u: rusty_h264_common::aligned::AlignedBytes,
pub v: rusty_h264_common::aligned::AlignedBytes,
pub poc: i32,
pub frame_num: u32,
pub mv: Vec<(i16, i16)>,
pub ref_idx: Vec<i8>,
pub mv1: Vec<(i16, i16)>,
pub ref_idx1: Vec<i8>,
pub w4: usize,
pub hpel: HpelOnce<alloc::sync::Arc<rusty_h264_common::inter::HpelPlanes>>,
}
#[cfg(feature = "std")]
pub type HpelOnce<T> = rusty_h264_common::once::OnceLock<T>;
#[cfg(not(feature = "std"))]
pub type HpelOnce<T> = core::cell::OnceCell<T>;
impl RefFrame {
pub(crate) fn hpel(&self, cw: usize, ch: usize) -> &rusty_h264_common::inter::HpelPlanes {
self.hpel.get_or_init(|| {
alloc::sync::Arc::new(rusty_h264_common::inter::build_hpel_planes(&self.y, cw, ch))
})
}
}
#[cfg(feature = "profile")]
pub fn satdpath_snapshot() -> Vec<u64> {
crate::mb16::satdpath::snapshot()
}
#[cfg(not(feature = "profile"))]
pub fn satdpath_snapshot() -> Vec<u64> {
Vec::new()
}
#[cfg(feature = "profile")]
pub fn satdpath_reset() {
crate::mb16::satdpath::reset()
}
#[cfg(not(feature = "profile"))]
pub fn satdpath_reset() {}
#[cfg(feature = "profile")]
pub fn spstats_redundant() -> u64 {
crate::mb16::spstats::redundant_count()
}
#[cfg(not(feature = "profile"))]
pub fn spstats_redundant() -> u64 {
0
}
#[cfg(feature = "profile")]
pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) {
crate::mb16::spstats::snapshot()
}
#[cfg(not(feature = "profile"))]
pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) {
(Vec::new(), Vec::new())
}
#[cfg(feature = "profile")]
pub fn spstats_reset() {
crate::mb16::spstats::reset()
}
#[cfg(not(feature = "profile"))]
pub fn spstats_reset() {}
pub const DIA_DEFAULT_MASK: u32 = crate::mb16::DIA_DEFAULT;
pub fn set_dia_mask(m: u32) {
crate::mb16::set_dia_mask(m)
}
pub fn set_me_sadfp(on: bool) {
crate::mb16::set_me_sadfp(on)
}
pub fn set_me_sadfp_mode(m: u32) {
crate::mb16::set_me_sadfp_mode(m)
}
pub fn set_me_fc(on: bool) {
crate::mb16::set_me_fc(on)
}
pub fn set_split_mg(milli: u32) {
crate::mb16::set_split_mg(milli)
}
pub fn set_mv_smooth(on: bool) {
crate::mb16::set_mv_smooth(on)
}
pub fn set_mv_smooth_mode(m: u32) {
crate::mb16::set_mv_smooth_mode(m)
}
pub fn set_sp_fc(on: bool) {
crate::mb16::set_sp_fc(on)
}
pub fn set_subme(level: u32) {
let (pat, cap) = match level {
1 => (3, 0),
2 => (2, 0),
3 => (0, 2),
4 => (0, 3),
_ => (0, 0),
};
set_subpel_pattern(pat);
crate::mb16::set_sp_maxit(cap);
}
pub fn set_turbo(on: bool) {
set_split_t(if on { 10_000_000 } else { 0 });
}
pub fn set_sp_maxit(n: u32) {
crate::mb16::set_sp_maxit(n)
}
#[cfg(feature = "profile")]
pub fn diastats_snapshot() -> Vec<(u64, u64)> {
crate::mb16::diastats::snapshot()
}
#[cfg(not(feature = "profile"))]
pub fn diastats_snapshot() -> Vec<(u64, u64)> {
Vec::new()
}
#[cfg(feature = "profile")]
pub fn diastats_reset() {
crate::mb16::diastats::reset()
}
#[cfg(not(feature = "profile"))]
pub fn diastats_reset() {}
pub fn set_defer_subpel(on: bool) {
crate::mb16::DEFER_SUBPEL.store(
if on { 1 } else { 0 },
core::sync::atomic::Ordering::Relaxed,
);
}
pub fn set_split_t(t: u32) {
crate::mb16::SPLIT_T.store(t, core::sync::atomic::Ordering::Relaxed);
}
pub fn set_subpel_dispatch(on: bool) {
crate::mb16::SP_DISPATCH.store(
if on { 1 } else { 0 },
core::sync::atomic::Ordering::Relaxed,
);
}
pub fn set_subpel_pattern(p: u32) {
crate::mb16::SUBPEL_PAT.store(p, core::sync::atomic::Ordering::Relaxed);
}
impl Encoder {
pub fn new(cfg: EncoderConfig) -> Result<Self, EncodeError> {
if !matches!(
cfg.profile,
Profile::ConstrainedBaseline | Profile::Baseline | Profile::Main | Profile::High
) {
return Err(EncodeError::Unsupported("unsupported profile"));
}
let mut cfg = cfg;
if cfg.transform_8x8 && !matches!(cfg.profile, Profile::High) {
cfg.transform_8x8 = false;
}
if cfg.bframes > 0 && !matches!(cfg.profile, Profile::Main | Profile::High) {
return Err(EncodeError::Unsupported(
"B-frames require Main or High profile",
));
}
if cfg.chroma != ChromaFormat::Yuv420 {
return Err(EncodeError::Unsupported("only 4:2:0 chroma"));
}
if cfg.width == 0 || cfg.height == 0 || cfg.width % 2 != 0 || cfg.height % 2 != 0 {
return Err(EncodeError::Unsupported(
"dimensions must be positive and even",
));
}
let sps = Sps::from_config(&cfg);
let pps = Pps::from_config(&cfg);
let rc = (cfg.bitrate > 0).then(|| RateControl::new(cfg.bitrate, cfg.framerate, cfg.qp));
let (cfg_w, cfg_h) = (cfg.width, cfg.height);
let sps_nal = sps.to_nal();
let pps_nal = pps.to_nal();
Ok(Self {
cfg,
sps,
pps,
frame_index: 0,
next_frame_num: 0,
gop_index: 0,
refs: Vec::new(),
rc,
pending_qpo: None,
pending_aq_probe: None,
la_queue: Vec::new(),
since_idr: 0,
force_idr: false,
keyframe_requested: false,
last_src: None,
last_prep: None,
cut_hist: [1.0, 1.0],
gather: None,
bw: BitWriter::with_capacity(cfg_w * cfg_h / 2 + 4096),
sps_nal,
pps_nal,
})
}
pub(crate) fn set_pending_qpo(&mut self, qpo: Vec<i32>) {
self.pending_qpo = Some(qpo);
}
pub(crate) fn set_aq_probe(&mut self, f: YuvFrame) {
self.pending_aq_probe = Some(f);
}
pub fn config(&self) -> &EncoderConfig {
&self.cfg
}
pub fn encode(&mut self, frame: &YuvFrame) -> Vec<u8> {
self.try_encode(frame).expect("frame matched config")
}
pub fn encode_planes(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
if frame.width != self.cfg.width || frame.height != self.cfg.height {
return Err(EncodeError::FrameMismatch);
}
if frame.is_valid() {
return self.try_encode_planes(frame);
}
let mut g = self
.gather
.take()
.unwrap_or_else(|| YuvFrame::black(frame.width, frame.height));
frame.copy_into(&mut g);
let r = self.try_encode_planes(&g.as_planes());
self.gather = Some(g);
r
}
pub fn encode_into(
&mut self,
frame: &YuvPlanes<'_>,
out: &mut [u8],
) -> Result<usize, EncodeError> {
if frame.width != self.cfg.width || frame.height != self.cfg.height {
return Err(EncodeError::FrameMismatch);
}
let buffered = self.lookahead_active() || (self.cfg.scenecut > 0 && self.cfg.gop_size > 1);
if buffered {
let au = self.encode_planes(frame)?;
return copy_out(&au, out);
}
if core::mem::take(&mut self.keyframe_requested) {
self.force_idr = true;
}
let mut sink = Sink::Slice {
buf: out,
len: 0,
needed: 0,
};
if frame.is_valid() {
self.encode_direct_into(frame, &mut sink)?;
} else {
let mut g = self
.gather
.take()
.unwrap_or_else(|| YuvFrame::black(frame.width, frame.height));
frame.copy_into(&mut g);
let r = self.encode_direct_into(&g.as_planes(), &mut sink);
self.gather = Some(g);
r?;
}
match sink {
Sink::Slice { len, needed, .. } if needed > len => {
Err(EncodeError::BufferTooSmall { needed })
}
Sink::Slice { len, .. } => Ok(len),
Sink::Vec(_) => unreachable!("the direct path writes to the slice sink"),
}
}
pub fn flush_into(&mut self, out: &mut [u8]) -> Result<usize, EncodeError> {
let tail = self.try_flush()?;
copy_out(&tail, out)
}
pub fn request_keyframe(&mut self) {
self.keyframe_requested = true;
}
pub fn try_encode(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
self.try_encode_planes(&frame.as_planes())
}
fn try_encode_planes(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
if core::mem::take(&mut self.keyframe_requested) {
let mut out = if self.la_queue.is_empty() {
Vec::new()
} else {
self.try_flush()?
};
self.force_idr = true;
out.extend_from_slice(&self.try_encode_planes(frame)?);
return Ok(out);
}
if self.cfg.scenecut > 0 && self.cfg.gop_size > 1 {
if self.last_src.is_some() {
let counter = self.since_idr + self.la_queue.len() as u32;
let minki = self.cfg.min_keyint.max(1);
let r = if counter + 2 >= minki {
let cur_prep = mbtree::pair_prep(&self.cfg, frame);
let prev_prep = match self.last_prep.take() {
Some(p) => p,
None => mbtree::pair_prep(
&self.cfg,
&self.last_src.as_ref().expect("guarded").as_planes(),
),
};
let r = mbtree::pair_ratio_prepped(&self.cfg, &cur_prep, &prev_prep);
self.last_prep = Some(cur_prep);
r
} else {
self.last_prep = None;
1.0 };
let (p1, p2) = (self.cut_hist[0], self.cut_hist[1]);
self.cut_hist = [r, p1];
if counter >= minki && lookahead::is_scene_cut(&self.cfg, r, p1, p2) {
let flushed = self.try_flush()?;
self.force_idr = true;
self.pending_aq_probe = self.last_src.take();
self.last_src = Some(frame.to_frame());
let mut out = flushed;
out.extend_from_slice(&self.try_encode_inner(frame)?);
return Ok(out);
}
}
self.last_src = Some(frame.to_frame());
}
self.try_encode_inner(frame)
}
fn try_encode_inner(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
if !self.lookahead_active() {
return self.encode_direct(frame);
}
if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
return Err(EncodeError::FrameMismatch);
}
self.la_queue.push(frame.to_frame());
let gop = self.cfg.gop_size.max(1);
let rem_to_idr = (gop - (self.since_idr % gop)) as usize;
let window = (self.cfg.lookahead.max(1) as usize).min(rem_to_idr);
if self.la_queue.len() >= window {
self.emit_lookahead_gop()
} else {
Ok(Vec::new())
}
}
fn lookahead_active(&self) -> bool {
self.cfg.mbtree && self.cfg.bframes == 0 && self.cfg.bitrate == 0
}
fn emit_lookahead_gop(&mut self) -> Result<Vec<u8>, EncodeError> {
let frames = core::mem::take(&mut self.la_queue);
let offs = mbtree::gop_qp_offsets(&self.cfg, &frames, self.cfg.mbtree_strength);
let mut out = Vec::new();
for (i, f) in frames.iter().enumerate() {
if let Some(o) = offs.get(i) {
self.pending_qpo = Some(o.clone());
}
out.extend_from_slice(&self.encode_direct(&f.as_planes())?);
}
Ok(out)
}
pub fn flush(&mut self) -> Vec<u8> {
self.try_flush()
.expect("buffered frames matched the config when accepted")
}
pub fn try_flush(&mut self) -> Result<Vec<u8>, EncodeError> {
if self.la_queue.is_empty() {
return Ok(Vec::new());
}
self.emit_lookahead_gop()
}
fn encode_direct(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
let mut out = Vec::new();
self.encode_direct_into(frame, &mut Sink::Vec(&mut out))?;
Ok(out)
}
fn encode_direct_into(
&mut self,
frame: &YuvPlanes<'_>,
out: &mut Sink<'_>,
) -> Result<(), EncodeError> {
let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
return Err(EncodeError::FrameMismatch);
}
if self.cfg.bframes > 0 {
return Err(EncodeError::Unsupported(
"B-frames need encode_all (lookahead)",
));
}
let forced = core::mem::take(&mut self.force_idr);
let is_idr = self.cfg.gop_size <= 1
|| forced
|| self.since_idr == 0
|| self.since_idr >= self.cfg.gop_size;
if is_idr {
self.gop_index = 0;
self.next_frame_num = 0;
self.refs.clear();
}
let frame_num = self.next_frame_num;
let poc_lsb = (2 * self.gop_index) % 256;
let qpo = self.pending_qpo.take().unwrap_or_default();
let aq_probe_owned = self.pending_aq_probe.take();
let aq_probe = aq_probe_owned.as_ref().map(YuvFrame::as_planes);
let complexity = if self.rc.is_some() {
lookahead::complexity(
&self.cfg,
frame,
if is_idr { None } else { self.refs.first() },
)
} else {
0.0
};
let qp = match &self.rc {
Some(rc) => rc.pick_qp(is_idr, complexity),
None if is_idr => (self.cfg.qp as i32 + self.cfg.i_qp_offset).clamp(0, 51) as u8,
None => self.cfg.qp,
};
let mut w = core::mem::take(&mut self.bw);
w.clear();
let (nal_type, mut reference) = if is_idr {
out.nal(
self.sps_nal.ref_idc,
self.sps_nal.nal_type,
&self.sps_nal.rbsp,
);
out.nal(
self.pps_nal.ref_idc,
self.pps_nal.nal_type,
&self.pps_nal.rbsp,
);
slice::write_idr_slice_header(&mut w, &self.cfg, qp);
let r = if self.cfg.cabac {
mb16::encode_slice_data_cabac_intra(
&mut w,
&self.cfg,
frame,
qp,
&qpo,
aq_probe.as_ref(),
)
} else {
mb16::encode_slice_data(
&mut w,
&self.cfg,
frame,
qp,
false,
&[],
&qpo,
aq_probe.as_ref(),
&[],
)
};
(NalUnitType::IdrSlice, r)
} else {
let wp: Vec<(i32, i32)> = if self.cfg.weightp {
estimate_luma_weights(&self.cfg, frame, &self.refs)
} else {
Vec::new()
};
let wp_hdr = if self.cfg.weightp {
Some(wp.as_slice())
} else {
None
};
slice::write_p_slice_header(
&mut w,
&self.cfg,
qp,
frame_num,
poc_lsb,
self.refs.len(),
wp_hdr,
);
let r = if self.cfg.cabac {
mb16::encode_slice_data_cabac_p(&mut w, &self.cfg, frame, qp, &self.refs, &qpo, &wp)
} else {
mb16::encode_slice_data(
&mut w, &self.cfg, frame, qp, true, &self.refs, &qpo, None, &wp,
)
};
(NalUnitType::NonIdrSlice, r)
};
reference.poc = 2 * self.gop_index as i32;
reference.frame_num = frame_num;
let slice_bytes = w.finish();
if let Some(rc) = &mut self.rc {
rc.update(is_idr, slice_bytes.len() * 8, qp, complexity);
}
{
let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::EncNal);
out.nal(3, nal_type, slice_bytes);
}
self.bw = w;
self.refs.insert(0, reference);
self.refs.truncate(self.cfg.num_ref_frames.max(1) as usize);
self.frame_index += 1;
self.gop_index += 1;
self.next_frame_num = (self.next_frame_num + 1) % 16;
self.since_idr = if is_idr { 1 } else { self.since_idr + 1 };
if self.cfg.gop_size > 1 && self.since_idr >= self.cfg.gop_size {
self.pending_aq_probe = Some(frame.to_frame());
}
Ok(())
}
pub fn encode_all(&self, frames: &[YuvFrame]) -> Result<Vec<Vec<u8>>, EncodeError> {
for f in frames {
if f.width != self.cfg.width || f.height != self.cfg.height || !f.is_valid() {
return Err(EncodeError::FrameMismatch);
}
}
let grain_seq = self.cfg.preset != Preset::Fast
&& rusty_h264_common::knob("RFF_GRAIN_SUBPEL")
.map(|v| v != "0")
.unwrap_or(true)
&& frames.len() >= 2
&& {
let cl0 = mbtree::coded_luma(&self.cfg, &frames[0].as_planes());
let cl1 = mbtree::coded_luma(&self.cfg, &frames[1].as_planes());
crate::signals::FrameSignals::new(
&cl1,
self.cfg.mb_width() * 16,
self.cfg.mb_width(),
self.cfg.mb_height(),
Some(&cl0),
)
.grain_signature()
};
let _guard = mb16::SeqFastPath::set(grain_seq);
if self.cfg.bframes > 0 {
let seg_starts = lookahead::segment_gops(&self.cfg, frames);
let seg_range = |k: usize| {
(
seg_starts[k],
seg_starts.get(k + 1).copied().unwrap_or(frames.len()),
)
};
let n_gops = seg_starts.len();
let (w, h) = (self.cfg.width, self.cfg.height);
let gop_sig: Vec<f64> = (0..n_gops)
.map(|g| {
let (s, e) = seg_range(g);
gop_bi_residual(&frames[s..e], w, h, 1)
})
.collect();
let gop_fav: Vec<bool> = if self.cfg.bframes_adaptive {
(0..n_gops)
.map(|g| {
let (s, e) = seg_range(g);
if e - s < 2 {
return bframes_favorable(gop_sig[g]);
}
let cl0 = mbtree::coded_luma(&self.cfg, &frames[s].as_planes());
let cl1 = mbtree::coded_luma(&self.cfg, &frames[s + 1].as_planes());
let cw = self.cfg.mb_width() * 16;
let sig = signals::FrameSignals::new(
&cl1,
cw,
self.cfg.mb_width(),
self.cfg.mb_height(),
Some(&cl0),
);
!sig.is_screen()
})
.collect()
} else {
vec![true; n_gops]
};
let gop_iqp: Vec<i32> = gop_sig
.iter()
.map(|&s| gop_iqp_offset(s, self.cfg.i_qp_offset))
.collect();
let gop_bqp: Vec<i32> = gop_sig
.iter()
.map(|&s| gop_bframe_qp_offset(s, self.cfg.bframe_qp_offset))
.collect();
let bcount = if self.cfg.bframes_adaptive {
adaptive_bcount(frames, w, h, self.cfg.bframes as usize)
} else {
self.cfg.bframes as usize
};
if gop_fav.iter().any(|&f| f) {
return Ok(self.encode_all_bframes(
frames,
bcount,
&seg_starts,
&gop_fav,
&gop_iqp,
&gop_bqp,
));
}
let mut pcfg = self.cfg.clone();
pcfg.bframes = 0;
return Encoder::new(pcfg)?.encode_all(frames);
}
if self.cfg.bitrate > 0 {
let mut enc = Encoder::new(self.cfg.clone())?;
let seg_starts = lookahead::segment_gops(&self.cfg, frames);
let offs: Vec<Vec<i32>> = if self.cfg.mbtree {
seg_starts
.iter()
.enumerate()
.flat_map(|(k, &s)| {
let end = seg_starts.get(k + 1).copied().unwrap_or(frames.len());
frames[s..end].chunks(self.cfg.lookahead.max(1) as usize)
})
.flat_map(|w| mbtree::gop_qp_offsets(&self.cfg, w, self.cfg.mbtree_strength))
.collect()
} else {
Vec::new()
};
let mut next_seg = 1usize; return frames
.iter()
.enumerate()
.map(|(i, f)| {
if seg_starts.get(next_seg) == Some(&i) {
enc.force_idr = true;
next_seg += 1;
}
if let Some(qpo) = offs.get(i) {
enc.pending_qpo = Some(qpo.clone());
}
enc.encode_direct(&f.as_planes())
})
.collect();
}
let seg_starts = lookahead::segment_gops(&self.cfg, frames);
let gops: Vec<&[YuvFrame]> = seg_starts
.iter()
.enumerate()
.map(|(k, &s)| &frames[s..seg_starts.get(k + 1).copied().unwrap_or(frames.len())])
.collect();
if gops.is_empty() || frames.is_empty() {
return Ok(Vec::new());
}
let mut out: Vec<Option<Vec<Vec<u8>>>> = (0..gops.len()).map(|_| None).collect();
let cfg = &self.cfg;
let gops_ref = &gops;
let encode_gop = |i: usize| -> Vec<Vec<u8>> {
let mut enc = Encoder::new(cfg.clone()).expect("config");
let offs: Vec<Vec<i32>> = if cfg.mbtree {
gops_ref[i]
.chunks(cfg.lookahead.max(1) as usize)
.flat_map(|w| mbtree::gop_qp_offsets(cfg, w, cfg.mbtree_strength))
.collect()
} else {
Vec::new()
};
if i > 0 {
if let Some(pf) = gops_ref[i - 1].last() {
enc.set_aq_probe(pf.clone());
}
}
let aus: Vec<Vec<u8>> = gops_ref[i]
.iter()
.enumerate()
.map(|(fi, f)| {
if let Some(o) = offs.get(fi) {
enc.set_pending_qpo(o.clone());
}
enc.encode_direct(&f.as_planes())
.expect("frame matched config")
})
.collect();
aus
};
#[cfg(feature = "std")]
{
let n = rusty_h264_common::knob("RUSTY_THREADS")
.and_then(|v| v.parse().ok())
.or_else(|| std::thread::available_parallelism().map(|n| n.get()).ok())
.unwrap_or(1)
.min(gops.len());
let encode_gop = &encode_gop;
std::thread::scope(|s| {
let handles: Vec<_> = (0..n)
.map(|t| {
s.spawn(move || {
let _fast = mb16::SeqFastPath::set(grain_seq);
let mut local = Vec::new();
let mut i = t;
while i < gops_ref.len() {
local.push((i, encode_gop(i)));
i += n;
}
local
})
})
.collect();
for h in handles {
for (i, aus) in h.join().expect("encode worker panicked") {
out[i] = Some(aus);
}
}
});
}
#[cfg(not(feature = "std"))]
{
for i in 0..gops.len() {
out[i] = Some(encode_gop(i));
}
}
Ok(out.into_iter().flatten().flatten().collect())
}
fn encode_all_bframes(
&self,
frames: &[YuvFrame],
bcount: usize,
seg_starts: &[usize],
gop_favorable: &[bool],
gop_iqp: &[i32],
gop_bqp: &[i32],
) -> Vec<Vec<u8>> {
let n = frames.len();
if n == 0 {
return Vec::new();
}
let step = bcount.max(1) + 1; let mut seg_of = vec![0usize; n];
let mut seg_start_of = vec![0usize; n];
{
let mut k = 0usize;
for d in 0..n {
if k + 1 < seg_starts.len() && seg_starts[k + 1] == d {
k += 1;
}
seg_of[d] = k;
seg_start_of[d] = seg_starts[k];
}
}
let mut cfg = self.cfg.clone();
cfg.num_ref_frames = cfg.num_ref_frames.max(if cfg.b_pyramid { 3 } else { 2 }); let sps = Sps::from_config(&cfg);
let pps = Pps::from_config(&cfg);
let mut is_anchor = vec![false; n];
for (k, &s) in seg_starts.iter().enumerate() {
let e = seg_starts.get(k + 1).copied().unwrap_or(n).min(n);
if !gop_favorable.get(k).copied().unwrap_or(true) {
is_anchor[s..e].fill(true);
continue;
}
let boundary = k + 1 < seg_starts.len();
let mut phase = 0usize;
for d in s..e {
is_anchor[d] = phase == 0 || (boundary && d + 1 == e);
phase += 1;
if phase == step {
phase = 0;
}
}
}
is_anchor[n - 1] = true;
if self.cfg.bframes_adaptive {
let (w, h) = (self.cfg.width, self.cfg.height);
fn memo_mean(frames: &[YuvFrame], means: &mut [f64], x: usize) -> f64 {
if means[x].is_nan() {
let f = &frames[x];
let mut s = 0u64;
let mut i = 0;
while i < f.y.len() {
s += f.y[i] as u64;
i += 64;
}
let c = f.y.len().div_ceil(64) as u64;
means[x] = s as f64 / c.max(1) as f64;
}
means[x]
}
let mut means = vec![f64::NAN; n];
let mut a = 0usize;
while a + 1 < n {
if !is_anchor[a] {
a += 1;
continue;
}
let mut b = a + 1;
while b < n && !is_anchor[b] {
b += 1;
}
if b > a + 1 && b < n {
let flash = (a..b).any(|x| {
(memo_mean(frames, &mut means, x) - memo_mean(frames, &mut means, x + 1))
.abs()
> 2.5
});
let res = gop_bi_residual(&frames[a..=b], w, h, 1);
if flash || !bframes_favorable(res) {
for x in a + 1..b {
is_anchor[x] = true;
}
}
}
a = b;
}
}
let mbtree_off: Vec<Vec<i32>> = if cfg.mbtree {
let mut off = vec![Vec::new(); n];
for (k, &s) in seg_starts.iter().enumerate() {
let seg_end = seg_starts.get(k + 1).copied().unwrap_or(n).min(n);
let anchors: Vec<usize> = (s..seg_end).filter(|&d| is_anchor[d]).collect();
for aw in anchors.chunks(cfg.lookahead.max(1) as usize) {
let aframes: Vec<&YuvFrame> = aw.iter().map(|&d| &frames[d]).collect();
let offs = mbtree::gop_qp_offsets_refs(&cfg, &aframes, cfg.mbtree_strength);
for (o, &d) in offs.into_iter().zip(aw.iter()) {
off[d] = o;
}
}
}
off
} else {
Vec::new()
};
let pyramid = cfg.b_pyramid && cfg.cabac && step >= 3;
let mut is_bref = vec![false; n];
let mut order: Vec<usize> = Vec::with_capacity(n);
let mut prev: Option<usize> = None;
for d in 0..n {
if !is_anchor[d] {
continue;
}
order.push(d);
if let Some(p) = prev {
if pyramid && d - p > 2 {
let m = (p + d) / 2;
is_bref[m] = true;
order.push(m);
order.extend(((p + 1)..d).filter(|&x| x != m));
} else {
order.extend((p + 1)..d);
}
}
prev = Some(d);
}
let mut dpb: Vec<RefFrame> = Vec::new();
let mut aus: Vec<Vec<u8>> = Vec::with_capacity(n);
let mut frame_num: u32 = 0;
for &d in &order {
let is_idr = d == seg_start_of[d];
if is_idr {
dpb.clear();
frame_num = 0;
}
let is_b = !is_anchor[d];
let bref = is_b && is_bref[d];
let poc = ((d - seg_start_of[d]) as i32) * 2; let iqp = gop_iqp.get(seg_of[d]).copied().unwrap_or(cfg.i_qp_offset);
let bqp_leaf = gop_bqp
.get(seg_of[d])
.copied()
.unwrap_or(cfg.bframe_qp_offset);
let bqp = if bref { (bqp_leaf + 1) / 2 } else { bqp_leaf };
let qpo: &[i32] = mbtree_off.get(d).map(|v| v.as_slice()).unwrap_or(&[]);
let aq_probe = if is_idr && d > 0 {
frames.get(d - 1).map(YuvFrame::as_planes)
} else {
None
};
let (au, recon) = code_picture(
&cfg,
&sps,
&pps,
&frames[d].as_planes(),
is_idr,
is_b,
bref,
poc,
frame_num,
&dpb,
iqp,
bqp,
qpo,
aq_probe.as_ref(),
);
aus.push(au);
if !is_b {
if let Some(r) = recon {
dpb.insert(0, r);
dpb.truncate(cfg.num_ref_frames as usize);
}
frame_num = (frame_num + 1) % 16;
} else if let Some(r) = recon {
dpb.insert(0, r);
dpb.truncate(cfg.num_ref_frames as usize);
frame_num = (frame_num + 1) % 16;
}
}
aus
}
}
fn code_picture(
cfg: &EncoderConfig,
sps: &Sps,
pps: &Pps,
frame: &YuvPlanes<'_>,
is_idr: bool,
is_b: bool,
b_is_ref: bool,
poc: i32,
frame_num: u32,
dpb: &[RefFrame],
i_qp_offset: i32,
b_qp_offset: i32,
qpo: &[i32],
aq_probe: Option<&YuvPlanes<'_>>,
) -> (Vec<u8>, Option<RefFrame>) {
let mut out = Vec::new();
let mut w = BitWriter::with_capacity(cfg.width * cfg.height / 2 + 4096);
let poc_lsb = (poc as u32) & 0xFF; let qp = if is_b {
(cfg.qp as i32 + b_qp_offset).clamp(0, 51) as u8
} else if is_idr {
(cfg.qp as i32 + i_qp_offset).clamp(0, 51) as u8
} else {
cfg.qp
};
let (nal_type, nal_ref_idc, recon) = if is_idr {
sps.to_nal().write_annex_b(&mut out);
pps.to_nal().write_annex_b(&mut out);
slice::write_idr_slice_header(&mut w, cfg, qp);
let mut r = if cfg.cabac {
mb16::encode_slice_data_cabac_intra(&mut w, cfg, frame, qp, qpo, aq_probe)
} else {
mb16::encode_slice_data(&mut w, cfg, frame, qp, false, &[], qpo, aq_probe, &[])
};
r.poc = poc;
r.frame_num = frame_num;
(NalUnitType::IdrSlice, 3u8, Some(r))
} else if is_b {
let l0 = dpb.iter().filter(|r| r.poc < poc).max_by_key(|r| r.poc);
let l1 = dpb.iter().filter(|r| r.poc > poc).min_by_key(|r| r.poc);
let as_ref = b_is_ref && cfg.cabac && l0.is_some() && l1.is_some();
slice::write_b_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, 1, 1, as_ref);
let mut b_recon: Option<RefFrame> = None;
match (l0, l1) {
(Some(l0), Some(l1)) if cfg.cabac => {
b_recon = mb16::encode_slice_data_cabac_b(
&mut w,
cfg,
frame,
qp,
poc,
l0,
l1,
&[],
as_ref,
);
}
(Some(l0), Some(l1)) => {
mb16::encode_slice_data_b(&mut w, cfg, frame, qp, poc, l0, l1, &[]);
}
_ => {
let n = cfg.mb_width() * cfg.mb_height();
if cfg.cabac {
mb16::encode_all_skip_b_cabac(&mut w, cfg, qp, n);
} else {
w.write_ue(n as u32);
w.rbsp_trailing_bits();
}
}
}
if as_ref {
if let Some(r) = &mut b_recon {
r.poc = poc;
r.frame_num = frame_num;
}
(NalUnitType::NonIdrSlice, 2u8, b_recon)
} else {
(NalUnitType::NonIdrSlice, 0u8, None)
}
} else {
let p_dpb: &[RefFrame] = dpb;
let wp: Vec<(i32, i32)> = if cfg.weightp {
estimate_luma_weights(cfg, frame, p_dpb)
} else {
Vec::new()
};
let wp_hdr = if cfg.weightp {
Some(wp.as_slice())
} else {
None
};
slice::write_p_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, p_dpb.len(), wp_hdr);
let mut r = if cfg.cabac {
mb16::encode_slice_data_cabac_p(&mut w, cfg, frame, qp, p_dpb, qpo, &wp)
} else {
mb16::encode_slice_data(&mut w, cfg, frame, qp, true, dpb, qpo, None, &wp)
};
r.poc = poc;
r.frame_num = frame_num;
(NalUnitType::NonIdrSlice, 3u8, Some(r))
};
let slice_bytes = w.into_bytes();
NalUnit::new(nal_ref_idc, nal_type, slice_bytes).write_annex_b(&mut out);
(out, recon)
}
const BI_THRESH: f64 = 4.0;
const B_GAP_THRESH: f64 = 8.2;
fn bframes_favorable(residual: f64) -> bool {
residual < B_GAP_THRESH
}
fn gop_iqp_offset(residual: f64, base: i32) -> i32 {
if base == 0 {
return 0;
}
let bonus =
rusty_h264_common::fmath::round(2.0 * ((BI_THRESH - residual) / BI_THRESH).clamp(0.0, 1.0))
as i32;
base - bonus
}
fn gop_bframe_qp_offset(residual: f64, base: i32) -> i32 {
const RAMP: f64 = 0.3; let boost =
rusty_h264_common::fmath::round(4.0 * ((RAMP - residual) / RAMP).clamp(0.0, 1.0)) as i32;
base + boost
}
pub fn temporal_decay_ratio(frames: &[YuvFrame], w: usize, h: usize) -> f64 {
let g1 = gop_bi_residual(frames, w, h, 1);
let g2 = gop_bi_residual(frames, w, h, 2);
if !g1.is_finite() || !g2.is_finite() {
return f64::INFINITY;
}
g2 / g1.max(1e-3)
}
fn adaptive_bcount(frames: &[YuvFrame], w: usize, h: usize, max_b: usize) -> usize {
let cap = max_b.clamp(1, 3);
let g1 = gop_bi_residual(frames, w, h, 1);
let g2 = gop_bi_residual(frames, w, h, 2);
if !g1.is_finite() || !g2.is_finite() {
return 1;
}
let ratio = g2 / g1.max(1e-3);
let c = if ratio >= 1.4 {
1
} else if ratio >= 1.3 {
2
} else {
3
};
c.clamp(1, cap)
}
fn gop_bi_residual(frames: &[YuvFrame], w: usize, h: usize, gap: usize) -> f64 {
let n = frames.len();
if n < 2 * gap + 1 || w < 48 || h < 48 {
return f64::INFINITY;
}
let sad = |cur: &[u8], rf: &[u8], dx: isize, dy: isize| -> u64 {
let mut s = 0u64;
let mut y = 16;
while y < h - 16 {
let cbase = (y * w) as isize;
let rbase = ((y as isize + dy) * w as isize) + dx;
let mut x = 16isize;
while x < (w - 16) as isize {
let c = cur[(cbase + x) as usize] as i32;
let r = rf[(rbase + x) as usize] as i32;
s += (c - r).unsigned_abs() as u64;
x += 8;
}
y += 8;
}
s
};
let global_me = |cur: &[u8], rf: &[u8]| -> (isize, isize) {
let (mut best, mut bc) = ((0isize, 0isize), u64::MAX);
let mut dy = -12;
while dy <= 12 {
let mut dx = -12;
while dx <= 12 {
let c = sad(cur, rf, dx, dy);
if c < bc {
bc = c;
best = (dx, dy);
}
dx += 4;
}
dy += 4;
}
let centre = best;
for dy in best.1 - 3..=best.1 + 3 {
for dx in best.0 - 3..=best.0 + 3 {
if (dx, dy) == centre {
continue;
}
let c = sad(cur, rf, dx, dy);
if c < bc {
bc = c;
best = (dx, dy);
}
}
}
best
};
let mut n_samp = 0usize;
{
let mut y = 16;
while y < h - 16 {
let mut x = 16;
while x < w - 16 {
n_samp += 1;
x += 8;
}
y += 8;
}
}
let step = (n / 5).max(1);
let (mut total, mut cnt) = (0f64, 0usize);
let mut d = gap;
while d < n - gap {
let (cur, past, fut) = (&frames[d].y, &frames[d - gap].y, &frames[d + gap].y);
let (mpx, mpy) = global_me(cur, past);
let (mfx, mfy) = global_me(cur, fut);
let mut bi = 0u64;
let mut y = 16;
while y < h - 16 {
let mut x = 16isize;
while x < (w - 16) as isize {
let c = cur[y * w + x as usize] as i32;
let p = past[((y as isize + mpy) * w as isize + x + mpx) as usize] as i32;
let f = fut[((y as isize + mfy) * w as isize + x + mfx) as usize] as i32;
bi += (c - ((p + f + 1) >> 1)).unsigned_abs() as u64;
x += 8;
}
y += 8;
}
total += bi as f64 / n_samp as f64;
cnt += 1;
d += step;
}
if cnt > 0 {
total / cnt as f64
} else {
f64::INFINITY
}
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(unused_imports)]
use alloc::{
boxed::Box,
format,
string::{String, ToString},
vec,
vec::Vec,
};
#[allow(unused_imports)]
use rusty_h264_common::once::OnceLock;
fn textured(w: usize, h: usize, t: usize) -> YuvFrame {
let mut f = YuvFrame::black(w, h);
for y in 0..h {
for x in 0..w {
let v = ((x + 3 * t) * 255 / w) as u32 + ((y * 7 + x * 3 + t) % 23) as u32;
f.y[y * w + x] = v.min(255) as u8;
}
}
f
}
#[test]
fn fast_preset_never_builds_the_half_pel_cache() {
let (w, h) = (64, 48);
let frames: Vec<YuvFrame> = (0..4).map(|t| textured(w, h, t)).collect();
for refs in [1, 3] {
let mut cfg = EncoderConfig::baseline(w, h);
cfg.gop_size = 8;
cfg.min_keyint = 8;
cfg.num_ref_frames = refs;
assert_eq!(cfg.preset, Preset::Fast);
let mut enc = Encoder::new(cfg).unwrap();
for f in &frames {
let _ = enc.encode(f);
}
assert!(!enc.refs.is_empty());
assert!(
enc.refs.iter().all(|r| r.hpel.get().is_none()),
"Fast built a half-pel cache (refs={refs})"
);
}
let mut cfg = EncoderConfig::baseline(w, h);
cfg.gop_size = 8;
cfg.min_keyint = 8;
cfg.num_ref_frames = 3;
cfg.preset = Preset::Balanced;
let mut enc = Encoder::new(cfg).unwrap();
for f in &frames {
let _ = enc.encode(f);
}
assert!(
enc.refs.iter().any(|r| r.hpel.get().is_some()),
"Balanced never built the half-pel cache"
);
}
#[test]
fn rejects_unsupported_config() {
let mut cfg = EncoderConfig::new(16, 16);
cfg.profile = Profile::High;
cfg.transform_8x8 = true;
cfg.cabac = true;
assert!(
Encoder::new(cfg).is_ok(),
"High + 8x8 + CABAC must be accepted"
);
let mut cfg = EncoderConfig::new(16, 16);
cfg.profile = Profile::Main;
cfg.transform_8x8 = true;
let enc = Encoder::new(cfg).expect("Main + 8x8 must clamp, not fail");
assert!(
!enc.cfg.transform_8x8,
"8x8 must be cleared when the profile cannot signal it"
);
}
#[test]
fn encodes_access_unit_with_sps_pps_idr() {
use rusty_h264_common::nal::split_annex_b;
let cfg = EncoderConfig::new(32, 32);
let mut enc = Encoder::new(cfg).unwrap();
let frame = YuvFrame::black(32, 32);
let mut au = enc.encode(&frame);
au.extend(enc.flush());
let nals = split_annex_b(&au);
assert_eq!(nals.len(), 3);
assert_eq!(NalUnitType::from_id(nals[0][0]), NalUnitType::Sps);
assert_eq!(NalUnitType::from_id(nals[1][0]), NalUnitType::Pps);
assert_eq!(NalUnitType::from_id(nals[2][0]), NalUnitType::IdrSlice);
}
#[test]
fn encode_all_matches_sequential_cqp() {
let (w, h) = (48usize, 32usize);
let mut cfg = EncoderConfig::new(w, h);
cfg.gop_size = 4; let frames: Vec<YuvFrame> = (0..10u8)
.map(|t| YuvFrame {
width: w,
height: h,
y: (0..w * h)
.map(|i| (i as u8).wrapping_add(t.wrapping_mul(7)))
.collect(),
u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
})
.collect();
let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
seq.extend_from_slice(&seq_enc.flush()); let par: Vec<u8> = Encoder::new(cfg)
.unwrap()
.encode_all(&frames)
.unwrap()
.concat();
assert_eq!(seq, par, "GOP-parallel must equal sequential+flush at CQP");
}
#[test]
fn encode_all_matches_sequential_quality_preset() {
let (w, h) = (48usize, 32usize);
let mut cfg = EncoderConfig::new(w, h);
cfg.gop_size = 3; cfg.preset = crate::config::Preset::Quality;
let frames: Vec<YuvFrame> = (0..12u8)
.map(|t| YuvFrame {
width: w,
height: h,
y: (0..w * h)
.map(|i| {
let base = (i as u8).wrapping_add(t.wrapping_mul(3));
if t % 2 == 0 {
base
} else {
base.wrapping_mul(37).wrapping_add(i as u8)
}
})
.collect(),
u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
})
.collect();
let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
seq.extend_from_slice(&seq_enc.flush());
let par: Vec<u8> = Encoder::new(cfg)
.unwrap()
.encode_all(&frames)
.unwrap()
.concat();
assert_eq!(
seq, par,
"quality-preset GOP-parallel must equal sequential+flush"
);
}
#[test]
fn rejects_mismatched_frame() {
let cfg = EncoderConfig::new(16, 16);
let mut enc = Encoder::new(cfg).unwrap();
let frame = YuvFrame::black(32, 16);
assert_eq!(enc.try_encode(&frame), Err(EncodeError::FrameMismatch));
}
}