avcodec/cbs: add API to discard units by AVDiscard

3 files changed, 44 insertions, 0 deletions

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 504197e06d..cf5211249b 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -1026,3 +1026,24 @@ int ff_cbs_make_unit_writable(CodedBitstreamContext *ctx,
     av_buffer_unref(&ref);
     return 0;
 }
+
+void ff_cbs_discard_units(CodedBitstreamContext *ctx,
+                          CodedBitstreamFragment *frag,
+                          enum AVDiscard skip,
+                          int flags)
+{
+    if (!ctx->codec->discarded_unit)
+        return;
+
+    for (int i = frag->nb_units - 1; i >= 0; i--) {
+        if (ctx->codec->discarded_unit(ctx, &frag->units[i], skip)) {
+            // discard all units
+            if (!(flags & DISCARD_FLAG_KEEP_NON_VCL)) {
+                ff_cbs_fragment_free(frag);
+                return;
+            }
+
+            ff_cbs_delete_unit(frag, i);
+        }
+    }
+}
diff --git a/libavcodec/cbs.h b/libavcodec/cbs.h
index ee21623dac..b4131db5fe 100644
--- a/libavcodec/cbs.h
+++ b/libavcodec/cbs.h
@@ -26,6 +26,7 @@
 
 #include "codec_id.h"
 #include "codec_par.h"
+#include "defs.h"
 #include "packet.h"
 
 
@@ -432,5 +433,21 @@ int ff_cbs_make_unit_refcounted(CodedBitstreamContext *ctx,
 int ff_cbs_make_unit_writable(CodedBitstreamContext *ctx,
                               CodedBitstreamUnit *unit);
 
+enum CbsDiscardFlags {
+    DISCARD_FLAG_NONE = 0,
+
+    /**
+     * keep non-vcl units even if the picture has been dropped.
+     */
+    DISCARD_FLAG_KEEP_NON_VCL = 0x01,
+};
+
+/**
+ * Discard units accroding to 'skip'.
+ */
+void ff_cbs_discard_units(CodedBitstreamContext *ctx,
+                          CodedBitstreamFragment *frag,
+                          enum AVDiscard skip,
+                          int flags);
 
 #endif /* AVCODEC_CBS_H */
diff --git a/libavcodec/cbs_internal.h b/libavcodec/cbs_internal.h
index e585c77934..077954eee5 100644
--- a/libavcodec/cbs_internal.h
+++ b/libavcodec/cbs_internal.h
@@ -133,6 +133,12 @@ typedef struct CodedBitstreamType {
                       CodedBitstreamUnit *unit,
                       PutBitContext *pbc);
 
+    // Return 1 when the unit should be dropped according to 'skip',
+    // 0 otherwise.
+    int (*discarded_unit)(CodedBitstreamContext *ctx,
+                          const CodedBitstreamUnit *unit,
+                          enum AVDiscard skip);
+
     // Read the data from all of frag->units and assemble it into
     // a bitstream for the whole fragment.
     int (*assemble_fragment)(CodedBitstreamContext *ctx,