diff --git a/crates/vp8/src/decoder.rs b/crates/vp8/src/decoder.rs index b20a89ed5..9e3f66667 100644 --- a/crates/vp8/src/decoder.rs +++ b/crates/vp8/src/decoder.rs @@ -8,7 +8,7 @@ use crate::boolcoder::BoolDecoder; use crate::error::{Error, Result, ensure}; use crate::header::{EntropyContext, FrameHeader, FrameTag, LfDeltas, Segmentation, read_header}; use crate::loopfilter::{self, Kind, Params}; -use crate::predict::{self, RefPlane}; +use crate::predict::{self, BlockRect, DstPlane, RefPlane}; use crate::spec_tables::{ AC_QLOOKUP, BILINEAR_FILTERS, BMODE_PROBS, COEFF_BANDS, DC_QLOOKUP, KF_BMODE_PROBS, KF_UV_MODE_PROBS, KF_YMODE_PROBS, MODE_CONTEXTS, MVPARTITION_PROBS, SIXTAP_FILTERS, SUB_MV_REF_PROBS, @@ -168,252 +168,58 @@ impl Decoder { refs_untouched && probs_kept && segments_kept && deltas_kept } - /// Decode one compressed frame. Returns the picture when the frame is meant to be shown - /// (`show_frame`), `None` for a hidden frame (an altref). - pub fn decode(&mut self, data: &[u8], pts: i64) -> Result> { - let tag = FrameTag::parse(data)?; - ensure!(tag.first_part_size > 0, "empty first partition"); - if tag.key_frame { - ensure!(tag.width > 0 && tag.height > 0, "key frame of size {}x{}", tag.width, tag.height); - let (w, h) = (usize::from(tag.width), usize::from(tag.height)); - let (mb_w, mb_h) = (w.div_ceil(16), h.div_ceil(16)); - let samples = mb_w.checked_mul(16).and_then(|w| w.checked_mul(mb_h)).and_then(|h| h.checked_mul(16)); - ensure!(samples.is_some_and(|n| n <= MAX_LUMA_SAMPLES), "frame size {}x{} exceeds the decoder memory limit", tag.width, tag.height); - if w != self.width || h != self.height { - self.width = w; - self.height = h; - self.mb_w = mb_w; - self.mb_h = mb_h; - self.seg_map = vec![0; self.mb_w * self.mb_h]; - } - // key frame defaults (9.3, 9.4, 9.7, 13.4, 16.1, 17.2) - self.ctx = EntropyContext::default(); - self.seg.absolute = false; - self.seg.quant = [0; 4]; - self.seg.lf_level = [0; 4]; - self.lf.ref_deltas = [0; 4]; - self.lf.mode_deltas = [0; 4]; - self.sign_bias = [false; 4]; - } else if self.refs[0].is_none() { - return Err(Error::MissingReference("inter frame before the first key frame".into())); - } - ensure!(self.mb_w > 0 && self.mb_w <= MAX_MBS && self.mb_h <= MAX_MBS, "no frame size"); - let first_start = tag.header_len; - let first_end = first_start.saturating_add(tag.first_part_size); - ensure!(first_end <= data.len(), "first partition ({} bytes) runs past the frame ({} bytes)", tag.first_part_size, data.len()); - let mut bd = BoolDecoder::new(&data[first_start..first_end]); - let mut saved = None; - let hdr = read_header(&mut bd, tag.key_frame, &mut self.seg, &mut self.lf, &mut self.ctx, &mut saved); - if !tag.key_frame { - self.sign_bias[2] = hdr.sign_bias_golden; - self.sign_bias[3] = hdr.sign_bias_altref; + /// Key frame setup: frame size (with its limits) and the default state of 9.3, 9.4, 9.7, + /// 13.4, 16.1 and 17.2. + fn begin_key_frame(&mut self, tag: &FrameTag) -> Result<()> { + ensure!(tag.width > 0 && tag.height > 0, "key frame of size {}x{}", tag.width, tag.height); + let (w, h) = (usize::from(tag.width), usize::from(tag.height)); + let (mb_w, mb_h) = (w.div_ceil(16), h.div_ceil(16)); + let samples = mb_w.checked_mul(16).and_then(|w| w.checked_mul(mb_h)).and_then(|h| h.checked_mul(16)); + ensure!(samples.is_some_and(|n| n <= MAX_LUMA_SAMPLES), "frame size {}x{} exceeds the decoder memory limit", tag.width, tag.height); + if w != self.width || h != self.height { + self.width = w; + self.height = h; + self.mb_w = mb_w; + self.mb_h = mb_h; + self.seg_map = vec![0; self.mb_w * self.mb_h]; } + self.ctx = EntropyContext::default(); + self.seg.absolute = false; + self.seg.quant = [0; 4]; + self.seg.lf_level = [0; 4]; + self.lf.ref_deltas = [0; 4]; + self.lf.mode_deltas = [0; 4]; + self.sign_bias = [false; 4]; + Ok(()) + } - // token partitions (9.5): sizes of all but the last, then the data - let n = hdr.partitions; - let sizes_start = first_end; - let mut pos = sizes_start.saturating_add(3 * (n - 1)).min(data.len()); - let mut parts: Vec<&[u8]> = Vec::with_capacity(n); - for i in 0..n { - let size = if i + 1 < n { - let s = sizes_start + 3 * i; - match data.get(s..s + 3) { - Some(b) => usize::from(b[0]) | usize::from(b[1]) << 8 | usize::from(b[2]) << 16, - None => 0, - } - } else { - data.len() - pos - }; - let end = pos.saturating_add(size).min(data.len()); - parts.push(&data[pos..end]); - pos = end; - } - let mut token_decoders: Vec = parts.into_iter().map(BoolDecoder::new).collect(); - - let mut frame = Frame::new(self.mb_w * 16, self.mb_h * 16); - let dequant: [Dequant; 4] = std::array::from_fn(|s| { + /// Dequantization factors of each of the four segments (9.6, 14.1). + fn segment_dequant(&self, hdr: &FrameHeader) -> [Dequant; 4] { + std::array::from_fn(|s| { let q = if self.seg.enabled { let v = i32::from(self.seg.quant[s]); if self.seg.absolute { v } else { hdr.quant.y_ac + v } } else { hdr.quant.y_ac }; - Dequant::new(q.clamp(0, 127), &hdr) - }); - let filters: &[i32; 48] = if tag.version == 0 || tag.version > 3 { &SIXTAP_FILTERS } else { &BILINEAR_FILTERS }; - let full_pixel = tag.version == 3; + Dequant::new(q.clamp(0, 127), hdr) + }) + } - let stride_info = self.mb_w + 1; - let mut infos = vec![MbInfo::default(); stride_info * (self.mb_h + 1)]; - let mut lf_info: Vec<(u8, bool)> = vec![(0, false); self.mb_w * self.mb_h]; - // contexts: B_PRED modes above / left (key frames), non-zero flags above / left - let mut above_bmodes = vec![[B_DC_PRED; 4]; self.mb_w]; - let mut above_nz = vec![[0u8; 9]; self.mb_w]; - let refs = self.refs.clone(); - - for mb_y in 0..self.mb_h { - let mut left_bmodes = [B_DC_PRED; 4]; - let mut left_nz = [0u8; 9]; - let part = mb_y % n; - for mb_x in 0..self.mb_w { - let idx = (mb_y + 1) * stride_info + mb_x + 1; - // ---- prediction record (first partition) - let seg_idx = mb_y * self.mb_w + mb_x; - if self.seg.update_map { - let s = bd.read_tree(&SEGMENT_TREE, &self.seg.tree_probs); - if let Some(m) = self.seg_map.get_mut(seg_idx) { - *m = s; - } - } - let segment = if self.seg.enabled { self.seg_map.get(seg_idx).copied().unwrap_or(0) as usize & 3 } else { 0 }; - let coeff_skip = hdr.skip_prob.is_some_and(|p| bd.read_bool(p)); - let mut bmodes = [B_DC_PRED; 16]; - let mut uv_mode = DC_PRED; - let mut info = MbInfo::default(); - if tag.key_frame { - info.mode = bd.read_tree(&KF_YMODE_TREE, &KF_YMODE_PROBS); - if info.mode == B_PRED { - for b in 0..16 { - let a = if b < 4 { above_bmodes[mb_x][b] } else { bmodes[b - 4] }; - let l = if b & 3 == 0 { left_bmodes[b >> 2] } else { bmodes[b - 1] }; - let o = (usize::from(a) * 10 + usize::from(l)) * 9; - bmodes[b] = bd.read_tree(&BMODE_TREE, &KF_BMODE_PROBS[o..o + 9]); - } - } else { - bmodes = [implied_bmode(info.mode); 16]; - } - uv_mode = bd.read_tree(&UV_MODE_TREE, &KF_UV_MODE_PROBS); - } else if bd.read_bool(hdr.prob_intra) { - info.ref_frame = if bd.read_bool(hdr.prob_last) { 2 + u8::from(bd.read_bool(hdr.prob_golden)) } else { 1 }; - self.read_inter_modes(&mut bd, &infos, idx, stride_info, mb_x, mb_y, &mut info); - } else { - info.mode = bd.read_tree(&YMODE_TREE, &self.ctx.ymode); - if info.mode == B_PRED { - for m in &mut bmodes { - *m = bd.read_tree(&BMODE_TREE, &BMODE_PROBS); - } - } - uv_mode = bd.read_tree(&UV_MODE_TREE, &self.ctx.uv_mode); - } - above_bmodes[mb_x] = [bmodes[12], bmodes[13], bmodes[14], bmodes[15]]; - left_bmodes = [bmodes[3], bmodes[7], bmodes[11], bmodes[15]]; - if let Some(slot) = infos.get_mut(idx) { - *slot = info; - } - - // ---- residual (token partition) - let has_y2 = info.mode != B_PRED && info.mode != MV_SPLIT; - let mut coeffs = [[0i16; 16]; 25]; - let mut nonzero = [false; 25]; - let td = &mut token_decoders[part]; - let mut has_coeffs = false; - if coeff_skip { - let keep_y2 = if has_y2 { 0 } else { left_nz[8] }; - let keep_y2_above = if has_y2 { 0 } else { above_nz[mb_x][8] }; - left_nz = [0; 9]; - above_nz[mb_x] = [0; 9]; - left_nz[8] = keep_y2; - above_nz[mb_x][8] = keep_y2_above; - } else { - let dq = &dequant[segment]; - let an = &mut above_nz[mb_x]; - let first = if has_y2 { - let ctx = usize::from(an[8] + left_nz[8]); - let eob = read_coefficients(td, &self.ctx.coeff[1], ctx, 0, dq.y2, &mut coeffs[24]); - let nz = u8::from(eob > 0); - an[8] = nz; - left_nz[8] = nz; - nonzero[24] = eob > 0; - has_coeffs |= eob > 0; - 1 - } else { - 0 - }; - let y_probs = if has_y2 { &self.ctx.coeff[0] } else { &self.ctx.coeff[3] }; - for b in 0..16 { - let (bx, by) = (b & 3, b >> 2); - let ctx = usize::from(an[bx] + left_nz[by]); - let eob = read_coefficients(td, y_probs, ctx, first, dq.y, &mut coeffs[b]); - let nz = eob > first; - an[bx] = u8::from(nz); - left_nz[by] = u8::from(nz); - nonzero[b] = nz; - has_coeffs |= nz; - } - for b in 16..24 { - let (off, k) = if b < 20 { (4, b - 16) } else { (6, b - 20) }; - let (bx, by) = (k & 1, k >> 1); - let ctx = usize::from(an[off + bx] + left_nz[off + by]); - let eob = read_coefficients(td, &self.ctx.coeff[2], ctx, 0, dq.uv, &mut coeffs[b]); - let nz = eob > 0; - an[off + bx] = u8::from(nz); - left_nz[off + by] = u8::from(nz); - nonzero[b] = nz; - has_coeffs |= nz; - } - if has_y2 && nonzero[24] { - let dc = transform::inverse_wht(&coeffs[24]); - for b in 0..16 { - coeffs[b][0] = dc[b]; - nonzero[b] |= dc[b] != 0; - } - } - } - - // ---- loop filter level (15.1, 9.3, 9.4) - let level = self.filter_level(&hdr, segment, &info); - let inner = !has_y2 || has_coeffs; - if let Some(l) = lf_info.get_mut(seg_idx) { - *l = (level, inner); - } - - // ---- reconstruction - let (px, py) = (mb_x * 16, mb_y * 16); - let ys = frame.w; - let cs = frame.w / 2; - if info.ref_frame == 0 { - if info.mode == B_PRED { - for b in 0..16 { - predict::predict_subblock(&mut frame.y, ys, px, py, b, bmodes[b], mb_x + 1 == self.mb_w); - if nonzero[b] { - let pos = (py + 4 * (b >> 2)) * ys + px + 4 * (b & 3); - transform::add_residue(&mut frame.y, ys, pos, &transform::inverse_dct(&coeffs[b])); - } - } - } else { - predict::predict_block(&mut frame.y, ys, px, py, 16, info.mode); - add_luma_residue(&mut frame.y, ys, px, py, &coeffs, &nonzero); - } - predict::predict_block(&mut frame.u, cs, px / 2, py / 2, 8, uv_mode); - predict::predict_block(&mut frame.v, cs, px / 2, py / 2, 8, uv_mode); - } else { - let Some(r) = refs.get(usize::from(info.ref_frame) - 1).and_then(Option::as_deref) else { - return Err(Error::MissingReference(format!("reference frame {}", info.ref_frame))); - }; - predict_inter_mb(r, &mut frame, mb_x, mb_y, &info, filters, full_pixel); - add_luma_residue(&mut frame.y, ys, px, py, &coeffs, &nonzero); - } - for b in 16..24 { - if nonzero[b] { - let k = (b - 16) & 3; - let pos = (py / 2 + 4 * (k >> 1)) * cs + px / 2 + 4 * (k & 1); - let plane = if b < 20 { &mut frame.u } else { &mut frame.v }; - transform::add_residue(plane, cs, pos, &transform::inverse_dct(&coeffs[b])); - } - } + /// Reads the segment id (when the map is updated) and returns the segment of the macroblock. + fn read_segment(&mut self, bd: &mut BoolDecoder, seg_idx: usize) -> usize { + if self.seg.update_map { + let s = bd.read_tree(&SEGMENT_TREE, &self.seg.tree_probs); + if let Some(m) = self.seg_map.get_mut(seg_idx) { + *m = s; } } + if self.seg.enabled { self.seg_map.get(seg_idx).copied().unwrap_or(0) as usize & 3 } else { 0 } + } - if hdr.filter_level > 0 { - loop_filter(&mut frame, self.mb_w, self.mb_h, &lf_info, &hdr, tag.key_frame); - } - if let Some(s) = saved { - self.ctx = s; - } - - // reference updates (9.7, 9.8): the copies first (altref, then golden from the possibly - // updated altref), then the refreshes with the new frame - let frame = Arc::new(frame); + /// Reference updates (9.7, 9.8): the copies first (altref, then golden from the possibly + /// updated altref), then the refreshes with the new frame. + fn update_references(&mut self, hdr: &FrameHeader, frame: &Arc) { let [last, golden, altref] = self.refs.clone(); let mut new_golden = golden; let mut new_altref = altref; @@ -435,6 +241,118 @@ impl Decoder { } let new_last = if hdr.refresh_last { Some(frame.clone()) } else { last }; self.refs = [new_last, new_golden, new_altref]; + } + + /// Prediction record of a macroblock of an inter frame (11.2, 16). + fn read_inter_frame_modes(&self, bd: &mut BoolDecoder, hdr: &FrameHeader, infos: &[MbInfo], pos: MbPos) -> MbModes { + let mut modes = MbModes::default(); + if bd.read_bool(hdr.prob_intra) { + modes.info.ref_frame = if bd.read_bool(hdr.prob_last) { 2 + u8::from(bd.read_bool(hdr.prob_golden)) } else { 1 }; + self.read_inter_modes(bd, infos, pos, &mut modes.info); + } else { + modes.info.mode = bd.read_tree(&YMODE_TREE, &self.ctx.ymode); + if modes.info.mode == B_PRED { + for m in &mut modes.bmodes { + *m = bd.read_tree(&BMODE_TREE, &BMODE_PROBS); + } + } + modes.uv_mode = bd.read_tree(&UV_MODE_TREE, &self.ctx.uv_mode); + } + modes + } + + /// Decode one compressed frame. Returns the picture when the frame is meant to be shown + /// (`show_frame`), `None` for a hidden frame (an altref). + pub fn decode(&mut self, data: &[u8], pts: i64) -> Result> { + let tag = FrameTag::parse(data)?; + ensure!(tag.first_part_size > 0, "empty first partition"); + if tag.key_frame { + self.begin_key_frame(&tag)?; + } else if self.refs[0].is_none() { + return Err(Error::MissingReference("inter frame before the first key frame".into())); + } + ensure!(self.mb_w > 0 && self.mb_w <= MAX_MBS && self.mb_h <= MAX_MBS, "no frame size"); + let first_start = tag.header_len; + let first_end = first_start.saturating_add(tag.first_part_size); + ensure!(first_end <= data.len(), "first partition ({} bytes) runs past the frame ({} bytes)", tag.first_part_size, data.len()); + let mut bd = BoolDecoder::new(&data[first_start..first_end]); + let mut saved = None; + let hdr = read_header(&mut bd, tag.key_frame, &mut self.seg, &mut self.lf, &mut self.ctx, &mut saved); + if !tag.key_frame { + self.sign_bias[2] = hdr.sign_bias_golden; + self.sign_bias[3] = hdr.sign_bias_altref; + } + + let n = hdr.partitions; + let mut token_decoders = split_token_partitions(data, first_end, n); + + let mut frame = Frame::new(self.mb_w * 16, self.mb_h * 16); + let dequant = self.segment_dequant(&hdr); + let recon = InterRecon { + refs: self.refs.clone(), + filters: if tag.version == 0 || tag.version > 3 { &SIXTAP_FILTERS } else { &BILINEAR_FILTERS }, + full_pixel: tag.version == 3, + }; + + let stride_info = self.mb_w + 1; + let mut infos = vec![MbInfo::default(); stride_info * (self.mb_h + 1)]; + let mut lf_info: Vec<(u8, bool)> = vec![(0, false); self.mb_w * self.mb_h]; + // contexts: B_PRED modes above / left (key frames), non-zero flags above / left + let mut above_bmodes = vec![[B_DC_PRED; 4]; self.mb_w]; + let mut above_nz = vec![[0u8; 9]; self.mb_w]; + + for mb_y in 0..self.mb_h { + let mut left_bmodes = [B_DC_PRED; 4]; + let mut left_nz = [0u8; 9]; + let part = mb_y % n; + for mb_x in 0..self.mb_w { + let pos = MbPos { idx: (mb_y + 1) * stride_info + mb_x + 1, stride: stride_info, mb_x, mb_y }; + // ---- prediction record (first partition) + let seg_idx = mb_y * self.mb_w + mb_x; + let segment = self.read_segment(&mut bd, seg_idx); + let coeff_skip = hdr.skip_prob.is_some_and(|p| bd.read_bool(p)); + let modes = if tag.key_frame { + read_key_frame_modes(&mut bd, &above_bmodes[mb_x], &left_bmodes) + } else { + self.read_inter_frame_modes(&mut bd, &hdr, &infos, pos) + }; + let info = modes.info; + above_bmodes[mb_x] = [modes.bmodes[12], modes.bmodes[13], modes.bmodes[14], modes.bmodes[15]]; + left_bmodes = [modes.bmodes[3], modes.bmodes[7], modes.bmodes[11], modes.bmodes[15]]; + if let Some(slot) = infos.get_mut(pos.idx) { + *slot = info; + } + + // ---- residual (token partition) + let has_y2 = info.mode != B_PRED && info.mode != MV_SPLIT; + let residual = if coeff_skip { + skip_residual(has_y2, &mut above_nz[mb_x], &mut left_nz); + MbResidual::default() + } else { + read_residual(&mut token_decoders[part], &self.ctx.coeff, &dequant[segment], has_y2, &mut above_nz[mb_x], &mut left_nz) + }; + + // ---- loop filter level (15.1, 9.3, 9.4) + let level = self.filter_level(&hdr, segment, &info); + let inner = !has_y2 || residual.has_coeffs; + if let Some(l) = lf_info.get_mut(seg_idx) { + *l = (level, inner); + } + + // ---- reconstruction + reconstruct_mb(&mut frame, (mb_x, mb_y), self.mb_w, &modes, &residual, &recon)?; + } + } + + if hdr.filter_level > 0 { + loop_filter(&mut frame, self.mb_w, self.mb_h, &lf_info, &hdr, tag.key_frame); + } + if let Some(s) = saved { + self.ctx = s; + } + + let frame = Arc::new(frame); + self.update_references(&hdr, &frame); self.stats.frames += 1; self.stats.key_frames += u64::from(tag.key_frame); @@ -495,7 +413,8 @@ impl Decoder { } /// Inter prediction record (16.2-16.4, 17) of the macroblock at `idx` in `infos`. - fn read_inter_modes(&self, bd: &mut BoolDecoder, infos: &[MbInfo], idx: usize, stride: usize, mb_x: usize, mb_y: usize, info: &mut MbInfo) { + fn read_inter_modes(&self, bd: &mut BoolDecoder, infos: &[MbInfo], pos: MbPos, info: &mut MbInfo) { + let MbPos { idx, stride, mb_x, mb_y } = pos; let get = |i: usize| infos.get(i).copied().unwrap_or_default(); let above = get(idx - stride); let left = get(idx - 1); @@ -619,6 +538,188 @@ impl Decoder { } } +/// Position of a macroblock: its index in the bordered `infos` array (row stride `stride`) and +/// its macroblock coordinates. +#[derive(Clone, Copy)] +struct MbPos { + idx: usize, + stride: usize, + mb_x: usize, + mb_y: usize, +} + +/// Prediction record of one macroblock. +struct MbModes { + info: MbInfo, + /// Subblock intra modes (`B_DC_PRED` unless the macroblock is `B_PRED`, or on key frames). + bmodes: [u8; 16], + uv_mode: u8, +} + +impl Default for MbModes { + fn default() -> Self { + MbModes { info: MbInfo::default(), bmodes: [B_DC_PRED; 16], uv_mode: DC_PRED } + } +} + +/// Dequantized coefficients of one macroblock: 16 Y, 4 U, 4 V and the Y2 block. +#[derive(Default)] +struct MbResidual { + coeffs: [[i16; 16]; 25], + nonzero: [bool; 25], + has_coeffs: bool, +} + +/// What inter prediction of a frame needs: the reference frames and the sub-pixel filter setup. +struct InterRecon { + refs: [Option>; 3], + filters: &'static [i32; 48], + full_pixel: bool, +} + +/// Token partitions (9.5): sizes of all but the last, then the data. +fn split_token_partitions(data: &[u8], sizes_start: usize, n: usize) -> Vec> { + let mut pos = sizes_start.saturating_add(3 * (n - 1)).min(data.len()); + let mut parts: Vec = Vec::with_capacity(n); + for i in 0..n { + let size = if i + 1 < n { + let s = sizes_start + 3 * i; + match data.get(s..s + 3) { + Some(b) => usize::from(b[0]) | usize::from(b[1]) << 8 | usize::from(b[2]) << 16, + None => 0, + } + } else { + data.len() - pos + }; + let end = pos.saturating_add(size).min(data.len()); + parts.push(BoolDecoder::new(&data[pos..end])); + pos = end; + } + parts +} + +/// Prediction record of a key frame macroblock (11.2, 11.3). +fn read_key_frame_modes(bd: &mut BoolDecoder, above_bmodes: &[u8; 4], left_bmodes: &[u8; 4]) -> MbModes { + let mut modes = MbModes::default(); + modes.info.mode = bd.read_tree(&KF_YMODE_TREE, &KF_YMODE_PROBS); + if modes.info.mode == B_PRED { + for b in 0..16 { + let a = if b < 4 { above_bmodes[b] } else { modes.bmodes[b - 4] }; + let l = if b & 3 == 0 { left_bmodes[b >> 2] } else { modes.bmodes[b - 1] }; + let o = (usize::from(a) * 10 + usize::from(l)) * 9; + modes.bmodes[b] = bd.read_tree(&BMODE_TREE, &KF_BMODE_PROBS[o..o + 9]); + } + } else { + modes.bmodes = [implied_bmode(modes.info.mode); 16]; + } + modes.uv_mode = bd.read_tree(&UV_MODE_TREE, &KF_UV_MODE_PROBS); + modes +} + +/// Non-zero context update for a macroblock without coefficients: everything is reset, except +/// the Y2 flags when the macroblock has no Y2 block. +fn skip_residual(has_y2: bool, above_nz: &mut [u8; 9], left_nz: &mut [u8; 9]) { + let keep_y2 = if has_y2 { 0 } else { left_nz[8] }; + let keep_y2_above = if has_y2 { 0 } else { above_nz[8] }; + *left_nz = [0; 9]; + *above_nz = [0; 9]; + left_nz[8] = keep_y2; + above_nz[8] = keep_y2_above; +} + +/// Tokens of all blocks of one macroblock (13), updating the non-zero contexts. +fn read_residual( + td: &mut BoolDecoder, + coeff_probs: &[[[[u8; 11]; 3]; 8]; 4], + dq: &Dequant, + has_y2: bool, + an: &mut [u8; 9], + left_nz: &mut [u8; 9], +) -> MbResidual { + let mut r = MbResidual::default(); + let first = if has_y2 { + let ctx = usize::from(an[8] + left_nz[8]); + let eob = read_coefficients(td, &coeff_probs[1], ctx, 0, dq.y2, &mut r.coeffs[24]); + let nz = u8::from(eob > 0); + an[8] = nz; + left_nz[8] = nz; + r.nonzero[24] = eob > 0; + r.has_coeffs |= eob > 0; + 1 + } else { + 0 + }; + let y_probs = if has_y2 { &coeff_probs[0] } else { &coeff_probs[3] }; + for b in 0..16 { + let (bx, by) = (b & 3, b >> 2); + let ctx = usize::from(an[bx] + left_nz[by]); + let eob = read_coefficients(td, y_probs, ctx, first, dq.y, &mut r.coeffs[b]); + let nz = eob > first; + an[bx] = u8::from(nz); + left_nz[by] = u8::from(nz); + r.nonzero[b] = nz; + r.has_coeffs |= nz; + } + for b in 16..24 { + let (off, k) = if b < 20 { (4, b - 16) } else { (6, b - 20) }; + let (bx, by) = (k & 1, k >> 1); + let ctx = usize::from(an[off + bx] + left_nz[off + by]); + let eob = read_coefficients(td, &coeff_probs[2], ctx, 0, dq.uv, &mut r.coeffs[b]); + let nz = eob > 0; + an[off + bx] = u8::from(nz); + left_nz[off + by] = u8::from(nz); + r.nonzero[b] = nz; + r.has_coeffs |= nz; + } + if has_y2 && r.nonzero[24] { + let dc = transform::inverse_wht(&r.coeffs[24]); + for b in 0..16 { + r.coeffs[b][0] = dc[b]; + r.nonzero[b] |= dc[b] != 0; + } + } + r +} + +/// Prediction plus residue of one macroblock into `frame`. +fn reconstruct_mb(frame: &mut Frame, (mb_x, mb_y): (usize, usize), mb_w: usize, modes: &MbModes, res: &MbResidual, recon: &InterRecon) -> Result<()> { + let info = &modes.info; + let (px, py) = (mb_x * 16, mb_y * 16); + let ys = frame.w; + let cs = frame.w / 2; + if info.ref_frame == 0 { + if info.mode == B_PRED { + for b in 0..16 { + predict::predict_subblock(&mut frame.y, ys, px, py, b, modes.bmodes[b], mb_x + 1 == mb_w); + if res.nonzero[b] { + let pos = (py + 4 * (b >> 2)) * ys + px + 4 * (b & 3); + transform::add_residue(&mut frame.y, ys, pos, &transform::inverse_dct(&res.coeffs[b])); + } + } + } else { + predict::predict_block(&mut frame.y, ys, px, py, 16, info.mode); + add_luma_residue(&mut frame.y, ys, px, py, &res.coeffs, &res.nonzero); + } + predict::predict_block(&mut frame.u, cs, px / 2, py / 2, 8, modes.uv_mode); + predict::predict_block(&mut frame.v, cs, px / 2, py / 2, 8, modes.uv_mode); + } else { + let Some(r) = recon.refs.get(usize::from(info.ref_frame) - 1).and_then(Option::as_deref) else { + return Err(Error::MissingReference(format!("reference frame {}", info.ref_frame))); + }; + predict_inter_mb(r, frame, mb_x, mb_y, info, recon.filters, recon.full_pixel); + add_luma_residue(&mut frame.y, ys, px, py, &res.coeffs, &res.nonzero); + } + for b in 16..24 { + if res.nonzero[b] { + let k = (b - 16) & 3; + let pos = (py / 2 + 4 * (k >> 1)) * cs + px / 2 + 4 * (k & 1); + let plane = if b < 20 { &mut frame.u } else { &mut frame.v }; + transform::add_residue(plane, cs, pos, &transform::inverse_dct(&res.coeffs[b])); + } + } + Ok(()) +} + /// `read_mvcomponent` (17.1), in quarter pixels. fn read_mv_component(bd: &mut BoolDecoder, p: &[u8; 19]) -> i32 { let a = if bd.read_bool(p[MVP_IS_SHORT]) { @@ -730,16 +831,17 @@ fn predict_inter_mb(r: &Frame, frame: &mut Frame, mb_x: usize, mb_y: usize, info let fp = |v: i32| if full_pixel { v & !7 } else { v }; if info.mode != MV_SPLIT { let m = info.mv; - predict::predict_inter(&ry, &mut frame.y, ys, px, py, 16, 16, m.col * 2, m.row * 2, filters); + predict::predict_inter(&ry, DstPlane { data: &mut frame.y, stride: ys }, BlockRect { x: px, y: py, w: 16, h: 16 }, (m.col * 2, m.row * 2), filters); // chroma: the luma vector in quarter pixels is the chroma vector in eighth pixels let (cx, cy) = (fp(m.col), fp(m.row)); - predict::predict_inter(&ru, &mut frame.u, cs, px / 2, py / 2, 8, 8, cx, cy, filters); - predict::predict_inter(&rv, &mut frame.v, cs, px / 2, py / 2, 8, 8, cx, cy, filters); + predict::predict_inter(&ru, DstPlane { data: &mut frame.u, stride: cs }, BlockRect { x: px / 2, y: py / 2, w: 8, h: 8 }, (cx, cy), filters); + predict::predict_inter(&rv, DstPlane { data: &mut frame.v, stride: cs }, BlockRect { x: px / 2, y: py / 2, w: 8, h: 8 }, (cx, cy), filters); return; } for b in 0..16 { let m = info.bmv[b]; - predict::predict_inter(&ry, &mut frame.y, ys, px + 4 * (b & 3), py + 4 * (b >> 2), 4, 4, m.col * 2, m.row * 2, filters); + let rect = BlockRect { x: px + 4 * (b & 3), y: py + 4 * (b >> 2), w: 4, h: 4 }; + predict::predict_inter(&ry, DstPlane { data: &mut frame.y, stride: ys }, rect, (m.col * 2, m.row * 2), filters); } for i in 0..2 { for j in 0..2 { @@ -750,8 +852,9 @@ fn predict_inter_mb(r: &Frame, frame: &mut Frame, mb_x: usize, mb_y: usize, info let cx = fp(avg(sum(|m| m.col))); let cy = fp(avg(sum(|m| m.row))); let (bx, by) = (px / 2 + 4 * j, py / 2 + 4 * i); - predict::predict_inter(&ru, &mut frame.u, cs, bx, by, 4, 4, cx, cy, filters); - predict::predict_inter(&rv, &mut frame.v, cs, bx, by, 4, 4, cx, cy, filters); + let rect = BlockRect { x: bx, y: by, w: 4, h: 4 }; + predict::predict_inter(&ru, DstPlane { data: &mut frame.u, stride: cs }, rect, (cx, cy), filters); + predict::predict_inter(&rv, DstPlane { data: &mut frame.v, stride: cs }, rect, (cx, cy), filters); } } } diff --git a/crates/vp8/src/header.rs b/crates/vp8/src/header.rs index 9c38c5fdb..0d9d60bc6 100644 --- a/crates/vp8/src/header.rs +++ b/crates/vp8/src/header.rs @@ -141,6 +141,110 @@ pub struct FrameHeader { pub prob_golden: u8, } +/// segmentation_header (9.3). +fn read_segmentation(d: &mut BoolDecoder, seg: &mut Segmentation) { + seg.enabled = d.read_flag(); + seg.update_map = false; + if !seg.enabled { + return; + } + seg.update_map = d.read_flag(); + let update_data = d.read_flag(); + if update_data { + seg.absolute = d.read_flag(); + for q in &mut seg.quant { + *q = d.read_optional_signed(7) as i8; + } + for l in &mut seg.lf_level { + *l = d.read_optional_signed(6) as i8; + } + } + if seg.update_map { + for p in &mut seg.tree_probs { + *p = if d.read_flag() { d.read_literal(8) as u8 } else { crate::tables::DEFAULT_SEGMENT_PROB }; + } + } +} + +/// Loop filter delta adjustments (9.6). +fn read_lf_deltas(d: &mut BoolDecoder, lf: &mut LfDeltas) { + lf.enabled = d.read_flag(); + if lf.enabled && d.read_flag() { + for v in lf.ref_deltas.iter_mut().chain(lf.mode_deltas.iter_mut()) { + if d.read_flag() { + *v = d.read_signed(6) as i8; + } + } + } +} + +/// Reference frame refresh / copy / sign bias flags (9.7, 9.8) and the entropy refresh flag. +fn read_reference_updates(d: &mut BoolDecoder, key_frame: bool, h: &mut FrameHeader) { + if key_frame { + h.refresh_golden = true; + h.refresh_altref = true; + h.refresh_entropy = d.read_flag(); + h.refresh_last = true; + return; + } + h.refresh_golden = d.read_flag(); + h.refresh_altref = d.read_flag(); + if !h.refresh_golden { + h.copy_to_golden = d.read_literal(2) as u8; + } + if !h.refresh_altref { + h.copy_to_altref = d.read_literal(2) as u8; + } + h.sign_bias_golden = d.read_flag(); + h.sign_bias_altref = d.read_flag(); + h.refresh_entropy = d.read_flag(); + h.refresh_last = d.read_flag(); +} + +/// token_prob_update (13.4). +fn read_token_prob_updates(d: &mut BoolDecoder, ctx: &mut EntropyContext) { + let mut i = 0; + for plane in ctx.coeff.iter_mut() { + for band in plane.iter_mut() { + for c in band.iter_mut() { + for p in c.iter_mut() { + if d.read_bool(crate::spec_tables::COEFF_UPDATE_PROBS[i]) { + *p = d.read_literal(8) as u8; + } + i += 1; + } + } + } + } +} + +/// Inter frame probabilities: intra/last/golden, mode probabilities and mv_prob_update (17.2). +fn read_inter_probs(d: &mut BoolDecoder, h: &mut FrameHeader, ctx: &mut EntropyContext) { + h.prob_intra = d.read_literal(8) as u8; + h.prob_last = d.read_literal(8) as u8; + h.prob_golden = d.read_literal(8) as u8; + if d.read_flag() { + for p in &mut ctx.ymode { + *p = d.read_literal(8) as u8; + } + } + if d.read_flag() { + for p in &mut ctx.uv_mode { + *p = d.read_literal(8) as u8; + } + } + let mut i = 0; + for comp in ctx.mv.iter_mut() { + for p in comp.iter_mut() { + if d.read_bool(crate::spec_tables::MV_UPDATE_PROBS[i]) { + let x = d.read_literal(7) as u8; + *p = if x != 0 { x << 1 } else { 1 }; + } + i += 1; + } + } +} + /// Read the frame header from the first partition, updating the persistent state. /// `saved` receives the entropy context to restore after the frame when /// `refresh_entropy_probs` is 0. @@ -157,42 +261,11 @@ pub fn read_header( h.color_space = d.read_literal(1) as u8; h.clamping_type = d.read_literal(1) as u8; } - seg.enabled = d.read_flag(); - seg.update_map = false; - if seg.enabled { - seg.update_map = d.read_flag(); - let update_data = d.read_flag(); - if update_data { - seg.absolute = d.read_flag(); - for q in &mut seg.quant { - *q = d.read_optional_signed(7) as i8; - } - for l in &mut seg.lf_level { - *l = d.read_optional_signed(6) as i8; - } - } - if seg.update_map { - for p in &mut seg.tree_probs { - *p = if d.read_flag() { d.read_literal(8) as u8 } else { crate::tables::DEFAULT_SEGMENT_PROB }; - } - } - } + read_segmentation(d, seg); h.filter_type = d.read_literal(1) as u8; h.filter_level = d.read_literal(6) as u8; h.sharpness = d.read_literal(3) as u8; - lf.enabled = d.read_flag(); - if lf.enabled && d.read_flag() { - for v in &mut lf.ref_deltas { - if d.read_flag() { - *v = d.read_signed(6) as i8; - } - } - for v in &mut lf.mode_deltas { - if d.read_flag() { - *v = d.read_signed(6) as i8; - } - } - } + read_lf_deltas(d, lf); h.partitions = 1 << d.read_literal(2); h.quant = QuantIndices { y_ac: d.read_literal(7) as i32, @@ -202,67 +275,13 @@ pub fn read_header( uv_dc_delta: d.read_optional_signed(4), uv_ac_delta: d.read_optional_signed(4), }; - if key_frame { - h.refresh_golden = true; - h.refresh_altref = true; - h.refresh_entropy = d.read_flag(); - h.refresh_last = true; - } else { - h.refresh_golden = d.read_flag(); - h.refresh_altref = d.read_flag(); - if !h.refresh_golden { - h.copy_to_golden = d.read_literal(2) as u8; - } - if !h.refresh_altref { - h.copy_to_altref = d.read_literal(2) as u8; - } - h.sign_bias_golden = d.read_flag(); - h.sign_bias_altref = d.read_flag(); - h.refresh_entropy = d.read_flag(); - h.refresh_last = d.read_flag(); - } + read_reference_updates(d, key_frame, &mut h); // The probabilities in effect before this frame's updates come back after it. *saved = if h.refresh_entropy { None } else { Some(ctx.clone()) }; - // token_prob_update (13.4) - let mut i = 0; - for plane in ctx.coeff.iter_mut() { - for band in plane.iter_mut() { - for c in band.iter_mut() { - for p in c.iter_mut() { - if d.read_bool(crate::spec_tables::COEFF_UPDATE_PROBS[i]) { - *p = d.read_literal(8) as u8; - } - i += 1; - } - } - } - } + read_token_prob_updates(d, ctx); h.skip_prob = if d.read_flag() { Some(d.read_literal(8) as u8) } else { None }; if !key_frame { - h.prob_intra = d.read_literal(8) as u8; - h.prob_last = d.read_literal(8) as u8; - h.prob_golden = d.read_literal(8) as u8; - if d.read_flag() { - for p in &mut ctx.ymode { - *p = d.read_literal(8) as u8; - } - } - if d.read_flag() { - for p in &mut ctx.uv_mode { - *p = d.read_literal(8) as u8; - } - } - // mv_prob_update (17.2) - let mut i = 0; - for comp in ctx.mv.iter_mut() { - for p in comp.iter_mut() { - if d.read_bool(crate::spec_tables::MV_UPDATE_PROBS[i]) { - let x = d.read_literal(7) as u8; - *p = if x != 0 { x << 1 } else { 1 }; - } - i += 1; - } - } + read_inter_probs(d, &mut h, ctx); } h } diff --git a/crates/vp8/src/predict.rs b/crates/vp8/src/predict.rs index 411eecc68..f4945b499 100644 --- a/crates/vp8/src/predict.rs +++ b/crates/vp8/src/predict.rs @@ -244,21 +244,27 @@ impl RefPlane<'_> { } } -/// Predict the `w`x`h` block at `(bx, by)` of a plane from `reference` displaced by the motion -/// vector `(mvx, mvy)` in eighth pixels, using `filters` (8 x 6 taps: the six-tap or bilinear -/// set, 18.3), writing into `dst` (same layout as the block's plane). -pub fn predict_inter( - reference: &RefPlane, - dst: &mut [u8], - dst_stride: usize, - bx: usize, - by: usize, - w: usize, - h: usize, - mvx: i32, - mvy: i32, - filters: &[i32; 48], -) { +/// Destination plane of an inter prediction (same layout as the block's plane). +pub struct DstPlane<'a> { + pub data: &'a mut [u8], + pub stride: usize, +} + +/// A `w`x`h` block at `(x, y)` of a plane. +#[derive(Clone, Copy)] +pub struct BlockRect { + pub x: usize, + pub y: usize, + pub w: usize, + pub h: usize, +} + +/// Predict the block `rect` of a plane from `reference` displaced by the motion vector +/// `(mvx, mvy)` in eighth pixels, using `filters` (8 x 6 taps: the six-tap or bilinear set, +/// 18.3), writing into `dst`. +pub fn predict_inter(reference: &RefPlane, dst: DstPlane, rect: BlockRect, (mvx, mvy): (i32, i32), filters: &[i32; 48]) { + let DstPlane { data: dst, stride: dst_stride } = dst; + let BlockRect { x: bx, y: by, w, h } = rect; let x0 = bx as isize + (mvx >> 3) as isize; let y0 = by as isize + (mvy >> 3) as isize; let (fx, fy) = ((mvx & 7) as usize, (mvy & 7) as usize);