From 6d3ca9eee89c9c9751b6e833b26fb899a581faf7 Mon Sep 17 00:00:00 2001
From: Rainbaby <rainbaby@outlook.jp>
Date: Sat, 2 Apr 2022 00:27:59 +0800
Subject: [PATCH] Fix RPU header parsing for Profile 7

Add the missing nlq_pred_pivot_value for it.
---
 dolby_vision/src/rpu/dovi_rpu.rs        |  3 ++-
 dolby_vision/src/rpu/rpu_data_header.rs | 28 +++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/dolby_vision/src/rpu/dovi_rpu.rs b/dolby_vision/src/rpu/dovi_rpu.rs
index d38b1d0..9d3a33b 100644
--- a/dolby_vision/src/rpu/dovi_rpu.rs
+++ b/dolby_vision/src/rpu/dovi_rpu.rs
@@ -326,7 +326,7 @@ impl DoviRpu {
 
         header.nlq_method_idc = Some(0);
         header.nlq_num_pivots_minus2 = Some(0);
-        header.num_x_partitions_minus1 = 2046;
+        header.nlq_pred_pivot_value = Some(vec![0, (1 << header.bl_bit_depth_minus8 + 8) - 1]);
 
         if let Some(ref mut rpu_data_nlq) = self.rpu_data_nlq {
             rpu_data_nlq.convert_to_mel();
@@ -348,6 +348,7 @@ impl DoviRpu {
 
         header.nlq_method_idc = None;
         header.nlq_num_pivots_minus2 = None;
+        header.nlq_pred_pivot_value = None;
 
         header.num_x_partitions_minus1 = 0;
         header.num_y_partitions_minus1 = 0;
diff --git a/dolby_vision/src/rpu/rpu_data_header.rs b/dolby_vision/src/rpu/rpu_data_header.rs
index b464ade..a8986b6 100644
--- a/dolby_vision/src/rpu/rpu_data_header.rs
+++ b/dolby_vision/src/rpu/rpu_data_header.rs
@@ -37,6 +37,7 @@ pub struct RpuDataHeader {
     pub pred_pivot_value: [Vec<u64>; NUM_COMPONENTS],
     pub nlq_method_idc: Option<u8>,
     pub nlq_num_pivots_minus2: Option<u8>,
+    pub nlq_pred_pivot_value: Option<Vec<u64>>,
     pub num_x_partitions_minus1: u64,
     pub num_y_partitions_minus1: u64,
 }
@@ -113,6 +114,15 @@ impl RpuDataHeader {
                     if rpu_nal.rpu_format & 0x700 == 0 && !rpu_nal.disable_residual_flag {
                         rpu_nal.nlq_method_idc = Some(reader.get_n(3));
                         rpu_nal.nlq_num_pivots_minus2 = Some(0);
+
+                        let mut nlq_pred_pivot_value = vec![0; 2];
+
+                        for nlq_pivot_idx in 0..2 {
+                            nlq_pred_pivot_value[nlq_pivot_idx] =
+                                reader.get_n((rpu_nal.bl_bit_depth_minus8 + 8) as usize);
+                        }
+
+                        rpu_nal.nlq_pred_pivot_value = Some(nlq_pred_pivot_value);
                     }
 
                     rpu_nal.num_x_partitions_minus1 = reader.get_ue()?;
@@ -145,6 +155,10 @@ impl RpuDataHeader {
                     self.nlq_num_pivots_minus2.is_none(),
                     "profile 5: nlq_num_pivots_minus2 should be undefined"
                 );
+                ensure!(
+                    self.nlq_pred_pivot_value.is_none(),
+                    "profile 5: nlq_pred_pivot_value should be undefined"
+                );
             }
             7 => {
                 ensure!(
@@ -165,6 +179,10 @@ impl RpuDataHeader {
                     self.nlq_num_pivots_minus2.is_none(),
                     "profile 8: nlq_num_pivots_minus2 should be undefined"
                 );
+                ensure!(
+                    self.nlq_pred_pivot_value.is_none(),
+                    "profile 8: nlq_pred_pivot_value should be undefined"
+                );
             }
             _ => (),
         };
@@ -285,6 +303,15 @@ impl RpuDataHeader {
                         if let Some(nlq_method_idc) = self.nlq_method_idc {
                             writer.write_n(&nlq_method_idc.to_be_bytes(), 3);
                         }
+
+                        if let Some(nlq_pred_pivot_value) = &self.nlq_pred_pivot_value {
+                            for nlq_pivot_idx in 0..2 {
+                                writer.write_n(
+                                    &nlq_pred_pivot_value[nlq_pivot_idx].to_be_bytes(),
+                                    (self.bl_bit_depth_minus8 + 8) as usize,
+                                )
+                            }
+                        }
                     }
 
                     writer.write_ue(self.num_x_partitions_minus1);
@@ -324,6 +351,7 @@ impl RpuDataHeader {
             pred_pivot_value: [vec![0, 1023], vec![0, 1023], vec![0, 1023]],
             nlq_method_idc: None,
             nlq_num_pivots_minus2: None,
+            nlq_pred_pivot_value: None,
             num_x_partitions_minus1: 0,
             num_y_partitions_minus1: 0,
         }