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