h264: introduce frame pool for average block buffers in MT decoder

[nihav.git] / nihav-itu / src / codecs / h264 / decoder_mt.rs

diff --git a/nihav-itu/src/codecs/h264/decoder_mt.rs b/nihav-itu/src/codecs/h264/decoder_mt.rs
index 69a335b45667441ec5aaf6c6062a33c10a9976b2..fac66c58bb0820bd1bc06de49dbce89c2af6643d 100644 (file)
--- a/nihav-itu/src/codecs/h264/decoder_mt.rs
+++ b/nihav-itu/src/codecs/h264/decoder_mt.rs
@@ -1,13 +1,11 @@
-/* TODO:
-  * buffer pool for DSP avg frames
-*/
-
 use nihav_core::codecs::*;
 use nihav_core::io::bitreader::*;
 
 use super::*;
 use super::dispatch::*;
 
+const AVG_BUF_VINFO: NAVideoInfo = NAVideoInfo { width: 32, height: 32, flipped: false, format: YUV420_FORMAT, bits: 12 };
+
 pub struct FrameDecoder {
     pub slices:         Vec<(SliceHeader, usize, SliceRefs, Vec<u8>)>,
     pub cur_pic:        PictureInfo,
@@ -61,7 +59,7 @@ impl FrameDecoder {
             17, 18, 20, 24, 19, 21, 26, 28, 23, 27, 29, 30, 22, 25, 38, 41
         ];
 
-        let mut mb_idx = slice_hdr.first_mb_in_slice as usize;
+        let mut mb_idx = slice_hdr.first_mb_in_slice;
         let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
         let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
         while br.tell() < full_size && mb_idx < self.num_mbs {
@@ -155,7 +153,7 @@ impl FrameDecoder {
         Ok(mb_idx)
     }
     fn decode_slice_cabac(&mut self, cabac: &mut CABAC, slice_hdr: &SliceHeader, refs: &SliceRefs) -> DecoderResult<usize> {
-        let mut mb_idx = slice_hdr.first_mb_in_slice as usize;
+        let mut mb_idx = slice_hdr.first_mb_in_slice;
         let mut prev_mb_skipped = false;
         let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
         let mut last_qp_diff = false;
@@ -472,6 +470,7 @@ struct H264MTDecoder {
     deblock_skip:   bool,
     max_last_poc:   u32,
     poc_base:       u32,
+    avg_pool:       NAVideoBufferPool<u8>,
 }
 
 impl H264MTDecoder {
@@ -491,6 +490,7 @@ impl H264MTDecoder {
             deblock_skip:   false,
             max_last_poc:   0,
             poc_base:       0,
+            avg_pool:       NAVideoBufferPool::new(8),
         }
     }
     fn handle_nal(&mut self, src: Vec<u8>, supp: &mut NADecoderSupport, skip_decoding: bool, user_id: u32, time: NATimeInfo) -> DecoderResult<()> {
@@ -558,8 +558,13 @@ impl H264MTDecoder {
                     let height = sps.pic_height_in_mbs << 4;
                     let num_mbs = sps.pic_width_in_mbs * sps.pic_height_in_mbs;
 
-                    let avg_vi = NAVideoInfo { width: 32, height: 32, flipped: false, format: YUV420_FORMAT, bits: 12 };
-                    let avg_buf = alloc_video_buffer(avg_vi, 4).unwrap().get_vbuf().unwrap();
+                    let avg_buf = if let Some(buf) = self.avg_pool.get_free() {
+                            buf
+                        } else {
+                            let new_avg_buf = alloc_video_buffer(AVG_BUF_VINFO, 4).unwrap().get_vbuf().unwrap();
+                            self.avg_pool.add_frame(new_avg_buf.clone());
+                            new_avg_buf
+                        };
                     let mut mc_dsp = H264MC::new(avg_buf);
                     mc_dsp.set_dimensions(width, height);
 
@@ -777,7 +782,7 @@ impl NADecoderMT for H264MTDecoder {
             }
 
             let num_bufs = if !self.sps.is_empty() {
-                    self.sps[0].num_ref_frames as usize + 1
+                    self.sps[0].num_ref_frames + 1
                 } else {
                     3
                 }.max(16 + 1);
@@ -789,6 +794,8 @@ impl NADecoderMT for H264MTDecoder {
             supp.pool_u8.set_dec_bufs(num_bufs + nthreads);
             supp.pool_u8.prealloc_video(NAVideoInfo::new(width, height, false, fmt), 4)?;
 
+            self.avg_pool.prealloc_video(AVG_BUF_VINFO, 4)?;
+
             Ok(())
         } else {
             Err(DecoderError::InvalidData)
@@ -923,6 +930,7 @@ impl NADecoderMT for H264MTDecoder {
     }
     fn flush(&mut self) {
         clear_threads(&mut self.dispatch);
+        self.frame_refs.clear_refs();
     }
 }