From 90e9d2feabbbff0611aa873c67f75b1dc919ea08 Mon Sep 17 00:00:00 2001
From: Moritz Bunkus <mo@bunkus.online>
Date: Wed, 3 Mar 2021 22:15:40 +0100
Subject: [PATCH] HEVC ES parser: optionally normalize VPS/SPS/PPS placement

---
 src/common/hevc_es_parser.cpp | 36 ++++++++++++++++++++++++++++++++---
 src/common/hevc_es_parser.h   |  4 +++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/common/hevc_es_parser.cpp b/src/common/hevc_es_parser.cpp
index b96d05c0e..0746b1d2d 100644
--- a/src/common/hevc_es_parser.cpp
+++ b/src/common/hevc_es_parser.cpp
@@ -106,6 +106,11 @@ es_parser_c::discard_actual_frames(bool discard) {
   m_discard_actual_frames = discard;
 }
 
+void
+es_parser_c::normalize_parameter_sets(bool normalize) {
+  m_normalize_parameter_sets = normalize;
+}
+
 void
 es_parser_c::add_bytes(unsigned char *buffer,
                        size_t size) {
@@ -226,6 +231,25 @@ es_parser_c::add_timestamp(int64_t timestamp) {
   ++m_stats.num_timestamps_in;
 }
 
+void
+es_parser_c::add_parameter_sets_to_extra_data() {
+  std::vector<memory_cptr> old_extra_data;
+
+  for (auto const &data : m_extra_data)
+    if (old_extra_data.empty() || (*data != *old_extra_data.back()))
+      old_extra_data.emplace_back(data);
+
+  m_extra_data.clear();
+  m_extra_data.reserve(m_vps_list.size() + m_sps_list.size() + m_pps_list.size() + old_extra_data.size());
+
+  auto create_nalu_and_add = [this](memory_cptr const &nalu) { m_extra_data.push_back(create_nalu_with_size(nalu)); };
+
+  std::for_each(m_vps_list.begin(),     m_vps_list.end(),     create_nalu_and_add);
+  std::for_each(m_sps_list.begin(),     m_sps_list.end(),     create_nalu_and_add);
+  std::for_each(m_pps_list.begin(),     m_pps_list.end(),     create_nalu_and_add);
+  std::copy(    old_extra_data.begin(), old_extra_data.end(), std::back_inserter(m_extra_data));
+}
+
 void
 es_parser_c::flush_incomplete_frame() {
   if (!m_have_incomplete_frame || !m_hevcc_ready)
@@ -287,6 +311,9 @@ es_parser_c::handle_slice_nalu(memory_cptr const &nalu,
     m_b_frames_since_keyframe = false;
     cleanup();
 
+    if (m_normalize_parameter_sets)
+      add_parameter_sets_to_extra_data();
+
   } else
     m_b_frames_since_keyframe |= is_b_slice;
 
@@ -344,7 +371,8 @@ es_parser_c::handle_vps_nalu(memory_cptr const &nalu) {
     m_codec_private.vps_data_id                = vps_info.id;
   }
 
-  m_extra_data.push_back(create_nalu_with_size(nalu));
+  if (!m_normalize_parameter_sets)
+    m_extra_data.push_back(create_nalu_with_size(nalu));
 }
 
 void
@@ -386,7 +414,8 @@ es_parser_c::handle_sps_nalu(memory_cptr const &nalu) {
   } else
     use_sps_info = false;
 
-  m_extra_data.push_back(create_nalu_with_size(parsed_nalu));
+  if (!m_normalize_parameter_sets)
+    m_extra_data.push_back(create_nalu_with_size(parsed_nalu));
 
   // Update codec private if needed
   if (-1 == m_codec_private.sps_data_id)
@@ -450,7 +479,8 @@ es_parser_c::handle_pps_nalu(memory_cptr const &nalu) {
     m_hevcc_changed    = true;
   }
 
-  m_extra_data.push_back(create_nalu_with_size(nalu));
+  if (!m_normalize_parameter_sets)
+    m_extra_data.push_back(create_nalu_with_size(nalu));
 }
 
 void
diff --git a/src/common/hevc_es_parser.h b/src/common/hevc_es_parser.h
index 6f849eef7..5ca7df93d 100644
--- a/src/common/hevc_es_parser.h
+++ b/src/common/hevc_es_parser.h
@@ -81,7 +81,7 @@ protected:
   bool m_have_incomplete_frame{};
   std::deque<std::pair<memory_cptr, uint64_t>> m_unhandled_nalus;
 
-  bool m_simple_picture_order{}, m_discard_actual_frames{};
+  bool m_simple_picture_order{}, m_discard_actual_frames{}, m_normalize_parameter_sets{};
 
   debugging_option_c m_debug_keyframe_detection{"hevc_parser|hevc_keyframe_detection"}, m_debug_nalu_types{"hevc_parser|hevc_nalu_types"}, m_debug_timestamps{"hevc_parser|hevc_timestamps"}, m_debug_sps_info{"hevc_parser|hevc_sps|hevc_sps_info"}, m_debug_parameter_sets{"hevc_parser|hevc_parameter_sets"};
   static std::unordered_map<int, std::string> ms_nalu_names_by_type;
@@ -175,6 +175,7 @@ public:
   }
 
   void discard_actual_frames(bool discard = true);
+  void normalize_parameter_sets(bool normalize = true);
 
   int get_num_skipped_frames() const {
     return m_num_skipped_frames;
@@ -218,6 +219,7 @@ protected:
   void calculate_frame_order();
   void calculate_frame_timestamps();
   void calculate_frame_references_and_update_stats();
+  void add_parameter_sets_to_extra_data();
   static void init_nalu_names();
 };