moq-vaapi 0.0.1

// Copyright 2023 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::borrow::Cow;
use std::fmt::Debug;
use std::io::Cursor;
use std::io::Seek;
use std::io::SeekFrom;

#[allow(clippy::len_without_is_empty)]
pub trait Header: Sized {
	/// Parse the NALU header, returning it.
	fn parse<T: AsRef<[u8]>>(cursor: &mut Cursor<T>) -> Result<Self, String>;
	/// Whether this header type indicates EOS.
	fn is_end(&self) -> bool;
	/// The length of the header.
	fn len(&self) -> usize;
}

#[derive(Debug)]
pub struct Nalu<'a, U> {
	pub header: U,
	/// The mapping that backs this NALU. Possibly shared with the other NALUs
	/// in the Access Unit.
	pub data: Cow<'a, [u8]>,

	pub size: usize,
	pub offset: usize,
}

impl<'a, U> Nalu<'a, U>
where
	U: Debug + Header,
{
	/// Find the next Annex B encoded NAL unit.
	pub fn next(cursor: &mut Cursor<&'a [u8]>) -> Result<Nalu<'a, U>, String> {
		let bitstream = cursor.clone().into_inner();
		let pos = usize::try_from(cursor.position()).map_err(|err| err.to_string())?;

		// Find the start code for this NALU
		let current_nalu_offset = match Nalu::<'a, U>::find_start_code(cursor, pos) {
			Some(offset) => offset,
			None => return Err("No NAL found".into()),
		};

		let mut start_code_offset = pos + current_nalu_offset;

		// If the preceding byte is 00, then we actually have a four byte SC,
		// i.e. 00 00 00 01 Where the first 00 is the "zero_byte()"
		if start_code_offset > 0 && cursor.get_ref()[start_code_offset - 1] == 00 {
			start_code_offset -= 1;
		}

		// The NALU offset is its offset + 3 bytes to skip the start code.
		let nalu_offset = pos + current_nalu_offset + 3;

		// Set the bitstream position to the start of the current NALU
		cursor.set_position(u64::try_from(nalu_offset).map_err(|err| err.to_string())?);

		let hdr = U::parse(cursor)?;

		// Find the start of the subsequent NALU.
		let mut next_nalu_offset = match Nalu::<'a, U>::find_start_code(cursor, nalu_offset) {
			Some(offset) => offset,
			None => {
				let cur_pos = cursor.position();
				let end_pos = cursor.seek(SeekFrom::End(0)).map_err(|err| err.to_string())?;
				let _ = cursor.seek(SeekFrom::Start(cur_pos)).map_err(|err| err.to_string())?;
				(end_pos - cur_pos) as usize
			} // Whatever data is left must be part of the current NALU
		};

		while next_nalu_offset > 0 && cursor.get_ref()[nalu_offset + next_nalu_offset - 1] == 00 {
			// Discard trailing_zero_8bits
			next_nalu_offset -= 1;
		}

		let nal_size = if hdr.is_end() {
			// the NALU is comprised of only the header
			hdr.len()
		} else {
			next_nalu_offset
		};

		Ok(Nalu {
			header: hdr,
			data: Cow::from(&bitstream[start_code_offset..nalu_offset + nal_size]),
			size: nal_size,
			offset: nalu_offset - start_code_offset,
		})
	}
}

impl<'a, U> Nalu<'a, U>
where
	U: Debug,
{
	fn find_start_code(data: &mut Cursor<&'a [u8]>, offset: usize) -> Option<usize> {
		// discard all zeroes until the start code pattern is found
		data.get_ref()[offset..]
			.windows(3)
			.position(|window| window == [0x00, 0x00, 0x01])
	}

	pub fn into_owned(self) -> Nalu<'static, U> {
		Nalu {
			header: self.header,
			size: self.size,
			offset: self.offset,
			data: Cow::Owned(self.data.into_owned()),
		}
	}
}

impl<'a, U> AsRef<[u8]> for Nalu<'a, U> {
	fn as_ref(&self) -> &[u8] {
		&self.data[self.offset..self.offset + self.size]
	}
}