security(bitstream): cache RBSP stop-bit lookup (#2880)

crates/bitstream/src/lib.rs

@@ -7,7 +7,7 @@
 
 #![cfg_attr(not(test), deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable))]
 
-use std::fmt;
+use std::{cell::Cell, fmt};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum BitError {
@@ -35,11 +35,13 @@ pub struct BitReader<'a> {
     data: &'a [u8],
     /// Position in bits.
     pos: usize,
+    /// Cached position of the RBSP stop bit, computed only by [`Self::more_rbsp_data`].
+    rbsp_stop: Cell<Option<Option<usize>>>,
 }
 
 impl<'a> BitReader<'a> {
     pub fn new(data: &'a [u8]) -> Self {
-        Self { data, pos: 0 }
+        Self { data, pos: 0, rbsp_stop: Cell::new(None) }
     }
     pub fn data(&self) -> &'a [u8] {
         self.data
@@ -157,14 +159,13 @@ impl<'a> BitReader<'a> {
         if self.pos >= total {
             return false;
         }
-        // find last set bit (the rbsp_stop_one_bit)
-        let mut last = None;
-        for (i, &b) in self.data.iter().enumerate().rev() {
-            if b != 0 {
-                last = Some(i * 8 + 7 - b.trailing_zeros() as usize);
-                break;
-            }
-        }
+        let last = self.rbsp_stop.get().unwrap_or_else(|| {
+            // The trailing stop bit is immutable for this reader.  Caching it avoids searching a
+            // long zero suffix after every macroblock.
+            let stop = self.data.iter().rposition(|&b| b != 0).map(|i| i * 8 + 7 - self.data[i].trailing_zeros() as usize);
+            self.rbsp_stop.set(Some(stop));
+            stop
+        });
         match last {
             Some(stop) => self.pos < stop,
             None => false,
@@ -443,7 +444,9 @@ mod tests {
     fn more_rbsp() {
         let d = [0b1100_0000];
         let mut r = BitReader::new(&d);
+        assert_eq!(r.rbsp_stop.get(), None);
         assert!(r.more_rbsp_data());
+        assert_eq!(r.rbsp_stop.get(), Some(Some(1)));
         r.read_bit().unwrap();
         assert!(!r.more_rbsp_data());
     }