use std::iter::FusedIterator;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameType {
Idr,
P,
B,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GopPattern {
Ipppp { gop: usize },
Ibpbp { gop: usize, b_count: usize },
}
impl GopPattern {
pub const fn iphone_default() -> Self {
Self::Ibpbp { gop: 30, b_count: 1 }
}
pub const fn legacy_ipppp(gop: usize) -> Self {
Self::Ipppp { gop }
}
pub fn has_b_frames(&self) -> bool {
matches!(self, Self::Ibpbp { .. })
}
pub fn gop_size(&self) -> usize {
match self {
Self::Ipppp { gop } | Self::Ibpbp { gop, .. } => *gop,
}
}
pub fn legacy_b_count(&self) -> usize {
match self {
Self::Ipppp { .. } => 0,
Self::Ibpbp { b_count, .. } => *b_count,
}
}
pub fn frame_type_at(&self, display_idx: usize) -> FrameType {
match self {
Self::Ipppp { gop } => {
if display_idx.is_multiple_of(*gop) {
FrameType::Idr
} else {
FrameType::P
}
}
Self::Ibpbp { gop, b_count } => {
let pos_in_gop = display_idx % gop;
if pos_in_gop == 0 {
return FrameType::Idr;
}
if pos_in_gop == gop - 1 {
return FrameType::P;
}
let m = b_count + 1;
let sub_gop_pos = (pos_in_gop - 1) % m;
if sub_gop_pos < *b_count {
FrameType::B
} else {
FrameType::P
}
}
}
}
}
impl Default for GopPattern {
fn default() -> Self {
Self::iphone_default()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EncodeOrderFrame {
pub encode_idx: u32,
pub display_idx: u32,
pub gop_idx: u32,
pub frame_type: FrameType,
}
pub fn iter_encode_order(
n_frames: usize,
pattern: GopPattern,
) -> EncodeOrderIter {
EncodeOrderIter::new(n_frames, pattern)
}
#[derive(Debug, Clone)]
pub struct EncodeOrderIter {
frames: std::vec::IntoIter<EncodeOrderFrame>,
}
impl EncodeOrderIter {
fn new(n_frames: usize, pattern: GopPattern) -> Self {
let mut frames = Vec::with_capacity(n_frames);
let mut encode_idx: u32 = 0;
let mut display_pos = 0usize;
let mut gop_idx: u32 = 0;
let gop_size = pattern.gop_size();
debug_assert!(gop_size > 0, "gop_size must be > 0");
while display_pos < n_frames {
let gop_start = display_pos;
let gop_end = (gop_start + gop_size).min(n_frames);
match pattern {
GopPattern::Ipppp { .. } => {
for d in gop_start..gop_end {
let ft = pattern.frame_type_at(d);
frames.push(EncodeOrderFrame {
encode_idx,
display_idx: d as u32,
gop_idx,
frame_type: ft,
});
encode_idx += 1;
}
}
GopPattern::Ibpbp { b_count, .. } => {
let m = b_count + 1;
frames.push(EncodeOrderFrame {
encode_idx,
display_idx: gop_start as u32,
gop_idx,
frame_type: FrameType::Idr,
});
encode_idx += 1;
let mut sub_start = gop_start + 1;
while sub_start < gop_end {
let sub_end = (sub_start + m).min(gop_end);
let anchor = sub_end - 1;
let anchor_ft = pattern.frame_type_at(anchor);
frames.push(EncodeOrderFrame {
encode_idx,
display_idx: anchor as u32,
gop_idx,
frame_type: anchor_ft,
});
encode_idx += 1;
for d in sub_start..anchor {
let ft = pattern.frame_type_at(d);
debug_assert_eq!(
ft,
FrameType::B,
"sub-GOP intermediate must be B (display_idx={d})",
);
frames.push(EncodeOrderFrame {
encode_idx,
display_idx: d as u32,
gop_idx,
frame_type: ft,
});
encode_idx += 1;
}
sub_start = sub_end;
}
}
}
display_pos = gop_end;
gop_idx += 1;
}
debug_assert_eq!(frames.len(), n_frames);
Self {
frames: frames.into_iter(),
}
}
}
impl Iterator for EncodeOrderIter {
type Item = EncodeOrderFrame;
fn next(&mut self) -> Option<Self::Item> {
self.frames.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.frames.size_hint()
}
}
impl ExactSizeIterator for EncodeOrderIter {}
impl FusedIterator for EncodeOrderIter {}
#[cfg(test)]
mod tests {
use super::*;
fn collect_types(n_frames: usize, pattern: GopPattern) -> Vec<FrameType> {
iter_encode_order(n_frames, pattern)
.map(|f| f.frame_type)
.collect()
}
fn collect_display(n_frames: usize, pattern: GopPattern) -> Vec<u32> {
iter_encode_order(n_frames, pattern)
.map(|f| f.display_idx)
.collect()
}
#[test]
fn ipppp_encode_equals_display_order() {
let pat = GopPattern::Ipppp { gop: 5 };
let display: Vec<u32> = collect_display(10, pat);
assert_eq!(display, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
}
#[test]
fn ipppp_idr_at_gop_boundary_only() {
let pat = GopPattern::Ipppp { gop: 5 };
let types = collect_types(10, pat);
assert_eq!(
types,
vec![
FrameType::Idr, FrameType::P, FrameType::P, FrameType::P, FrameType::P,
FrameType::Idr, FrameType::P, FrameType::P, FrameType::P, FrameType::P,
],
);
}
#[test]
fn ibpbp_m2_encode_order_is_anchor_first() {
let pat = GopPattern::Ibpbp { gop: 5, b_count: 1 };
let display = collect_display(5, pat);
assert_eq!(display, vec![0, 2, 1, 4, 3]);
}
#[test]
fn ibpbp_m2_frame_types_match_encode_order() {
let pat = GopPattern::Ibpbp { gop: 5, b_count: 1 };
let types = collect_types(5, pat);
assert_eq!(
types,
vec![FrameType::Idr, FrameType::P, FrameType::B, FrameType::P, FrameType::B],
);
}
#[test]
fn ibpbp_m2_closed_gop_last_is_p() {
let pat = GopPattern::Ibpbp { gop: 10, b_count: 1 };
let frames: Vec<_> = iter_encode_order(10, pat).collect();
let last = frames.iter().find(|f| f.display_idx == 9).unwrap();
assert_eq!(last.frame_type, FrameType::P);
}
#[test]
fn ibpbp_m3_encode_order_two_bs_per_subgop() {
let pat = GopPattern::Ibpbp { gop: 7, b_count: 2 };
let display = collect_display(7, pat);
assert_eq!(display, vec![0, 3, 1, 2, 6, 4, 5]);
}
#[test]
fn multi_gop_idr_periodicity() {
let pat = GopPattern::Ibpbp { gop: 5, b_count: 1 };
let frames: Vec<_> = iter_encode_order(15, pat).collect();
let gop_indices: Vec<u32> = frames.iter().map(|f| f.gop_idx).collect();
assert_eq!(gop_indices.iter().max(), Some(&2));
let idr_displays: Vec<u32> = frames
.iter()
.filter(|f| f.frame_type == FrameType::Idr)
.map(|f| f.display_idx)
.collect();
assert_eq!(idr_displays, vec![0, 5, 10]);
}
#[test]
fn encode_idx_is_monotonic_from_zero() {
let pat = GopPattern::Ibpbp { gop: 5, b_count: 1 };
let frames: Vec<_> = iter_encode_order(10, pat).collect();
for (i, f) in frames.iter().enumerate() {
assert_eq!(f.encode_idx as usize, i);
}
}
#[test]
fn display_indices_form_a_permutation() {
for &(gop, b_count) in &[(5usize, 1usize), (7, 2), (10, 1), (30, 1)] {
let pat = GopPattern::Ibpbp { gop, b_count };
for n in [gop, 2 * gop, 3 * gop, 5, 10, 15, 30] {
let frames: Vec<_> = iter_encode_order(n, pat).collect();
assert_eq!(frames.len(), n);
let mut display: Vec<u32> = frames.iter().map(|f| f.display_idx).collect();
display.sort();
let expected: Vec<u32> = (0..n as u32).collect();
assert_eq!(display, expected,
"display permutation broken for gop={gop} b_count={b_count} n={n}");
}
}
}
#[test]
fn iphone_default_is_m2() {
let pat = GopPattern::iphone_default();
assert!(pat.has_b_frames());
assert_eq!(pat.gop_size(), 30);
if let GopPattern::Ibpbp { gop, b_count } = pat {
assert_eq!(gop, 30);
assert_eq!(b_count, 1);
} else {
panic!("iphone_default should be Ibpbp");
}
}
#[test]
fn legacy_ipppp_has_no_b_frames() {
let pat = GopPattern::legacy_ipppp(10);
assert!(!pat.has_b_frames());
assert_eq!(pat.gop_size(), 10);
}
#[test]
fn b_frames_never_first_or_last_in_gop() {
for &(gop, b_count) in &[(5usize, 1usize), (7, 2), (10, 1)] {
let pat = GopPattern::Ibpbp { gop, b_count };
for d in [0, gop - 1, gop, 2 * gop - 1, 2 * gop] {
let ft = pat.frame_type_at(d);
assert_ne!(
ft,
FrameType::B,
"B at display_idx={d} for gop={gop}, b_count={b_count}",
);
}
}
}
}