moq-mux 0.7.2

//! AAC.
//!
//! ISO 14496-3 AudioSpecificConfig parse and encode lives in [`Config`].
//! [`Import`] publishes raw AAC frames (not ADTS) to a moq broadcast.

mod import;

pub use import::*;

use bytes::{Buf, Bytes};

/// AAC parsing errors.
#[derive(Debug, Clone, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
	#[error("AudioSpecificConfig must be at least 2 bytes")]
	ConfigTooShort,

	#[error("extended audioObjectType requires 2 additional bytes")]
	ExtendedConfigTooShort,

	#[error("AudioSpecificConfig incomplete")]
	IncompleteConfig,

	#[error("explicit sample rate requires 3 additional bytes")]
	ExplicitSampleRateTooShort,

	#[error("unsupported sample rate index: {0}")]
	UnsupportedSampleRateIndex(u8),
}

pub type Result<T> = std::result::Result<T, Error>;

/// Typed AAC configuration mirroring the relevant fields of an
/// AudioSpecificConfig.
pub struct Config {
	pub profile: u8,
	pub sample_rate: u32,
	pub channel_count: u32,
}

impl Config {
	/// Parse an AudioSpecificConfig buffer.
	///
	/// Handles basic formats (object_type < 31), extended formats
	/// (object_type == 31), and explicit sample rates (freq_index == 15). The
	/// fields are bit-packed and not byte-aligned, so a bit reader is required:
	/// with an explicit 24-bit rate the channelConfiguration lands mid-byte after
	/// it. Any SBR/PS extension bits after the core fields are consumed.
	pub fn parse<T: Buf>(buf: &mut T) -> Result<Self> {
		if buf.remaining() < 2 {
			return Err(Error::ConfigTooShort);
		}

		let mut reader = BitReader::new(buf);

		// audioObjectType: 5 bits, escaped to 6 more when it reads 31.
		let mut object_type = reader.read(5, Error::ConfigTooShort)? as u8;
		if object_type == 31 {
			object_type = 32 + reader.read(6, Error::ExtendedConfigTooShort)? as u8;
		}

		// samplingFrequencyIndex: 4 bits; index 15 means an explicit 24-bit rate follows.
		let freq_index = reader.read(4, Error::IncompleteConfig)? as u8;
		let sample_rate = if freq_index == 15 {
			reader.read(24, Error::ExplicitSampleRateTooShort)?
		} else {
			*SAMPLE_RATES
				.get(freq_index as usize)
				.ok_or(Error::UnsupportedSampleRateIndex(freq_index))?
		};

		// channelConfiguration: 4 bits, immediately after the (possibly explicit) rate.
		let channel_config = reader.read(4, Error::IncompleteConfig)? as u8;
		let channel_count = channel_count_from_config(channel_config);

		// AudioSpecificConfig can carry variable-length extensions (SBR, PS, etc.).
		// We've extracted the essential fields; drain the rest so the buffer is advanced.
		if buf.remaining() > 0 {
			buf.advance(buf.remaining());
		}

		Ok(Self {
			profile: object_type,
			sample_rate,
			channel_count,
		})
	}

	/// Encode this configuration as an AudioSpecificConfig (ISO 14496-3 §1.6.2.1).
	///
	/// Standard sample rates produce 2 bytes; non-standard rates fall back to
	/// the 5-byte form with an explicit 24-bit frequency.
	pub fn encode(&self) -> Bytes {
		// audioObjectType is a 5-bit field; mask to prevent shift overflow.
		let profile = self.profile & 0x1F;

		let freq_index: u8 = match self.sample_rate {
			96000 => 0,
			88200 => 1,
			64000 => 2,
			48000 => 3,
			44100 => 4,
			32000 => 5,
			24000 => 6,
			22050 => 7,
			16000 => 8,
			12000 => 9,
			11025 => 10,
			8000 => 11,
			7350 => 12,
			_ => 0xF, // explicit 24-bit frequency follows
		};

		let channel_config = channel_config_from_count(self.channel_count) as u64;

		if freq_index != 0xF {
			// 5 + 4 + 4 = 13 bits → 2 bytes (3 bits padding)
			let b0 = (profile << 3) | (freq_index >> 1);
			let b1 = ((freq_index & 1) << 7) | ((channel_config as u8 & 0x0F) << 3);
			Bytes::from(vec![b0, b1])
		} else {
			// 5 + 4 + 24 + 4 = 37 bits → 5 bytes (3 bits padding)
			let mut bits: u64 = 0;
			bits |= (profile as u64) << 35;
			bits |= 0xF_u64 << 31;
			bits |= (self.sample_rate as u64) << 7;
			bits |= (channel_config & 0xF) << 3;
			let all = bits.to_be_bytes();
			Bytes::copy_from_slice(&all[3..8])
		}
	}
}

/// The 13 standard AAC sampling frequencies, indexed by samplingFrequencyIndex
/// (ISO 14496-3 Table 1.18). Index 15 is the escape for an explicit 24-bit rate.
const SAMPLE_RATES: [u32; 13] = [
	96000, 88200, 64000, 48000, 44100, 32000, 24000, 22050, 16000, 12000, 11025, 8000, 7350,
];

/// MSB-first bit reader that pulls bytes from a [`Buf`] on demand.
///
/// AudioSpecificConfig is bit-packed: an explicit 24-bit sample rate pushes the
/// following channelConfiguration off byte boundaries, so the fields can't be
/// read a whole byte at a time.
struct BitReader<'a, T: Buf> {
	buf: &'a mut T,
	current: u8,
	bits_left: u8,
}

impl<'a, T: Buf> BitReader<'a, T> {
	fn new(buf: &'a mut T) -> Self {
		Self {
			buf,
			current: 0,
			bits_left: 0,
		}
	}

	/// Read `n` bits (n <= 32) MSB-first, returning `short` if the buffer runs dry.
	fn read(&mut self, n: u8, short: Error) -> Result<u32> {
		let mut value = 0u32;
		for _ in 0..n {
			if self.bits_left == 0 {
				if !self.buf.has_remaining() {
					return Err(short);
				}
				self.current = self.buf.get_u8();
				self.bits_left = 8;
			}
			self.bits_left -= 1;
			value = (value << 1) | u32::from((self.current >> self.bits_left) & 1);
		}
		Ok(value)
	}
}

/// Map an AAC `channel_config` (ISO 14496-3 Table 1.19) to its real channel count.
/// Configs 1..=6 happen to be identity (5.1 has config=6 and 6 channels). Config
/// 7 is 7.1 = 8 channels. Config 0 means "described elsewhere" — we default to
/// stereo.
fn channel_count_from_config(channel_config: u8) -> u32 {
	match channel_config {
		1..=6 => channel_config as u32,
		7 => 8,
		0 => {
			tracing::warn!("channel_config=0 (program config element) unsupported, defaulting to stereo");
			2
		}
		_ => {
			tracing::warn!(channel_config, "unsupported channel config, defaulting to stereo");
			2
		}
	}
}

/// Inverse of [`channel_count_from_config`]. Defaults to stereo for unsupported
/// counts (channel configs > 7 are reserved).
fn channel_config_from_count(channel_count: u32) -> u8 {
	match channel_count {
		1..=6 => channel_count as u8,
		8 => 7,
		_ => {
			tracing::warn!(channel_count, "unsupported channel count, defaulting to stereo");
			2
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn parses_standard_2_byte_config() {
		// AAC-LC (profile=2), 44100 Hz (freq_index=4), stereo (channels=2).
		// b0 = 0x12 (00010 0100b → object_type=2, freq_index high 3 bits=010)
		// b1 = 0x10 (0001 0000b → freq_index low bit=0, channel_config=2, padding=000)
		let buf = vec![0x12, 0x10];
		let cfg = Config::parse(&mut buf.as_slice()).unwrap();
		assert_eq!(cfg.profile, 2);
		assert_eq!(cfg.sample_rate, 44100);
		assert_eq!(cfg.channel_count, 2);
	}

	#[test]
	fn round_trip_explicit_sample_rate() {
		// A non-standard rate (no freq_index) forces the explicit 24-bit form, where
		// channelConfiguration lands mid-byte after the rate. A byte-aligned parser
		// misreads both fields; the bit reader round-trips them.
		let cfg = Config {
			profile: 2,
			sample_rate: 44_056, // not in the standard table
			channel_count: 2,
		};
		let encoded = cfg.encode();
		assert_eq!(encoded.len(), 5, "explicit-rate config is 5 bytes");

		let parsed = Config::parse(&mut encoded.as_ref()).unwrap();
		assert_eq!(parsed.profile, 2);
		assert_eq!(parsed.sample_rate, 44_056);
		assert_eq!(parsed.channel_count, 2);
	}

	#[test]
	fn parses_extended_object_type() {
		// audioObjectType 31 escapes to a 6-bit extended type. Bytes encode
		// AOT=31, ext=4 (-> object_type 36), freq_index=3 (48000), channel_config=2,
		// which straddle byte boundaries: 11111 000100 0011 0010 + padding.
		let buf: [u8; 3] = [0xF8, 0x86, 0x40];
		let cfg = Config::parse(&mut buf.as_slice()).unwrap();
		assert_eq!(cfg.profile, 36);
		assert_eq!(cfg.sample_rate, 48_000);
		assert_eq!(cfg.channel_count, 2);
	}

	#[test]
	fn round_trip_5_1_channels() {
		// 5.1 surround: config=6, 6 channels.
		let cfg = Config {
			profile: 2,
			sample_rate: 48000,
			channel_count: 6,
		};
		let encoded = cfg.encode();
		let parsed = Config::parse(&mut encoded.as_ref()).unwrap();
		assert_eq!(parsed.channel_count, 6);
	}

	#[test]
	fn round_trip_7_1_channels() {
		// 7.1 surround: config=7, but 8 channels.
		let cfg = Config {
			profile: 2,
			sample_rate: 48000,
			channel_count: 8,
		};
		let encoded = cfg.encode();
		let parsed = Config::parse(&mut encoded.as_ref()).unwrap();
		assert_eq!(parsed.channel_count, 8, "7.1 surround should round-trip as 8 channels");
	}

	#[test]
	fn channel_config_zero_falls_back_to_stereo() {
		// Config 0 means "described in PCE" which we don't implement.
		assert_eq!(channel_count_from_config(0), 2);
	}

	#[test]
	fn unsupported_channel_count_falls_back_to_stereo_config() {
		assert_eq!(channel_config_from_count(9), 2);
	}
}