vp6: handle partial golden frame update in alpha case
[nihav.git] / nihav-duck / src / codecs / vp56.rs
index de3a8498c939517553e38996b5a83026070c1d07..807756a495687533836950daa513237d528c2c14 100644 (file)
@@ -1,5 +1,6 @@
 use nihav_core::codecs::*;
 use nihav_core::io::bitreader::*;
+use nihav_codec_support::codecs::{MV, ZERO_MV};
 use super::vpcommon::*;
 
 pub const TOKEN_LARGE: u8 = 5;
@@ -69,7 +70,7 @@ struct Node {
 }
 
 fn prob2weight(a: u8, b: u8) -> u8 {
-    let w = (((a as u16) * (b as u16)) >> 8) as u8;
+    let w = ((u16::from(a) * u16::from(b)) >> 8) as u8;
     if w == 0 {
         1
     } else {
@@ -131,7 +132,7 @@ impl VP6Huff {
         let mut nlen = 0;
 
         for w in weights.iter().rev() {
-            let weight = *w as u16;
+            let weight = u16::from(*w);
             let mut pos = nlen;
             for i in 0..nlen {
                 if nodes[i].weight > weight {
@@ -191,8 +192,8 @@ impl<'a> ReadHuff for BitReader<'a> {
     fn read_huff(&mut self, huff: &VP6Huff) -> DecoderResult<u8> {
         let peekval                             = self.peek(16);
         for (i, (code, bit)) in huff.codes.iter().zip(huff.bits.iter()).enumerate() {
-            if (peekval >> (16 - *bit)) == (*code as u32) {
-                self.skip(*bit as u32)?;
+            if (peekval >> (16 - *bit)) == u32::from(*code) {
+                self.skip(u32::from(*bit))?;
                 return Ok(i as u8);
             }
         }
@@ -351,6 +352,7 @@ pub struct VP56Decoder {
     mb_w:       usize,
     mb_h:       usize,
     models:     VP56Models,
+    amodels:    VP56Models,
     coeffs:     [[i16; 64]; 6],
     last_mbt:   VPMBType,
 
@@ -393,7 +395,7 @@ pub fn expand_token_bc(bc: &mut BoolCoder, val_probs: &[u8; 11], token: u8, vers
         if token != 0 {
             sign                                = bc.read_bool();
         }
-        level = token as i16;
+        level = i16::from(token);
     } else {
         let cat: usize = vp_tree!(bc, val_probs[6],
                                   vp_tree!(bc, val_probs[7], 0, 1),
@@ -433,6 +435,7 @@ impl VP56Decoder {
             mb_w:       0,
             mb_h:       0,
             models:     VP56Models::new(),
+            amodels:    VP56Models::new(),
             coeffs:     [[0; 64]; 6],
             last_mbt:   VPMBType::InterNoMV,
 
@@ -458,7 +461,7 @@ impl VP56Decoder {
         self.top_ctx = [vec![0; self.mb_w * 2], vec![0; self.mb_w], vec![0; self.mb_w], vec![0; self.mb_w * 2]];
     }
     pub fn init(&mut self, supp: &mut NADecoderSupport, vinfo: NAVideoInfo) -> DecoderResult<()> {
-        supp.pool_u8.set_dec_bufs(3);
+        supp.pool_u8.set_dec_bufs(3 + if vinfo.get_format().has_alpha() { 1 } else { 0 });
         supp.pool_u8.prealloc_video(NAVideoInfo::new(vinfo.get_width(), vinfo.get_height(), false, vinfo.get_format()), 4)?;
         self.set_dimensions(vinfo.get_width(), vinfo.get_height());
         self.dc_pred.resize(self.mb_w);
@@ -517,14 +520,62 @@ impl VP56Decoder {
             }
         }
 
+        let psrc = &src[if self.has_alpha { 3 } else { 0 }..aoffset];
+        self.decode_planes(br, &mut dframe, &mut bc, &hdr, psrc, false)?;
+
+        if self.has_alpha {
+            let asrc = &src[aoffset + 3..];
+            let mut bc = BoolCoder::new(asrc)?;
+            let ahdr = br.parse_header(&mut bc)?;
+            validate!(ahdr.mb_w == hdr.mb_w && ahdr.mb_h == hdr.mb_h);
+            std::mem::swap(&mut self.models, &mut self.amodels);
+            let ret = self.decode_planes(br, &mut dframe, &mut bc, &ahdr, asrc, true);
+            std::mem::swap(&mut self.models, &mut self.amodels);
+            if let Err(err) = ret {
+                return Err(err);
+            }
+            match (hdr.is_golden, ahdr.is_golden) {
+                (true, true) => { self.shuf.add_golden_frame(buf.clone()); },
+                (true, false) => {
+                    let cur_golden = self.shuf.get_golden().unwrap();
+                    let off    = cur_golden.get_offset(3);
+                    let stride = cur_golden.get_stride(3);
+                    let mut new_golden = supp.pool_u8.get_copy(&buf).unwrap();
+                    let dst = new_golden.get_data_mut().unwrap();
+                    let src = cur_golden.get_data();
+                    dst[off..][..stride * self.mb_h * 16].copy_from_slice(&src[off..][..stride * self.mb_h * 16]);
+                    self.shuf.add_golden_frame(new_golden);
+                },
+                (false, true) => {
+                    let cur_golden = self.shuf.get_golden().unwrap();
+                    let off    = cur_golden.get_offset(3);
+                    let stride = cur_golden.get_stride(3);
+                    let mut new_golden = supp.pool_u8.get_copy(&cur_golden).unwrap();
+                    let dst = new_golden.get_data_mut().unwrap();
+                    let src = buf.get_data();
+                    dst[off..][..stride * self.mb_h * 16].copy_from_slice(&src[off..][..stride * self.mb_h * 16]);
+                    self.shuf.add_golden_frame(new_golden);
+                },
+                _ => {},
+            };
+        }
+
+        if hdr.is_golden && !self.has_alpha {
+            self.shuf.add_golden_frame(buf.clone());
+        }
+        self.shuf.add_frame(buf.clone());
+
+        Ok((NABufferType::Video(buf), if hdr.is_intra { FrameType::I } else { FrameType::P }))
+    }
+    fn decode_planes(&mut self, br: &mut dyn VP56Parser, dframe: &mut NASimpleVideoFrame<u8>, bc: &mut BoolCoder, hdr: &VP56Header, src: &[u8], alpha: bool) -> DecoderResult<()> {
         let mut cr;
         if hdr.multistream {
-            let off = (if self.has_alpha { 3 } else { 0 }) + (hdr.offset as usize);
+            let off = hdr.offset as usize;
             if !hdr.use_huffman {
                 let bc2 = BoolCoder::new(&src[off..])?;
                 cr = CoeffReader::Bool(bc2);
             } else {
-                let br = BitReader::new(&src[off..aoffset], BitReaderMode::BE);
+                let br = BitReader::new(&src[off..], BitReaderMode::BE);
                 cr = CoeffReader::Huff(br);
             }
         } else {
@@ -535,10 +586,10 @@ impl VP56Decoder {
             br.reset_models(&mut self.models);
             self.reset_mbtype_models();
         } else {
-            self.decode_mode_prob_models(&mut bc)?;
-            br.decode_mv_models(&mut bc, &mut self.models.mv_models)?;
+            self.decode_mode_prob_models(bc)?;
+            br.decode_mv_models(bc, &mut self.models.mv_models)?;
         }
-        br.decode_coeff_models(&mut bc, &mut self.models, hdr.is_intra)?;
+        br.decode_coeff_models(bc, &mut self.models, hdr.is_intra)?;
         if hdr.use_huffman {
             for i in 0..2 {
                 self.models.vp6huff.dc_token_tree[i].build_codes(&self.models.coeff_models[i].dc_value_probs);
@@ -562,7 +613,7 @@ impl VP56Decoder {
         self.fstate = FrameState::new();
         self.fstate.dc_quant = VP56_DC_QUANTS[hdr.quant as usize] * 4;
         self.fstate.ac_quant = VP56_AC_QUANTS[hdr.quant as usize] * 4;
-        self.loop_thr = VP56_FILTER_LIMITS[hdr.quant as usize] as i16;
+        self.loop_thr = i16::from(VP56_FILTER_LIMITS[hdr.quant as usize]);
 
         self.last_mbt = VPMBType::InterNoMV;
         for vec in self.top_ctx.iter_mut() {
@@ -582,25 +633,12 @@ impl VP56Decoder {
             self.fstate.last_idx = [24; 4];
             for mb_x in 0..self.mb_w {
                 self.fstate.mb_x = mb_x;
-                self.decode_mb(&mut dframe, &mut bc, &mut cr, br, &hdr, false)?;
+                self.decode_mb(dframe, bc, &mut cr, br, &hdr, alpha)?;
                 self.dc_pred.next_mb();
             }
             self.dc_pred.update_row();
         }
-
-        if self.has_alpha {
-            let asrc = &src[aoffset + 3..];
-            let mut bc = BoolCoder::new(asrc)?;
-            let ahdr = br.parse_header(&mut bc)?;
-            validate!(ahdr.mb_w == hdr.mb_w && ahdr.mb_h == hdr.mb_h);
-        }
-
-        if hdr.is_golden {
-            self.shuf.add_golden_frame(buf.clone());
-        }
-        self.shuf.add_frame(buf.clone());
-
-        Ok((NABufferType::Video(buf), if hdr.is_intra { FrameType::I } else { FrameType::P }))
+        Ok(())
     }
     fn reset_mbtype_models(&mut self) {
         const DEFAULT_XMITTED_PROBS: [[u8; 20]; 3] = [
@@ -655,11 +693,11 @@ impl VP56Decoder {
                 let mut total = 0;
                 for i in 0..10 {
                     if i == mode { continue; }
-                    cnt[i] = 100 * (prob_xmitted[i * 2] as u32);
+                    cnt[i] = 100 * u32::from(prob_xmitted[i * 2]);
                     total += cnt[i];
                 }
-                let sum = (prob_xmitted[mode * 2] as u32) + (prob_xmitted[mode * 2 + 1] as u32);
-                mdl.probs[9] = 255 - rescale_mb_mode_prob(prob_xmitted[mode * 2 + 1] as u32, sum);
+                let sum = u32::from(prob_xmitted[mode * 2]) + u32::from(prob_xmitted[mode * 2 + 1]);
+                mdl.probs[9] = 255 - rescale_mb_mode_prob(u32::from(prob_xmitted[mode * 2 + 1]), sum);
 
                 let inter_mv0_weight = (cnt[0] as u32) + (cnt[2] as u32);
                 let inter_mv1_weight = (cnt[3] as u32) + (cnt[4] as u32);
@@ -750,6 +788,7 @@ impl VP56Decoder {
         }
         Ok(self.last_mbt)
     }
+    #[allow(clippy::cognitive_complexity)]
     fn decode_mb(&mut self, frm: &mut NASimpleVideoFrame<u8>, bc: &mut BoolCoder, cr: &mut CoeffReader, br: &mut dyn VP56Parser, hdr: &VP56Header, alpha: bool) -> DecoderResult<()> {
         const FOURMV_SUB_TYPE: [VPMBType; 4] = [ VPMBType::InterNoMV, VPMBType::InterMV, VPMBType::InterNearest, VPMBType::InterNear ];
 
@@ -765,7 +804,7 @@ impl VP56Decoder {
             let prob = if mb_x == 0 {
                     iprob
                 } else if !self.ilace_mb {
-                    iprob + (((256 - (iprob as u16)) >> 1) as u8)
+                    iprob + (((256 - u16::from(iprob)) >> 1) as u8)
                 } else {
                     iprob - (iprob >> 1)
                 };
@@ -889,25 +928,24 @@ impl VP56Decoder {
                 }
             }
         }
-        if !alpha {
-            for blk_no in 4..6 {
-                self.fstate.plane = blk_no - 3;
-                self.fstate.ctx_idx = blk_no - 2;
-                self.fstate.top_ctx = self.top_ctx[self.fstate.plane][mb_x];
-                match cr {
-                    CoeffReader::None              => {
-                        br.decode_block(bc, &mut self.coeffs[blk_no], &self.models.coeff_models[1], &self.models.vp6models, &mut self.fstate)?;
-                    },
-                    CoeffReader::Bool(ref mut bcc) => {
-                        br.decode_block(bcc, &mut self.coeffs[blk_no], &self.models.coeff_models[1], &self.models.vp6models, &mut self.fstate)?;
-                    },
-                    CoeffReader::Huff(ref mut brc) => {
-                        br.decode_block_huff(brc, &mut self.coeffs[blk_no], &self.models.vp6models, &self.models.vp6huff, &mut self.fstate)?;
-                    },
-                };
-                self.top_ctx[self.fstate.plane][mb_x] = self.fstate.top_ctx;
-                self.predict_dc(mb_type, mb_pos, blk_no, alpha);
-
+        for blk_no in 4..6 {
+            self.fstate.plane = blk_no - 3;
+            self.fstate.ctx_idx = blk_no - 2;
+            self.fstate.top_ctx = self.top_ctx[self.fstate.plane][mb_x];
+            match cr {
+                CoeffReader::None              => {
+                    br.decode_block(bc, &mut self.coeffs[blk_no], &self.models.coeff_models[1], &self.models.vp6models, &mut self.fstate)?;
+                },
+                CoeffReader::Bool(ref mut bcc) => {
+                    br.decode_block(bcc, &mut self.coeffs[blk_no], &self.models.coeff_models[1], &self.models.vp6models, &mut self.fstate)?;
+                },
+                CoeffReader::Huff(ref mut brc) => {
+                    br.decode_block_huff(brc, &mut self.coeffs[blk_no], &self.models.vp6models, &self.models.vp6huff, &mut self.fstate)?;
+                },
+            };
+            self.top_ctx[self.fstate.plane][mb_x] = self.fstate.top_ctx;
+            self.predict_dc(mb_type, mb_pos, blk_no, alpha);
+            if !alpha {
                 let has_ac = self.fstate.last_idx[self.fstate.ctx_idx] > 0;
                 if mb_type.is_intra() {
                     if has_ac {
@@ -944,7 +982,7 @@ impl VP56Decoder {
             let x = self.fstate.mb_x * 8;
             let y = self.fstate.mb_y * 8;
             br.mc_block(frm, self.mc_buf.clone(), src.clone(), 1, x, y, mv, self.loop_thr);
-            br.mc_block(frm, self.mc_buf.clone(), src.clone(), 2, x, y, mv, self.loop_thr);
+            br.mc_block(frm, self.mc_buf.clone(), src,         2, x, y, mv, self.loop_thr);
         }
     }
     fn do_fourmv(&mut self, br: &dyn VP56Parser, frm: &mut NASimpleVideoFrame<u8>, mvs: &[MV; 4], alpha: bool) {
@@ -963,7 +1001,7 @@ impl VP56Decoder {
             let sum = mvs[0] + mvs[1] + mvs[2] + mvs[3];
             let mv = MV { x: sum.x / 4, y: sum.y / 4 };
             br.mc_block(frm, self.mc_buf.clone(), src.clone(), 1, x, y, mv, self.loop_thr);
-            br.mc_block(frm, self.mc_buf.clone(), src.clone(), 2, x, y, mv, self.loop_thr);
+            br.mc_block(frm, self.mc_buf.clone(), src,         2, x, y, mv, self.loop_thr);
         }
     }
     fn predict_dc(&mut self, mb_type: VPMBType, _mb_pos: usize, blk_no: usize, _alpha: bool) {