1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use anyhow::{Result, bail};

use self::slice::SliceNAL;

use super::BitVecReader;

pub(crate) mod hrd_parameters;
pub(crate) mod pps;
pub(crate) mod profile_tier_level;
pub(crate) mod scaling_list_data;
pub(crate) mod short_term_rps;
pub(crate) mod slice;
pub(crate) mod sps;
pub(crate) mod vps;
pub(crate) mod vui_parameters;

// https://github.com/virinext/hevcesbrowser/blob/master/hevcparser/include/Hevc.h
pub const NAL_TRAIL_N: u8 = 0;
pub const NAL_TRAIL_R: u8 = 1;
pub const NAL_TSA_N: u8 = 2;
pub const NAL_TSA_R: u8 = 3;
pub const NAL_STSA_N: u8 = 4;
pub const NAL_STSA_R: u8 = 5;
pub const NAL_RADL_N: u8 = 6;
pub const NAL_RADL_R: u8 = 7;
pub const NAL_RASL_N: u8 = 8;
pub const NAL_RASL_R: u8 = 9;
pub const NAL_BLA_W_LP: u8 = 16;
pub const NAL_BLA_W_RADL: u8 = 17;
pub const NAL_BLA_N_LP: u8 = 18;
pub const NAL_IDR_W_RADL: u8 = 19;
pub const NAL_IDR_N_LP: u8 = 20;
pub const NAL_CRA_NUT: u8 = 21;
pub const NAL_IRAP_VCL23: u8 = 23;
pub const NAL_VPS: u8 = 32;
pub const NAL_SPS: u8 = 33;
pub const NAL_PPS: u8 = 34;
pub const NAL_AUD: u8 = 35;
pub const NAL_EOS_NUT: u8 = 36;
pub const NAL_EOB_NUT: u8 = 37;
pub const NAL_FD_NUT: u8 = 38;
pub const NAL_SEI_PREFIX: u8 = 39;
pub const NAL_SEI_SUFFIX: u8 = 40;
pub const NAL_UNSPEC62: u8 = 62;
pub const NAL_UNSPEC63: u8 = 63;

pub const USER_DATA_REGISTERED_ITU_T_35: u8 = 4;

#[derive(Default, Debug, Clone)]
pub struct NALUnit {
    pub start: usize,
    pub end: usize,

    pub nal_type: u8,
    pub nuh_layer_id: u8,
    pub temporal_id: u8,
    pub start_code_len: u8,

    pub decoded_frame_index: u64,
}

#[derive(Default, Debug, Clone)]
pub struct Frame {
    pub decoded_number: u64,
    pub presentation_number: u64,
    pub frame_type: u64,

    pub nals: Vec<NALUnit>,
    pub first_slice: SliceNAL,
}

#[derive(Default, Debug, Clone)]
pub struct SeiMessage {
    num_payload_type_ff_bytes: usize,
    last_payload_type_byte: u8,

    num_payload_size_ff_bytes: usize,
    last_payload_size_byte: u8,

    payload_type: u8,
    payload_size: usize,
}

impl SeiMessage {
    pub fn from_bytes(data: &[u8]) -> Result<SeiMessage> {
        let mut reader = BitVecReader::new(data.to_vec());

        SeiMessage::parse(&mut reader)
    }

    pub fn parse(reader: &mut BitVecReader) -> Result<SeiMessage> {
        // forbidden_zero_bit
        reader.skip_n(1);

        let nal_type = reader.get_n::<u8>(6);

        if nal_type != NAL_SEI_PREFIX {
            bail!("NAL type {} is not SEI_PREFIX", nal_type);
        }

        if reader.available() < 9 && matches!(nal_type, NAL_EOS_NUT | NAL_EOB_NUT) {
        } else {
            reader.skip_n(6); // nuh_layer_id
            reader.skip_n(3); // temporal_id
        }

        let mut msg;

        loop {
            msg = SeiMessage::default();

            msg.last_payload_type_byte = reader.get_n(8);
            while msg.last_payload_type_byte == 0xFF {
                msg.num_payload_type_ff_bytes += 1;
                msg.last_payload_type_byte = reader.get_n(8);

                msg.payload_type += 255;
            }

            msg.payload_type += msg.last_payload_type_byte;

            msg.last_payload_size_byte = reader.get_n(8);
            while msg.last_payload_size_byte == 0xFF {
                msg.num_payload_size_ff_bytes += 1;
                msg.last_payload_size_byte = reader.get_n(8);
                
                msg.payload_size += 255;
            }
            
            msg.payload_size += msg.last_payload_size_byte as usize;

            reader.skip_n(msg.payload_size * 8);

            if reader.available() <= 8 || reader.get_n::<u8>(8) == 0x80 {
                break;
            }
        }

        Ok(msg)
    }
}