h264: add multi-threaded decoder
[nihav.git] / nihav-itu / src / codecs / h264 / decoder_mt.rs
CommitLineData
11d7aef2
KS
1/* TODO:
2 * buffer pool for DSP avg frames
3*/
4
5use nihav_core::codecs::*;
6use nihav_core::io::bitreader::*;
7
8use super::*;
9use super::dispatch::*;
10
11pub struct FrameDecoder {
12 pub slices: Vec<(SliceHeader, usize, SliceRefs, Vec<u8>)>,
13 pub cur_pic: PictureInfo,
14 sps: Arc<SeqParameterSet>,
15 pps: Arc<PicParameterSet>,
16 pub num_mbs: usize,
17 mc_dsp: H264MC,
18 dispatch: Shareable<ThreadDispatcher>,
19 sstate: SliceState,
20 cavlc_cb: Arc<CAVLCTables>,
21 ipcm_buf: [u8; 256 + 64 + 64],
22 is_mbaff: bool,
23 deblock_skip: bool,
24}
25
26impl FrameDecoder {
27 pub fn decode_slice(&mut self, hdr: &SliceHeader, hdr_size: usize, refs: &SliceRefs, nal: &[u8]) -> DecoderResult<usize> {
28 self.sstate.reset(self.sps.pic_width_in_mbs, self.sps.pic_height_in_mbs, hdr.first_mb_in_slice);
29
30 let mut full_size = nal.len() * 8;
31 for &byte in nal.iter().rev() {
32 if byte == 0 {
33 full_size -= 8;
34 } else {
35 full_size -= (byte.trailing_zeros() + 1) as usize;
36 break;
37 }
38 }
39 validate!(full_size > 0);
40
41 let mut br = BitReader::new(&nal[hdr_size / 8..], BitReaderMode::BE);
42 if !self.pps.entropy_coding_mode {
43 br.skip((hdr_size & 7) as u32)?;
44 self.decode_slice_cavlc(&mut br, full_size - (hdr_size & !7), hdr, refs)
45 } else {
46 let csrc = &nal[(hdr_size + 7) / 8..];
47 validate!(csrc.len() >= 2);
48 let mut cabac = CABAC::new(csrc, hdr.slice_type, hdr.slice_qp, hdr.cabac_init_idc as usize)?;
49 self.decode_slice_cabac(&mut cabac, hdr, refs)
50 }
51 }
52 fn decode_slice_cavlc(&mut self, br: &mut BitReader, full_size: usize, slice_hdr: &SliceHeader, refs: &SliceRefs) -> DecoderResult<usize> {
53 const INTRA_CBP: [u8; 48] = [
54 47, 31, 15, 0, 23, 27, 29, 30, 7, 11, 13, 14, 39, 43, 45, 46,
55 16, 3, 5, 10, 12, 19, 21, 26, 28, 35, 37, 42, 44, 1, 2, 4,
56 8, 17, 18, 20, 24, 6, 9, 22, 25, 32, 33, 34, 36, 40, 38, 41
57 ];
58 const INTER_CBP: [u8; 48] = [
59 0, 16, 1, 2, 4, 8, 32, 3, 5, 10, 12, 15, 47, 7, 11, 13,
60 14, 6, 9, 31, 35, 37, 42, 44, 33, 34, 36, 40, 39, 43, 45, 46,
61 17, 18, 20, 24, 19, 21, 26, 28, 23, 27, 29, 30, 22, 25, 38, 41
62 ];
63
64 let mut mb_idx = slice_hdr.first_mb_in_slice as usize;
65 let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
66 let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
67 while br.tell() < full_size && mb_idx < self.num_mbs {
68 mb_info.coded = [false; 25];
69 mb_info.ref_l0 = [ZERO_REF; 4];
70 mb_info.ref_l1 = [ZERO_REF; 4];
71 mb_info.mv_l0 = [ZERO_MV; 16];
72 mb_info.mv_l1 = [ZERO_MV; 16];
73 mb_info.chroma_dc = [[0; 4]; 2];
74 mb_info.cbpy = 0;
75 mb_info.cbpc = 0;
76
77 if !slice_hdr.slice_type.is_intra() {
78 let mb_skip_run = br.read_ue()? as usize;
79 validate!(mb_idx + mb_skip_run <= self.num_mbs);
80 mb_info.mb_type = skip_type;
81 for _ in 0..mb_skip_run {
82 self.handle_macroblock(slice_hdr, &mut mb_info, refs)?;
83 mb_idx += 1;
84 }
85 if mb_idx == self.num_mbs || br.tell() >= full_size {
86 break;
87 }
88 }
89 if br.tell() < full_size {
90 if self.is_mbaff && ((mb_idx & 1) == 0) {
91 let _mb_field_decoding = br.read_bool()?;
92 }
93 let mut mb_type = decode_mb_type_cavlc(br, slice_hdr)?;
94 mb_info.mb_type = mb_type;
95 mb_info.transform_size_8x8 = false;
96 if mb_type == MBType::PCM {
97 br.align();
98 for pix in self.ipcm_buf[..256 + 64 + 64].iter_mut() {
99 *pix = br.read(8)? as u8;
100 }
101 self.sstate.fill_ncoded(16);
102 } else {
103 if self.pps.transform_8x8_mode && mb_type == MBType::Intra4x4 {
104 mb_info.transform_size_8x8 = br.read_bool()?;
105 if mb_info.transform_size_8x8 {
106 mb_type = MBType::Intra8x8;
107 mb_info.mb_type = MBType::Intra8x8;
108 }
109 }
110 decode_mb_pred_cavlc(br, slice_hdr, mb_type, &mut self.sstate, &mut mb_info)?;
111 let (cbpy, cbpc) = if let MBType::Intra16x16(_, cbpy, cbpc) = mb_type {
112 (cbpy, cbpc)
113 } else {
114 let cbp_id = br.read_ue()? as usize;
115 validate!(cbp_id < INTRA_CBP.len());
116 let cbp = if mb_type == MBType::Intra4x4 || mb_type == MBType::Intra8x8 {
117 INTRA_CBP[cbp_id]
118 } else {
119 INTER_CBP[cbp_id]
120 };
121 if self.pps.transform_8x8_mode && (cbp & 0xF) != 0 && mb_info.can_have_8x8_tx(self.sps.direct_8x8_inference) {
122 mb_info.transform_size_8x8 = br.read_bool()?;
123 }
124 ((cbp & 0xF), (cbp >> 4))
125 };
126 mb_info.cbpy = cbpy;
127 mb_info.cbpc = cbpc;
128 self.sstate.get_cur_mb().cbp = (cbpc << 4) | cbpy;
129 if cbpy != 0 || cbpc != 0 || mb_type.is_intra16x16() {
130 let mb_qp_delta = br.read_se()?;
131 validate!(mb_qp_delta >= -26 && mb_qp_delta <= 25);
132 let new_qp = mb_qp_delta + i32::from(mb_info.qp_y);
133 mb_info.qp_y = if new_qp < 0 {
134 (new_qp + 52) as u8
135 } else if new_qp >= 52 {
136 (new_qp - 52) as u8
137 } else {
138 new_qp as u8
139 };
140 mb_info.coeffs = [[0; 16]; 25];
141 if self.pps.transform_8x8_mode {
142 mb_info.clear_coeffs8x8();
143 }
144 mb_info.chroma_dc = [[0; 4]; 2];
145 decode_residual_cavlc(br, &mut self.sstate, &mut mb_info, &self.cavlc_cb)?;
146 }
147 }
148 self.handle_macroblock(slice_hdr, &mut mb_info, refs)?;
149 }
150 mb_idx += 1;
151 if let Ok(disp) = self.dispatch.read() {
152 disp.update_pos(self.cur_pic.full_id, mb_idx);
153 }
154 }
155 Ok(mb_idx)
156 }
157 fn decode_slice_cabac(&mut self, cabac: &mut CABAC, slice_hdr: &SliceHeader, refs: &SliceRefs) -> DecoderResult<usize> {
158 let mut mb_idx = slice_hdr.first_mb_in_slice as usize;
159 let mut prev_mb_skipped = false;
160 let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
161 let mut last_qp_diff = false;
162
163 let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
164
165 while mb_idx < self.num_mbs {
166 mb_info.coded = [false; 25];
167 mb_info.ref_l0 = [ZERO_REF; 4];
168 mb_info.ref_l1 = [ZERO_REF; 4];
169 mb_info.mv_l0 = [ZERO_MV; 16];
170 mb_info.mv_l1 = [ZERO_MV; 16];
171 mb_info.chroma_dc = [[0; 4]; 2];
172 mb_info.cbpy = 0;
173 mb_info.cbpc = 0;
174 let mb_skip = cabac_decode_mbskip(cabac, &self.sstate, slice_hdr);
175 if !mb_skip {
176 if self.is_mbaff && (((mb_idx & 1) == 0) || (prev_mb_skipped && ((mb_idx & 1) == 1))) {
177 let _mb_field_decoding = cabac.decode_bit(70);
178 }
179 let mut mb_type = cabac_decode_mb_type(cabac, slice_hdr, &self.sstate);
180 mb_info.mb_type = mb_type;
181 mb_info.transform_size_8x8 = false;
182 if mb_type == MBType::PCM {
183 let ipcm_size = 256 + 64 + 64;
184 validate!(cabac.pos + ipcm_size <= cabac.src.len());
185 self.ipcm_buf[..ipcm_size].copy_from_slice(&cabac.src[cabac.pos..][..ipcm_size]);
186 cabac.pos += ipcm_size;
187 cabac.reinit()?;
188 last_qp_diff = false;
189 } else {
190 if self.pps.transform_8x8_mode && mb_type == MBType::Intra4x4 {
191 let mut ctx = 0;
192 if self.sstate.get_top_mb().transform_8x8 {
193 ctx += 1;
194 }
195 if self.sstate.get_left_mb().transform_8x8 {
196 ctx += 1;
197 }
198 mb_info.transform_size_8x8 = cabac.decode_bit(399 + ctx);
199 if mb_info.transform_size_8x8 {
200 mb_type = MBType::Intra8x8;
201 mb_info.mb_type = MBType::Intra8x8;
202 }
203 }
204 decode_mb_pred_cabac(cabac, slice_hdr, mb_type, &mut self.sstate, &mut mb_info);
205 let (cbpy, cbpc) = if let MBType::Intra16x16(_, cbpy, cbpc) = mb_type {
206 (cbpy, cbpc)
207 } else {
208 decode_cbp_cabac(cabac, &self.sstate)
209 };
210 if self.pps.transform_8x8_mode && cbpy != 0 && mb_info.can_have_8x8_tx(self.sps.direct_8x8_inference) {
211 let mut ctx = 0;
212 if self.sstate.get_top_mb().transform_8x8 {
213 ctx += 1;
214 }
215 if self.sstate.get_left_mb().transform_8x8 {
216 ctx += 1;
217 }
218 mb_info.transform_size_8x8 = cabac.decode_bit(399 + ctx);
219 }
220 if mb_type.is_intra() {
221 self.sstate.get_cur_mb().cmode = mb_info.chroma_ipred;
222 }
223 mb_info.cbpy = cbpy;
224 mb_info.cbpc = cbpc;
225 self.sstate.get_cur_mb().cbp = (cbpc << 4) | cbpy;
226 if cbpy != 0 || cbpc != 0 || mb_type.is_intra16x16() {
227 let mb_qp_delta = decode_mb_qp_delta_cabac(cabac, last_qp_diff as usize);
228 validate!(mb_qp_delta >= -26 && mb_qp_delta <= 25);
229 last_qp_diff = mb_qp_delta != 0;
230 let new_qp = mb_qp_delta + i32::from(mb_info.qp_y);
231 mb_info.qp_y = if new_qp < 0 {
232 (new_qp + 52) as u8
233 } else if new_qp >= 52 {
234 (new_qp - 52) as u8
235 } else {
236 new_qp as u8
237 };
238 mb_info.coeffs = [[0; 16]; 25];
239 if self.pps.transform_8x8_mode {
240 mb_info.clear_coeffs8x8();
241 }
242 mb_info.chroma_dc = [[0; 4]; 2];
243 decode_residual_cabac(cabac, &mut self.sstate, &mut mb_info);
244 } else {
245 last_qp_diff = false;
246 }
247 }
248 } else {
249 mb_info.mb_type = skip_type;
250 mb_info.transform_size_8x8 = false;
251 last_qp_diff = false;
252 }
253 self.handle_macroblock(slice_hdr, &mut mb_info, refs)?;
254 prev_mb_skipped = mb_skip;
255 if !(self.is_mbaff && ((mb_idx & 1) == 0)) && cabac.decode_terminate() {
256 if let Ok(disp) = self.dispatch.read() {
257 disp.update_pos(self.cur_pic.full_id, mb_idx + 1);
258 }
259 return Ok(mb_idx + 1);
260 }
261 mb_idx += 1;
262 if let Ok(disp) = self.dispatch.read() {
263 disp.update_pos(self.cur_pic.full_id, mb_idx);
264 }
265 }
266 Err(DecoderError::InvalidData)
267 }
268 #[allow(clippy::cognitive_complexity)]
269 fn handle_macroblock(&mut self, slice_hdr: &SliceHeader, mb_info: &mut CurrentMBInfo, refs: &SliceRefs) -> DecoderResult<()> {
270 let qp_y = mb_info.qp_y;
271 let qpr = ((qp_y as i8) + self.pps.chroma_qp_index_offset).max(0).min(51) as usize;
272 let qp_u = CHROMA_QUANTS[qpr];
273 let qpb = ((qp_y as i8) + self.pps.second_chroma_qp_index_offset).max(0).min(51) as usize;
274 let qp_v = CHROMA_QUANTS[qpb];
275
276 let tx_bypass = qp_y == 0 && self.sps.qpprime_y_zero_transform_bypass;
277
278 self.sstate.get_cur_mb().mb_type = mb_info.mb_type.into();
279 if mb_info.mb_type != MBType::PCM {
280 self.sstate.get_cur_mb().qp_y = qp_y;
281 self.sstate.get_cur_mb().qp_u = qp_u;
282 self.sstate.get_cur_mb().qp_v = qp_v;
283 self.sstate.get_cur_mb().transform_8x8 = mb_info.transform_size_8x8;
284 }
285 let has_dc = mb_info.mb_type.is_intra16x16() && mb_info.coded[24];
286 if has_dc {
287 idct_luma_dc(&mut mb_info.coeffs[24], qp_y);
288 for i in 0..16 {
289 mb_info.coeffs[i][0] = mb_info.coeffs[24][i];
290 }
291 }
292 if !mb_info.transform_size_8x8 {
293 let quant_dc = !mb_info.mb_type.is_intra16x16();
294 for i in 0..16 {
295 if mb_info.coded[i] {
296 if !tx_bypass {
297 idct(&mut mb_info.coeffs[i], qp_y, quant_dc);
298 }
299 } else if has_dc {
300 if !tx_bypass {
301 idct_dc(&mut mb_info.coeffs[i], qp_y, quant_dc);
302 }
303 mb_info.coded[i] = true;
304 }
305 }
306 } else {
307 for i in 0..4 {
308 if mb_info.coded[(i & 1) * 2 + (i & 2) * 4] && !tx_bypass {
309 dequant8x8(&mut mb_info.coeffs8x8[i].coeffs, &self.pps.scaling_list_8x8[!mb_info.mb_type.is_intra() as usize]);
310 idct8x8(&mut mb_info.coeffs8x8[i].coeffs, qp_y);
311 }
312 }
313 }
314 for chroma in 0..2 {
315 let qp_c = if chroma == 0 { qp_u } else { qp_v };
316 if mb_info.cbpc != 0 {
317 chroma_dc_transform(&mut mb_info.chroma_dc[chroma], qp_c);
318 }
319 for i in 0..4 {
320 let blk_no = 16 + chroma * 4 + i;
321 mb_info.coeffs[blk_no][0] = mb_info.chroma_dc[chroma][i];
322 if mb_info.coded[blk_no] {
323 idct(&mut mb_info.coeffs[blk_no], qp_c, false);
324 } else if mb_info.coeffs[blk_no][0] != 0 {
325 idct_dc(&mut mb_info.coeffs[blk_no], qp_c, false);
326 mb_info.coded[blk_no] = true;
327 }
328 }
329 }
330 if !self.pps.entropy_coding_mode || mb_info.mb_type.is_skip() || mb_info.mb_type.is_intra() {
331 self.sstate.reset_mb_mv();
332 }
333 if !mb_info.mb_type.is_intra() {
334 let temporal_mv = !slice_hdr.direct_spatial_mv_pred;
335 let cur_id = self.cur_pic.full_id as u16;
336 // wait for the reference macroblock MV to be available
337 if matches!(mb_info.mb_type, MBType::Direct | MBType::BSkip | MBType::B8x8) {
338 if let Some(ref_id) = refs.get_ref_id(0, mb_info.ref_l1[0].index()) {
339 wait_for_mb(&self.dispatch, &self.sstate, self.sstate.mb_x * 16, self.sstate.mb_y * 16, ZERO_MV, ref_id)?;
340 }
341 }
342 Self::pred_mv(&mut self.sstate, refs, mb_info, cur_id, temporal_mv, self.sps.direct_8x8_inference);
343 }
344 if !self.pps.constrained_intra_pred && mb_info.mb_type != MBType::Intra4x4 && mb_info.mb_type != MBType::Intra8x8 {
345 self.sstate.fill_ipred(IntraPredMode::DC);
346 }
347
348 let xpos = self.sstate.mb_x * 16;
349 let ypos = self.sstate.mb_y * 16;
350 let mut frm = NASimpleVideoFrame::from_video_buf(&mut self.cur_pic.buf).unwrap();
351 if mb_info.mb_type != MBType::PCM {
352 let weight_mode = if self.pps.weighted_pred && slice_hdr.slice_type.is_p() {
353 1
354 } else if slice_hdr.slice_type.is_b() {
355 self.pps.weighted_bipred_idc
356 } else {
357 0
358 };
359 recon_mb_mt(&mut frm, slice_hdr, mb_info, &mut self.sstate, refs, &mut self.mc_dsp, weight_mode, &self.dispatch)?;
360 } else {
361 for (dline, src) in frm.data[frm.offset[0] + xpos + ypos * frm.stride[0]..].chunks_mut(frm.stride[0]).take(16).zip(self.ipcm_buf.chunks(16)) {
362 dline[..16].copy_from_slice(src);
363 }
364 for (dline, src) in frm.data[frm.offset[1] + xpos/2 + ypos/2 * frm.stride[1]..].chunks_mut(frm.stride[1]).take(8).zip(self.ipcm_buf[256..].chunks(8)) {
365 dline[..8].copy_from_slice(src);
366 }
367 for (dline, src) in frm.data[frm.offset[2] + xpos/2 + ypos/2 * frm.stride[2]..].chunks_mut(frm.stride[2]).take(8).zip(self.ipcm_buf[256 + 64..].chunks(8)) {
368 dline[..8].copy_from_slice(src);
369 }
370 }
371 self.sstate.save_ipred_context(&frm);
372
373 let mv_info = &mut self.cur_pic.mv_info;
374 let mb_pos = self.sstate.mb_x + self.sstate.mb_y * mv_info.mb_stride;
375 let mut mb = FrameMBInfo::new();
376 mb.mb_type = mb_info.mb_type.into();
377 for blk4 in 0..16 {
378 mb.mv[blk4] = self.sstate.get_cur_blk4(blk4).mv;
379 }
380 for blk8 in 0..4 {
381 mb.ref_poc[blk8] = refs.map_refs(self.sstate.get_cur_blk8(blk8).ref_idx);
382 mb.ref_idx[blk8] = self.sstate.get_cur_blk8(blk8).ref_idx;
383 }
384 mv_info.mbs[mb_pos] = mb;
385
386 let deblock_mode = slice_hdr.disable_deblocking_filter_idc;
387 if !self.deblock_skip && deblock_mode != 1 {
388 let is_s = slice_hdr.slice_type == SliceType::SI || slice_hdr.slice_type == SliceType::SP;
389 self.sstate.fill_deblock(refs, deblock_mode, is_s);
390 let mut frm = NASimpleVideoFrame::from_video_buf(&mut self.cur_pic.buf).unwrap();
391 let lf_alpha = slice_hdr.slice_alpha_c0_offset;
392 let lf_beta = slice_hdr.slice_beta_offset;
393 loop_filter_mb(&mut frm, &self.sstate, lf_alpha, lf_beta);
394 }
395 self.sstate.next_mb();
396 Ok(())
397 }
398
399 fn pred_mv(sstate: &mut SliceState, frame_refs: &SliceRefs, mb_info: &mut CurrentMBInfo, cur_id: u16, temporal_mv: bool, direct_8x8: bool) {
400 let mb_type = mb_info.mb_type;
401 if !mb_type.is_4x4() {
402 let (pw, ph) = mb_type.size();
403 let mut xoff = 0;
404 let mut yoff = 0;
405 if mb_type == MBType::Direct || mb_type == MBType::BSkip {
406 sstate.predict_direct_mb(frame_refs, temporal_mv, direct_8x8, cur_id);
407 }
408 for part in 0..mb_type.num_parts() {
409 if !mb_type.is_l1(part) {
410 match mb_type {
411 MBType::PSkip => sstate.predict_pskip(),
412 MBType::BSkip | MBType::Direct => {
413 },
414 _ => {
415 sstate.predict(xoff, yoff, pw, ph, 0,
416 mb_info.mv_l0[part], mb_info.ref_l0[part]);
417 },
418 };
419 }
420 if !mb_type.is_l0(part) && mb_type != MBType::BSkip && mb_type != MBType::Direct {
421 sstate.predict(xoff, yoff, pw, ph, 1, mb_info.mv_l1[part], mb_info.ref_l1[part]);
422 }
423 if pw != 16 {
424 xoff += pw;
425 } else {
426 yoff += ph;
427 }
428 }
429 } else {
430 for part in 0..4 {
431 let sub_type = mb_info.sub_mb_type[part];
432 let mut xoff = (part & 1) * 8;
433 let mut yoff = (part & 2) * 4;
434 let orig_x = xoff;
435 let (pw, ph) = sub_type.size();
436 for subpart in 0..sub_type.num_parts() {
437 if sub_type != SubMBType::Direct8x8 {
438 if !sub_type.is_l1() {
439 sstate.predict(xoff, yoff, pw, ph, 0, mb_info.mv_l0[part * 4 + subpart], mb_info.ref_l0[part]);
440 }
441 if !sub_type.is_l0() {
442 sstate.predict(xoff, yoff, pw, ph, 1, mb_info.mv_l1[part * 4 + subpart], mb_info.ref_l1[part]);
443 }
444 } else {
445 for sblk in 0..4 {
446 sstate.predict_direct_sub(frame_refs, temporal_mv, direct_8x8, cur_id, (xoff / 4) + (sblk & 1) + (yoff / 4) * 4 + (sblk & 2) * 2);
447 }
448 }
449 xoff += pw;
450 if xoff == orig_x + 8 {
451 xoff -= 8;
452 yoff += ph;
453 }
454 }
455 }
456 }
457 }
458}
459
460struct H264MTDecoder {
461 info: NACodecInfoRef,
462 nal_len: u8,
463 dispatch: Shareable<ThreadDispatcher>,
464 frame_refs: FrameRefs,
465 skip_mode: FrameSkipMode,
466 sps: Vec<Arc<SeqParameterSet>>,
467 cur_sps: usize,
468 pps: Vec<Arc<PicParameterSet>>,
469 cur_pps: usize,
470 cur_fdec: Option<FrameDecoder>,
471 cavlc_cb: Arc<CAVLCTables>,
472 deblock_skip: bool,
473 max_last_poc: u32,
474 poc_base: u32,
475}
476
477impl H264MTDecoder {
478 fn new() -> Self {
479 Self {
480 info: NACodecInfoRef::default(),
481 nal_len: 0,
482 dispatch: Arc::new(RwLock::new(ThreadDispatcher::new())),
483 frame_refs: FrameRefs::new(),
484 skip_mode: FrameSkipMode::default(),
485 sps: Vec::new(),
486 cur_sps: 0,
487 pps: Vec::new(),
488 cur_pps: 0,
489 cur_fdec: None,
490 cavlc_cb: Arc::new(CAVLCTables::new()),
491 deblock_skip: false,
492 max_last_poc: 0,
493 poc_base: 0,
494 }
495 }
496 fn handle_nal(&mut self, src: Vec<u8>, supp: &mut NADecoderSupport, skip_decoding: bool, user_id: u32, time: NATimeInfo) -> DecoderResult<()> {
497 validate!(!src.is_empty());
498 validate!((src[0] & 0x80) == 0);
499 let nal_ref_idc = src[0] >> 5;
500 let nal_unit_type = src[0] & 0x1F;
501
502 let mut full_size = src.len() * 8;
503 for &byte in src.iter().rev() {
504 if byte == 0 {
505 full_size -= 8;
506 } else {
507 full_size -= (byte.trailing_zeros() + 1) as usize;
508 break;
509 }
510 }
511 validate!(full_size > 0);
512 match nal_unit_type {
513 1 | 5 if !skip_decoding => {
514 let is_idr = nal_unit_type == 5;
515 let mut br = BitReader::new(&src[..(full_size + 7)/8], BitReaderMode::BE);
516 br.skip(8)?;
517
518 let slice_hdr = parse_slice_header(&mut br, self.sps.as_slice(), self.pps.as_slice(), is_idr, nal_ref_idc)?;
519 let hdr_size = br.tell();
520 validate!(br.tell() < full_size);
521 let full_id;
522 if slice_hdr.first_mb_in_slice == 0 {
523 validate!(self.cur_fdec.is_none());
524 for (i, pps) in self.pps.iter().enumerate() {
525 if pps.pic_parameter_set_id == slice_hdr.pic_parameter_set_id {
526 self.cur_pps = i;
527 break;
528 }
529 }
530 for (i, sps) in self.sps.iter().enumerate() {
531 if sps.seq_parameter_set_id == self.pps[self.cur_pps].seq_parameter_set_id {
532 self.cur_sps = i;
533 break;
534 }
535 }
536
537 let mut cur_full_id = self.frame_refs.calc_picture_num(&slice_hdr, is_idr, nal_ref_idc, &self.sps[self.cur_sps]) + self.poc_base;
538 if is_idr {
539 if cur_full_id <= self.max_last_poc {
540 self.poc_base = self.max_last_poc + 2 - (cur_full_id - self.poc_base);
541 cur_full_id = self.max_last_poc + 2;
542 }
543 }
544 self.max_last_poc = self.max_last_poc.max(cur_full_id);
545 full_id = cur_full_id;
546
547 let sps = &self.sps[self.cur_sps];
548 if sps.chroma_format_idc != 1 || sps.bit_depth_luma != 8 || sps.bit_depth_chroma != 8 {
549 println!(" chroma fmt {} bits {}/{}", sps.chroma_format_idc, sps.bit_depth_luma, sps.bit_depth_chroma);
550 return Err(DecoderError::NotImplemented);
551 }
552
553 if is_idr {
554 self.frame_refs.clear_refs();
555 }
556
557 let width = sps.pic_width_in_mbs << 4;
558 let height = sps.pic_height_in_mbs << 4;
559 let num_mbs = sps.pic_width_in_mbs * sps.pic_height_in_mbs;
560
561 let avg_vi = NAVideoInfo { width: 32, height: 32, flipped: false, format: YUV420_FORMAT, bits: 12 };
562 let avg_buf = alloc_video_buffer(avg_vi, 4).unwrap().get_vbuf().unwrap();
563 let mut mc_dsp = H264MC::new(avg_buf);
564 mc_dsp.set_dimensions(width, height);
565
566 let is_mbaff = sps.mb_adaptive_frame_field && !slice_hdr.field_pic;
567 if is_mbaff {
568 println!("MBAFF");
569 return Err(DecoderError::NotImplemented);
570 }
571 if !sps.frame_mbs_only {
572 println!("PAFF?");
573 return Err(DecoderError::NotImplemented);
574 }
575
576 let cur_vinfo = supp.pool_u8.get_info();
577 let tmp_vinfo = NAVideoInfo::new(width, height, false, YUV420_FORMAT);
578 if cur_vinfo != Some(tmp_vinfo) {
579 supp.pool_u8.reset();
580 supp.pool_u8.prealloc_video(tmp_vinfo, 4)?;
581 }
582
583 let buf = if let Some(pic) = supp.pool_u8.get_free() {
584 pic
585 } else {
586 if supp.pool_u8.get_num_used() > 256 {
587 return Err(DecoderError::AllocError);
588 }
589 if let Ok(nbuf) = alloc_video_buffer(tmp_vinfo, 4) {
590 let vbuf = nbuf.get_vbuf().unwrap();
591 supp.pool_u8.add_frame(vbuf.clone());
592 vbuf
593 } else {
594 return Err(DecoderError::AllocError);
595 }
596 };
597
598 let cur_pic = PictureInfo {
599 id: slice_hdr.frame_num,
600 full_id, user_id, time,
601 pic_type: slice_hdr.slice_type.to_frame_type(),
602 buf,
603 cur_mb: 0,
604 is_ref: nal_ref_idc != 0,
605 is_idr,
606 long_term: get_long_term_id(is_idr, &slice_hdr),
607 mv_info: NABufferRef::new(FrameMV::new(sps.pic_width_in_mbs, sps.pic_height_in_mbs)),
608 };
609
610 self.cur_fdec = Some(FrameDecoder{
611 slices: Vec::new(),
612 sstate: SliceState::new(),
613 ipcm_buf: [0; 256 + 64 + 64],
614 //width, height,
615 num_mbs,
616 sps: Arc::clone(sps),
617 pps: Arc::clone(&self.pps[self.cur_pps]),
618 dispatch: Arc::clone(&self.dispatch),
619 cavlc_cb: Arc::clone(&self.cavlc_cb),
620 mc_dsp,
621 cur_pic,
622 is_mbaff,
623 deblock_skip: self.deblock_skip,
624 });
625 } else {
626 if let Some(ref mut fdec) = self.cur_fdec {
627 let new_type = slice_hdr.slice_type.to_frame_type();
628 let pic = &mut fdec.cur_pic;
629 pic.pic_type = match (pic.pic_type, new_type) {
630 (FrameType::I, _) => new_type,
631 (_, FrameType::B) => FrameType::B,
632 _ => pic.pic_type,
633 };
634 full_id = pic.full_id;
635 } else {
636 return Ok(());
637 }
638 }
639
640 let sps = &self.sps[self.cur_sps];
641
642 self.frame_refs.select_refs(sps, &slice_hdr, full_id);
643
644 if slice_hdr.adaptive_ref_pic_marking_mode {
645 self.frame_refs.apply_adaptive_marking(&slice_hdr.adaptive_ref_pic_marking, slice_hdr.frame_num, 1 << self.sps[self.cur_sps].log2_max_frame_num)?;
646 }
647 if let Some(ref mut fdec) = self.cur_fdec {
648 fdec.slices.push((slice_hdr, hdr_size, self.frame_refs.cur_refs.clone(), src));
649 }
650 },
651 2 => { // slice data partition A
652 //slice header
653 //slice id = read_ue()
654 //cat 2 slice data (all but MB layer residual)
655 return Err(DecoderError::NotImplemented);
656 },
657 3 => { // slice data partition B
658 //slice id = read_ue()
659 //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
660 //cat 3 slice data (MB layer residual)
661 return Err(DecoderError::NotImplemented);
662 },
663 4 => { // slice data partition C
664 //slice id = read_ue()
665 //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
666 //cat 4 slice data (MB layer residual)
667 return Err(DecoderError::NotImplemented);
668 },
669 6 => {}, //SEI
670 7 => {
671 let sps = parse_sps(&src[1..])?;
672 self.sps.push(Arc::new(sps));
673 },
674 8 => {
675 validate!(full_size >= 8 + 16);
676 let pps = parse_pps(&src[1..], self.sps.as_slice(), full_size - 8)?;
677 let mut found = false;
678 for stored_pps in self.pps.iter_mut() {
679 if stored_pps.pic_parameter_set_id == pps.pic_parameter_set_id {
680 *stored_pps = Arc::clone(&pps);
681 found = true;
682 break;
683 }
684 }
685 if !found {
686 self.pps.push(pps);
687 }
688 },
689 9 => { // access unit delimiter
690 },
691 10 => {}, //end of sequence
692 11 => {}, //end of stream
693 12 => {}, //filler
694 _ => {},
695 };
696
697 Ok(())
698 }
699}
700
701impl NADecoderMT for H264MTDecoder {
702 fn init(&mut self, supp: &mut NADecoderSupport, info: NACodecInfoRef, nthreads: usize) -> DecoderResult<()> {
703 if let NACodecTypeInfo::Video(vinfo) = info.get_properties() {
704 let fmt = YUV420_FORMAT;
705 let myinfo = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, false, fmt));
706 self.info = NACodecInfo::new_ref(info.get_name(), myinfo, info.get_extradata()).into_ref();
707
708 let edata = info.get_extradata().unwrap();
709//print!("edata:"); for &el in edata.iter() { print!(" {:02X}", el); } println!();
710 if edata.len() > 11 && &edata[0..4] == b"avcC" {
711 let mut mr = MemoryReader::new_read(edata.as_slice());
712 let mut br = ByteReader::new(&mut mr);
713
714 br.read_skip(4)?;
715 let version = br.read_byte()?;
716 validate!(version == 1);
717 let profile = br.read_byte()?;
718 let _compatibility = br.read_byte()?;
719 let _level = br.read_byte()?;
720 let b = br.read_byte()?;
721 validate!((b & 0xFC) == 0xFC);
722 self.nal_len = (b & 3) + 1;
723 let b = br.read_byte()?;
724 validate!((b & 0xE0) == 0xE0);
725 let num_sps = (b & 0x1F) as usize;
726 for _ in 0..num_sps {
727 let len = br.read_u16be()? as usize;
728 let offset = br.tell() as usize;
729 validate!((br.peek_byte()? & 0x1F) == 7);
730 let mut nal_buf = Vec::new();
731 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
732 self.handle_nal(nal_buf, supp, true, 0, NATimeInfo::new(None, None, None, 0, 0))?;
733 br.read_skip(len)?;
734 }
735 let num_pps = br.read_byte()? as usize;
736 for _ in 0..num_pps {
737 let len = br.read_u16be()? as usize;
738 let offset = br.tell() as usize;
739 validate!((br.peek_byte()? & 0x1F) == 8);
740 let mut nal_buf = Vec::new();
741 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
742 self.handle_nal(nal_buf, supp, true, 0, NATimeInfo::new(None, None, None, 0, 0))?;
743 br.read_skip(len)?;
744 }
745 if br.left() > 0 {
746 match profile {
747 100 | 110 | 122 | 144 => {
748 let b = br.read_byte()?;
749 validate!((b & 0xFC) == 0xFC);
750 // b & 3 -> chroma format
751 let b = br.read_byte()?;
752 validate!((b & 0xF8) == 0xF8);
753 // b & 7 -> luma depth minus 8
754 let b = br.read_byte()?;
755 validate!((b & 0xF8) == 0xF8);
756 // b & 7 -> chroma depth minus 8
757 let num_spsext = br.read_byte()? as usize;
758 for _ in 0..num_spsext {
759 let len = br.read_u16be()? as usize;
760 // parse spsext
761 br.read_skip(len)?;
762 }
763 },
764 _ => {},
765 };
766 }
767 } else {
768 return Err(DecoderError::NotImplemented);
769 }
770
771 let mut width = vinfo.get_width();
772 let mut height = vinfo.get_height();
773
774 if (width == 0 || height == 0) && !self.sps.is_empty() {
775 width = self.sps[0].pic_width_in_mbs * 16;
776 height = self.sps[0].pic_height_in_mbs * 16;
777 }
778
779 let num_bufs = if !self.sps.is_empty() {
780 self.sps[0].num_ref_frames as usize + 1
781 } else {
782 3
783 }.max(16 + 1);
784 if let Ok(ref mut sd) = self.dispatch.write() {
785 sd.max_threads = nthreads;
786 } else {
787 return Err(DecoderError::Bug);
788 }
789 supp.pool_u8.set_dec_bufs(num_bufs + nthreads);
790 supp.pool_u8.prealloc_video(NAVideoInfo::new(width, height, false, fmt), 4)?;
791
792 Ok(())
793 } else {
794 Err(DecoderError::InvalidData)
795 }
796 }
797 fn can_take_input(&mut self) -> bool {
798 if let Ok(ref sd) = self.dispatch.read() {
799 sd.can_decode_more()
800 } else {
801 false
802 }
803 }
804 fn queue_pkt(&mut self, supp: &mut NADecoderSupport, pkt: &NAPacket, user_id: u32) -> DecoderResult<bool> {
805 if !self.can_take_input() {
806 return Ok(false);
807 }
808
809 let src = pkt.get_buffer();
810
811 let mut mr = MemoryReader::new_read(&src);
812 let mut br = ByteReader::new(&mut mr);
813 let mut nal_buf = Vec::with_capacity(src.len());
814
815 if self.nal_len > 0 {
816 let mut skip_decoding = false;
817 if self.skip_mode != FrameSkipMode::None {
818 let mut pic_type = FrameType::I;
819 let mut is_ref = false;
820 while br.left() > 0 {
821 let size = match self.nal_len {
822 1 => br.read_byte()? as usize,
823 2 => br.read_u16be()? as usize,
824 3 => br.read_u24be()? as usize,
825 4 => br.read_u32be()? as usize,
826 _ => unreachable!(),
827 };
828 validate!(br.left() >= (size as i64));
829 let offset = br.tell() as usize;
830 let size = unescape_nal(&src[offset..][..size], &mut nal_buf);
831 validate!(size > 0);
832 let nal_ref_idc = nal_buf[0] >> 5;
833 let nal_unit_type = nal_buf[0] & 0x1F;
834 if nal_unit_type == 1 || nal_unit_type == 5 {
835 let mut bitr = BitReader::new(&nal_buf[1..], BitReaderMode::BE);
836 let (first_mb, slice_type) = parse_slice_header_minimal(&mut bitr)?;
837 if first_mb == 0 && nal_ref_idc != 0 {
838 is_ref = true;
839 }
840 let new_type = slice_type.to_frame_type();
841 pic_type = match (pic_type, new_type) {
842 (FrameType::I, _) => new_type,
843 (_, FrameType::B) => FrameType::B,
844 _ => pic_type,
845 };
846 }
847 br.read_skip(size)?;
848 }
849 match self.skip_mode {
850 FrameSkipMode::IntraOnly => {
851 skip_decoding = pic_type != FrameType::I;
852 },
853 FrameSkipMode::KeyframesOnly => {
854 if !is_ref {
855 skip_decoding = true;
856 }
857 },
858 _ => {},
859 };
860 br.seek(SeekFrom::Start(0))?;
861 }
862
863 let mut initial_ref_frames = Vec::new();
864 self.frame_refs.fill_ref_nums(&mut initial_ref_frames);
865
866 while br.left() > 0 {
867 let size = match self.nal_len {
868 1 => br.read_byte()? as usize,
869 2 => br.read_u16be()? as usize,
870 3 => br.read_u24be()? as usize,
871 4 => br.read_u32be()? as usize,
872 _ => unreachable!(),
873 };
874 validate!(br.left() >= (size as i64));
875 let offset = br.tell() as usize;
876 let mut cur_nal_buf = Vec::with_capacity(size);
877 let _size = unescape_nal(&src[offset..][..size], &mut cur_nal_buf);
878 self.handle_nal(cur_nal_buf, supp, skip_decoding, user_id, pkt.ts)?;
879 br.read_skip(size)?;
880 }
881 let mut fdec = None;
882 std::mem::swap(&mut fdec, &mut self.cur_fdec);
883 if let Some(fdc) = fdec {
884 let cpic = &fdc.cur_pic;
885 if cpic.is_ref {
886 self.frame_refs.add_short_term(cpic.clone(), self.sps[self.cur_sps].num_ref_frames);
887 }
888 if let Some(lt_idx) = cpic.long_term {
889 self.frame_refs.add_long_term(lt_idx, cpic.clone());
890 }
891 let mut ref_frames = Vec::new();
892 self.frame_refs.fill_ref_nums(&mut ref_frames);
893 queue_decoding(&mut self.dispatch, fdc, &initial_ref_frames, &ref_frames);
894 }
895 } else {
896//todo NAL detection
897 unimplemented!();
898 }
899 Ok(true)
900 }
901 fn has_output(&mut self) -> bool {
902 if let Ok(ref ds) = self.dispatch.read() {
903 ds.has_output()
904 } else {
905 panic!("can't peek into status");
906 }
907 }
908 fn get_frame(&mut self) -> (DecoderResult<NAFrameRef>, u32) {
909 match wait_for_one(&mut self.dispatch) {
910 Ok(cpic) => {
911 let bufinfo = NABufferType::Video(cpic.buf.clone());
912 let ftype = cpic.pic_type;
913 let dts = Some(u64::from(cpic.full_id));
914 let mut frm = NAFrame::new(cpic.time, ftype, cpic.is_idr, self.info.clone(), bufinfo);
915 if let (Some(mydts), None) = (dts, frm.get_dts()) {
916 frm.set_dts(Some(mydts));
917 }
918 frm.set_id(cpic.user_id as i64);
919 (Ok(frm.into_ref()), cpic.user_id)
920 },
921 Err((err, id)) => (Err(err), id),
922 }
923 }
924 fn flush(&mut self) {
925 clear_threads(&mut self.dispatch);
926 }
927}
928
929impl NAOptionHandler for H264MTDecoder {
930 fn get_supported_options(&self) -> &[NAOptionDefinition] { DECODER_OPTIONS }
931 fn set_options(&mut self, options: &[NAOption]) {
932 for option in options.iter() {
933 for opt_def in DECODER_OPTIONS.iter() {
934 if opt_def.check(option).is_ok() {
935 match (option.name, &option.value) {
936 (FRAME_SKIP_OPTION, NAValue::String(ref strval)) => {
937 if let Ok(smode) = FrameSkipMode::from_str(strval) {
938 self.skip_mode = smode;
939 }
940 },
941 (DEBLOCK_SKIP_OPTION, NAValue::Bool(val)) => {
942 self.deblock_skip = *val;
943 },
944 _ => {},
945 }
946 }
947 }
948 }
949 }
950 fn query_option_value(&self, name: &str) -> Option<NAValue> {
951 match name {
952 FRAME_SKIP_OPTION => Some(NAValue::String(self.skip_mode.to_string())),
953 DEBLOCK_SKIP_OPTION => Some(NAValue::Bool(self.deblock_skip)),
954 _ => None,
955 }
956 }
957}
958
959pub fn get_decoder_mt() -> Box<dyn NADecoderMT + Send> {
960 Box::new(H264MTDecoder::new())
961}