[nihav.git] / nihav-core / src / reorder.rs

//! Output frame reordering.
//!
//! NihAV decoders output frames in the same order as they are put in.
//! In result if you want to have frames in display order you might need some frame reorderer.
//! This module provides such functionality depending on codec type: audio codecs and video codecs without B-frames do not need any reorderer and can use `NoReorderer` if the common interface is required. Codecs with B-frames should use `IPBReorderer`. For codecs with very complex reordering rules like H.264 or H.256 `PictureIDReorderer` will be added eventually.
//!
//! You can find out required reorderer by quering codec properties using `nihav_core::register` module.
use std::mem::swap;
use std::collections::VecDeque;
pub use crate::frame::{FrameType, NAFrameRef};

/// A trait for frame reorderer.
pub trait FrameReorderer {
    /// Stores a newly decoded frame.
    fn add_frame(&mut self, fref: NAFrameRef) -> bool;
    /// Gets the next frame to be displayed (or `None` if that is not possible).
    fn get_frame(&mut self) -> Option<NAFrameRef>;
    /// Clears all stored frames.
    fn flush(&mut self);
    /// Retrieves the last frames stored by the reorderer.
    fn get_last_frames(&mut self) -> Option<NAFrameRef>;
}

/// Zero reorderer.
pub struct NoReorderer {
    fref:   Option<NAFrameRef>,
}

impl NoReorderer {
    /// Constructs a new instance of `NoReorderer`.
    pub fn new() -> Self {
        Self { fref: None }
    }
}

impl Default for NoReorderer {
    fn default() -> Self {
        Self::new()
    }
}

impl FrameReorderer for NoReorderer {
    fn add_frame(&mut self, fref: NAFrameRef) -> bool {
        if self.fref.is_none() {
            self.fref = Some(fref);
            true
        } else {
            false
        }
    }
    fn get_frame(&mut self) -> Option<NAFrameRef> {
        let mut ret = None;
        swap(&mut ret, &mut self.fref);
        ret
    }
    fn flush(&mut self) { self.fref = None; }
    fn get_last_frames(&mut self) -> Option<NAFrameRef> { None }
}

/// Frame reorderer for codecs with I/P/B frames.
#[derive(Default)]
pub struct IPBReorderer {
    rframe:     Option<NAFrameRef>,
    bframe:     Option<NAFrameRef>,
}

impl IPBReorderer {
    /// Constructs a new instance of `IPBReorderer`.
    pub fn new() -> Self { Self::default() }
}

impl FrameReorderer for IPBReorderer {
    fn add_frame(&mut self, fref: NAFrameRef) -> bool {
        if self.rframe.is_some() && self.bframe.is_some() { return false; }
        let is_b = fref.get_frame_type() == FrameType::B;
        if is_b && self.bframe.is_some() { return false; }
        if is_b {
            self.bframe = Some(fref);
        } else {
            std::mem::swap(&mut self.bframe, &mut self.rframe);
            self.rframe = Some(fref);
        }
        true
    }
    fn get_frame(&mut self) -> Option<NAFrameRef> {
        let mut ret = None;
        if self.bframe.is_some() {
            std::mem::swap(&mut ret, &mut self.bframe);
        }
        ret
    }
    fn flush(&mut self) {
        self.rframe = None;
        self.bframe = None;
    }
    fn get_last_frames(&mut self) -> Option<NAFrameRef> {
        let mut ret = None;
        if self.bframe.is_some() {
            std::mem::swap(&mut ret, &mut self.bframe);
        } else if self.rframe.is_some() {
            std::mem::swap(&mut ret, &mut self.rframe);
        }
        ret
    }
}

/// Frame reorderer for codecs with complex I/P/B frame structure like ITU H.26x.
#[derive(Default)]
pub struct ComplexReorderer {
    last_ref_dts:   Option<u64>,
    ready_idx:      usize,
    frames:         Vec<NAFrameRef>,
}

impl ComplexReorderer {
    /// Constructs a new instance of `IPBReorderer`.
    pub fn new() -> Self { Self::default() }
}

impl FrameReorderer for ComplexReorderer {
    fn add_frame(&mut self, fref: NAFrameRef) -> bool {
        if self.frames.len() >= 64 {
            return false;
        }
        let is_ref = fref.frame_type == FrameType::I || fref.frame_type == FrameType::P;
        if !is_ref {
            if self.frames.is_empty() || fref.get_dts().is_none() {
                self.frames.push(fref);
            } else if let Some(new_dts) = fref.get_dts() {
                let mut idx = 0;
                for (i, frm) in self.frames.iter().enumerate() {
                    idx = i;
                    if let Some(dts) = frm.get_dts() {
                        if dts > new_dts {
                            break;
                        }
                    }
                }
                self.frames.insert(idx, fref);
            }
        } else {
            for (i, frm) in self.frames.iter().enumerate() {
                if frm.get_dts() == self.last_ref_dts {
                    self.ready_idx = i + 1;
                }
            }
            self.last_ref_dts = fref.get_dts();
            self.frames.push(fref);
        }
        true
    }
    fn get_frame(&mut self) -> Option<NAFrameRef> {
        if self.ready_idx > 0 {
            self.ready_idx -= 1;
            Some(self.frames.remove(0))
        } else {
            None
        }
    }
    fn flush(&mut self) {
        self.last_ref_dts = None;
        self.ready_idx = 0;
        self.frames.clear();
    }
    fn get_last_frames(&mut self) -> Option<NAFrameRef> {
        if !self.frames.is_empty() {
            Some(self.frames.remove(0))
        } else {
            None
        }
    }
}

/// A generic reorderer for a multi-threaded decoder.
#[derive(Default)]
pub struct MTFrameReorderer {
    ids:            Vec<u32>,
    frames:         VecDeque<(u32, NAFrameRef)>,
    cur_id:         u32,
    flush_mode:     bool,
    output_to:      Option<u32>,
    last_ts:        Option<u64>,
}

impl MTFrameReorderer {
    /// Constructs a new instance of `MTFrameReorderer`.
    pub fn new() -> Self { Self::default() }
    /// Registers the fact that a new frame is queued for decoding and returns an internal ID for it.
    pub fn register_frame(&mut self) -> u32 {
        if self.flush_mode {
            self.flush_mode = false;
        }
        let ret = self.cur_id;
        self.cur_id += 1;
        self.ids.push(ret);
        ret
    }
    /// Puts a newly decoded frame into the internal queue.
    pub fn add_frame(&mut self, frm: NAFrameRef, id: u32) {
        //let ftype = frm.get_frame_type();
        let frm_id = if let Some(ts) = frm.ts.dts { ts } else { u64::from(id) };
        let mut idx = 0;
        for (_, frm) in self.frames.iter() {
            let cur_id = if let Some(ts) = frm.ts.dts { ts } else { frm.id as u64 };
            if frm_id < cur_id {
                break;
            }
            idx += 1;
        }
        self.frames.insert(idx, (id, frm));
        /*if self.frames.len() > 48 {
            for (id, frm) in self.frames.iter() { print!(" {}{}({})", frm.get_frame_type(), frm.get_dts().unwrap_or(0), *id); } println!();
            print!("reg IDs:"); for &id in self.ids.iter() { print!(" {}", id); } println!();
            panic!("too many frames in the queue");
        }*/
    }
    /// Removes the registered frame (e.g. in case of a decoding error).
    pub fn drop_frame(&mut self, id: u32) {
        self.ids.retain(|&el| el != id);
    }
    fn get_first_frame(&mut self) -> Option<NAFrameRef> {
        let (id, frm) = self.frames.pop_front().unwrap();
        self.drop_frame(id);
        self.last_ts = frm.get_dts();
        Some(frm)
    }
    /// Gets the next frame to be displayed (or `None` if that is not possible).
    #[allow(clippy::collapsible_if)]
    #[allow(clippy::collapsible_else_if)]
    pub fn get_frame(&mut self) -> Option<NAFrameRef> {
        // check if we have consequent timestamps that we can output
        if !self.frames.is_empty() {
            if let Some(dts) = self.frames[0].1.get_dts() {
                let last_ts = self.last_ts.unwrap_or(0);
                if self.last_ts.is_none() || (dts == last_ts + 1) {
                    self.output_to = None;
                    return self.get_first_frame();
                }
            }
        }
        if !self.flush_mode {
            'out_loop: loop {
                if let Some(last_id) = self.output_to {
                    if self.frames[0].0 != last_id {
                        return self.get_first_frame();
                    } else {
                        self.output_to = None;
                    }
                }
                for (pos, (id, frm)) in self.frames.iter().enumerate() {
                    if frm.is_keyframe() || (self.frames.len() > 32 && matches!(frm.get_frame_type(), FrameType::I | FrameType::P)) {
                        let kf_id = *id;
                        self.ids.sort();
                        if pos == 0 && kf_id == self.ids[0] {
                            return self.get_first_frame();
                        }
                        let end = self.ids.iter().position(|&id| id == kf_id).unwrap();
                        for ref_id in self.ids[..end].iter() {
                            if !self.frames.iter().any(|(id, _)| id == ref_id) {
                                return None;
                            }
                        }
                        self.output_to = if pos < self.frames.len() - 1 {
                                Some(self.frames[pos + 1].0)
                            } else {
                                Some(kf_id)
                            };
                        continue 'out_loop;
                    }
                }
                return None;
            }
        } else {
            if !self.frames.is_empty() {
                Some(self.frames.pop_front().unwrap().1)
            } else {
                None
            }
        }
    }
    /// Retrieves the last frames stored by the reorderer.
    pub fn get_last_frames(&mut self) -> Option<NAFrameRef> {
        self.flush_mode = true;
        self.get_frame()
    }
    /// Clears all stored frames.
    pub fn flush(&mut self) {
        self.flush_mode = false;
        self.frames.clear();
        self.ids.clear();
        self.output_to = None;
        self.last_ts = None;
    }
}
Commit	Line	Data
6c9578be KS	1	//! Output frame reordering.
	2	//!
	3	//! NihAV decoders output frames in the same order as they are put in.
	4	//! In result if you want to have frames in display order you might need some frame reorderer.
	5	//! This module provides such functionality depending on codec type: audio codecs and video codecs without B-frames do not need any reorderer and can use `NoReorderer` if the common interface is required. Codecs with B-frames should use `IPBReorderer`. For codecs with very complex reordering rules like H.264 or H.256 `PictureIDReorderer` will be added eventually.
	6	//!
	7	//! You can find out required reorderer by quering codec properties using `nihav_core::register` module.
ca8452f0	8	use std::mem::swap;
4213ad59	9	use std::collections::VecDeque;
ca8452f0 KS	10	pub use crate::frame::{FrameType, NAFrameRef};
ca8452f0 KS	11
6c9578be	12	/// A trait for frame reorderer.
ca8452f0	13	pub trait FrameReorderer {
6c9578be	14	/// Stores a newly decoded frame.
ca8452f0	15	fn add_frame(&mut self, fref: NAFrameRef) -> bool;
6c9578be	16	/// Gets the next frame to be displayed (or `None` if that is not possible).
ca8452f0	17	fn get_frame(&mut self) -> Option<NAFrameRef>;
6c9578be	18	/// Clears all stored frames.
ca8452f0	19	fn flush(&mut self);
6c9578be	20	/// Retrieves the last frames stored by the reorderer.
ca8452f0 KS	21	fn get_last_frames(&mut self) -> Option<NAFrameRef>;
	22	}
	23
6c9578be	24	/// Zero reorderer.
ca8452f0 KS	25	pub struct NoReorderer {
	26	fref: Option<NAFrameRef>,
	27	}
	28
	29	impl NoReorderer {
6c9578be	30	/// Constructs a new instance of `NoReorderer`.
ca8452f0 KS	31	pub fn new() -> Self {
	32	Self { fref: None }
	33	}
	34	}
	35
b36f412c KS	36	impl Default for NoReorderer {
	37	fn default() -> Self {
	38	Self::new()
	39	}
	40	}
	41
ca8452f0 KS	42	impl FrameReorderer for NoReorderer {
	43	fn add_frame(&mut self, fref: NAFrameRef) -> bool {
	44	if self.fref.is_none() {
	45	self.fref = Some(fref);
	46	true
	47	} else {
	48	false
	49	}
	50	}
	51	fn get_frame(&mut self) -> Option<NAFrameRef> {
	52	let mut ret = None;
	53	swap(&mut ret, &mut self.fref);
	54	ret
	55	}
	56	fn flush(&mut self) { self.fref = None; }
	57	fn get_last_frames(&mut self) -> Option<NAFrameRef> { None }
	58	}
	59
6c9578be	60	/// Frame reorderer for codecs with I/P/B frames.
f9c35397	61	#[derive(Default)]
ca8452f0	62	pub struct IPBReorderer {
f9c35397 KS	63	rframe: Option<NAFrameRef>,
f9c35397 KS	64	bframe: Option<NAFrameRef>,
ca8452f0 KS	65	}
	66
	67	impl IPBReorderer {
6c9578be	68	/// Constructs a new instance of `IPBReorderer`.
f9c35397	69	pub fn new() -> Self { Self::default() }
ca8452f0 KS	70	}
	71
	72	impl FrameReorderer for IPBReorderer {
	73	fn add_frame(&mut self, fref: NAFrameRef) -> bool {
f9c35397 KS	74	if self.rframe.is_some() && self.bframe.is_some() { return false; }
	75	let is_b = fref.get_frame_type() == FrameType::B;
	76	if is_b && self.bframe.is_some() { return false; }
	77	if is_b {
	78	self.bframe = Some(fref);
ca8452f0	79	} else {
f9c35397 KS	80	std::mem::swap(&mut self.bframe, &mut self.rframe);
f9c35397 KS	81	self.rframe = Some(fref);
ca8452f0	82	}
f9c35397	83	true
ca8452f0 KS	84	}
ca8452f0 KS	85	fn get_frame(&mut self) -> Option<NAFrameRef> {
f9c35397 KS	86	let mut ret = None;
	87	if self.bframe.is_some() {
	88	std::mem::swap(&mut ret, &mut self.bframe);
ca8452f0	89	}
f9c35397	90	ret
ca8452f0 KS	91	}
ca8452f0 KS	92	fn flush(&mut self) {
f9c35397 KS	93	self.rframe = None;
f9c35397 KS	94	self.bframe = None;
ca8452f0 KS	95	}
ca8452f0 KS	96	fn get_last_frames(&mut self) -> Option<NAFrameRef> {
f9c35397 KS	97	let mut ret = None;
	98	if self.bframe.is_some() {
	99	std::mem::swap(&mut ret, &mut self.bframe);
	100	} else if self.rframe.is_some() {
	101	std::mem::swap(&mut ret, &mut self.rframe);
	102	}
	103	ret
ca8452f0 KS	104	}
	105	}
	106
8480761d KS	107	/// Frame reorderer for codecs with complex I/P/B frame structure like ITU H.26x.
	108	#[derive(Default)]
	109	pub struct ComplexReorderer {
	110	last_ref_dts: Option<u64>,
	111	ready_idx: usize,
	112	frames: Vec<NAFrameRef>,
	113	}
	114
	115	impl ComplexReorderer {
	116	/// Constructs a new instance of `IPBReorderer`.
	117	pub fn new() -> Self { Self::default() }
	118	}
	119
	120	impl FrameReorderer for ComplexReorderer {
	121	fn add_frame(&mut self, fref: NAFrameRef) -> bool {
	122	if self.frames.len() >= 64 {
	123	return false;
	124	}
	125	let is_ref = fref.frame_type == FrameType::I \|\| fref.frame_type == FrameType::P;
	126	if !is_ref {
	127	if self.frames.is_empty() \|\| fref.get_dts().is_none() {
	128	self.frames.push(fref);
	129	} else if let Some(new_dts) = fref.get_dts() {
	130	let mut idx = 0;
	131	for (i, frm) in self.frames.iter().enumerate() {
	132	idx = i;
	133	if let Some(dts) = frm.get_dts() {
	134	if dts > new_dts {
	135	break;
	136	}
	137	}
	138	}
	139	self.frames.insert(idx, fref);
	140	}
	141	} else {
	142	for (i, frm) in self.frames.iter().enumerate() {
	143	if frm.get_dts() == self.last_ref_dts {
	144	self.ready_idx = i + 1;
	145	}
	146	}
	147	self.last_ref_dts = fref.get_dts();
	148	self.frames.push(fref);
	149	}
	150	true
	151	}
	152	fn get_frame(&mut self) -> Option<NAFrameRef> {
	153	if self.ready_idx > 0 {
	154	self.ready_idx -= 1;
	155	Some(self.frames.remove(0))
	156	} else {
	157	None
	158	}
	159	}
	160	fn flush(&mut self) {
	161	self.last_ref_dts = None;
	162	self.ready_idx = 0;
b191eef3	163	self.frames.clear();
8480761d KS	164	}
	165	fn get_last_frames(&mut self) -> Option<NAFrameRef> {
	166	if !self.frames.is_empty() {
	167	Some(self.frames.remove(0))
	168	} else {
	169	None
	170	}
	171	}
	172	}
	173
4213ad59 KS	174	/// A generic reorderer for a multi-threaded decoder.
	175	#[derive(Default)]
	176	pub struct MTFrameReorderer {
	177	ids: Vec<u32>,
	178	frames: VecDeque<(u32, NAFrameRef)>,
	179	cur_id: u32,
	180	flush_mode: bool,
	181	output_to: Option<u32>,
	182	last_ts: Option<u64>,
	183	}
	184
	185	impl MTFrameReorderer {
	186	/// Constructs a new instance of `MTFrameReorderer`.
	187	pub fn new() -> Self { Self::default() }
	188	/// Registers the fact that a new frame is queued for decoding and returns an internal ID for it.
	189	pub fn register_frame(&mut self) -> u32 {
	190	if self.flush_mode {
	191	self.flush_mode = false;
	192	}
	193	let ret = self.cur_id;
	194	self.cur_id += 1;
	195	self.ids.push(ret);
	196	ret
	197	}
	198	/// Puts a newly decoded frame into the internal queue.
	199	pub fn add_frame(&mut self, frm: NAFrameRef, id: u32) {
	200	//let ftype = frm.get_frame_type();
	201	let frm_id = if let Some(ts) = frm.ts.dts { ts } else { u64::from(id) };
	202	let mut idx = 0;
	203	for (_, frm) in self.frames.iter() {
	204	let cur_id = if let Some(ts) = frm.ts.dts { ts } else { frm.id as u64 };
	205	if frm_id < cur_id {
	206	break;
	207	}
	208	idx += 1;
	209	}
	210	self.frames.insert(idx, (id, frm));
	211	/*if self.frames.len() > 48 {
	212	for (id, frm) in self.frames.iter() { print!(" {}{}({})", frm.get_frame_type(), frm.get_dts().unwrap_or(0), *id); } println!();
	213	print!("reg IDs:"); for &id in self.ids.iter() { print!(" {}", id); } println!();
	214	panic!("too many frames in the queue");
	215	}*/
	216	}
	217	/// Removes the registered frame (e.g. in case of a decoding error).
	218	pub fn drop_frame(&mut self, id: u32) {
	219	self.ids.retain(\|&el\| el != id);
	220	}
	221	fn get_first_frame(&mut self) -> Option<NAFrameRef> {
	222	let (id, frm) = self.frames.pop_front().unwrap();
	223	self.drop_frame(id);
	224	self.last_ts = frm.get_dts();
	225	Some(frm)
	226	}
	227	/// Gets the next frame to be displayed (or `None` if that is not possible).
	228	#[allow(clippy::collapsible_if)]
e6aaad5c	229	#[allow(clippy::collapsible_else_if)]
4213ad59 KS	230	pub fn get_frame(&mut self) -> Option<NAFrameRef> {
	231	// check if we have consequent timestamps that we can output
	232	if !self.frames.is_empty() {
	233	if let Some(dts) = self.frames[0].1.get_dts() {
	234	let last_ts = self.last_ts.unwrap_or(0);
	235	if self.last_ts.is_none() \|\| (dts == last_ts + 1) {
	236	self.output_to = None;
	237	return self.get_first_frame();
	238	}
	239	}
	240	}
	241	if !self.flush_mode {
	242	'out_loop: loop {
	243	if let Some(last_id) = self.output_to {
	244	if self.frames[0].0 != last_id {
	245	return self.get_first_frame();
	246	} else {
	247	self.output_to = None;
	248	}
	249	}
	250	for (pos, (id, frm)) in self.frames.iter().enumerate() {
	251	if frm.is_keyframe() \|\| (self.frames.len() > 32 && matches!(frm.get_frame_type(), FrameType::I \| FrameType::P)) {
	252	let kf_id = *id;
	253	self.ids.sort();
	254	if pos == 0 && kf_id == self.ids[0] {
	255	return self.get_first_frame();
	256	}
	257	let end = self.ids.iter().position(\|&id\| id == kf_id).unwrap();
	258	for ref_id in self.ids[..end].iter() {
e6aaad5c	259	if !self.frames.iter().any(\|(id, _)\| id == ref_id) {
4213ad59 KS	260	return None;
	261	}
	262	}
	263	self.output_to = if pos < self.frames.len() - 1 {
	264	Some(self.frames[pos + 1].0)
	265	} else {
	266	Some(kf_id)
	267	};
	268	continue 'out_loop;
	269	}
	270	}
	271	return None;
	272	}
	273	} else {
	274	if !self.frames.is_empty() {
	275	Some(self.frames.pop_front().unwrap().1)
	276	} else {
	277	None
	278	}
	279	}
	280	}
	281	/// Retrieves the last frames stored by the reorderer.
	282	pub fn get_last_frames(&mut self) -> Option<NAFrameRef> {
	283	self.flush_mode = true;
	284	self.get_frame()
	285	}
	286	/// Clears all stored frames.
	287	pub fn flush(&mut self) {
	288	self.flush_mode = false;
	289	self.frames.clear();
1bb151a6	290	self.ids.clear();
4213ad59 KS	291	self.output_to = None;
	292	self.last_ts = None;
	293	}
	294	}