transotf/src/av/mov_avc.rs

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::fmt::{Debug, Error, Formatter};
use std::io::{Cursor, Seek, SeekFrom};

#[derive(Debug)]
pub struct AvcDecoderConfigurationRecord {
    pub profile_idc: u8,
    pub constraint_set_flag: u8,
    pub level_idc: u8,
    pub sequence_parameter_set: Vec<u8>,
    pub picture_parameter_set: Vec<u8>,
}

impl AvcDecoderConfigurationRecord {
    pub fn from_parameter_sets(
        sequence_parameter_set: Vec<u8>,
        picture_parameter_set: Vec<u8>,
    ) -> AvcDecoderConfigurationRecord {
        AvcDecoderConfigurationRecord {
            profile_idc: sequence_parameter_set[1],
            constraint_set_flag: sequence_parameter_set[2],
            level_idc: sequence_parameter_set[3],
            sequence_parameter_set,
            picture_parameter_set,
        }
    }

    pub fn to_bytes(&self) -> Result<Vec<u8>, String> {
        let mut buffer = vec![];
        // configuration_version
        buffer.push(1);
        buffer.push(self.profile_idc);
        buffer.push(self.constraint_set_flag);
        buffer.push(self.level_idc);
        // reserved and length_size_minus_one
        buffer.push(0b1111_1100 | 0b0000_0011);
        // reserved and num_of_sequence_parameter_set_ext
        buffer.push(0b1110_0000 | 0b0000_0001);

        buffer
            .write_u16::<BigEndian>(self.sequence_parameter_set.len() as u16)
            .map_err(|x| format!("Failed to write u16 to vec ({})", x))?;

        buffer.append(&mut self.sequence_parameter_set.clone());
        // num_of_picture_parameter_set_ext
        buffer.push(0b0000_0001);
        buffer
            .write_u16::<BigEndian>(self.picture_parameter_set.len() as u16)
            .map_err(|x| format!("Failed to write u16 to vec ({})", x))?;
        buffer.append(&mut self.picture_parameter_set.clone());

        Ok(buffer)
    }
}

#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
#[repr(u8)]
pub enum NalUnitType {
    CodedSliceOfANonIdrPicture = 1,
    CodedSliceDataPartitionA = 2,
    CodedSliceDataPartitionB = 3,
    CodedSliceDataPartitionC = 4,
    CodedSliceOfAnIdrPicture = 5,
    SupplementalEnhancementInformation = 6,
    SequenceParameterSet = 7,
    PictureParameterSet = 8,
    AccessUnitDelimiter = 9,
    EndOfSequence = 10,
    EndOfStream = 11,
    FilterData = 12,
    SequenceParameterSetExtension = 13,
    PrefixNalUnit = 14,
    SubsetSequenceParameterSet = 15,
    CodedSliceOfAnAuxiliaryCodedPictureWithoutPartitioning = 19,
    CodedSliceExtension = 20,
}

impl NalUnitType {
    fn from_u8(n: u8) -> Result<Self, String> {
        Ok(match n {
            1 => NalUnitType::CodedSliceOfANonIdrPicture,
            2 => NalUnitType::CodedSliceDataPartitionA,
            3 => NalUnitType::CodedSliceDataPartitionB,
            4 => NalUnitType::CodedSliceDataPartitionC,
            5 => NalUnitType::CodedSliceOfAnIdrPicture,
            6 => NalUnitType::SupplementalEnhancementInformation,
            7 => NalUnitType::SequenceParameterSet,
            8 => NalUnitType::PictureParameterSet,
            9 => NalUnitType::AccessUnitDelimiter,
            10 => NalUnitType::EndOfSequence,
            11 => NalUnitType::EndOfStream,
            12 => NalUnitType::FilterData,
            13 => NalUnitType::SequenceParameterSetExtension,
            14 => NalUnitType::PrefixNalUnit,
            15 => NalUnitType::SubsetSequenceParameterSet,
            19 => NalUnitType::CodedSliceOfAnAuxiliaryCodedPictureWithoutPartitioning,
            20 => NalUnitType::CodedSliceExtension,
            _ => return Err(format!("Couldn't find NalUnitType for {:?}", n)),
        })
    }
}

#[derive(Debug, Copy, Clone)]
pub struct NalUnit {
    pub nal_ref_idc: u8,
    pub nal_unit_type: NalUnitType,
}

impl NalUnit {
    fn from_byte(byte: u8) -> Result<Self, String> {
        let nal_ref_idc = (byte >> 5) & 0b11;
        let nal_unit_type = NalUnitType::from_u8(byte & 0b1_1111)?;

        Ok(NalUnit {
            nal_ref_idc,
            nal_unit_type,
        })
    }

    fn to_byte(&self) -> u8 {
        let nal_ref_idc = self.nal_ref_idc << 5;
        let nal_unit_type = self.nal_unit_type as u8;
        return nal_unit_type | nal_ref_idc;
    }
}

#[derive(Clone)]
pub struct NalBuffer {
    pub nal_unit: NalUnit,
    pub buffer: Vec<u8>,
}

struct StringyBoi(String);

impl Debug for StringyBoi {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        f.write_str(&self.0)
    }
}

impl Debug for NalBuffer {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        if f.alternate() {
            f.debug_tuple("NalBuffer")
                .field(&self.nal_unit)
                .field(&self.buffer)
                .finish()
        } else {
            f.debug_tuple("NalBuffer")
                .field(&self.nal_unit)
                .field(&StringyBoi(format!("len={}", self.buffer.len())))
                .finish()
        }
    }
}

fn find_mdat(data: &Vec<u8>) -> Option<usize> {
    for i in 0..=data.len() - 3 {
        if data[i] == 'm' as u8
            && data.get(i + 3).copied() == Some('t' as u8)
            && &data[i..i + 4] == b"mdat"
        {
            return Some(i + 4);
        }
    }

    None
}

#[cfg(debug_assertions)]
pub fn verify_data(data: Vec<u8>) {
    let index = find_mdat(&data).expect("no mdat found") as u64;
    let mut cursor = Cursor::new(&data);
    cursor.set_position(index - 8);
    let size = cursor.read_u32::<BigEndian>().expect(":O") as u64;
    cursor.set_position(index);
    while cursor.position() < (size + index - 8) {
        let skip = cursor.read_u32::<BigEndian>().expect(":C") as i64;
        cursor
            .seek(SeekFrom::Current(skip))
            .expect("Failed to seek");
    }

    assert_eq!(cursor.position() - index, size - 8);
}

pub fn annex_b_to_avc(mut data: Vec<u8>, nal_length: u8) -> Result<Vec<u8>, String> {
    let nal_length = nal_length as usize;
    let mut result: Vec<u8> = vec![];
    let index = find_mdat(&data).ok_or("No mdat found".to_string())?;
    let end = {
        let mut cursor = Cursor::new(&data);
        cursor.set_position((index - 8) as u64);
        index + (cursor.read_u32::<BigEndian>().expect(":O") as usize) - 8
    };

    result.extend_from_slice(&data[0..index]);
    let matcher = vec![0u8; nal_length - 1 as usize];
    let mut last_position = index + nal_length - 1;
    for i in (index + matcher.len() + nal_length)..end {
        if data[i] == 1 && &data[1 + i - nal_length..i] == &matcher[..] {
            let block_length = (i - last_position - 4) as u32;
            let mut header = &mut data[last_position - 3..last_position + 1];
            header
                .write_u32::<BigEndian>(block_length)
                .map_err(|x| format!("Failed writing to buffer: {}", x))?;

            last_position = i;
        }
    }

    let data_len = data.len();
    let mut header = &mut data[last_position - 3..=last_position];
    header
        .write_u32::<BigEndian>(((data_len - 1) - last_position) as u32)
        .map_err(|x| format!("Failed writing to buffer: {}", x))?;

    #[cfg(debug_assertions)]
    verify_data(data.clone());
    Ok(data)
}

impl NalBuffer {
    pub fn to_vec(&self) -> Vec<u8> {
        let mut bytes = vec![0; 1 + self.buffer.len()];
        {
            let (left, right) = bytes.split_at_mut(1);
            left[0] = self.nal_unit.to_byte();
            right.copy_from_slice(&self.buffer);
        }

        bytes
    }

    pub fn from_stream(stream: Vec<u8>) -> Result<Vec<NalBuffer>, String> {
        let mut i = 0;
        let mut buffers = vec![];
        let mut byte_buffer = vec![];
        let mut current: Option<NalBuffer> = None;
        while i < stream.len() {
            if stream.len() > i + 3
                && stream[i] == 0
                && stream[i + 1] == 0
                && (stream[i + 2] == 1
                    || (stream.len() > i + 4 && stream[i + 2] == 0 && stream[i + 3] == 1))
            {
                if let Some(mut buffer) = current.take() {
                    buffer.buffer = byte_buffer;
                    byte_buffer = vec![];
                    buffers.push(buffer);
                }

                i += if stream[i + 2] == 1 { 3 } else { 4 };
                if i >= stream.len() {
                    break;
                }

                current = Some(NalBuffer {
                    nal_unit: NalUnit::from_byte(stream[i])?,
                    buffer: vec![],
                });
                i += 1;
                if i >= stream.len() {
                    break;
                }
            }

            if stream.len() > i + 4
                && stream[i] == 0
                && stream[i + 1] == 0
                && stream[i + 2] == 0
                && stream[i + 3] == 3
            {
                byte_buffer.append(&mut (&stream[i..i + 2]).to_vec());
                i += 4;
                if i >= stream.len() {
                    break;
                }
            }

            byte_buffer.push(stream[i]);
            i += 1;
        }

        if let Some(mut buffer) = current.take() {
            buffer.buffer = byte_buffer;
            buffers.push(buffer);
        }

        Ok(buffers)
    }
}

#[cfg(test)]
mod test {
    use crate::av::mov_avc::{annex_b_to_avc, verify_data};
    use std::fs::File;
    use std::io::Read;

    #[test]
    fn test_annex_b_to_avc() {
        let mdat = b"mdat".to_vec();
        let mut insert: Vec<u8> = vec![
            0, 0, 0, 1, 1, 2, 3, 4, 0, 0, 0, 1, 1, 2, 3, 4, 5, 0, 0, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8,
            9,
        ];
        let mut buffer: Vec<u8> = vec![0, 0, 0, insert.len() as u8 + 8];
        let mut result: Vec<u8> = vec![0, 0, 0, insert.len() as u8 + 8];
        buffer.extend_from_slice(&mdat);
        buffer.append(&mut insert);
        result.extend(mdat);
        result.append(&mut vec![
            0, 0, 0, 4, 1, 2, 3, 4, 0, 0, 0, 5, 1, 2, 3, 4, 5, 0, 0, 0, 9, 1, 2, 3, 4, 5, 6, 7, 8,
            9,
        ]);
        assert_eq!(annex_b_to_avc(buffer, 4), Ok(result.clone()));

        verify_data(result);
    }

    #[test]
    fn bs() {
        let mut x = File::open("/home/eater/shit/dash-exmp/chunk-stream0-00001.m4s").expect("ok");
        let mut xxx = vec![];
        x.read_to_end(&mut xxx).unwrap();
        verify_data(xxx);
        let mut x = File::open("/tmp/transotf/segment-00000.m4s").expect("ok");
        let mut xxx = vec![];
        x.read_to_end(&mut xxx).unwrap();
        verify_data(xxx);
    }
}