arcly-stream 0.1.6

//! AV1 OBU (Open Bitstream Unit) framing: LEB128 + OBU iteration.
//!
//! Supports the "low-overhead bitstream format" (size-delimited OBUs via
//! `obu_has_size_field`), which is what MP4/RTP/WebM carriage uses.

/// Decode an unsigned LEB128 integer at `*pos`, advancing it. Up to 8 bytes.
pub fn leb128(data: &[u8], pos: &mut usize) -> Option<u64> {
    let mut value: u64 = 0;
    for i in 0..8 {
        let byte = *data.get(*pos)?;
        *pos += 1;
        value |= ((byte & 0x7f) as u64) << (i * 7);
        if byte & 0x80 == 0 {
            return Some(value);
        }
    }
    None
}

/// One OBU: its type and payload (header bytes stripped).
pub struct Obu<'a> {
    /// `obu_type` (e.g. 1 = sequence header, 6 = frame).
    pub obu_type: u8,
    /// OBU payload, excluding the OBU header and size field.
    pub payload: &'a [u8],
    /// The complete OBU bytes (header + optional size field + payload), as they
    /// appeared in the input — used by RTP carriage to re-frame OBU elements.
    pub raw: &'a [u8],
    /// Whether the OBU header had its `obu_extension_flag` set (a 2-byte header).
    pub has_extension: bool,
}

/// Iterate the OBUs of an AV1 temporal unit in low-overhead format.
pub fn iter_obus(data: &[u8]) -> ObuIter<'_> {
    ObuIter { data, pos: 0 }
}

/// Iterator over [`Obu`]s. See [`iter_obus`].
pub struct ObuIter<'a> {
    data: &'a [u8],
    pos: usize,
}

impl<'a> Iterator for ObuIter<'a> {
    type Item = Obu<'a>;

    fn next(&mut self) -> Option<Obu<'a>> {
        if self.pos >= self.data.len() {
            return None;
        }
        let start = self.pos;
        let header = self.data[self.pos];
        self.pos += 1;
        let obu_type = (header >> 3) & 0x0F;
        let extension_flag = (header >> 2) & 1;
        let has_size_field = (header >> 1) & 1;
        if extension_flag == 1 {
            self.pos += 1; // extension header byte
        }
        let size = if has_size_field == 1 {
            leb128(self.data, &mut self.pos)? as usize
        } else {
            self.data.len().saturating_sub(self.pos)
        };
        let end = self.pos.checked_add(size)?;
        if end > self.data.len() {
            return None; // truncated
        }
        let payload = &self.data[self.pos..end];
        let raw = &self.data[start..end];
        self.pos = end;
        Some(Obu {
            obu_type,
            payload,
            raw,
            has_extension: extension_flag == 1,
        })
    }
}

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

    #[test]
    fn leb128_roundtrip() {
        let mut pos = 0;
        assert_eq!(leb128(&[0x00], &mut pos), Some(0));
        pos = 0;
        assert_eq!(leb128(&[0x7f], &mut pos), Some(127));
        pos = 0;
        // 300 = 0b1_00101100 → bytes 0xAC 0x02
        assert_eq!(leb128(&[0xAC, 0x02], &mut pos), Some(300));
        assert_eq!(pos, 2);
    }

    #[test]
    fn iterates_size_delimited_obus() {
        // Temporal delimiter (type 2, has_size, size 0) + a 3-byte type-6 OBU.
        let data = [0x12, 0x00, 0x32, 0x03, 0xAA, 0xBB, 0xCC];
        let obus: Vec<(u8, usize)> = iter_obus(&data)
            .map(|o| (o.obu_type, o.payload.len()))
            .collect();
        assert_eq!(obus, vec![(2, 0), (6, 3)]);
    }
}