use crate::cdf::{
FrameInterOrderHints, InterFrameContext, InterWalkPixels, MotionFieldMvs, PartitionWalker,
QuantizerParams, TileCdfContext, TileDecodeParams, TileGeometry,
};
use crate::encoder::ivf::IvfReader;
use crate::encoder::tile_group_obu::{parse_tile_group_obu_body, TilePayload};
use crate::film_grain::film_grain_synthesis;
use crate::frame_header::{
parse_frame_header_with_refs, FrameHeader, RefInfo, NUM_REF_FRAMES, PRIMARY_REF_NONE,
};
use crate::inter_pred::get_relative_dist;
use crate::loop_filter::PlaneBuffer;
use crate::obu::{ObuIter, ObuType};
use crate::sequence_header::{parse_sequence_header, SequenceHeader};
use crate::symbol_decoder::SymbolDecoder;
use crate::uncompressed_header_tail::{
SegmentationParams, ALTREF_FRAME, LAST_FRAME, MAX_SEGMENTS, SEG_LVL_ALT_Q, SEG_LVL_SKIP,
};
use crate::Error;
use crate::{PlaneRefSpec, RefFrameStoreEntry};
use std::sync::Arc;
use std::time::{Duration, Instant};
#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct DecodeStageTimings {
pub parse: Duration,
pub tiles: Duration,
pub deblock: Duration,
pub cdef: Duration,
pub superres: Duration,
pub restoration: Duration,
pub ref_store: Duration,
pub grain: Duration,
pub output: Duration,
}
impl DecodeStageTimings {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn total(&self) -> Duration {
self.parse
+ self.tiles
+ self.deblock
+ self.cdef
+ self.superres
+ self.restoration
+ self.ref_store
+ self.grain
+ self.output
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SpecFrame {
pub width: u32,
pub height: u32,
pub planes: Vec<Vec<u8>>,
pub plane_dims: Vec<(u32, u32)>,
pub bit_depth: u8,
}
#[derive(Debug, Clone)]
struct SpecRefSlot {
planes: Vec<Vec<u16>>,
plane_dims: Vec<(u32, u32)>,
mf_mvs: Vec<i16>,
mf_ref_frames: Vec<i8>,
saved_order_hints: [i32; ALTREF_FRAME + 1],
mi_rows: u32,
mi_cols: u32,
frame_is_intra: bool,
frame_type_is_key: bool,
bit_depth: u8,
cdfs: Box<TileCdfContext>,
gm_params: [[i32; 6]; 8],
lf_ref_deltas: [i8; 8],
lf_mode_deltas: [i8; 2],
segment_ids: Vec<i32>,
grain_params: Option<crate::uncompressed_header_tail::FilmGrainParams>,
}
#[derive(Debug, Clone)]
struct SpecRefState {
info: RefInfo,
slots: [Option<Arc<SpecRefSlot>>; NUM_REF_FRAMES as usize],
}
impl SpecRefState {
fn new() -> Self {
Self {
info: RefInfo::default(),
slots: Default::default(),
}
}
}
struct DecodedFrameInternal {
frame: SpecFrame,
ref_planes: Vec<Vec<u16>>,
ref_plane_dims: Vec<(u32, u32)>,
mf_mvs: Vec<i16>,
mf_ref_frames: Vec<i8>,
order_hints_by_ref: [i32; ALTREF_FRAME + 1],
end_cdfs: Box<TileCdfContext>,
mi_rows: u32,
mi_cols: u32,
segment_ids: Vec<i32>,
grain_params: Option<crate::uncompressed_header_tail::FilmGrainParams>,
}
pub(crate) fn lossless_array(
qp: &crate::uncompressed_header_tail::QuantizationParams,
sp: &SegmentationParams,
) -> [bool; MAX_SEGMENTS] {
let deltas_all_zero = qp.delta_q_y_dc == 0
&& qp.delta_q_u_dc == 0
&& qp.delta_q_u_ac == 0
&& qp.delta_q_v_dc == 0
&& qp.delta_q_v_ac == 0;
let mut out = [false; MAX_SEGMENTS];
for (segment_id, slot) in out.iter_mut().enumerate() {
let qindex = if sp.enabled && sp.segment_feature_active[segment_id][SEG_LVL_ALT_Q] {
let data = i32::from(sp.segment_feature_data[segment_id][SEG_LVL_ALT_Q]);
(i32::from(qp.base_q_idx) + data).clamp(0, 255)
} else {
i32::from(qp.base_q_idx)
};
*slot = qindex == 0 && deltas_all_zero;
}
out
}
pub fn decode_frame_spec(
seq: &SequenceHeader,
fh: &FrameHeader,
tile_group_body: &[u8],
) -> Result<SpecFrame, Error> {
let mut stages = DecodeStageTimings::default();
Ok(decode_frame_spec_full(seq, fh, tile_group_body, None, &mut stages, 1)?.frame)
}
fn parse_tile_group_body_with_ti(
ti: &crate::tile_info::TileInfo,
body: &[u8],
) -> Result<crate::encoder::tile_group_obu::ParsedTileGroup, Error> {
let ceil_log2 = |v: u32| -> u32 {
if v <= 1 {
0
} else {
32 - (v - 1).leading_zeros()
}
};
if ti.mi_col_starts.len() != (ti.tile_cols as usize) + 1
|| ti.mi_row_starts.len() != (ti.tile_rows as usize) + 1
{
return Err(Error::PartitionWalkOutOfRange);
}
parse_tile_group_obu_body(
body,
ti.tile_cols * ti.tile_rows,
ceil_log2(ti.tile_cols),
ceil_log2(ti.tile_rows),
u32::from(ti.tile_size_bytes),
)
}
fn parse_whole_frame_tile_group(
ti: &crate::tile_info::TileInfo,
body: &[u8],
require_flag_zero: bool,
) -> Result<Vec<TilePayload>, Error> {
let parsed = parse_tile_group_body_with_ti(ti, body)?;
let num_tiles = ti.tile_cols * ti.tile_rows;
if parsed.tg_start != 0 || parsed.tg_end + 1 != num_tiles {
return Err(Error::TileGroupInvalid);
}
if require_flag_zero && parsed.tile_start_and_end_present_flag {
return Err(Error::TileGroupInvalid);
}
Ok(parsed.tiles)
}
fn decode_frame_spec_full(
seq: &SequenceHeader,
fh: &FrameHeader,
tile_group_body: &[u8],
refs: Option<&SpecRefState>,
stages: &mut DecodeStageTimings,
threads: usize,
) -> Result<DecodedFrameInternal, Error> {
let ti = fh
.tile_info
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let t0 = Instant::now();
let tiles = parse_whole_frame_tile_group(ti, tile_group_body, false)?;
stages.parse += t0.elapsed();
decode_frame_spec_tiles(seq, fh, &tiles, refs, stages, threads)
}
fn decode_frame_spec_tiles(
seq: &SequenceHeader,
fh: &FrameHeader,
tiles: &[TilePayload],
refs: Option<&SpecRefState>,
stages: &mut DecodeStageTimings,
threads: usize,
) -> Result<DecodedFrameInternal, Error> {
if fh.show_existing_frame {
return Err(Error::PartitionWalkOutOfRange);
}
if !fh.frame_is_intra && refs.is_none() {
return Err(Error::PartitionWalkOutOfRange);
}
let cc = &seq.color_config;
if !matches!(cc.bit_depth, 8 | 10 | 12) {
return Err(Error::PartitionWalkOutOfRange);
}
let fs = fh
.frame_size
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
if fs.use_superres && fs.upscaled_width <= fs.frame_width {
return Err(Error::PartitionWalkOutOfRange);
}
let ti = fh
.tile_info
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let qp = fh
.quantization_params
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let default_seg = SegmentationParams::disabled();
let sp = fh.segmentation_params.as_ref().unwrap_or(&default_seg);
let dq = fh.delta_q_params.unwrap_or_default();
let dlf = fh.delta_lf_params.unwrap_or_default();
let lr = fh
.lr_params
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let lr_walk = crate::cdf::LrParams {
num_planes: cc.num_planes as usize,
frame_restoration_type: [
lr.frame_restoration_type[0] as u8,
lr.frame_restoration_type[1] as u8,
lr.frame_restoration_type[2] as u8,
],
loop_restoration_size: lr.loop_restoration_size,
subsampling_x: u8::from(cc.subsampling_x),
subsampling_y: u8::from(cc.subsampling_y),
frame_height: fs.frame_height,
upscaled_width: fs.upscaled_width,
use_superres: fs.use_superres,
superres_denom: fs.superres_denom,
allow_intrabc: fh.allow_intrabc,
};
let lossless = lossless_array(qp, sp);
let coded_lossless = lossless.iter().all(|&l| l);
let seg_alt_q_active = {
let mut a = [false; MAX_SEGMENTS];
for (i, slot) in a.iter_mut().enumerate() {
*slot = sp.segment_feature_active[i][SEG_LVL_ALT_Q];
}
a
};
let seg_alt_q_data = {
let mut a = [0i16; MAX_SEGMENTS];
for (i, slot) in a.iter_mut().enumerate() {
*slot = sp.segment_feature_data[i][SEG_LVL_ALT_Q];
}
a
};
let mut seg_qm_level = [[15u8; MAX_SEGMENTS]; 3];
if qp.using_qmatrix {
for segment_id in 0..MAX_SEGMENTS {
if !lossless[segment_id] {
seg_qm_level[0][segment_id] = qp.qm_y;
seg_qm_level[1][segment_id] = qp.qm_u;
seg_qm_level[2][segment_id] = qp.qm_v;
}
}
}
let quant = QuantizerParams {
base_q_idx: qp.base_q_idx,
delta_q_y_dc: qp.delta_q_y_dc,
delta_q_u_dc: qp.delta_q_u_dc,
delta_q_u_ac: qp.delta_q_u_ac,
delta_q_v_dc: qp.delta_q_v_dc,
delta_q_v_ac: qp.delta_q_v_ac,
using_qmatrix: qp.using_qmatrix,
bit_depth: cc.bit_depth,
delta_q_present: dq.delta_q_present,
current_q_index: qp.base_q_idx,
segmentation_enabled: sp.enabled,
seg_alt_q_active,
seg_alt_q_data,
seg_qm_level,
};
let seg_skip: [bool; MAX_SEGMENTS] = if sp.enabled {
core::array::from_fn(|s| sp.segment_feature_active[s][SEG_LVL_SKIP])
} else {
[false; MAX_SEGMENTS]
};
let cdef_bits = fh
.cdef_params
.as_ref()
.map_or(0, |c| u32::from(c.cdef_bits));
let tx_mode_select = matches!(
fh.tx_mode,
Some(crate::uncompressed_header_tail::TxMode::TxModeSelect)
);
let params = TileDecodeParams {
frame_is_intra: fh.frame_is_intra,
subsampling_x: u8::from(cc.subsampling_x),
subsampling_y: u8::from(cc.subsampling_y),
num_planes: cc.num_planes,
seg_id_pre_skip: sp.seg_id_pre_skip,
segmentation_enabled: sp.enabled,
seg_skip,
last_active_seg_id: sp.last_active_seg_id,
lossless_array: &lossless,
coded_lossless,
enable_cdef: seq.enable_cdef,
allow_intrabc: fh.allow_intrabc,
cdef_bits,
use_128x128_superblock: seq.use_128x128_superblock,
delta_q_res: dq.delta_q_res,
delta_lf_present: dlf.delta_lf_present,
delta_lf_multi: dlf.delta_lf_multi,
mono_chrome: cc.mono_chrome,
delta_lf_res: dlf.delta_lf_res,
allow_screen_content_tools: fh.allow_screen_content_tools,
enable_filter_intra: seq.enable_filter_intra,
bit_depth: cc.bit_depth,
tx_mode_select,
reduced_tx_set: fh.reduced_tx_set.unwrap_or(false),
enable_intra_edge_filter: seq.enable_intra_edge_filter,
};
let num_tiles = ti.tile_cols * ti.tile_rows;
if ti.mi_col_starts.len() != (ti.tile_cols as usize) + 1
|| ti.mi_row_starts.len() != (ti.tile_rows as usize) + 1
{
return Err(Error::PartitionWalkOutOfRange);
}
if tiles.len() != num_tiles as usize {
return Err(Error::TileGroupInvalid);
}
let mi_rows = fs.mi_rows;
let mi_cols = fs.mi_cols;
let mut walker = PartitionWalker::new(
mi_rows,
mi_cols,
TileGeometry {
mi_row_start: 0,
mi_row_end: mi_rows,
mi_col_start: 0,
mi_col_end: mi_cols,
},
)
.ok_or(Error::PartitionWalkOutOfRange)?;
let is_inter_frame = !fh.frame_is_intra;
let sub_x_u8 = u8::from(cc.subsampling_x);
let sub_y_u8 = u8::from(cc.subsampling_y);
let num_planes_usize = cc.num_planes as usize;
let mfmvs = if is_inter_frame
&& fh
.inter_refs
.as_ref()
.is_some_and(|ir| ir.use_ref_frame_mvs)
{
motion_field_estimation(
refs.ok_or(Error::PartitionWalkOutOfRange)?,
fh,
seq,
mi_rows,
mi_cols,
)?
} else {
MotionFieldMvs::new_invalid(mi_rows, mi_cols)
};
let prev_segment_ids: Option<Vec<i32>> =
if is_inter_frame && sp.enabled && fh.primary_ref_frame != PRIMARY_REF_NONE {
let st = refs.ok_or(Error::PartitionWalkOutOfRange)?;
let ir = fh
.inter_refs
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let slot = ir.ref_frame_idx[fh.primary_ref_frame as usize] as usize;
match st.slots.get(slot).and_then(|s| s.as_ref()) {
Some(prev)
if prev.mi_rows == mi_rows
&& prev.mi_cols == mi_cols
&& prev.segment_ids.len() == (mi_rows as usize) * (mi_cols as usize) =>
{
Some(prev.segment_ids.clone())
}
_ => Some(vec![0i32; (mi_rows as usize) * (mi_cols as usize)]),
}
} else if is_inter_frame && sp.enabled {
Some(vec![0i32; (mi_rows as usize) * (mi_cols as usize)])
} else {
None
};
let dummy_plane: Vec<u16> = vec![0u16; 64];
let dummy_entry = || RefFrameStoreEntry {
plane: &dummy_plane,
stride: 8,
upscaled_width: 8,
width: 8,
height: 8,
};
let mut plane_stores: Vec<[RefFrameStoreEntry<'_>; NUM_REF_FRAMES as usize]> = Vec::new();
let mut ref_frame_idx = [0u8; 7];
let mut order_hints_by_ref = [0i32; ALTREF_FRAME + 1];
let mut sign_bias = [0i32; 8];
let mut is_scaled_per_ref = [false; 7];
if is_inter_frame {
let st = refs.ok_or(Error::PartitionWalkOutOfRange)?;
let ir = fh
.inter_refs
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
for plane in 0..num_planes_usize {
let arr: [RefFrameStoreEntry<'_>; NUM_REF_FRAMES as usize] =
core::array::from_fn(|slot| match st.slots[slot].as_ref() {
Some(s) if plane < s.planes.len() => {
let (w, _h) = s.plane_dims[plane];
let (lw, lh) = s.plane_dims[0];
RefFrameStoreEntry {
plane: &s.planes[plane],
stride: w as usize,
upscaled_width: lw,
width: lw,
height: lh,
}
}
_ => dummy_entry(),
});
plane_stores.push(arr);
}
for i in 0..7 {
let slot = ir.ref_frame_idx[i] as usize;
if slot >= NUM_REF_FRAMES as usize {
return Err(Error::PartitionWalkOutOfRange);
}
ref_frame_idx[i] = slot as u8;
let hint = st.info.order_hint[slot] as i32;
order_hints_by_ref[LAST_FRAME + i] = hint;
if seq.enable_order_hint {
sign_bias[LAST_FRAME + i] = i32::from(
get_relative_dist(hint, fh.order_hint as i32, u32::from(seq.order_hint_bits))
> 0,
);
}
let no_scale = 1u64 << crate::inter_pred::REF_SCALE_SHIFT;
let fw = u64::from(fs.frame_width.max(1));
let fhh = u64::from(fs.frame_height.max(1));
let x_scale = ((u64::from(st.info.upscaled_width[slot])
<< crate::inter_pred::REF_SCALE_SHIFT)
+ fw / 2)
/ fw;
let y_scale = ((u64::from(st.info.frame_height[slot])
<< crate::inter_pred::REF_SCALE_SHIFT)
+ fhh / 2)
/ fhh;
is_scaled_per_ref[i] = x_scale != no_scale || y_scale != no_scale;
}
}
let order_hints = FrameInterOrderHints {
order_hint_bits: if seq.enable_order_hint {
u32::from(seq.order_hint_bits)
} else {
0
},
current_order_hint: fh.order_hint as i32,
order_hints_by_ref,
};
let plane_ref_specs: Vec<PlaneRefSpec<'_>> = (0..plane_stores.len())
.map(|p| PlaneRefSpec {
plane: p as u8,
subsampling_x: if p > 0 { sub_x_u8 } else { 0 },
subsampling_y: if p > 0 { sub_y_u8 } else { 0 },
frame_store: &plane_stores[p],
frame_width: fs.frame_width,
frame_height: fs.frame_height,
})
.collect();
let pixels = InterWalkPixels {
ref_frame_idx,
bit_depth: cc.bit_depth,
plane_refs: &plane_ref_specs,
order_hints,
};
let ictx: Option<InterFrameContext<'_>> = if is_inter_frame {
let ir = fh
.inter_refs
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut c = InterFrameContext::identity_default(&mfmvs);
c.segmentation_update_map = sp.update_map;
c.segmentation_temporal_update = sp.temporal_update;
c.seg_feature_active = sp.segment_feature_active;
c.seg_feature_data = sp.segment_feature_data;
c.prev_segment_ids = prev_segment_ids.as_deref();
c.skip_mode_present = fh.skip_mode_present.unwrap_or(false);
if let Some(smf) = fh.skip_mode_frame {
c.skip_mode_frame = [i32::from(smf[0]), i32::from(smf[1])];
}
c.reference_select = fh.reference_select.unwrap_or(false);
if let Some(g) = fh.global_motion_params.as_ref() {
for r in 0..8 {
c.gm_type[r] = g.gm_type[r] as i32;
c.gm_params[r] = g.gm_params[r];
}
}
c.ref_frame_sign_bias = sign_bias;
c.allow_high_precision_mv = ir.allow_high_precision_mv;
c.force_integer_mv = fh.force_integer_mv;
c.use_ref_frame_mvs = ir.use_ref_frame_mvs;
c.is_motion_mode_switchable = ir.is_motion_mode_switchable;
c.allow_warped_motion = fh.allow_warped_motion.unwrap_or(false);
c.is_scaled_per_ref = is_scaled_per_ref;
c.enable_interintra_compound = seq.enable_interintra_compound;
c.enable_masked_compound = seq.enable_masked_compound;
c.enable_jnt_comp = seq.enable_jnt_comp;
c.order_hints = order_hints;
c.interpolation_filter = ir.interpolation_filter as u8;
c.enable_dual_filter = seq.enable_dual_filter;
c.pixels = Some(&pixels);
Some(c)
} else {
None
};
let frame_start_cdfs: Box<TileCdfContext> = if fh.primary_ref_frame != PRIMARY_REF_NONE {
let st = refs.ok_or(Error::PartitionWalkOutOfRange)?;
let ir = fh
.inter_refs
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let slot = ir.ref_frame_idx[fh.primary_ref_frame as usize] as usize;
let slot_state = st
.slots
.get(slot)
.and_then(|s| s.as_ref())
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut loaded = slot_state.cdfs.clone();
loaded.zero_counts();
loaded
} else {
let mut c = TileCdfContext::new_from_defaults();
c.init_coeff_cdfs(qp.base_q_idx);
Box::new(c)
};
let mut end_cdfs: Option<Box<TileCdfContext>> = None;
let t_tiles = Instant::now();
let tile_parallel =
threads > 1 && tiles.len() > 1 && fh.frame_is_intra && !fh.allow_intrabc && !is_inter_frame;
if tile_parallel {
use std::sync::atomic::{AtomicUsize, Ordering};
let worker_count = threads.min(tiles.len());
let next = AtomicUsize::new(0);
let lr_walk_ref = if lr.uses_lr { Some(&lr_walk) } else { None };
let decode_one =
|tile_num: usize| -> Result<(PartitionWalker, Box<TileCdfContext>), Error> {
let tile = &tiles[tile_num];
let tile_row = (tile_num as u32) / ti.tile_cols;
let tile_col = (tile_num as u32) % ti.tile_cols;
let (r0, r1) = (
ti.mi_row_starts[tile_row as usize],
ti.mi_row_starts[tile_row as usize + 1],
);
let (c0, c1) = (
ti.mi_col_starts[tile_col as usize],
ti.mi_col_starts[tile_col as usize + 1],
);
let (h4, w4) = (r1 - r0, c1 - c0);
let mut local = PartitionWalker::new(
h4,
w4,
TileGeometry {
mi_row_start: 0,
mi_row_end: h4,
mi_col_start: 0,
mi_col_end: w4,
},
)
.ok_or(Error::PartitionWalkOutOfRange)?;
local.set_lr_origin(r0, c0);
local.begin_tile(TileGeometry {
mi_row_start: 0,
mi_row_end: h4,
mi_col_start: 0,
mi_col_end: w4,
});
local.set_current_q_index(i32::from(qp.base_q_idx));
let mut decoder = SymbolDecoder::init_symbol(
&tile.bytes,
tile.bytes.len(),
fh.disable_cdf_update,
)?;
let mut cdfs = frame_start_cdfs.as_ref().clone();
local.decode_tile_syntax_with_lr(
&mut decoder,
&mut cdfs,
¶ms,
lr_walk_ref,
None,
&quant,
dq.delta_q_present,
)?;
Ok((local, Box::new(cdfs)))
};
let mut first_err: Option<Error> = None;
std::thread::scope(|scope| {
let (tx, rx) = std::sync::mpsc::channel();
for _ in 0..worker_count {
let tx = tx.clone();
let next = &next;
let decode_one = &decode_one;
scope.spawn(move || loop {
let t = next.fetch_add(1, Ordering::SeqCst);
if t >= tiles.len() {
break;
}
let out = decode_one(t);
if tx.send((t, out)).is_err() {
break;
}
});
}
drop(tx);
for (t, out) in rx {
match out {
Ok((local, cdfs)) => {
let tile_row = (t as u32) / ti.tile_cols;
let tile_col = (t as u32) % ti.tile_cols;
let (r0, r1) = (
ti.mi_row_starts[tile_row as usize],
ti.mi_row_starts[tile_row as usize + 1],
);
let (c0, c1) = (
ti.mi_col_starts[tile_col as usize],
ti.mi_col_starts[tile_col as usize + 1],
);
walker.import_rect_from(
&local,
0,
0,
r0,
c0,
r1 - r0,
c1 - c0,
sub_x_u8,
sub_y_u8,
);
walker.import_curr_frame_rect_from(&local, r0, c0, sub_x_u8, sub_y_u8);
if lr.uses_lr {
walker.import_lr_units_from(&local, r0, r1, c0, c1, &lr_walk);
}
if !fh.disable_frame_end_update_cdf && t as u32 == ti.context_update_tile_id
{
end_cdfs = Some(cdfs);
}
}
Err(e) => {
if first_err.is_none() {
first_err = Some(e);
}
}
}
}
});
if let Some(e) = first_err {
return Err(e);
}
}
for (tile_num, tile) in tiles.iter().enumerate() {
if tile_parallel {
break;
}
let tile_row = (tile_num as u32) / ti.tile_cols;
let tile_col = (tile_num as u32) % ti.tile_cols;
walker.begin_tile(TileGeometry {
mi_row_start: ti.mi_row_starts[tile_row as usize],
mi_row_end: ti.mi_row_starts[tile_row as usize + 1],
mi_col_start: ti.mi_col_starts[tile_col as usize],
mi_col_end: ti.mi_col_starts[tile_col as usize + 1],
});
walker.set_current_q_index(i32::from(qp.base_q_idx));
let mut decoder =
SymbolDecoder::init_symbol(&tile.bytes, tile.bytes.len(), fh.disable_cdf_update)?;
let mut cdfs = frame_start_cdfs.as_ref().clone();
walker.decode_tile_syntax_with_lr(
&mut decoder,
&mut cdfs,
¶ms,
if lr.uses_lr { Some(&lr_walk) } else { None },
ictx.as_ref(),
&quant,
dq.delta_q_present,
)?;
if !fh.disable_frame_end_update_cdf && tile_num as u32 == ti.context_update_tile_id {
end_cdfs = Some(Box::new(cdfs));
}
}
let end_cdfs = end_cdfs.unwrap_or(frame_start_cdfs);
stages.tiles += t_tiles.elapsed();
let sub_x = u8::from(cc.subsampling_x) as u32;
let sub_y = u8::from(cc.subsampling_y) as u32;
let num_planes = cc.num_planes as usize;
let mut plane_bufs: Vec<Vec<i32>> = Vec::with_capacity(num_planes);
let mut plane_dims: Vec<(u32, u32)> = Vec::with_capacity(num_planes);
for plane in 0..num_planes {
let (rows, cols, src) = walker
.take_curr_frame_plane(plane)
.ok_or(Error::PartitionWalkOutOfRange)?;
let (pw, ph) = if plane == 0 {
(fs.frame_width, fs.frame_height)
} else {
(
(fs.frame_width + sub_x) >> sub_x,
(fs.frame_height + sub_y) >> sub_y,
)
};
if rows < ph || cols < pw {
return Err(Error::PartitionWalkOutOfRange);
}
plane_bufs.push(src);
plane_dims.push((cols, rows));
}
let t_deblock = Instant::now();
if let Some(lf) = fh.loop_filter_params.as_ref() {
if (lf.loop_filter_level[0] != 0 || lf.loop_filter_level[1] != 0)
&& !coded_lossless
&& !fh.allow_intrabc
{
let mut bufs: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in plane_bufs.iter_mut().zip(plane_dims.iter()) {
bufs.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
if threads > 1 {
walker.loop_filter_frame_parallel_from_grid(
lf,
sp,
dlf.delta_lf_multi,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
fs.frame_width,
fs.frame_height,
&mut bufs,
threads,
);
} else {
walker.loop_filter_frame_from_grid(
lf,
sp,
dlf.delta_lf_multi,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
fs.frame_width,
fs.frame_height,
&mut bufs,
);
}
}
}
stages.deblock += t_deblock.elapsed();
let t_cdef = Instant::now();
let cdef_active = fh.cdef_params.as_ref().is_some_and(|cdef| {
!coded_lossless && !fh.allow_intrabc && seq.enable_cdef && !cdef.short_circuited
});
let superres_active = fs.use_superres && fs.upscaled_width > fs.frame_width;
if !superres_active {
let boundary: Vec<crate::loop_restoration::StripeBoundaryRows> = if lr.uses_lr {
plane_bufs
.iter_mut()
.zip(plane_dims.iter())
.enumerate()
.map(|(plane, (buf, &(pw, ph)))| {
let sy = if plane == 0 { 0 } else { sub_y };
let plane_h = (fs.frame_height + sy) >> sy;
let pb = PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
};
crate::loop_restoration::StripeBoundaryRows::capture(
&pb,
plane_h as i32 - 1,
sy as u8,
)
})
.collect()
} else {
Vec::new()
};
let mut bufs: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in plane_bufs.iter_mut().zip(plane_dims.iter()) {
bufs.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
if let (true, Some(cdef)) = (cdef_active, fh.cdef_params.as_ref()) {
if threads > 1 {
walker.cdef_frame_parallel_from_idx(
cdef,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
&mut bufs,
threads,
);
} else {
walker.cdef_frame_in_place_from_idx(
cdef,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
&mut bufs,
);
}
}
stages.cdef += t_cdef.elapsed();
let t_lr = Instant::now();
if lr.uses_lr {
if threads > 1 {
walker.loop_restore_frame_parallel_from_grid(
lr,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
mi_rows,
mi_cols,
fs.frame_height,
fs.upscaled_width,
&mut bufs,
&boundary,
threads,
);
} else {
walker.loop_restore_frame_in_place_from_grid(
lr,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
mi_rows,
mi_cols,
fs.frame_height,
fs.upscaled_width,
&mut bufs,
&boundary,
);
}
}
stages.restoration += t_lr.elapsed();
} else {
let mut deblocked: Option<Vec<Vec<i32>>> = if lr.uses_lr {
Some(plane_bufs.clone())
} else {
None
};
if let Some(cdef) = fh.cdef_params.as_ref() {
if !coded_lossless && !fh.allow_intrabc && seq.enable_cdef && !cdef.short_circuited {
let mut fallback: Option<Vec<Vec<i32>>> = None;
let src_owned: &mut Vec<Vec<i32>> = match deblocked {
Some(ref mut d) => d,
None => fallback.insert(plane_bufs.clone()),
};
let mut src: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in src_owned.iter_mut().zip(plane_dims.iter()) {
src.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
let mut dst: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in plane_bufs.iter_mut().zip(plane_dims.iter()) {
dst.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
walker.cdef_frame_from_idx(
cdef,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
&src,
&mut dst,
);
}
}
stages.cdef += t_cdef.elapsed();
let t_superres = Instant::now();
if fs.use_superres && fs.upscaled_width > fs.frame_width {
let sr_ctx = crate::superres::SuperresFrameContext {
use_superres: true,
frame_width: fs.frame_width,
upscaled_width: fs.upscaled_width,
frame_height: fs.frame_height,
mi_cols: fs.mi_cols,
num_planes: cc.num_planes,
bit_depth: cc.bit_depth,
subsampling_x: u8::from(cc.subsampling_x),
subsampling_y: u8::from(cc.subsampling_y),
};
let mut new_dims: Vec<(u32, u32)> = Vec::with_capacity(num_planes);
for plane in 0..num_planes {
let (out_w, out_h) = if plane == 0 {
(fs.upscaled_width, fs.frame_height)
} else {
(
(fs.upscaled_width + sub_x) >> sub_x,
(fs.frame_height + sub_y) >> sub_y,
)
};
new_dims.push((out_w, out_h));
}
let upscale_set = |bufs: &mut Vec<Vec<i32>>| -> Result<(), Error> {
let mut inputs_owned = std::mem::take(bufs);
let mut inputs: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in inputs_owned.iter_mut().zip(plane_dims.iter()) {
inputs.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
let mut outputs_owned: Vec<Vec<i32>> = new_dims
.iter()
.map(|&(pw, ph)| vec![0i32; (pw as usize) * (ph as usize)])
.collect();
{
let mut outputs: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in outputs_owned.iter_mut().zip(new_dims.iter()) {
outputs.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
crate::superres::upscale_frame(&sr_ctx, &inputs, &mut outputs)
.map_err(|_| Error::PartitionWalkOutOfRange)?;
}
*bufs = outputs_owned;
Ok(())
};
upscale_set(&mut plane_bufs)?;
if let Some(d) = deblocked.as_mut() {
upscale_set(d)?;
}
plane_dims = new_dims;
}
stages.superres += t_superres.elapsed();
let t_lr = Instant::now();
if let Some(mut curr_owned) = deblocked {
let mut cdef_owned: Vec<Vec<i32>> = plane_bufs.clone();
let mut curr: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in curr_owned.iter_mut().zip(plane_dims.iter()) {
curr.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
let mut cdef_bufs: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in cdef_owned.iter_mut().zip(plane_dims.iter()) {
cdef_bufs.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
let mut lr_out: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in plane_bufs.iter_mut().zip(plane_dims.iter()) {
lr_out.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
walker.loop_restore_frame_from_grid(
lr,
cc.num_planes,
cc.bit_depth,
sub_x as u8,
sub_y as u8,
mi_rows,
mi_cols,
fs.frame_height,
fs.upscaled_width,
&curr,
&cdef_bufs,
&mut lr_out,
);
}
stages.restoration += t_lr.elapsed();
}
let t_ref = Instant::now();
for plane in 0..num_planes {
let (out_w, out_h) = if plane == 0 {
(fs.upscaled_width, fs.frame_height)
} else {
(
(fs.upscaled_width + sub_x) >> sub_x,
(fs.frame_height + sub_y) >> sub_y,
)
};
let (pw, ph) = plane_dims[plane];
if (pw, ph) == (out_w, out_h) {
continue;
}
if pw < out_w || ph < out_h {
return Err(Error::PartitionWalkOutOfRange);
}
let src = &plane_bufs[plane];
let mut buf = vec![0i32; (out_w as usize) * (out_h as usize)];
for y in 0..out_h as usize {
let row = &src[y * pw as usize..y * pw as usize + out_w as usize];
buf[y * out_w as usize..(y + 1) * out_w as usize].copy_from_slice(row);
}
plane_bufs[plane] = buf;
plane_dims[plane] = (out_w, out_h);
}
const REFMVS_LIMIT: i16 = (1 << 12) - 1;
let cells = (mi_rows as usize) * (mi_cols as usize);
let mut mf_ref_frames: Vec<i8> = vec![-1; cells];
let mut mf_mvs: Vec<i16> = vec![0; cells * 2];
if is_inter_frame {
let st = refs.ok_or(Error::PartitionWalkOutOfRange)?;
let raw_refs = walker.ref_frames();
let raw_mvs = walker.mvs();
let hint_bits = if seq.enable_order_hint {
u32::from(seq.order_hint_bits)
} else {
0
};
for cell in 0..cells {
for list in 0..2usize {
let r = raw_refs[cell * 2 + list];
if r > 0 {
let slot = ref_frame_idx[(r - 1) as usize] as usize;
let dist = get_relative_dist(
st.info.order_hint[slot] as i32,
fh.order_hint as i32,
hint_bits,
);
if dist < 0 {
let mv_row = raw_mvs[(cell * 2 + list) * 2];
let mv_col = raw_mvs[(cell * 2 + list) * 2 + 1];
if mv_row.abs() <= REFMVS_LIMIT && mv_col.abs() <= REFMVS_LIMIT {
mf_ref_frames[cell] = r;
mf_mvs[cell * 2] = mv_row;
mf_mvs[cell * 2 + 1] = mv_col;
}
}
}
}
}
}
let segment_ids: Vec<i32> = walker.segment_ids().iter().map(|&v| v.max(0)).collect();
drop(walker);
let store_planes = !seq.reduced_still_picture_header;
let ref_plane_dims: Vec<(u32, u32)> = plane_dims.clone();
let to_u16 = |buf: &[i32]| -> Vec<u16> { buf.iter().map(|&v| v.max(0) as u16).collect() };
let max_val: i32 = (1 << cc.bit_depth) - 1;
let narrow = |buf: Vec<i32>| -> Vec<u8> {
if cc.bit_depth == 8 {
buf.into_iter().map(|v| v.clamp(0, 255) as u8).collect()
} else {
let mut out = Vec::with_capacity(buf.len() * 2);
for v in buf {
let v = v.clamp(0, max_val) as u16;
out.extend_from_slice(&v.to_le_bytes());
}
out
}
};
let effective_grain: Option<crate::uncompressed_header_tail::FilmGrainParams> =
match fh.film_grain_params.as_ref() {
Some(fg) if fg.apply_grain && fg.predicted => {
let slot = fg.film_grain_params_ref_idx as usize;
refs.and_then(|st| st.slots.get(slot))
.and_then(|s| s.as_ref())
.and_then(|s| s.grain_params.clone())
.map(|mut g| {
g.grain_seed = fg.grain_seed;
g.apply_grain = true;
g
})
}
Some(fg) if fg.apply_grain => Some(fg.clone()),
_ => None,
};
let (ref_planes, planes): (Vec<Vec<u16>>, Vec<Vec<u8>>) = match effective_grain.as_ref() {
Some(fg) if fg.apply_grain => {
let ref_planes: Vec<Vec<u16>> = if store_planes {
plane_bufs.iter().map(|b| to_u16(b)).collect()
} else {
Vec::new()
};
stages.ref_store += t_ref.elapsed();
let t_grain = Instant::now();
let mut bufs: Vec<PlaneBuffer<'_>> = Vec::with_capacity(num_planes);
for (buf, &(pw, ph)) in plane_bufs.iter_mut().zip(plane_dims.iter()) {
bufs.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
film_grain_synthesis(
fg,
cc.bit_depth,
cc.num_planes,
u8::from(cc.subsampling_x),
u8::from(cc.subsampling_y),
cc.matrix_coefficients,
&mut bufs,
);
stages.grain += t_grain.elapsed();
let t_out = Instant::now();
let planes: Vec<Vec<u8>> = plane_bufs.into_iter().map(narrow).collect();
stages.output += t_out.elapsed();
(ref_planes, planes)
}
_ => {
let t_out = Instant::now();
let mut ref_planes: Vec<Vec<u16>> = Vec::with_capacity(num_planes);
let mut planes: Vec<Vec<u8>> = Vec::with_capacity(num_planes);
for buf in plane_bufs {
if store_planes {
ref_planes.push(to_u16(&buf));
}
planes.push(narrow(buf));
}
stages.ref_store += t_ref.elapsed();
stages.output += t_out.elapsed();
(ref_planes, planes)
}
};
Ok(DecodedFrameInternal {
frame: SpecFrame {
width: fs.upscaled_width,
height: fs.frame_height,
planes,
plane_dims,
bit_depth: cc.bit_depth,
},
ref_planes,
ref_plane_dims,
mf_mvs,
mf_ref_frames,
order_hints_by_ref,
end_cdfs,
mi_rows,
mi_cols,
segment_ids,
grain_params: effective_grain,
})
}
fn motion_field_estimation(
refs: &SpecRefState,
fh: &FrameHeader,
seq: &SequenceHeader,
mi_rows: u32,
mi_cols: u32,
) -> Result<MotionFieldMvs, Error> {
use crate::inter_pred::{motion_field_estimation_core, MotionFieldSlot};
let ir = fh
.inter_refs
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let hint_bits = if seq.enable_order_hint {
u32::from(seq.order_hint_bits)
} else {
0
};
let order_hint = fh.order_hint as i32;
let mut order_hints = [0i32; ALTREF_FRAME + 1];
for i in 0..7 {
order_hints[LAST_FRAME + i] = refs.info.order_hint[ir.ref_frame_idx[i] as usize] as i32;
}
let mut slots: [Option<MotionFieldSlot<'_>>; 8] = [None; 8];
for (i, slot) in refs.slots.iter().enumerate() {
slots[i] = slot.as_ref().map(|st| MotionFieldSlot {
mf_mvs: &st.mf_mvs,
mf_ref_frames: &st.mf_ref_frames,
saved_order_hints: st.saved_order_hints,
mi_rows: st.mi_rows,
mi_cols: st.mi_cols,
frame_is_intra: st.frame_is_intra,
});
}
Ok(motion_field_estimation_core(
&slots,
&ir.ref_frame_idx,
&order_hints,
order_hint,
hint_bits,
mi_rows,
mi_cols,
))
}
pub fn decode_av1_spec(input: &[u8]) -> Result<Vec<SpecFrame>, Error> {
SpecDecodeSession::new().decode_ivf(input)
}
pub const MAX_PICTURE_SIZE: u32 = 35_651_584;
fn check_picture_size(fh: &FrameHeader, cap: u32) -> Result<(), Error> {
if let Some(fs) = fh.frame_size.as_ref() {
let samples = u64::from(fs.upscaled_width) * u64::from(fs.frame_height);
if samples > u64::from(cap) {
return Err(Error::PictureSizeExceedsLimit);
}
}
Ok(())
}
pub fn decode_av1_spec_at_operating_point(
input: &[u8],
operating_point: u8,
) -> Result<Vec<SpecFrame>, Error> {
let mut session = SpecDecodeSession::new();
session.set_operating_point(operating_point)?;
session.decode_ivf(input)
}
#[derive(Debug)]
pub struct SpecDecodeSession {
seq: Option<SequenceHeader>,
refs: SpecRefState,
chosen_op: u8,
op_idc: u16,
max_picture_size: u32,
stages: DecodeStageTimings,
threads: usize,
}
impl Default for SpecDecodeSession {
fn default() -> Self {
Self::new()
}
}
impl SpecDecodeSession {
#[must_use]
pub fn new() -> Self {
Self {
seq: None,
refs: SpecRefState::new(),
chosen_op: 0,
op_idc: 0,
max_picture_size: MAX_PICTURE_SIZE,
stages: DecodeStageTimings::default(),
threads: 1,
}
}
pub fn set_threads(&mut self, threads: usize) {
self.threads = threads.max(1);
}
#[must_use]
pub fn threads(&self) -> usize {
self.threads
}
pub fn set_max_picture_size(&mut self, max_samples: u32) {
self.max_picture_size = max_samples;
}
#[must_use]
pub fn max_picture_size(&self) -> u32 {
self.max_picture_size
}
pub fn decode_ivf(&mut self, input: &[u8]) -> Result<Vec<SpecFrame>, Error> {
let reader = IvfReader::new(input).map_err(|_| Error::UnexpectedEnd)?;
let records = reader.read_all().map_err(|_| Error::UnexpectedEnd)?;
let mut out = Vec::new();
for record in records {
out.extend(self.decode_temporal_unit(&record.payload)?);
}
Ok(out)
}
pub fn set_operating_point(&mut self, operating_point: u8) -> Result<(), Error> {
self.chosen_op = operating_point;
if let Some(s) = self.seq.as_ref() {
let idx = usize::from(operating_point);
if idx >= s.operating_points.len() {
return Err(Error::OperatingPointOutOfRange);
}
self.op_idc = s.operating_points[idx].operating_point_idc;
}
Ok(())
}
#[must_use]
pub fn operating_point_idc(&self) -> u16 {
self.op_idc
}
pub fn decode_temporal_unit(&mut self, payload: &[u8]) -> Result<Vec<SpecFrame>, Error> {
let mut out = Vec::new();
self.stages = DecodeStageTimings::default();
decode_temporal_unit_spec(
payload,
&mut self.seq,
&mut self.refs,
&mut out,
self.chosen_op,
&mut self.op_idc,
self.max_picture_size,
&mut self.stages,
self.threads,
)?;
Ok(out)
}
#[doc(hidden)]
#[must_use]
pub fn last_stage_timings(&self) -> DecodeStageTimings {
self.stages
}
pub fn reset_references(&mut self) {
self.refs = SpecRefState::new();
}
}
fn reference_frame_update(
refs: &mut SpecRefState,
fh: &FrameHeader,
decoded: DecodedFrameInternal,
) -> Result<SpecFrame, Error> {
if fh.refresh_frame_flags == 0 {
return Ok(decoded.frame);
}
let fs = fh
.frame_size
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut gm_params = crate::uncompressed_header_tail::prev_gm_params_default();
if let Some(g) = fh.global_motion_params.as_ref() {
gm_params = g.gm_params;
}
let (lf_ref_deltas, lf_mode_deltas) = fh
.loop_filter_params
.as_ref()
.map(|lf| (lf.loop_filter_ref_deltas, lf.loop_filter_mode_deltas))
.unwrap_or((
crate::uncompressed_header_tail::LOOP_FILTER_REF_DELTAS_DEFAULT,
[0i8; 2],
));
let (seg_active, seg_data) = fh
.segmentation_params
.as_ref()
.map(|sp| (sp.segment_feature_active, sp.segment_feature_data))
.unwrap_or((
[[false; crate::uncompressed_header_tail::SEG_LVL_MAX]
as [bool; crate::uncompressed_header_tail::SEG_LVL_MAX];
crate::uncompressed_header_tail::MAX_SEGMENTS],
[[0i16; crate::uncompressed_header_tail::SEG_LVL_MAX];
crate::uncompressed_header_tail::MAX_SEGMENTS],
));
let DecodedFrameInternal {
frame,
ref_planes,
ref_plane_dims,
mf_mvs,
mf_ref_frames,
order_hints_by_ref,
end_cdfs,
mi_rows,
mi_cols,
segment_ids,
grain_params,
} = decoded;
let payload = Arc::new(SpecRefSlot {
planes: ref_planes,
plane_dims: ref_plane_dims,
mf_mvs,
mf_ref_frames,
saved_order_hints: order_hints_by_ref,
mi_rows,
mi_cols,
frame_is_intra: fh.frame_is_intra,
frame_type_is_key: matches!(fh.frame_type, crate::frame_header::FrameType::Key),
bit_depth: frame.bit_depth,
cdfs: end_cdfs,
gm_params,
lf_ref_deltas,
lf_mode_deltas,
segment_ids,
grain_params,
});
for i in 0..NUM_REF_FRAMES as usize {
if (fh.refresh_frame_flags >> i) & 1 != 0 {
refs.info.valid[i] = true;
refs.info.order_hint[i] = fh.order_hint;
refs.info.frame_id[i] = fh.current_frame_id;
refs.info.upscaled_width[i] = fs.upscaled_width;
refs.info.frame_height[i] = fs.frame_height;
refs.info.render_width[i] = fs.render_width;
refs.info.render_height[i] = fs.render_height;
refs.info.frame_type_is_key[i] = payload.frame_type_is_key;
refs.info.saved_gm_params[i] = payload.gm_params;
refs.info.saved_lf_ref_deltas[i] = payload.lf_ref_deltas;
refs.info.saved_lf_mode_deltas[i] = payload.lf_mode_deltas;
refs.info.saved_seg_feature_active[i] = seg_active;
refs.info.saved_seg_feature_data[i] = seg_data;
refs.slots[i] = Some(Arc::clone(&payload));
}
}
Ok(frame)
}
fn output_existing_frame(
refs: &SpecRefState,
fh: &FrameHeader,
seq: &SequenceHeader,
) -> Result<SpecFrame, Error> {
let idx = fh
.frame_to_show_map_idx
.ok_or(Error::PartitionWalkOutOfRange)? as usize;
if idx >= NUM_REF_FRAMES as usize {
return Err(Error::PartitionWalkOutOfRange);
}
let slot = refs.slots[idx]
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let mut grain_bufs: Option<Vec<Vec<i32>>> = None;
if seq.film_grain_params_present {
if let Some(fg) = slot.grain_params.as_ref() {
if fg.apply_grain {
let mut bufs: Vec<Vec<i32>> = slot
.planes
.iter()
.map(|p| p.iter().map(|&v| v as i32).collect())
.collect();
let cc = &seq.color_config;
let mut views: Vec<PlaneBuffer<'_>> = Vec::with_capacity(bufs.len());
for (buf, &(pw, ph)) in bufs.iter_mut().zip(slot.plane_dims.iter()) {
views.push(PlaneBuffer {
rows: ph,
cols: pw,
samples: buf,
});
}
film_grain_synthesis(
fg,
slot.bit_depth,
cc.num_planes,
u8::from(cc.subsampling_x),
u8::from(cc.subsampling_y),
cc.matrix_coefficients,
&mut views,
);
grain_bufs = Some(bufs);
}
}
}
let max_val: i32 = (1i32 << slot.bit_depth) - 1;
let planes: Vec<Vec<u8>> = match grain_bufs {
Some(bufs) => bufs
.into_iter()
.map(|p| {
if slot.bit_depth == 8 {
p.into_iter().map(|v| v.clamp(0, 255) as u8).collect()
} else {
let mut out = Vec::with_capacity(p.len() * 2);
for v in p {
out.extend_from_slice(&(v.clamp(0, max_val) as u16).to_le_bytes());
}
out
}
})
.collect(),
None => slot
.planes
.iter()
.map(|p| {
if slot.bit_depth == 8 {
p.iter().map(|&v| v.min(255) as u8).collect()
} else {
let mut out = Vec::with_capacity(p.len() * 2);
for &v in p {
out.extend_from_slice(&v.to_le_bytes());
}
out
}
})
.collect(),
};
let (w, h) = *slot
.plane_dims
.first()
.ok_or(Error::PartitionWalkOutOfRange)?;
Ok(SpecFrame {
width: w,
height: h,
planes,
plane_dims: slot.plane_dims.clone(),
bit_depth: slot.bit_depth,
})
}
#[allow(clippy::too_many_arguments)]
fn decode_temporal_unit_spec(
payload: &[u8],
seq: &mut Option<SequenceHeader>,
refs: &mut SpecRefState,
out: &mut Vec<SpecFrame>,
chosen_op: u8,
op_idc: &mut u16,
max_picture_size: u32,
stages: &mut DecodeStageTimings,
threads: usize,
) -> Result<(), Error> {
struct PendingFrame {
fh: FrameHeader,
fh_payload: Vec<u8>,
tiles: Vec<TilePayload>,
next_tg_start: u32,
}
let mut pending: Option<PendingFrame> = None;
for desc in ObuIter::new(payload) {
let desc = desc?;
if *op_idc != 0
&& desc.extension_flag
&& !matches!(
desc.obu_type,
ObuType::SequenceHeader | ObuType::TemporalDelimiter
)
{
let in_temporal = (*op_idc >> desc.temporal_id) & 1 != 0;
let in_spatial = (*op_idc >> (desc.spatial_id + 8)) & 1 != 0;
if !in_temporal || !in_spatial {
continue;
}
}
match desc.obu_type {
ObuType::TemporalDelimiter | ObuType::Padding | ObuType::Metadata => {}
ObuType::TileList => {}
ObuType::SequenceHeader => {
let sh = parse_sequence_header(desc.payload)?;
let idx = usize::from(chosen_op);
if idx >= sh.operating_points.len() {
return Err(Error::OperatingPointOutOfRange);
}
*op_idc = sh.operating_points[idx].operating_point_idc;
*seq = Some(sh);
}
ObuType::FrameHeader => {
if pending.as_ref().is_some_and(|p| !p.tiles.is_empty()) {
return Err(Error::TileGroupInvalid);
}
let s = seq.as_ref().ok_or(Error::PartitionWalkOutOfRange)?;
let t_parse = Instant::now();
let fh = parse_frame_header_with_refs(desc.payload, s, &refs.info)?;
stages.parse += t_parse.elapsed();
check_picture_size(&fh, max_picture_size)?;
if fh.show_existing_frame {
out.push(output_existing_frame(refs, &fh, s)?);
if fh.refresh_frame_flags != 0 {
let idx = fh
.frame_to_show_map_idx
.ok_or(Error::PartitionWalkOutOfRange)?
as usize;
let payload = refs.slots[idx]
.as_ref()
.cloned()
.ok_or(Error::PartitionWalkOutOfRange)?;
let loaded_order_hint = refs.info.order_hint[idx];
let loaded_frame_id = refs.info.frame_id[idx];
let loaded_uw = refs.info.upscaled_width[idx];
let loaded_fh = refs.info.frame_height[idx];
let loaded_rw = refs.info.render_width[idx];
let loaded_rh = refs.info.render_height[idx];
for i in 0..NUM_REF_FRAMES as usize {
if (fh.refresh_frame_flags >> i) & 1 != 0 {
refs.info.valid[i] = true;
refs.info.order_hint[i] = loaded_order_hint;
refs.info.frame_id[i] = loaded_frame_id;
refs.info.upscaled_width[i] = loaded_uw;
refs.info.frame_height[i] = loaded_fh;
refs.info.render_width[i] = loaded_rw;
refs.info.render_height[i] = loaded_rh;
refs.info.frame_type_is_key[i] = payload.frame_type_is_key;
refs.info.saved_gm_params[i] = payload.gm_params;
refs.info.saved_lf_ref_deltas[i] = payload.lf_ref_deltas;
refs.info.saved_lf_mode_deltas[i] = payload.lf_mode_deltas;
refs.info.saved_seg_feature_active[i] =
refs.info.saved_seg_feature_active[idx];
refs.info.saved_seg_feature_data[i] =
refs.info.saved_seg_feature_data[idx];
refs.slots[i] = Some(payload.clone());
}
}
}
pending = None;
} else {
pending = Some(PendingFrame {
fh,
fh_payload: desc.payload.to_vec(),
tiles: Vec::new(),
next_tg_start: 0,
});
}
}
ObuType::RedundantFrameHeader => {
let p = pending.as_ref().ok_or(Error::TileGroupInvalid)?;
if desc.payload != &p.fh_payload[..] {
return Err(Error::TileGroupInvalid);
}
}
ObuType::TileGroup => {
let s = seq.as_ref().ok_or(Error::PartitionWalkOutOfRange)?;
let p = pending.as_mut().ok_or(Error::TileGroupInvalid)?;
let ti =
p.fh.tile_info
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let num_tiles = ti.tile_cols * ti.tile_rows;
let t_parse = Instant::now();
let parsed = parse_tile_group_body_with_ti(ti, desc.payload)?;
stages.parse += t_parse.elapsed();
if parsed.tg_start != p.next_tg_start {
return Err(Error::TileGroupInvalid);
}
p.next_tg_start = parsed.tg_end + 1;
p.tiles.extend(parsed.tiles);
if parsed.tg_end + 1 == num_tiles {
let pf = pending.take().expect("pending frame checked above");
let decoded =
decode_frame_spec_tiles(s, &pf.fh, &pf.tiles, Some(refs), stages, threads)?;
let frame = reference_frame_update(refs, &pf.fh, decoded)?;
if pf.fh.show_frame {
out.push(frame);
}
}
}
ObuType::Frame => {
if pending.as_ref().is_some_and(|p| !p.tiles.is_empty()) {
return Err(Error::TileGroupInvalid);
}
pending = None;
let s = seq.as_ref().ok_or(Error::PartitionWalkOutOfRange)?;
let t_parse = Instant::now();
let fh = parse_frame_header_with_refs(desc.payload, s, &refs.info)?;
check_picture_size(&fh, max_picture_size)?;
let tg_offset = fh.bits_consumed.div_ceil(8);
if tg_offset > desc.payload.len() {
return Err(Error::UnexpectedEnd);
}
let tg_body = &desc.payload[tg_offset..];
let ti = fh
.tile_info
.as_ref()
.ok_or(Error::PartitionWalkOutOfRange)?;
let tiles = parse_whole_frame_tile_group(ti, tg_body, true)?;
stages.parse += t_parse.elapsed();
let decoded = decode_frame_spec_tiles(s, &fh, &tiles, Some(refs), stages, threads)?;
let frame = reference_frame_update(refs, &fh, decoded)?;
if fh.show_frame {
out.push(frame);
}
}
_ => return Err(Error::PartitionWalkOutOfRange),
}
}
if pending.as_ref().is_some_and(|p| !p.tiles.is_empty()) {
return Err(Error::TileGroupInvalid);
}
Ok(())
}