[FFmpeg-devel] [PATCH] avcodec/cbs_h2645: fix storing Picture Header references in the context

[James Almer]

Signed-off-by: James Almer <jamrial at gmail.com>
---
 libavcodec/cbs_h2645.c                | 31 +++++++++++++++------------
 libavcodec/cbs_h266.h                 |  7 ++----
 libavcodec/cbs_h266_syntax_template.c |  7 +++---
 3 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index cdd7901518..cf480d71f6 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -525,12 +525,6 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
     if (frag->data_size == 0)
         return 0;
 
-    if (codec_id == AV_CODEC_ID_VVC) {
-        //we deactive picture header here to avoid reuse previous au's ph.
-        CodedBitstreamH266Context *h266 = ctx->priv_data;
-        h266->priv.ph = NULL;
-    }
-
     if (header && frag->data[0] && codec_id == AV_CODEC_ID_H264) {
         // AVCC header.
         size_t size, start, end;
@@ -793,19 +787,20 @@ cbs_h266_replace_ps(6, SPS, sps, sps_seq_parameter_set_id)
 cbs_h266_replace_ps(6, PPS, pps, pps_pic_parameter_set_id)
 
 static int cbs_h266_replace_ph(CodedBitstreamContext *ctx,
-                               CodedBitstreamUnit *unit)
+                               CodedBitstreamUnit *unit,
+                               H266RawPH *ph)
 {
     CodedBitstreamH266Context *h266 = ctx->priv_data;
     int err;
 
-    h266->priv.ph = NULL;
     err = ff_cbs_make_unit_refcounted(ctx, unit);
     if (err < 0)
         return err;
-    err = av_buffer_replace(&h266->priv.ph_ref, unit->content_ref);
+    av_assert0(unit->content_ref);
+    err = av_buffer_replace(&h266->ph_ref, unit->content_ref);
     if (err < 0)
         return err;
-    h266->priv.ph = (H266RawPH*)h266->priv.ph_ref->data;
+    h266->ph = ph;
     return 0;
 }
 
@@ -1111,7 +1106,7 @@ static int cbs_h266_read_nal_unit(CodedBitstreamContext *ctx,
             err = cbs_h266_read_ph(ctx, &gbc, ph);
             if (err < 0)
                 return err;
-            err = cbs_h266_replace_ph(ctx, unit);
+            err = cbs_h266_replace_ph(ctx, unit, ph);
             if (err < 0)
                 return err;
         }
@@ -1139,6 +1134,10 @@ static int cbs_h266_read_nal_unit(CodedBitstreamContext *ctx,
             pos = get_bits_count(&gbc);
             len = unit->data_size;
 
+            err = cbs_h266_replace_ph(ctx, unit, &slice->header.sh_picture_header);
+            if (err < 0)
+                return err;
+
             slice->data_size = len - pos / 8;
             slice->data_ref  = av_buffer_ref(unit->data_ref);
             if (!slice->data_ref)
@@ -1640,7 +1639,7 @@ static int cbs_h266_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h266_replace_ph(ctx, unit);
+            err = cbs_h266_replace_ph(ctx, unit, ph);
             if (err < 0)
                 return err;
         }
@@ -1661,6 +1660,10 @@ static int cbs_h266_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
+            err = cbs_h266_replace_ph(ctx, unit, &slice->header.sh_picture_header);
+            if (err < 0)
+                return err;
+
             if (slice->data) {
                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
                                                  slice->data_size,
@@ -1884,8 +1887,8 @@ static void cbs_h266_flush(CodedBitstreamContext *ctx)
         av_buffer_unref(&h266->pps_ref[i]);
         h266->pps[i] = NULL;
     }
-    av_buffer_unref(&h266->priv.ph_ref);
-    h266->priv.ph = NULL;
+    av_buffer_unref(&h266->ph_ref);
+    h266->ph = NULL;
 }
 
 static void cbs_h266_close(CodedBitstreamContext *ctx)
diff --git a/libavcodec/cbs_h266.h b/libavcodec/cbs_h266.h
index 03dfd4a954..460db7ab62 100644
--- a/libavcodec/cbs_h266.h
+++ b/libavcodec/cbs_h266.h
@@ -770,14 +770,11 @@ typedef struct CodedBitstreamH266Context {
     AVBufferRef *vps_ref[VVC_MAX_VPS_COUNT];
     AVBufferRef *sps_ref[VVC_MAX_SPS_COUNT];
     AVBufferRef *pps_ref[VVC_MAX_PPS_COUNT];
+    AVBufferRef *ph_ref;
     H266RawVPS  *vps[VVC_MAX_SPS_COUNT];
     H266RawSPS  *sps[VVC_MAX_SPS_COUNT];
     H266RawPPS  *pps[VVC_MAX_PPS_COUNT];
-
-    struct {
-        AVBufferRef *ph_ref;
-        H266RawPH   *ph;
-    } priv;
+    H266RawPH   *ph;
 } CodedBitstreamH266Context;
 
 #endif /* AVCODEC_CBS_H266_H */
diff --git a/libavcodec/cbs_h266_syntax_template.c b/libavcodec/cbs_h266_syntax_template.c
index 06f9f29e08..2861e6f223 100644
--- a/libavcodec/cbs_h266_syntax_template.c
+++ b/libavcodec/cbs_h266_syntax_template.c
@@ -2675,12 +2675,11 @@ static int FUNC(slice_header) (CodedBitstreamContext *ctx, RWContext *rw,
 
     flag(sh_picture_header_in_slice_header_flag);
     if (current->sh_picture_header_in_slice_header_flag) {
+        //7.4.8 if sh_picture_header_in_slice_header_flag is true, we do not have a PH NAL unit
         CHECK(FUNC(picture_header) (ctx, rw, &current->sh_picture_header));
         ph = &current->sh_picture_header;
-        //7.4.8 if sh_picture_header_in_slice_header_flag is true, we do not have PH NAL unit
-        h266->priv.ph = NULL;
     } else {
-        ph = h266->priv.ph;
+        ph = h266->ph;
         if (!ph) {
             av_log(ctx->log_ctx, AV_LOG_ERROR,
                    "Picture header not available.\n");
@@ -2822,7 +2821,7 @@ static int FUNC(slice_header) (CodedBitstreamContext *ctx, RWContext *rw,
               (ctx, rw, sps, pps, &current->sh_ref_pic_lists));
         ref_pic_lists = &current->sh_ref_pic_lists;
     } else {
-        ref_pic_lists = &h266->priv.ph->ph_ref_pic_lists;
+        ref_pic_lists = &ph->ph_ref_pic_lists;
     }
     if ((current->sh_slice_type != VVC_SLICE_TYPE_I &&
          ref_pic_lists->rpl_ref_list[0].num_ref_entries > 1) ||
-- 
2.41.0