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