h264: add multi-threaded decoder
[nihav.git] / nihav-itu / src / codecs / h264 / pic_ref.rs
CommitLineData
696e4e20 1use nihav_core::codecs::DecoderResult;
11d7aef2 2use nihav_core::frame::{FrameType, NAVideoBufferRef, NATimeInfo};
93839abd 3use nihav_core::refs::*;
696e4e20
KS
4use nihav_codec_support::codecs::MV;
5use super::sets::SeqParameterSet;
6use super::slice::*;
7use super::types::*;
8
9#[derive(Clone)]
10pub struct PictureInfo {
11 pub id: u16,
12 pub full_id: u32,
11d7aef2
KS
13 pub time: NATimeInfo,
14 pub user_id: u32,
696e4e20
KS
15 pub pic_type: FrameType,
16 pub buf: NAVideoBufferRef<u8>,
17 pub cur_mb: usize,
18 pub is_ref: bool,
11d7aef2 19 pub is_idr: bool,
696e4e20
KS
20 pub long_term: Option<usize>,
21
93839abd 22 pub mv_info: NABufferRef<FrameMV>,
696e4e20
KS
23}
24
25#[derive(Clone,Copy,Default, Debug)]
26pub struct FrameMBInfo {
27 pub mb_type: CompactMBType,
28 pub ref_poc: [[u16; 2]; 4],
29 pub ref_idx: [[PicRef; 2]; 4],
30 pub mv: [[MV; 2]; 16],
31}
32
33impl FrameMBInfo {
34 pub fn new() -> Self { Self::default() }
35}
36
37#[derive(Clone)]
38pub struct FrameMV {
39 pub mbs: Vec<FrameMBInfo>,
40 pub mb_stride: usize,
41}
42
43impl FrameMV {
44 pub fn new(mb_w: usize, mb_h: usize) -> Self {
45 Self {
46 mbs: vec![FrameMBInfo::default(); mb_w * mb_h],
47 mb_stride: mb_w,
48 }
49 }
50}
51
56a17e69
KS
52#[derive(Clone)]
53pub struct SliceRefs {
696e4e20
KS
54 pub ref_list0: Vec<Option<PictureInfo>>,
55 pub ref_list1: Vec<Option<PictureInfo>>,
56a17e69
KS
56 pub cur_id: u32,
57}
58
59impl SliceRefs {
11d7aef2
KS
60 pub fn get_ref_id(&self, list_id: u8, ref_id: usize) -> Option<u32> {
61 let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
62 if ref_list.len() > ref_id {
63 ref_list[ref_id].as_ref().map(|pic| pic.full_id)
64 } else {
65 None
66 }
67 }
56a17e69
KS
68 pub fn select_ref_pic(&self, list_id: u8, ref_id: usize) -> Option<NAVideoBufferRef<u8>> {
69 let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
70 if ref_list.len() > ref_id {
71 ref_list[ref_id].as_ref().map(|pic| pic.buf.clone())
72 } else {
73 None
74 }
75 }
76 pub fn get_colocated_info(&self, mb_x: usize, mb_y: usize) -> (FrameMBInfo, u16, bool) {
77 if let Some(ref ref_pic) = &self.ref_list1[0] {
78 let mv_info = &ref_pic.mv_info;
79 let mb = mv_info.mbs[mb_x + mb_y * mv_info.mb_stride];
80 (mb, ref_pic.full_id as u16, ref_pic.long_term.is_some())
81 } else {
82 (FrameMBInfo::default(), 0, false)
83 }
84 }
85 pub fn map_ref0(&self, ref0_id: u16) -> (PicRef, bool) {
86 let mut r0_idx = 0;
87 let mut long = false;
88 for (i, rpic0) in self.ref_list0.iter().enumerate() {
89 if let Some(ref pic) = rpic0 {
90 if (pic.full_id as u16) == ref0_id {
91 r0_idx = i as u8;
92 long = pic.long_term.is_some();
93 break;
94 }
95 }
96 }
97 (PicRef::new(r0_idx), long)
98 }
99 pub fn map_refs(&self, ref_idx: [PicRef; 2]) -> [u16; 2] {
100 let r0 = ref_idx[0].index();
101 let r1 = ref_idx[1].index();
102 let ref0 = if r0 < self.ref_list0.len() {
103 if let Some(ref pic) = self.ref_list0[r0] {
104 pic.full_id as u16
105 } else {
106 MISSING_POC
107 }
108 } else {
109 MISSING_POC
110 };
111 let ref1 = if r1 < self.ref_list1.len() {
112 if let Some(ref pic) = self.ref_list1[r1] {
113 pic.full_id as u16
114 } else {
115 MISSING_POC
116 }
117 } else {
118 MISSING_POC
119 };
120 [ref0, ref1]
121 }
122 pub fn cmp_refs(&self, ref1: [PicRef; 2], ref2: [PicRef; 2]) -> bool {
123 if ref1 != ref2 {
124 self.cmp_ref(ref1[0], ref2[0], 0) && self.cmp_ref(ref1[1], ref2[1], 1)
125 } else {
126 true
127 }
128 }
129 fn cmp_ref(&self, ref1: PicRef, ref2: PicRef, list: u8) -> bool {
130 if ref1 == ref2 {
131 true
132 } else {
133 let idx0 = ref1.index();
134 let idx1 = ref2.index();
135 if idx0 == idx1 {
136 return true;
137 }
138 let src = if list == 0 { &self.ref_list0 } else { &self.ref_list1 };
139 if idx0 >= src.len() || idx1 >= src.len() {
140//panic!("wrong refs");
141 return false;
142 }
143 if let (Some(ref pic0), Some(ref pic1)) = (&src[idx0], &src[idx1]) {
144 pic0.full_id == pic1.full_id
145 } else {
146//panic!("missing pics");
147 false
148 }
149 }
150 }
151}
152
153pub struct FrameRefs {
154 pub ref_pics: Vec<PictureInfo>,
155 pub cur_refs: SliceRefs,
696e4e20
KS
156 pub long_term: Vec<Option<PictureInfo>>,
157
158 prev_poc_msb: u32,
159 prev_poc_lsb: u16,
160 prev_ref_poc_lsb: u16,
161 prev_frame_num: u16,
162 frame_num_offset: u32,
02cfd8de 163 max_frame_num: i32,
696e4e20
KS
164}
165
166impl FrameRefs {
167 pub fn new() -> Self {
168 Self {
169 ref_pics: Vec::with_capacity(16),
56a17e69
KS
170 cur_refs: SliceRefs {
171 ref_list0: Vec::with_capacity(3),
172 ref_list1: Vec::with_capacity(3),
173 cur_id: 0,
174 },
696e4e20
KS
175 long_term: Vec::new(),
176
177 prev_poc_msb: 0,
178 prev_poc_lsb: 0,
179 prev_ref_poc_lsb: 0,
180 prev_frame_num: 0,
181 frame_num_offset: 0,
02cfd8de 182 max_frame_num: 0,
696e4e20
KS
183 }
184 }
11d7aef2
KS
185 pub fn fill_ref_nums(&self, dst: &mut Vec<u32>) {
186 for pic in self.ref_pics.iter() {
187 if !dst.contains(&pic.full_id) {
188 dst.push(pic.full_id);
189 }
190 }
191 for ltpic in self.long_term.iter() {
192 if let Some(ref pic) = ltpic {
193 if !dst.contains(&pic.full_id) {
194 dst.push(pic.full_id);
195 }
196 }
197 }
198 }
696e4e20 199 pub fn calc_picture_num(&mut self, slice_hdr: &SliceHeader, is_idr: bool, ref_id: u8, sps: &SeqParameterSet) -> u32 {
02cfd8de 200 self.max_frame_num = 1 << sps.log2_max_frame_num;
696e4e20
KS
201 match sps.pic_order_cnt_type {
202 0 => {
203 if is_idr {
204 //self.prev_poc_msb = 0;
205 self.prev_poc_lsb = 0;
206 } else {
207 self.prev_poc_lsb = self.prev_ref_poc_lsb;
208 }
209 let max_poc_lsb = 1 << sps.log2_max_pic_order_cnt_lsb;
210 let half_max_poc_lsb = 1 << (sps.log2_max_pic_order_cnt_lsb - 1);
211 let cur_lsb = slice_hdr.pic_order_cnt_lsb;
212 let poc_msb = if cur_lsb < self.prev_poc_lsb && (self.prev_poc_lsb - cur_lsb >= half_max_poc_lsb) {
213 self.prev_poc_msb + max_poc_lsb
214 } else if cur_lsb > self.prev_poc_lsb && (cur_lsb - self.prev_poc_lsb > half_max_poc_lsb) {
215 self.prev_poc_msb.wrapping_sub(max_poc_lsb)
216 } else {
217 self.prev_poc_msb
218 };
219 let poc = poc_msb + u32::from(cur_lsb);
220 if ref_id != 0 {
221 self.prev_ref_poc_lsb = slice_hdr.pic_order_cnt_lsb;
222 self.prev_poc_msb = poc_msb;
223 }
224 poc
225 },
226 1 => {
227 let off = if self.prev_frame_num > slice_hdr.frame_num {
228 self.frame_num_offset + (1 << sps.log2_max_frame_num)
229 } else {
230 self.frame_num_offset
231 };
232 let mut anum = if sps.num_ref_frames_in_pic_order_cnt_cycle != 0 {
233 (off as i32) + i32::from(slice_hdr.frame_num)
234 } else {
235 0
236 };
237 if ref_id == 0 && anum > 0 {
238 anum -= 1;
239 }
240 let (poc_cycle_cnt, fno_in_poc_cycle) = if anum > 0 {
241 let nrf = sps.num_ref_frames_in_pic_order_cnt_cycle as i32;
242 ((anum - 1) / nrf, (anum - 1) % nrf)
243 } else {
244 (0, 0)
245 };
246 let mut expected_delta = 0;
247 for &offset in sps.offset_for_ref_frame[..sps.num_ref_frames_in_pic_order_cnt_cycle].iter() {
248 expected_delta += offset;
249 }
250 let mut expected_poc = if anum > 0 {
251 let mut sum = poc_cycle_cnt * expected_delta;
252 for &offset in sps.offset_for_ref_frame[..=fno_in_poc_cycle as usize].iter() {
253 sum += offset;
254 }
255 sum
256 } else {
257 0
258 };
259 if ref_id == 0 {
260 expected_poc += sps.offset_for_non_ref_pic;
261 }
262 let (top_id, _bottom_id) = if !slice_hdr.field_pic {
263 let top_id = expected_poc + slice_hdr.delta_pic_order_cnt[0];
264 let bot_id = top_id + sps.offset_for_top_to_bottom_field + slice_hdr.delta_pic_order_cnt[1];
265 (top_id, bot_id)
266 } else if !slice_hdr.bottom_field {
267 (expected_poc + slice_hdr.delta_pic_order_cnt[0], 0)
268 } else {
269 (0, sps.offset_for_top_to_bottom_field + slice_hdr.delta_pic_order_cnt[1])
270 };
271 self.prev_frame_num = slice_hdr.frame_num;
272 self.frame_num_offset = off;
273 top_id as u32
274 },
275 _ => {
276 if slice_hdr.frame_num < self.prev_frame_num {
277 self.frame_num_offset += 1 << sps.log2_max_frame_num;
278 }
279 self.prev_frame_num = slice_hdr.frame_num;
280 self.frame_num_offset + u32::from(slice_hdr.frame_num)
281 },
282 }
283 }
284 pub fn apply_adaptive_marking(&mut self, marking: &AdaptiveMarking, cur_id: u16, max_id: u16) -> DecoderResult<()> {
285 let all_ref_pics = self.ref_pics.clone();
286
287 for (&op, (&arg1, &arg2)) in marking.memory_management_control_op.iter().zip(marking.operation_arg.iter().zip(marking.operation_arg2.iter())).take(marking.num_ops) {
288 match op {
289 1 => {
290 let src_id = cur_id.wrapping_sub(arg1) & (max_id - 1);
291 let mut found = false;
292 let mut idx = 0;
293 for (i, pic) in self.ref_pics.iter().enumerate() {
294 if pic.id == src_id {
295 found = true;
296 idx = i;
297 break;
298 }
299 }
300 if found {
301 self.ref_pics.remove(idx);
302 }
303 },
304 2 => { // mark long term picture as unused
305 let idx = arg1 as usize;
306 if idx < self.long_term.len() {
307 self.long_term[idx] = None;
308 }
309 },
310 3 => {
311 let src_id = cur_id.wrapping_sub(arg1) & (max_id - 1);
312
313 let didx = arg2 as usize;
314 for pic in all_ref_pics.iter() {
315 if pic.id == src_id {
316 if didx < self.long_term.len() {
317 self.long_term[didx] = Some(pic.clone());
318 }
319 break;
320 }
321 }
322 },
323 4 => {
324 self.long_term.resize(arg1 as usize, None);
325 },
326 5 => {
37952415
KS
327 self.ref_pics.clear();
328 self.long_term.clear();
696e4e20
KS
329 },
330 6 => {
331 // assign an long term index to current pic - done elsewhere
332 },
333 _ => {},
334 };
335 }
336 Ok(())
337 }
338 pub fn clear_refs(&mut self) {
37952415
KS
339 self.ref_pics.clear();
340 self.long_term.clear();
696e4e20 341 }
b7c882c1 342 #[allow(clippy::cognitive_complexity)]
696e4e20 343 pub fn select_refs(&mut self, sps: &SeqParameterSet, slice_hdr: &SliceHeader, cur_id: u32) {
56a17e69
KS
344 self.cur_refs.cur_id = cur_id;
345 self.cur_refs.ref_list0.clear();
346 self.cur_refs.ref_list1.clear();
696e4e20
KS
347 let pic_num_mask = if sps.log2_max_frame_num == 16 {
348 0xFFFF
349 } else {
350 (1 << sps.log2_max_frame_num) - 1
351 };
56a17e69 352
696e4e20
KS
353 if !slice_hdr.slice_type.is_intra() {
354 let has_reordering = slice_hdr.ref_pic_list_reordering_l0;
355 if !has_reordering {
356 let num_ref = slice_hdr.num_ref_idx_l0_active;
357 if slice_hdr.slice_type.is_p() {
358 if !self.ref_pics.is_empty() {
359 for pic in self.ref_pics.iter().rev().take(num_ref) {
56a17e69 360 self.cur_refs.ref_list0.push(Some(pic.clone()));
696e4e20
KS
361 }
362 }
363 } else {
364 let mut pivot = 0;
365 for (i, pic) in self.ref_pics.iter().enumerate() {
366 pivot = i;
367 if pic.full_id > cur_id {
368 break;
369 }
370 }
371 for pic in self.ref_pics[..pivot].iter().rev() {
56a17e69 372 if self.cur_refs.ref_list0.len() >= num_ref {
696e4e20
KS
373 break;
374 }
56a17e69 375 self.cur_refs.ref_list0.push(Some(pic.clone()));
696e4e20
KS
376 }
377 for pic in self.ref_pics.iter().skip(pivot) {
56a17e69 378 if self.cur_refs.ref_list0.len() >= num_ref {
696e4e20
KS
379 break;
380 }
56a17e69 381 self.cur_refs.ref_list0.push(Some(pic.clone()));
696e4e20
KS
382 }
383 }
56a17e69
KS
384 if !self.long_term.is_empty() && self.cur_refs.ref_list0.len() < num_ref {
385 let copy_size = num_ref - self.cur_refs.ref_list0.len();
696e4e20 386 for ltpic in self.long_term.iter().take(copy_size) {
56a17e69 387 self.cur_refs.ref_list0.push(ltpic.clone());
696e4e20
KS
388 }
389 }
390 } else {
56a17e69 391 form_ref_list(&mut self.cur_refs.ref_list0,
696e4e20
KS
392 &self.ref_pics, &self.long_term,
393 &slice_hdr.reordering_list_l0,
394 slice_hdr.frame_num, pic_num_mask);
395 }
396 if slice_hdr.slice_type.is_b() {
397 let has_reordering = slice_hdr.ref_pic_list_reordering_l1;
398 if !has_reordering {
399 let num_ref = slice_hdr.num_ref_idx_l1_active;
400 let mut pivot = 0;
401 for (i, pic) in self.ref_pics.iter().enumerate() {
402 pivot = i;
403 if pic.full_id > cur_id {
404 break;
405 }
406 }
407 for pic in self.ref_pics.iter().skip(pivot) {
56a17e69 408 if self.cur_refs.ref_list1.len() >= num_ref {
696e4e20
KS
409 break;
410 }
56a17e69 411 self.cur_refs.ref_list1.push(Some(pic.clone()));
696e4e20
KS
412 }
413 for pic in self.ref_pics[..pivot].iter().rev() {
56a17e69 414 if self.cur_refs.ref_list1.len() >= num_ref {
696e4e20
KS
415 break;
416 }
56a17e69 417 self.cur_refs.ref_list1.push(Some(pic.clone()));
696e4e20 418 }
56a17e69
KS
419 if !self.long_term.is_empty() && self.cur_refs.ref_list1.len() < num_ref {
420 let copy_size = num_ref - self.cur_refs.ref_list1.len();
696e4e20 421 for ltpic in self.long_term.iter().take(copy_size) {
56a17e69 422 self.cur_refs.ref_list1.push(ltpic.clone());
696e4e20
KS
423 }
424 }
56a17e69 425 if self.cur_refs.ref_list1.len() > 1 && self.cur_refs.ref_list0.len() == self.cur_refs.ref_list1.len() {
696e4e20 426 let mut equal = true;
56a17e69 427 for (pic1, pic2) in self.cur_refs.ref_list0.iter().zip(self.cur_refs.ref_list1.iter()) {
696e4e20
KS
428 match (pic1, pic2) {
429 (Some(p1), Some(p2)) => {
430 if p1.full_id != p2.full_id {
431 equal = false;
432 break;
433 }
434 },
435 (None, None) => {},
436 _ => {
437 equal = false;
438 break;
439 },
440 };
441 }
442 if equal {
56a17e69 443 self.cur_refs.ref_list1.swap(0, 1);
696e4e20
KS
444 }
445 }
446 } else {
56a17e69 447 form_ref_list(&mut self.cur_refs.ref_list1,
696e4e20
KS
448 &self.ref_pics, &self.long_term,
449 &slice_hdr.reordering_list_l1,
450 slice_hdr.frame_num, pic_num_mask);
451 }
452 }
453 }
454 }
455 pub fn add_short_term(&mut self, cpic: PictureInfo, num_ref_frames: usize) {
456 if !self.ref_pics.is_empty() && self.ref_pics.len() >= num_ref_frames {
02cfd8de
KS
457 let base_id = i32::from(cpic.id);
458 let mut min_id = base_id;
459 let mut min_idx = 0;
460 for (i, pic) in self.ref_pics.iter().enumerate() {
461 let mut pic_id = i32::from(pic.id);
462 if pic_id > base_id {
463 pic_id -= self.max_frame_num;
464 }
465 if pic_id < min_id {
466 min_id = pic_id;
467 min_idx = i;
468 }
469 }
470 self.ref_pics.remove(min_idx);
696e4e20
KS
471 }
472 if self.ref_pics.is_empty() || self.ref_pics.last().unwrap().full_id < cpic.full_id {
473 self.ref_pics.push(cpic);
474 } else {
475 let mut idx = 0;
476 for (i, pic) in self.ref_pics.iter().enumerate() {
477 if pic.full_id < cpic.full_id {
478 idx = i;
479 } else {
480 break;
481 }
482 }
483 self.ref_pics.insert(idx + 1, cpic);
484 }
485 }
486 pub fn add_long_term(&mut self, lt_idx: usize, cpic: PictureInfo) {
487 if lt_idx < self.long_term.len() {
488 self.long_term[lt_idx] = Some(cpic);
489 }
490 }
696e4e20
KS
491}
492
4a1ca15c 493fn form_ref_list(ref_list: &mut Vec<Option<PictureInfo>>, ref_pics: &[PictureInfo], long_term: &[Option<PictureInfo>], reord_info: &ReorderingInfo, cur_id: u16, pic_num_mask: u16) {
696e4e20
KS
494 let mut ref_pic_id = cur_id;
495 for (&op, &num) in reord_info.reordering_of_pic_nums_idc.iter().zip(reord_info.abs_diff_or_num.iter()).take(reord_info.num_ops) {
496 if op < 2 {
497 if op == 0 {
498 ref_pic_id = ref_pic_id.wrapping_sub(num) & pic_num_mask;
499 } else {
500 ref_pic_id = ref_pic_id.wrapping_add(num) & pic_num_mask;
501 }
502 let mut found = false;
503 for pic in ref_pics.iter() {
504 if pic.id == ref_pic_id {
505 ref_list.push(Some(pic.clone()));
506 found = true;
507 break;
508 }
509 }
510 if !found {
511 ref_list.push(None);
512 }
513 } else {
514 let idx = num as usize;
515 if idx < long_term.len() {
516 ref_list.push(long_term[idx].clone());
517 } else {
518 ref_list.push(None);
519 }
520 }
521 }
522}