Create simpler facades for decoding steps about h264-reader HOT 10 OPEN

pub trait RbspHandler {
    type Read: RbspBitRead;
    type Ctx;

    /// Processes a partially buffered NAL unit. `read` always reads from the beginning.
    /// Returns false if no additional calls are desired for this NAL unit.
    fn partial(&mut self, ctx: &mut Context<Self::Ctx>, header: NalHeader, read: Self::Read) -> bool {
        true
    }

    /// Processes a fully buffered NAL unit. `read` always reads from the beginning.
    /// Called iff every prior `partial` call for this NAL unit (zero or more) returned true.
    fn complete(&mut self, ctx: &mut Context<Self::Ctx>, header: NalHeader, read: Self::Read) {}
}

impl<Read: RbspBitRead, Ctx> RbspHandler for PicParameterSetNalHandler<Ctx> {
    type Read = Read;
    type Ctx = Ctx;

    fn complete(&mut self, ctx: &mut Context<Self::Ctx>, header: NalHeader, read: Self::Read) {
        match PicParameterSet::read(ctx, read) {
            Ok(pps) => ctx.put_pic_param_set(pps),
            Err(e) => error!("pps: {:?}", e),
        }
    }
}

impl<Read: RbspBitRead, Ctx> RbspHandler for SliceNalHandler<Ctx> {
    type Read = Read;
    type Ctx = Ctx;

    fn partial(&mut self, ctx: &mut Context<Self::Ctx>, header: NalHeader, read: Self::Read) -> bool {
        match SliceHeader::read(ctx, read) {
            Ok(header) => {
                info!("TODO: expose to caller: {:#?}", header);
                false
            },
            Err(SliceHeaderError::Incomplete) => true,
            Err(e) => {
                error!("slice_header() error: SliceHeaderError::{:?}", e);
                false
            },
        }
    }

    fn complete(&mut self, ctx: &mut Context<Self::Ctx>, header: NalHeader, read: Self::Read) -> bool {
        match SliceHeader::read(ctx, read) {
            Ok(pps) => info!("TODO: expose to caller: {:#?}", header),
            Err(e) => error!("slice_header() error: SliceHeaderError::{:?}", e),
        }
    }
}

mod nal {
    /// A single, partially- or completely-buffered NAL.
    /// The push module uses a [ByteSliceNal], but it's possible to implement
    /// this interface with any type that can produce a `std::io::BufRead`.
    pub trait Nal {
        type BufRead: std::io::BufRead;

        /// Returns if the NAL is completely buffered.
        pub fn complete(&self) -> bool;

        /// Returns the NAL header or error if missing/corrupt.
        pub fn header(&self) -> Result<NalHeader, NalHeaderError> { ... }

        /// Returns the bytes in NAL form (including the header byte and
        /// any emulation-prevention-three-bytes). If the NAL is incomplete,
	/// reads may fail with `ErrorKind::WouldBlock`.
        pub fn reader(&self) -> Self::BufRead;

        /// Returns the bytes in RBSP form (skipping header byte and
        /// emulation-prevention-three-bytes).
        pub fn rbsp_bytes(&self) -> rbsp::ByteReader<Self::BufRead> { ... }

        /// Returns a helper to access bits within the RBSP.
        pub fn rbsp_bits(&self) -> rbsp::BitReader<rbsp::ByteReader<Self::BufRead>> { ... }
    }

    /// A NAL which stores its bytes contiguously.
    pub struct ByteSliceNal<'a> { buf: &'a [u8], complete: bool }
    impl<'a> Nal for ByteSliceNal<'a> { ... }
    impl<'a> AsRef<u8> for ByteSliceNal<'a> { ... }
}

mod push {
    /// The results of a NAL handler during push parsing.
    pub enum HandlerResult {
        /// If this NAL is incomplete, continue buffering it for a following call.
        ReadMore,

        /// Stop buffering this NAL; move on to the next one immediately.
        /// This is no different than `ReadMore` when the NAL is complete.
        Skip,

        /// Abort the push call with the given error. Following pushes will also fail.
        AbortPush(Box<dyn Error>),
    }

    /// Accumulator for NAL push parsers.
    /// This is meant to be used by parsers for a specific format: Annex B, AVC, MPEG-TS, RTP, etc.
    /// Accumulates NALs in an internal buffer and delegates to a caller-supplied handler.
    pub struct Accumulator<'a> {
        buf: Vec<u8>,
        state: AccumulatorState,
        handler: &'a mut FnMut(ByteSliceNal),
    }

    impl Accumulator {
        pub fn new(nal_handler: &mut FnMut(ByteSliceNal) -> HandlerResult) -> Self;

        pub fn start(&mut self);
        // TODO: maybe should be push(..., end: bool) instead to avoid an extra !complete call to nal_handler.
        pub fn push(&mut self, buf: &[u8]) -> Result<(), Error>;
        pub fn end(&mut self);
    }

    /// Push parser for an Annex B form bytestream
    /// (one with NALs separated by `00 00 01` or `00 00 00 01`).
    pub struct AnnexBParser<'a> { accumulator: Accumulator<'a>, state: AnnexBParseState }

    impl<'a> AnnexBParser<'a> {
        pub fn new(nal_handler: &mut FnMut(nal: ByteSliceNal) -> HandlerResult) -> Self;
        pub fn push(&mut self, buf: &[u8]) -> Result<(), Error>;
        pub fn end_units(&mut self) -> Result<(), Error>;
    }
}

/// RBSP stuff, independent of push or pull parsing.
mod rbsp {
    /// Strips header byte and `emulation-prevention-three-byte` elements from the single-NAL input
    /// to yield the RBSP.
    pub struct ByteReader<N: std::io::BufRead> { ... }
    impl<N: std::io::BufRead> std::io::Read for ByteReader<N> { ... }
    impl<N: std::io::BufRead> std::io::BufRead for ByteReader<N> { ... }

    /// Reads the RBSP as bits given a byte-level reader.
    // (This only needs `BufRead` rather than `Read` for `has_more_rbsp_data`.)
    pub struct BitReader<R: std::io::BufRead>(bitstream_io::Reader<R, bitstream_io::BigEndian>);
    impl<R: std::io::BufRead> BitReader<R> {
        ...
    }
}

// Example usage
fn my_annexb_reader() {
    let mut ctx = Context::new();
    let mut parser = h264_reader::push::AnnexBParser::new(|nal| {
        match nal.header().unit_type() {
            Ok(UnitType::PicParameterSet) if !nal.complete() => return NalResult::ReadMore,
            Ok(UnitType::PicParameterSet) => {
                match PicParameterSet::read(ctx, nal.rbsp_bits()) {
                    Ok(pps) => ctx.put_pic_param_set(pps),
                    Err(e) => error!("Bad pps: {:?}", e),
                }
            },
            Ok(UnitType::SliceWithPartitioningIdr) => {
                match SliceHeader::read(ctx, nal.rbsp_bits()) {
                    Ok(hdr) => info!("Slice header: {:#?}", hdr),
                    Err(SliceHeaderError::NeedMoreData) => return NalResult::ReadMore,
                    Err(e) => error!("Bad slice header: {:?}", e),
                }
            },
            Ok(_) => {},
            Err(e) => error!("bad NAL header: {:#?}", e),
        }
        NalResult::Skip
    });
    let f = File::open("foo.h264").unwrap();
    let mut reader = BufReader::new(f);
    loop {
        let buf = reader.fill_buf().unwrap();
        if buf.is_empty() {
            break;
        }
        parser.push(buf);
        read.consume(buf.len());
    }
}