API overhaul

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+*   BREAKING CHANGE: major simplification of the push API.
 *   Annex B parser bugfixes.
 
 ## 0.5.0 - 2021-06-09

Cargo.toml

@@ -12,13 +12,13 @@ edition = "2018"
 
 [dependencies]
 bitstream-io = "1.1"
+hex-slice = "0.1.4"
 memchr = "2.1.1"
 rfc6381-codec = "0.1"
 log = "0.4"
 
 [dev-dependencies]
 hex-literal = "0.3.1"
-hex-slice = "0.1.4"
 criterion = "0.3"
 
 [[bench]]

benches/bench.rs

@@ -8,140 +8,146 @@
 
 #[macro_use]
 extern crate criterion;
-extern crate h264_reader;
 
-use criterion::Criterion;
+use criterion::{Bencher, Criterion, Throughput};
+use hex_literal::hex;
+use h264_reader::nal::{RefNal, Nal};
+use h264_reader::nal::UnitType;
+use h264_reader::nal::sei::SeiReader;
+use h264_reader::nal::slice::{SliceHeader, SliceHeaderError};
 use h264_reader::nal::sps::SeqParameterSet;
-use std::fs::File;
-use criterion::Throughput;
+use h264_reader::push::NalInterest;
+use h264_reader::rbsp::{self, BitReaderError};
+use std::io::{BufRead, ErrorKind};
 use std::convert::TryFrom;
-use std::io::Read;
-use hex_literal::hex;
 use h264_reader::annexb::AnnexBReader;
-use h264_reader::annexb::NalReader;
-use h264_reader::Context;
-use h264_reader::rbsp::RbspDecoder;
-use h264_reader::nal::NalHandler;
-use h264_reader::nal::NalHeader;
+use h264_reader::push::NalFragmentHandler;
 
-struct NullNalHandler {
-    start: u64,
+/// A NAL handler that does nothing, except maintain counters to limit optimization.
+#[derive(Default)]
+struct NullNalReader {
     push: u64,
     end: u64,
 }
-impl NalHandler for NullNalHandler {
-    type Ctx = ();
-
-    fn start(&mut self, _ctx: &mut Context<Self::Ctx>, _header: NalHeader) {
-        self.start += 1;
-    }
-
-    fn push(&mut self, _ctx: &mut Context<Self::Ctx>, _buf: &[u8]) {
+impl NalFragmentHandler for NullNalReader {
+    fn nal_fragment(&mut self, _bufs: &[&[u8]], end: bool) {
         self.push += 1;
-    }
-
-    fn end(&mut self, _ctx: &mut Context<Self::Ctx>) {
-        self.end += 1;
-    }
-}
-
-struct NullRbspNalReader {
-    decoder: RbspDecoder<NullNalHandler>,
-    decoder_started: bool,
-}
-impl NalReader for NullRbspNalReader {
-    type Ctx = ();
-
-    fn start(&mut self, _ctx: &mut Context<Self::Ctx>) {
-        assert!(!self.decoder_started);
-    }
-    fn push(&mut self, ctx: &mut Context<Self::Ctx>, mut buf: &[u8]) {
-        if !self.decoder_started && !buf.is_empty() {
-            let hdr = NalHeader::new(buf[0]).unwrap();
-            self.decoder.start(ctx, hdr);
-            buf = &buf[1..];
-            self.decoder_started = true;
-        }
-        if self.decoder_started {
-            self.decoder.push(ctx, buf);
+        if end {
+            self.end += 1;
         }
     }
-    fn end(&mut self, ctx: &mut Context<Self::Ctx>) {
-        assert!(self.decoder_started);
-        self.decoder.end(ctx);
-        self.decoder_started = false;
-    }
-}
-
-struct NullNalReader {
-    start: u64,
-    push: u64,
-    end: u64,
 }
-impl NalReader for NullNalReader {
-    type Ctx = ();
 
-    fn start(&mut self, _ctx: &mut Context<Self::Ctx>) {
-        self.start += 1;
-    }
-    fn push(&mut self, _ctx: &mut Context<Self::Ctx>, _buf: &[u8]) {
-        self.push += 1;
-    }
-    fn end(&mut self, _ctx: &mut Context<Self::Ctx>) {
-        self.end += 1;
-    }
+fn bench_annexb<'a, H, P>(mut r: AnnexBReader<H>, b: &mut Bencher, pushes: P)
+where H: NalFragmentHandler, P: Iterator<Item = &'a [u8]> + Clone {
+    b.iter(|| {
+        for p in pushes.clone() {
+            r.push(p);
+        }
+        r.reset();
+    })
 }
 
 fn h264_reader(c: &mut Criterion) {
-    let mut f = File::open("big_buck_bunny_1080p.h264").expect("file not found");
-    let len = f.metadata().unwrap().len();
-    let mut buf = vec![0; usize::try_from(len).unwrap()];
-    f.read(&mut buf[..]).unwrap();
-    let mut ctx = Context::default();
-    let nal_handler = NullNalHandler {
-        start: 0,
-        push: 0,
-        end: 0,
-    };
-    let rbsp_nal_reader = NullRbspNalReader {
-        decoder: RbspDecoder::new(nal_handler),
-        decoder_started: false,
+    let buf = std::fs::read("big_buck_bunny_1080p.h264").expect("reading h264 file failed");
+    let mut rbsp_len = 0;
+    let mut rbsp_len_nal_handler = |nal: RefNal<'_>| {
+        if nal.is_complete() {
+            let mut r = nal.rbsp_bytes();
+            loop {
+                let buf = r.fill_buf().unwrap();
+                let len = buf.len();
+                if len == 0 {
+                    break;
+                }
+                rbsp_len += u64::try_from(buf.len()).unwrap();
+                r.consume(len);
+            }
+        }
+        NalInterest::Buffer
     };
-    let nal_reader = NullNalReader {
-        start: 0,
-        push: 0,
-        end: 0,
+    let mut parsing_ctx = h264_reader::Context::default();
+    let mut scratch = Vec::new();
+    let mut parsing_nal_handler = |nal: RefNal<'_>| {
+        let nal_hdr = nal.header().unwrap();
+        match nal_hdr.nal_unit_type() {
+            UnitType::SeqParameterSet if nal.is_complete() => {
+                let sps = h264_reader::nal::sps::SeqParameterSet::from_bits(nal.rbsp_bits()).unwrap();
+                parsing_ctx.put_seq_param_set(sps);
+            },
+            UnitType::PicParameterSet if nal.is_complete() => {
+                let pps = h264_reader::nal::pps::PicParameterSet::from_bits(&parsing_ctx, nal.rbsp_bits()).unwrap();
+                parsing_ctx.put_pic_param_set(pps);
+            },
+            UnitType::SEI if nal.is_complete() => {
+                let mut r = SeiReader::from_rbsp_bytes(nal.rbsp_bytes(), &mut scratch);
+                while let Some(msg) = r.next().unwrap() {
+                    match msg.payload_type {
+                        h264_reader::nal::sei::HeaderType::BufferingPeriod => {}, // todo
+                        h264_reader::nal::sei::HeaderType::UserDataUnregistered => {}, // todo
+                        _ => panic!("unknown SEI payload type {:?}", msg.payload_type),
+                    }
+                }
+            },
+            UnitType::SliceLayerWithoutPartitioningIdr
+            | UnitType::SliceLayerWithoutPartitioningNonIdr => {
+                match SliceHeader::from_bits(&parsing_ctx, &mut nal.rbsp_bits(), nal.header().unwrap()) {
+                    Err(SliceHeaderError::RbspError(BitReaderError::ReaderErrorFor(_, e))) => {
+                        assert_eq!(e.kind(), ErrorKind::WouldBlock);
+                    },
+                    Err(e) => panic!("{:?}", e),
+                    Ok(_) => return NalInterest::Ignore,
+                }
+            },
+            _ => if nal.is_complete() { panic!("unknown slice type {:?}", nal_hdr) },
+        }
+        NalInterest::Buffer
     };
-    let mut annexb_rbsp_reader = AnnexBReader::new(rbsp_nal_reader);
-    let mut annexb_reader = AnnexBReader::new(nal_reader);
 
     let mut group = c.benchmark_group("parse_annexb");
-    group.throughput(Throughput::Bytes(len));
-    group.bench_function("annexb_only", |b| {
-        b.iter(|| {
-            annexb_reader.start(&mut ctx);
-            annexb_reader.push(&mut ctx, &buf[..]);
-            annexb_reader.end_units(&mut ctx);
-        })
-    });
-    group.bench_function("annexb_rbsp", |b| {
-        b.iter(|| {
-            annexb_rbsp_reader.start(&mut ctx);
-            annexb_rbsp_reader.push(&mut ctx, &buf[..]);
-            annexb_rbsp_reader.end_units(&mut ctx);
-        })
-    });
+    group.throughput(Throughput::Bytes(u64::try_from(buf.len()).unwrap()));
+
+    // Benchmark parsing in one big push (as when reading H.264 with a large buffer size),
+    // 184-byte pushes (like MPEG-TS), and 1440-byte pushes (~typical for RTP). RTP doesn't
+    // use Annex B encoding, but it does use the RBSP decoding and NAL parsing layers, so this
+    // is still informative.
+    group.bench_function("onepush_null", |b| bench_annexb(
+        AnnexBReader::for_fragment_handler(NullNalReader::default()), b, std::iter::once(&buf[..])));
+    group.bench_function("chunksize184_null", |b| bench_annexb(
+        AnnexBReader::for_fragment_handler(NullNalReader::default()), b, buf.chunks(184)));
+    group.bench_function("chunksize1440_null", |b| bench_annexb(
+        AnnexBReader::for_fragment_handler(NullNalReader::default()), b, buf.chunks(1440)));
+    group.bench_function("onepush_rbsp", |b| bench_annexb(
+        AnnexBReader::accumulate(&mut rbsp_len_nal_handler), b, std::iter::once(&buf[..])));
+    group.bench_function("chunksize184_rbsp", |b| bench_annexb(
+        AnnexBReader::accumulate(&mut rbsp_len_nal_handler), b, buf.chunks(184)));
+    group.bench_function("chunksize1440_rbsp", |b| bench_annexb(
+        AnnexBReader::accumulate(&mut rbsp_len_nal_handler), b, buf.chunks(1440)));
+    group.bench_function("onepush_parse", |b| bench_annexb(
+        AnnexBReader::accumulate(&mut parsing_nal_handler), b, std::iter::once(&buf[..])));
+    group.bench_function("chunksize184_parse", |b| bench_annexb(
+        AnnexBReader::accumulate(&mut parsing_nal_handler), b, buf.chunks(184)));
+    group.bench_function("chunksize1440_parse", |b| bench_annexb(
+        AnnexBReader::accumulate(&mut parsing_nal_handler), b, buf.chunks(1440)));
 }
 
 fn parse_nal(c: &mut Criterion) {
+    // Note this has no emulation prevention three bytes in it; sps[1..] can be
+    // passed directly to code that expects RBSP.
     let sps = hex!(
-        "64 00 16 AC 1B 1A 80 B0 3D FF FF
+        "67 64 00 16 AC 1B 1A 80 B0 3D FF FF
         00 28 00 21 6E 0C 0C 0C 80 00 01
         F4 00 00 27 10 74 30 07 D0 00 07
         A1 25 DE 5C 68 60 0F A0 00 0F 42
         4B BC B8 50");
+    let nal = RefNal::new(&sps[..], &[], true);
     let mut group = c.benchmark_group("parse_nal");
-    group.bench_function("sps", |b| b.iter(|| SeqParameterSet::from_bytes(&sps[..]).unwrap()));
+    group.bench_function("rbsp_sps", |b| b.iter(|| {
+        SeqParameterSet::from_bits(rbsp::BitReader::new(&sps[1..])).unwrap()
+    }));
+    group.bench_function("nal_sps", |b| b.iter(|| {
+        SeqParameterSet::from_bits(nal.rbsp_bits()).unwrap()
+    }));
 }
 
 criterion_group!(benches, h264_reader, parse_nal);

fuzz/fuzz_targets/annexb.rs

@@ -2,68 +2,60 @@
 //! single push call and a pair of push split at each possible byte location.
 
 #![no_main]
+use h264_reader::annexb::AnnexBReader;
 use hex_slice::AsHex;
-use h264_reader::Context;
 use libfuzzer_sys::fuzz_target;
 use std::convert::TryFrom;
 
 /// Encodes the stream as (4-byte length prefix, NAL)*, as commonly seen in AVC files.
 #[derive(Default)]
 struct AvcBuilder {
-    started: bool,
     cur: Vec<u8>,
+    all: Vec<u8>,
 }
 
-impl h264_reader::annexb::NalReader for AvcBuilder {
-    type Ctx = Vec<u8>;
-
-    fn start(&mut self, _ctx: &mut Context<Self::Ctx>) {
-        assert!(!self.started);
-        self.started = true;
-    }
-    fn push(&mut self, _ctx: &mut Context<Self::Ctx>, buf: &[u8]) {
-        assert!(self.started);
-        assert!(!buf.is_empty()); // useless empty push.
-        self.cur.extend_from_slice(buf);
-    }
-    fn end(&mut self, ctx: &mut Context<Self::Ctx>) {
-        assert!(self.started);
-        self.started = false;
-        let len = u32::try_from(self.cur.len()).unwrap();
-        ctx.user_context.extend_from_slice(&len.to_be_bytes()[..]);
-        ctx.user_context.extend_from_slice(&self.cur[..]);
-        self.cur.clear();
+impl h264_reader::push::NalFragmentHandler for AvcBuilder {
+    fn nal_fragment(&mut self, bufs: &[&[u8]], end: bool) {
+        assert!(!bufs.is_empty() || (!self.cur.is_empty() || end));
+        for buf in bufs {
+            assert!(!buf.is_empty());
+            self.cur.extend_from_slice(buf);
+        }
+        if end {
+            let len = u32::try_from(self.cur.len()).unwrap();
+            self.all.extend_from_slice(&len.to_be_bytes()[..]);
+            self.all.extend_from_slice(&self.cur[..]);
+            self.cur.clear();
+        }
     }
 }
 
 fuzz_target!(|data: &[u8]| {
     // Parse in a single push.
-    let mut single_push_ctx = h264_reader::Context::new(Vec::new());
-    let mut single_push = h264_reader::annexb::AnnexBReader::new(AvcBuilder::default());
-    single_push.start(&mut single_push_ctx);
-    single_push.push(&mut single_push_ctx, data);
-    single_push.end_units(&mut single_push_ctx);
+    let mut single_push = AnnexBReader::for_fragment_handler(AvcBuilder::default());
+    single_push.push(data);
+    single_push.reset();
+    let single_avc = single_push.into_fragment_handler();
 
     for i in 0..data.len() {
         // Parse in a split push.
-        let mut split_push_ctx = h264_reader::Context::new(Vec::new());
-        let mut split_push = h264_reader::annexb::AnnexBReader::new(AvcBuilder::default());
-        split_push.start(&mut split_push_ctx);
+        let mut split_push = AnnexBReader::for_fragment_handler(AvcBuilder::default());
         let (head, tail) = data.split_at(i);
-        split_push.push(&mut split_push_ctx, head);
-        split_push.push(&mut split_push_ctx, &[]); // also ensure empty pushes don't break.
-        split_push.push(&mut split_push_ctx, tail);
-        split_push.end_units(&mut split_push_ctx);
+        split_push.push(head);
+        split_push.push(&[]); // also ensure empty pushes don't break.
+        split_push.push(tail);
+        split_push.reset();
+        let split_avc = split_push.into_fragment_handler();
 
-        assert!(single_push_ctx.user_context.as_slice() == split_push_ctx.user_context.as_slice(),
+        assert!(single_avc.all.as_slice() == split_avc.all.as_slice(),
                 "inconsistent output.\n\
                 split point: {}\n\
                 input:       {:02x}\n\
                 single push: {:02x}\n\
                 split push:  {:02x}",
                 i,
                 data.as_hex(),
-                single_push_ctx.user_context.as_hex(),
-                split_push_ctx.user_context.as_hex());
+                single_avc.all.as_hex(),
+                split_avc.all.as_hex());
     }
 });