cros-codecs 0.0.6

// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::cell::RefCell;
use std::ops::Deref;
use std::rc::Rc;
use std::rc::Weak;

use log::debug;

use crate::codec::h264::parser::MaxLongTermFrameIdx;
use crate::codec::h264::parser::RefPicMarking;
use crate::codec::h264::parser::Slice;
use crate::codec::h264::parser::SliceType;
use crate::codec::h264::parser::Sps;
use crate::Resolution;

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum Field {
    #[default]
    Frame,
    Top,
    Bottom,
}

impl Field {
    /// Returns the field of opposite parity.
    pub fn opposite(&self) -> Self {
        match *self {
            Field::Frame => Field::Frame,
            Field::Top => Field::Bottom,
            Field::Bottom => Field::Top,
        }
    }
}

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum Reference {
    #[default]
    None,
    ShortTerm,
    LongTerm,
}

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum IsIdr {
    #[default]
    No,
    Yes {
        idr_pic_id: u16,
    },
}

/// The rank of a field, i.e. whether it is the first or second one to be parsed from the stream.
/// This is unrelated to the `Field` type, as the first field can be either `Top` or `Bottom`.
#[derive(Default, Debug)]
pub enum FieldRank {
    /// Frame has a single field.
    #[default]
    Single,
    /// Frame is interlaced, and this is the first field (with a reference to the second one).
    First(Weak<RefCell<PictureData>>),
    /// Frame is interlaced, and this is the second field (with a reference to the first one).
    Second(Rc<RefCell<PictureData>>),
}

#[derive(Default)]
pub struct PictureData {
    pub pic_order_cnt_type: u8,
    pub top_field_order_cnt: i32,
    pub bottom_field_order_cnt: i32,
    pub pic_order_cnt: i32,
    pub pic_order_cnt_msb: i32,
    pub pic_order_cnt_lsb: i32,
    pub delta_pic_order_cnt_bottom: i32,
    pub delta_pic_order_cnt0: i32,
    pub delta_pic_order_cnt1: i32,

    pub pic_num: i32,
    pub long_term_pic_num: u32,
    pub frame_num: u32,
    pub frame_num_offset: u32,
    pub frame_num_wrap: i32,
    pub long_term_frame_idx: u32,

    pub coded_resolution: Resolution,
    pub display_resolution: Resolution,

    pub type_: SliceType,
    pub nal_ref_idc: u8,
    pub is_idr: IsIdr,
    reference: Reference,
    pub ref_pic_list_modification_flag_l0: i32,
    pub abs_diff_pic_num_minus1: i32,

    // Does memory management op 5 needs to be executed after this
    // picture has finished decoding?
    pub has_mmco_5: bool,

    // Created by the decoding process for gaps in frame_num.
    // Not for decode or output.
    pub nonexisting: bool,

    pub field: Field,

    // Values from slice_hdr to be used during reference marking and
    // memory management after finishing this picture.
    pub ref_pic_marking: RefPicMarking,

    field_rank: FieldRank,

    pub timestamp: u64,
}

/// A `PictureData` within a `Rc<RefCell>` which field rank is guaranteed to be correct.
///
/// The field rank of `PictureData` is only final after both fields have been constructed - namely,
/// the first field can only point to the second one after the latter is available as a Rc. Methods
/// [`PictureData::into_rc`] and [`PictureData::split_frame`] take care of this, and is this only
/// producer of this type, ensuring all instances are correct.
#[derive(Default, Debug, Clone)]
pub struct RcPictureData {
    pic: Rc<RefCell<PictureData>>,
}

impl Deref for RcPictureData {
    type Target = Rc<RefCell<PictureData>>;

    fn deref(&self) -> &Self::Target {
        &self.pic
    }
}

impl PictureData {
    pub fn new_non_existing(frame_num: u32, timestamp: u64) -> Self {
        PictureData {
            frame_num,
            nonexisting: true,
            nal_ref_idc: 1,
            field: Field::Frame,
            pic_num: frame_num as i32,
            reference: Reference::ShortTerm,
            timestamp,
            ..Default::default()
        }
    }

    /// Create a new picture from a `slice`, `sps`, and `timestamp`.
    ///
    /// `first_field` is set if this picture is the second field of a frame.
    pub fn new_from_slice(
        slice: &Slice,
        sps: &Sps,
        timestamp: u64,
        first_field: Option<&RcPictureData>,
    ) -> Self {
        let hdr = &slice.header;
        let nalu_hdr = &slice.nalu.header;

        let is_idr = if nalu_hdr.idr_pic_flag {
            IsIdr::Yes { idr_pic_id: hdr.idr_pic_id }
        } else {
            IsIdr::No
        };

        let field = if hdr.field_pic_flag {
            if hdr.bottom_field_flag {
                Field::Bottom
            } else {
                Field::Top
            }
        } else {
            Field::Frame
        };

        let reference = if nalu_hdr.ref_idc != 0 { Reference::ShortTerm } else { Reference::None };

        let pic_num = if !hdr.field_pic_flag { hdr.frame_num } else { 2 * hdr.frame_num + 1 };

        let (
            pic_order_cnt_lsb,
            delta_pic_order_cnt_bottom,
            delta_pic_order_cnt0,
            delta_pic_order_cnt1,
        ) = match sps.pic_order_cnt_type {
            0 => (
                hdr.pic_order_cnt_lsb,
                hdr.delta_pic_order_cnt_bottom,
                Default::default(),
                Default::default(),
            ),
            1 => (
                Default::default(),
                Default::default(),
                hdr.delta_pic_order_cnt[0],
                hdr.delta_pic_order_cnt[1],
            ),
            _ => (Default::default(), Default::default(), Default::default(), Default::default()),
        };

        let coded_resolution = Resolution::from((sps.width(), sps.height()));

        let visible_rect = sps.visible_rectangle();

        let display_resolution = Resolution {
            width: visible_rect.max.x - visible_rect.min.x,
            height: visible_rect.max.y - visible_rect.min.y,
        };

        let mut pic = PictureData {
            pic_order_cnt_type: sps.pic_order_cnt_type,
            pic_order_cnt_lsb: i32::from(pic_order_cnt_lsb),
            delta_pic_order_cnt_bottom,
            delta_pic_order_cnt0,
            delta_pic_order_cnt1,
            pic_num: i32::from(pic_num),
            frame_num: u32::from(hdr.frame_num),
            nal_ref_idc: nalu_hdr.ref_idc,
            is_idr,
            reference,
            field,
            ref_pic_marking: hdr.dec_ref_pic_marking.clone(),
            coded_resolution,
            display_resolution,
            timestamp,
            ..Default::default()
        };

        if let Some(first_field) = first_field {
            pic.set_first_field_to(first_field);
        }

        pic
    }

    /// Whether the current picture is a reference, either ShortTerm or LongTerm.
    pub fn is_ref(&self) -> bool {
        !matches!(self.reference, Reference::None)
    }

    /// Whether this picture is a second field.
    pub fn is_second_field(&self) -> bool {
        matches!(self.field_rank, FieldRank::Second(..))
    }

    /// Returns the field rank of this picture, including a reference to its other field.
    pub fn field_rank(&self) -> &FieldRank {
        &self.field_rank
    }

    /// Returns a reference to the picture's Reference
    pub fn reference(&self) -> &Reference {
        &self.reference
    }

    /// Mark the picture as a reference picture.
    pub fn set_reference(&mut self, reference: Reference, apply_to_other_field: bool) {
        log::debug!("Set reference of {:#?} to {:?}", self, reference);
        self.reference = reference;

        if apply_to_other_field {
            if let Some(other_field) = self.other_field() {
                log::debug!(
                    "other_field: Set reference of {:#?} to {:?}",
                    &other_field.borrow(),
                    reference
                );
                other_field.borrow_mut().reference = reference;
            }
        }
    }

    /// Get a reference to the picture's other field, if there is any
    /// and its reference is still valid.
    pub fn other_field(&self) -> Option<Rc<RefCell<PictureData>>> {
        match &self.field_rank {
            FieldRank::Single => None,
            FieldRank::First(other_field) => other_field.upgrade(),
            FieldRank::Second(other_field) => Some(other_field.clone()),
        }
    }

    /// Set this picture's second field.
    fn set_second_field_to(&mut self, other_field: &Rc<RefCell<Self>>) {
        self.field_rank = FieldRank::First(Rc::downgrade(other_field));
    }

    /// Whether the current picture is the second field of a complementary ref pair.
    pub fn is_second_field_of_complementary_ref_pair(&self) -> bool {
        self.is_ref()
            && matches!(self.field_rank(), FieldRank::Second(first_field) if first_field.borrow().is_ref())
    }

    /// Set this picture's first field.
    fn set_first_field_to(&mut self, other_field: &Rc<RefCell<Self>>) {
        self.field_rank = FieldRank::Second(other_field.clone());
    }

    pub fn pic_num_f(&self, max_pic_num: i32) -> i32 {
        if !matches!(self.reference(), Reference::LongTerm) {
            self.pic_num
        } else {
            max_pic_num
        }
    }

    pub fn long_term_pic_num_f(&self, max_long_term_frame_idx: MaxLongTermFrameIdx) -> u32 {
        if matches!(self.reference(), Reference::LongTerm) {
            self.long_term_pic_num
        } else {
            2 * max_long_term_frame_idx.to_value_plus1()
        }
    }

    /// Consume this picture and return a Rc'd version.
    ///
    /// If the picture was a second field, adjust the field of the first field to point to this
    /// one.
    pub fn into_rc(self) -> RcPictureData {
        let self_rc = Rc::new(RefCell::new(self));

        if let FieldRank::Second(first_field) = self_rc.borrow().field_rank() {
            first_field.borrow_mut().set_second_field_to(&self_rc);
        }

        RcPictureData { pic: self_rc }
    }

    /// Split a frame into two complementary fields that reference one another.
    pub fn split_frame(mut self) -> (RcPictureData, RcPictureData) {
        assert!(matches!(self.field, Field::Frame));
        assert!(matches!(self.field_rank, FieldRank::Single));

        debug!(
            "Splitting picture (frame_num, POC) ({:?}, {:?})",
            self.frame_num, self.pic_order_cnt
        );

        let second_pic_order_cnt = if self.top_field_order_cnt < self.bottom_field_order_cnt {
            self.field = Field::Top;
            self.pic_order_cnt = self.top_field_order_cnt;

            self.bottom_field_order_cnt
        } else {
            self.field = Field::Bottom;
            self.pic_order_cnt = self.bottom_field_order_cnt;

            self.top_field_order_cnt
        };

        let second_field = PictureData {
            top_field_order_cnt: self.top_field_order_cnt,
            bottom_field_order_cnt: self.bottom_field_order_cnt,
            frame_num: self.frame_num,
            reference: self.reference,
            nonexisting: self.nonexisting,
            pic_order_cnt: second_pic_order_cnt,
            field: self.field.opposite(),
            ..Default::default()
        };

        debug!(
            "Split into picture (frame_num, POC) ({:?}, {:?}), field: {:?}",
            self.frame_num, self.pic_order_cnt, self.field
        );
        debug!(
            "Split into picture (frame_num, POC) ({:?}, {:?}), field {:?}",
            second_field.frame_num, second_field.pic_order_cnt, second_field.field
        );

        let first_field = Rc::new(RefCell::new(self));
        let second_field = Rc::new(RefCell::new(second_field));

        first_field.borrow_mut().set_second_field_to(&second_field);
        second_field.borrow_mut().set_first_field_to(&first_field);

        (RcPictureData { pic: first_field }, RcPictureData { pic: second_field })
    }
}

impl std::fmt::Debug for PictureData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PictureData")
            .field("pic_order_cnt_type", &self.pic_order_cnt_type)
            .field("top_field_order_cnt", &self.top_field_order_cnt)
            .field("bottom_field_order_cnt", &self.bottom_field_order_cnt)
            .field("pic_order_cnt", &self.pic_order_cnt)
            .field("pic_order_cnt_msb", &self.pic_order_cnt_msb)
            .field("pic_order_cnt_lsb", &self.pic_order_cnt_lsb)
            .field("delta_pic_order_cnt_bottom", &self.delta_pic_order_cnt_bottom)
            .field("delta_pic_order_cnt0", &self.delta_pic_order_cnt0)
            .field("delta_pic_order_cnt1", &self.delta_pic_order_cnt1)
            .field("pic_num", &self.pic_num)
            .field("long_term_pic_num", &self.long_term_pic_num)
            .field("frame_num", &self.frame_num)
            .field("frame_num_offset", &self.frame_num_offset)
            .field("frame_num_wrap", &self.frame_num_wrap)
            .field("long_term_frame_idx", &self.long_term_frame_idx)
            .field("coded_resolution", &self.coded_resolution)
            .field("display_resolution", &self.display_resolution)
            .field("type_", &self.type_)
            .field("nal_ref_idc", &self.nal_ref_idc)
            .field("is_idr", &self.is_idr)
            .field("reference", &self.reference)
            .field("ref_pic_list_modification_flag_l0", &self.ref_pic_list_modification_flag_l0)
            .field("abs_diff_pic_num_minus1", &self.abs_diff_pic_num_minus1)
            .field("has_mmco_5", &self.has_mmco_5)
            .field("nonexisting", &self.nonexisting)
            .field("field", &self.field)
            .field("ref_pic_marking", &self.ref_pic_marking)
            .field("field_rank", &self.field_rank)
            .finish()
    }
}