rpu/parser: Refactor handling of trailing bytes

quietvoid · Mar 19, 2024 · fd097b9 · fd097b9
1 parent ab27880
commit fd097b9
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 28 deletions.
diff --git a/dolby_vision/src/rpu/dovi_rpu.rs b/dolby_vision/src/rpu/dovi_rpu.rs
@@ -23,6 +23,7 @@ use crate::utils::{
 };
 
 pub(crate) const FINAL_BYTE: u8 = 0x80;
+const CRC32_TERMINATOR_BITS: u64 = 40;
 
 #[derive(Debug, Default, Clone)]
 #[cfg_attr(feature = "serde", derive(Serialize))]
@@ -97,6 +98,7 @@ impl DoviRpu {
 
     #[inline(always)]
     pub(crate) fn parse(data: &[u8]) -> Result<DoviRpu> {
+        let original_payload_size = data.len();
         let trailing_zeroes = data.iter().rev().take_while(|b| **b == 0).count();
 
         // Ignore trailing bytes
@@ -113,7 +115,7 @@ impl DoviRpu {
             bail!("Invalid RPU last byte: {}", last_byte);
         }
 
-        let dovi_rpu = DoviRpu::read_rpu_data(data, trailing_zeroes)?;
+        let mut dovi_rpu = DoviRpu::read_rpu_data(&data[..rpu_end])?;
 
         if received_crc32 != dovi_rpu.rpu_data_crc32 {
             bail!(
@@ -123,16 +125,19 @@ impl DoviRpu {
             );
         }
 
+        dovi_rpu.trailing_zeroes = trailing_zeroes;
+        dovi_rpu.original_payload_size = original_payload_size;
+
+        // Validate
+        dovi_rpu.validate()?;
+
         Ok(dovi_rpu)
     }
 
     #[inline(always)]
-    fn read_rpu_data(bytes: &[u8], trailing_zeroes: usize) -> Result<DoviRpu> {
+    fn read_rpu_data(bytes: &[u8]) -> Result<DoviRpu> {
         let mut reader = BsIoSliceReader::from_slice(bytes);
 
-        // CRC32 + 0x80 + trailing
-        let final_length = (32 + 8 + (trailing_zeroes * 8)) as u64;
-
         let rpu_prefix = reader.get_n(8)?;
         ensure!(rpu_prefix == 25, "rpu_nal_prefix should be 25");
 
@@ -160,7 +165,7 @@ impl DoviRpu {
             .unwrap_or(None);
 
         let vdr_dm_data = if header.vdr_dm_metadata_present_flag {
-            Some(vdr_dm_data_payload(&mut reader, &header, final_length)?)
+            Some(vdr_dm_data_payload(&mut reader, &header)?)
         } else {
             None
         };
@@ -171,11 +176,10 @@ impl DoviRpu {
         }
 
         // CRC32 is at the end, there can be more data in between
-
-        let remaining = if reader.available()? != final_length {
+        let remaining = if reader.available()? > CRC32_TERMINATOR_BITS {
             let mut remaining: BitVec<u8, Msb0> = BitVec::new();
 
-            while reader.available()? != final_length {
+            while reader.available()? != CRC32_TERMINATOR_BITS {
                 remaining.push(reader.get()?);
             }
 
@@ -188,23 +192,16 @@ impl DoviRpu {
         let last_byte: u8 = reader.get_n(8)?;
         ensure!(last_byte == FINAL_BYTE, "last byte should be 0x80");
 
-        let dovi_rpu = DoviRpu {
+        Ok(DoviRpu {
             dovi_profile: header.get_dovi_profile(),
             el_type,
             header,
             rpu_data_mapping,
             vdr_dm_data,
             remaining,
             rpu_data_crc32,
-            modified: false,
-            trailing_zeroes,
-            original_payload_size: bytes.len(),
-        };
-
-        // Validate
-        dovi_rpu.validate()?;
-
-        Ok(dovi_rpu)
+            ..Default::default()
+        })
     }
 
     pub fn write_hevc_unspec62_nalu(&self) -> Result<Vec<u8>> {

diff --git a/dolby_vision/src/rpu/mod.rs b/dolby_vision/src/rpu/mod.rs
@@ -29,10 +29,7 @@ pub enum ConversionMode {
 
 #[inline(always)]
 fn compute_crc32(data: &[u8]) -> u32 {
-    let mut digest = CRC32_INSTANCE.digest();
-    digest.update(data);
-
-    digest.finalize()
+    CRC32_INSTANCE.checksum(data)
 }
 
 impl From<u8> for ConversionMode {

diff --git a/dolby_vision/src/rpu/rpu_data_mapping.rs b/dolby_vision/src/rpu/rpu_data_mapping.rs
@@ -115,8 +115,7 @@ impl RpuDataMapping {
             curve.num_pivots_minus2 = reader.get_ue()?;
             let num_pivots = (curve.num_pivots_minus2 + 2) as usize;
 
-            curve.pivots = Vec::with_capacity(num_pivots);
-            curve.pivots.resize(num_pivots, Default::default());
+            curve.pivots = vec![Default::default(); num_pivots];
 
             for i in 0..num_pivots {
                 curve.pivots[i] = reader.get_n(bl_bit_depth)?;

diff --git a/dolby_vision/src/rpu/vdr_dm_data.rs b/dolby_vision/src/rpu/vdr_dm_data.rs
@@ -19,6 +19,9 @@ use super::profiles::DoviProfile;
 use super::extension_metadata::WithExtMetadataBlocks;
 use super::rpu_data_header::RpuDataHeader;
 
+// 16 bits min for required level 254 + CRC32 + 0x80
+const DM_DATA_PAYLOAD2_MIN_BITS: u64 = 56;
+
 #[derive(Debug, Default, Clone)]
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
 pub struct VdrDmData {
@@ -77,7 +80,6 @@ pub enum CmVersion {
 pub(crate) fn vdr_dm_data_payload(
     reader: &mut BsIoSliceReader,
     header: &RpuDataHeader,
-    final_length: u64,
 ) -> Result<VdrDmData> {
     let compressed_dm_data = header.reserved_zero_3bits == 1;
 
@@ -98,8 +100,7 @@ pub(crate) fn vdr_dm_data_payload(
         vdr_dm_data.cmv29_metadata = Some(DmData::V29(cmv29_dm_data));
     }
 
-    // 16 bits min for required level 254
-    if reader.available()? >= final_length + 16 {
+    if reader.available()? >= DM_DATA_PAYLOAD2_MIN_BITS {
         if let Some(cmv40_dm_data) = DmData::parse::<CmV40DmData>(reader)? {
             vdr_dm_data.cmv40_metadata = Some(DmData::V40(cmv40_dm_data));
         }