hevc_ps: expose SPS and VPS headers

2 files changed, 93 insertions, 48 deletions

diff --git a/libavcodec/hevc_ps.c b/libavcodec/hevc_ps.c
index 043e1bf308..aaaca5e6c7 100644
--- a/libavcodec/hevc_ps.c
+++ b/libavcodec/hevc_ps.c
@@ -357,81 +357,84 @@ static int parse_ptl(GetBitContext *gb, AVCodecContext *avctx,
 }
 
 static void decode_sublayer_hrd(GetBitContext *gb, unsigned int nb_cpb,
-                                int subpic_params_present)
+                                HEVCSublayerHdrParams *par, int subpic_params_present)
 {
     int i;
 
     for (i = 0; i < nb_cpb; i++) {
-        get_ue_golomb_long(gb); // bit_rate_value_minus1
-        get_ue_golomb_long(gb); // cpb_size_value_minus1
+        par->bit_rate_value_minus1[i] = get_ue_golomb_long(gb);
+        par->cpb_size_value_minus1[i] = get_ue_golomb_long(gb);
 
         if (subpic_params_present) {
-            get_ue_golomb_long(gb); // cpb_size_du_value_minus1
-            get_ue_golomb_long(gb); // bit_rate_du_value_minus1
+            par->cpb_size_du_value_minus1[i] = get_ue_golomb_long(gb);
+            par->bit_rate_du_value_minus1[i] = get_ue_golomb_long(gb);
         }
-        skip_bits1(gb); // cbr_flag
+
+        par->cbr_flag = get_bits1(gb);
     }
 }
 
 static int decode_hrd(GetBitContext *gb, int common_inf_present,
-                       int max_sublayers)
+                      HEVCHdrParams *hdr, int max_sublayers)
 {
-    int nal_params_present = 0, vcl_params_present = 0;
-    int subpic_params_present = 0;
-    int i;
-
     if (common_inf_present) {
-        nal_params_present = get_bits1(gb);
-        vcl_params_present = get_bits1(gb);
-
-        if (nal_params_present || vcl_params_present) {
-            subpic_params_present = get_bits1(gb);
-
-            if (subpic_params_present) {
-                skip_bits(gb, 8); // tick_divisor_minus2
-                skip_bits(gb, 5); // du_cpb_removal_delay_increment_length_minus1
-                skip_bits(gb, 1); // sub_pic_cpb_params_in_pic_timing_sei_flag
-                skip_bits(gb, 5); // dpb_output_delay_du_length_minus1
+        hdr->flags.nal_hrd_parameters_present_flag = get_bits1(gb);
+        hdr->flags.vcl_hrd_parameters_present_flag = get_bits1(gb);
+
+        if (hdr->flags.nal_hrd_parameters_present_flag ||
+            hdr->flags.vcl_hrd_parameters_present_flag) {
+            hdr->flags.sub_pic_hrd_params_present_flag = get_bits1(gb);
+
+            if (hdr->flags.sub_pic_hrd_params_present_flag) {
+                hdr->tick_divisor_minus2 = get_bits(gb, 8);
+                hdr->du_cpb_removal_delay_increment_length_minus1 = get_bits(gb, 5);
+                hdr->flags.sub_pic_cpb_params_in_pic_timing_sei_flag = get_bits1(gb);
+                hdr->dpb_output_delay_du_length_minus1 = get_bits(gb, 5);
             }
 
-            skip_bits(gb, 4); // bit_rate_scale
-            skip_bits(gb, 4); // cpb_size_scale
+            hdr->bit_rate_scale = get_bits(gb, 4);
+            hdr->cpb_size_scale = get_bits(gb, 4);
 
-            if (subpic_params_present)
-                skip_bits(gb, 4);  // cpb_size_du_scale
+            if (hdr->flags.sub_pic_hrd_params_present_flag)
+                hdr->cpb_size_du_scale = get_bits(gb, 4);
 
-            skip_bits(gb, 5); // initial_cpb_removal_delay_length_minus1
-            skip_bits(gb, 5); // au_cpb_removal_delay_length_minus1
-            skip_bits(gb, 5); // dpb_output_delay_length_minus1
+            hdr->initial_cpb_removal_delay_length_minus1 = get_bits(gb, 5);
+            hdr->au_cpb_removal_delay_length_minus1 = get_bits(gb, 5);
+            hdr->dpb_output_delay_length_minus1 = get_bits(gb, 5);
         }
     }
 
-    for (i = 0; i < max_sublayers; i++) {
-        int low_delay = 0;
-        unsigned int nb_cpb = 1;
-        int fixed_rate = get_bits1(gb);
+    for (int i = 0; i < max_sublayers; i++) {
+        hdr->flags.fixed_pic_rate_general_flag = get_bits1(gb);
+
+        hdr->cpb_cnt_minus1[i] = 1;
 
-        if (!fixed_rate)
-            fixed_rate = get_bits1(gb);
+        if (!hdr->flags.fixed_pic_rate_general_flag)
+            hdr->flags.fixed_pic_rate_within_cvs_flag = get_bits1(gb);
 
-        if (fixed_rate)
-            get_ue_golomb_long(gb);  // elemental_duration_in_tc_minus1
+        if (hdr->flags.fixed_pic_rate_within_cvs_flag)
+            hdr->elemental_duration_in_tc_minus1[i] = get_ue_golomb_long(gb);
         else
-            low_delay = get_bits1(gb);
+            hdr->flags.low_delay_hrd_flag = get_bits1(gb);
 
-        if (!low_delay) {
-            nb_cpb = get_ue_golomb_long(gb) + 1;
-            if (nb_cpb < 1 || nb_cpb > 32) {
-                av_log(NULL, AV_LOG_ERROR, "nb_cpb %d invalid\n", nb_cpb);
+        if (!hdr->flags.low_delay_hrd_flag) {
+            hdr->cpb_cnt_minus1[i] = get_ue_golomb_long(gb);
+            if (hdr->cpb_cnt_minus1[i] > 31) {
+                av_log(NULL, AV_LOG_ERROR, "nb_cpb %d invalid\n",
+                       hdr->cpb_cnt_minus1[i]);
                 return AVERROR_INVALIDDATA;
             }
         }
 
-        if (nal_params_present)
-            decode_sublayer_hrd(gb, nb_cpb, subpic_params_present);
-        if (vcl_params_present)
-            decode_sublayer_hrd(gb, nb_cpb, subpic_params_present);
+        if (hdr->flags.nal_hrd_parameters_present_flag)
+            decode_sublayer_hrd(gb, hdr->cpb_cnt_minus1[i], &hdr->nal_params[i],
+                                hdr->flags.sub_pic_hrd_params_present_flag);
+
+        if (hdr->flags.vcl_hrd_parameters_present_flag)
+            decode_sublayer_hrd(gb, hdr->cpb_cnt_minus1[i], &hdr->vcl_params[i],
+                                hdr->flags.sub_pic_hrd_params_present_flag);
     }
+
     return 0;
 }
 
@@ -538,7 +541,8 @@ int ff_hevc_decode_nal_vps(GetBitContext *gb, AVCodecContext *avctx,
             get_ue_golomb_long(gb); // hrd_layer_set_idx
             if (i)
                 common_inf_present = get_bits1(gb);
-            decode_hrd(gb, common_inf_present, vps->vps_max_sub_layers);
+            decode_hrd(gb, common_inf_present, &vps->hdr[i],
+                       vps->vps_max_sub_layers);
         }
     }
     get_bits1(gb); /* vps_extension_flag */
@@ -657,7 +661,7 @@ timing_info:
             vui->vui_num_ticks_poc_diff_one_minus1 = get_ue_golomb_long(gb);
         vui->vui_hrd_parameters_present_flag = get_bits1(gb);
         if (vui->vui_hrd_parameters_present_flag)
-            decode_hrd(gb, 1, sps->max_sub_layers);
+            decode_hrd(gb, 1, &sps->hdr, sps->max_sub_layers);
     }
 
     vui->bitstream_restriction_flag = get_bits1(gb);
diff --git a/libavcodec/hevc_ps.h b/libavcodec/hevc_ps.h
index 2124deb953..0124f5d37f 100644
--- a/libavcodec/hevc_ps.h
+++ b/libavcodec/hevc_ps.h
@@ -32,6 +32,43 @@
 #include "h2645_vui.h"
 #include "hevc.h"
 
+typedef struct HEVCSublayerHdrParams {
+    uint32_t bit_rate_value_minus1[HEVC_MAX_CPB_CNT];
+    uint32_t cpb_size_value_minus1[HEVC_MAX_CPB_CNT];
+    uint32_t cpb_size_du_value_minus1[HEVC_MAX_CPB_CNT];
+    uint32_t bit_rate_du_value_minus1[HEVC_MAX_CPB_CNT];
+    uint32_t cbr_flag;
+} HEVCSublayerHdrParams;
+
+typedef struct HEVCHdrFlagParams {
+    uint32_t nal_hrd_parameters_present_flag;
+    uint32_t vcl_hrd_parameters_present_flag;
+    uint32_t sub_pic_hrd_params_present_flag;
+    uint32_t sub_pic_cpb_params_in_pic_timing_sei_flag;
+    uint32_t fixed_pic_rate_general_flag;
+    uint32_t fixed_pic_rate_within_cvs_flag;
+    uint32_t low_delay_hrd_flag;
+} HEVCHdrFlagParams;
+
+typedef struct HEVCHdrParams {
+    HEVCHdrFlagParams flags;
+
+    uint8_t tick_divisor_minus2;
+    uint8_t du_cpb_removal_delay_increment_length_minus1;
+    uint8_t dpb_output_delay_du_length_minus1;
+    uint8_t bit_rate_scale;
+    uint8_t cpb_size_scale;
+    uint8_t cpb_size_du_scale;
+    uint8_t initial_cpb_removal_delay_length_minus1;
+    uint8_t au_cpb_removal_delay_length_minus1;
+    uint8_t dpb_output_delay_length_minus1;
+    uint8_t cpb_cnt_minus1[HEVC_MAX_SUB_LAYERS];
+    uint16_t elemental_duration_in_tc_minus1[HEVC_MAX_SUB_LAYERS];
+
+    HEVCSublayerHdrParams nal_params[HEVC_MAX_SUB_LAYERS];
+    HEVCSublayerHdrParams vcl_params[HEVC_MAX_SUB_LAYERS];
+} HEVCHdrParams;
+
 typedef struct ShortTermRPS {
     unsigned int num_negative_pics;
     int num_delta_pocs;
@@ -108,6 +145,8 @@ typedef struct PTL {
 } PTL;
 
 typedef struct HEVCVPS {
+    HEVCHdrParams hdr[HEVC_MAX_LAYER_SETS];
+
     uint8_t vps_temporal_id_nesting_flag;
     int vps_max_layers;
     int vps_max_sub_layers; ///< vps_max_temporal_layers_minus1 + 1
@@ -146,6 +185,8 @@ typedef struct HEVCSPS {
 
     HEVCWindow pic_conf_win;
 
+    HEVCHdrParams hdr;
+
     int bit_depth;
     int bit_depth_chroma;
     int pixel_shift;