decoder: make decode() return processed length

Require decode() to return how much of the input buffer has actually been consumed. This will be useful to signal the caller that the input was not correctly split and to simplify our decoding loop. For now, just return the size of the input buffer since that's what we are effectively consuming.
dabrain34 · Jul 20, 2023 · 36d403a · 36d403a
1 parent 2381715
commit 36d403a
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 17 deletions.
diff --git a/src/decoder/stateless.rs b/src/decoder/stateless.rs
@@ -195,7 +195,17 @@ pub trait StatelessVideoDecoder<M> {
     /// been detected that the client should acknowledge, or because there are no available output
     /// resources and dequeueing and returning pending frames will fix that. After the cause has
     /// been addressed, the client is responsible for calling this method again with the same data.
-    fn decode(&mut self, timestamp: u64, bitstream: &[u8]) -> std::result::Result<(), DecodeError>;
+    ///
+    /// The return value is the number of bytes in `bitstream` that have been processed. Usually
+    /// this will be equal to the length of `bitstream`, but some codecs may only do partial
+    /// processing if e.g. several units are sent at the same time. It is the responsibility of the
+    /// caller to check that all submitted input has been processed, and to resubmit the
+    /// unprocessed part if it hasn't. See the documentation of each codec for their expectations.
+    fn decode(
+        &mut self,
+        timestamp: u64,
+        bitstream: &[u8],
+    ) -> std::result::Result<usize, DecodeError>;
 
     /// Flush the decoder i.e. finish processing all pending decode requests and make sure the
     /// resulting frames are ready to be retrieved via [`next_event`].

diff --git a/src/decoder/stateless/h264.rs b/src/decoder/stateless/h264.rs
@@ -1802,7 +1802,11 @@ where
         None
     }
 
-    fn decode_access_unit(&mut self, timestamp: u64, bitstream: &[u8]) -> Result<(), DecodeError> {
+    fn decode_access_unit(
+        &mut self,
+        timestamp: u64,
+        bitstream: &[u8],
+    ) -> Result<(), DecodeError> {
         if self.backend.surface_pool().num_free_surfaces() == 0 {
             return Err(DecodeError::NotEnoughOutputBuffers(1));
         }
@@ -1870,7 +1874,7 @@ where
     B: StatelessH264DecoderBackend,
     B::Handle: Clone + 'static,
 {
-    fn decode(&mut self, timestamp: u64, mut bitstream: &[u8]) -> Result<(), DecodeError> {
+    fn decode(&mut self, timestamp: u64, mut bitstream: &[u8]) -> Result<usize, DecodeError> {
         let sps = Self::peek_sps(&mut self.codec.parser, bitstream);
 
         if let Some(sps) = sps {
@@ -1896,13 +1900,17 @@ where
             }
         }
 
+        let frame_len = bitstream.len();
+
         match &mut self.decoding_state {
             // Skip input until we get information from the stream.
-            DecodingState::AwaitingStreamInfo | DecodingState::Reset => Ok(()),
+            DecodingState::AwaitingStreamInfo | DecodingState::Reset => (),
             // Ask the client to confirm the format before we can process this.
-            DecodingState::AwaitingFormat(_) => Err(DecodeError::CheckEvents),
-            DecodingState::Decoding => self.decode_access_unit(timestamp, bitstream),
-        }
+            DecodingState::AwaitingFormat(_) => return Err(DecodeError::CheckEvents),
+            DecodingState::Decoding => self.decode_access_unit(timestamp, bitstream)?,
+        };
+
+        Ok(frame_len)
     }
 
     fn flush(&mut self) {

diff --git a/src/decoder/stateless/h265.rs b/src/decoder/stateless/h265.rs
@@ -848,7 +848,11 @@ where
         Ok(())
     }
 
-    fn decode_access_unit(&mut self, timestamp: u64, bitstream: &[u8]) -> Result<(), DecodeError> {
+    fn decode_access_unit(
+        &mut self,
+        timestamp: u64,
+        bitstream: &[u8],
+    ) -> Result<usize, DecodeError> {
         if self.backend.surface_pool().num_free_surfaces() == 0 {
             return Err(DecodeError::NotEnoughOutputBuffers(1));
         }
@@ -964,7 +968,7 @@ where
         let (picture, handle) = self.submit_picture()?;
         self.finish_picture(picture, handle)?;
 
-        Ok(())
+        Ok(bitstream.len())
     }
 
     /// Submits the picture to the accelerator.
@@ -987,7 +991,7 @@ impl<T, P> StatelessVideoDecoder<T::Descriptor> for Decoder<T, P>
 where
     T: DecodedHandle + Clone + 'static,
 {
-    fn decode(&mut self, timestamp: u64, mut bitstream: &[u8]) -> Result<(), DecodeError> {
+    fn decode(&mut self, timestamp: u64, mut bitstream: &[u8]) -> Result<usize, DecodeError> {
         let sps = Self::peek_sps(&mut self.parser, bitstream)?;
 
         if let Some(sps) = sps {
@@ -1013,7 +1017,7 @@ where
 
         match &mut self.decoding_state {
             // Skip input until we get information from the stream.
-            DecodingState::AwaitingStreamInfo | DecodingState::Reset => Ok(()),
+            DecodingState::AwaitingStreamInfo | DecodingState::Reset => Ok(bitstream.len()),
             // Ask the client to confirm the format before we can process this.
             DecodingState::AwaitingFormat(_) => Err(DecodeError::CheckEvents),
             DecodingState::Decoding => self.decode_access_unit(timestamp, bitstream),

diff --git a/src/decoder/stateless/vp8.rs b/src/decoder/stateless/vp8.rs
@@ -204,7 +204,7 @@ where
     B: StatelessVp8DecoderBackend,
     B::Handle: Clone + 'static,
 {
-    fn decode(&mut self, timestamp: u64, bitstream: &[u8]) -> Result<(), DecodeError> {
+    fn decode(&mut self, timestamp: u64, bitstream: &[u8]) -> Result<usize, DecodeError> {
         let frame = self.codec.parser.parse_frame(bitstream)?;
 
         if frame.header.key_frame {
@@ -219,10 +219,14 @@ where
 
         match &mut self.decoding_state {
             // Skip input until we get information from the stream.
-            DecodingState::AwaitingStreamInfo | DecodingState::Reset => Ok(()),
+            DecodingState::AwaitingStreamInfo | DecodingState::Reset => Ok(bitstream.len()),
             // Ask the client to confirm the format before we can process this.
             DecodingState::AwaitingFormat(_) => Err(DecodeError::CheckEvents),
-            DecodingState::Decoding => self.handle_frame(frame, timestamp),
+            DecodingState::Decoding => {
+                let len = frame.bitstream.len();
+                self.handle_frame(frame, timestamp)?;
+                Ok(len)
+            }
         }
     }
 

diff --git a/src/decoder/stateless/vp9.rs b/src/decoder/stateless/vp9.rs
@@ -179,7 +179,7 @@ where
     B: StatelessVp9DecoderBackend,
     B::Handle: Clone + 'static,
 {
-    fn decode(&mut self, timestamp: u64, bitstream: &[u8]) -> Result<(), DecodeError> {
+    fn decode(&mut self, timestamp: u64, bitstream: &[u8]) -> Result<usize, DecodeError> {
         let frames = self.codec.parser.parse_chunk(bitstream)?;
 
         let num_free_surfaces = self.backend.surface_pool().num_free_surfaces();
@@ -232,7 +232,7 @@ where
             }
         }
 
-        Ok(())
+        Ok(bitstream.len())
     }
 
     fn flush(&mut self) {

diff --git a/src/utils.rs b/src/utils.rs
@@ -413,7 +413,7 @@ where
     for (frame_num, packet) in stream_iter.enumerate() {
         loop {
             match decoder.decode(frame_num as u64, packet.as_ref()) {
-                Ok(()) => {
+                Ok(_) => {
                     if blocking_mode == BlockingMode::Blocking {
                         check_events(decoder)?;
                     }