hevc_parser: Set pict_type, key_frame and output_picture_number.

Conflicts:
	libavcodec/hevc.c

3 files changed, 233 insertions, 9 deletions

diff --git a/libavcodec/hevc.c b/libavcodec/hevc.c
index 9a8886735e..aa174b227e 100644
--- a/libavcodec/hevc.c
+++ b/libavcodec/hevc.c
@@ -2192,8 +2192,8 @@ static int decode_nal_unit(HEVCContext *s, const uint8_t *nal, int length)
 
 /* FIXME: This is adapted from ff_h264_decode_nal, avoiding duplication
    between these functions would be nice. */
-static int extract_rbsp(HEVCContext *s, const uint8_t *src, int length,
-                        HEVCNAL *nal)
+int ff_hevc_extract_rbsp(HEVCContext *s, const uint8_t *src, int length,
+                         HEVCNAL *nal)
 {
     int i, si, di;
     uint8_t *dst;
@@ -2279,7 +2279,8 @@ static int extract_rbsp(HEVCContext *s, const uint8_t *src, int length,
                     if (!s->skipped_bytes_pos)
                         return AVERROR(ENOMEM);
                 }
-                s->skipped_bytes_pos[s->skipped_bytes-1] = di - 1;
+                if (s->skipped_bytes_pos)
+                    s->skipped_bytes_pos[s->skipped_bytes-1] = di - 1;
                 continue;
             } else // next start code
                 goto nsc;
@@ -2363,7 +2364,7 @@ static int decode_nal_units(HEVCContext *s, const uint8_t *buf, int length)
         s->skipped_bytes_pos = s->skipped_bytes_pos_nal[s->nb_nals];
         nal = &s->nals[s->nb_nals];
 
-        consumed = extract_rbsp(s, buf, extract_length, nal);
+        consumed = ff_hevc_extract_rbsp(s, buf, extract_length, nal);
 
         s->skipped_bytes_nal[s->nb_nals] = s->skipped_bytes;
         s->skipped_bytes_pos_size_nal[s->nb_nals] = s->skipped_bytes_pos_size;
diff --git a/libavcodec/hevc.h b/libavcodec/hevc.h
index 01db69b74b..0975131b49 100644
--- a/libavcodec/hevc.h
+++ b/libavcodec/hevc.h
@@ -872,6 +872,9 @@ int ff_hevc_decode_nal_sps(HEVCContext *s);
 int ff_hevc_decode_nal_pps(HEVCContext *s);
 int ff_hevc_decode_nal_sei(HEVCContext *s);
 
+int ff_hevc_extract_rbsp(HEVCContext *s, const uint8_t *src, int length,
+                         HEVCNAL *nal);
+
 /**
  * Mark all frames in DPB as unused for reference.
  */
diff --git a/libavcodec/hevc_parser.c b/libavcodec/hevc_parser.c
index 7e038bea0c..642f9752a4 100644
--- a/libavcodec/hevc_parser.c
+++ b/libavcodec/hevc_parser.c
@@ -23,9 +23,15 @@
 #include "libavutil/common.h"
 #include "parser.h"
 #include "hevc.h"
+#include "golomb.h"
 
 #define START_CODE 0x000001 ///< start_code_prefix_one_3bytes
 
+typedef struct HEVCParseContext {
+    HEVCContext  h;
+    ParseContext pc;
+} HEVCParseContext;
+
 /**
  * Find the end of the current frame in the bitstream.
  * @return the position of the first byte of the next frame, or END_NOT_FOUND
@@ -33,7 +39,7 @@
 static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf, int buf_size)
 {
     int i;
-    ParseContext *pc = s->priv_data;
+    ParseContext *pc = &((HEVCParseContext *)s->priv_data)->pc;
 
     for (i = 0; i < buf_size; i++) {
         int nut;
@@ -57,7 +63,6 @@ static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf, int
             if (first_slice_segment_in_pic_flag) {
                 if (!pc->frame_start_found) {
                     pc->frame_start_found = 1;
-                    s->key_frame = nut >= NAL_BLA_W_LP && nut <= NAL_CRA_NUT;
                 } else { // First slice of next frame found
                     pc->frame_start_found = 0;
                     return i - 5;
@@ -69,13 +74,193 @@ static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf, int
     return END_NOT_FOUND;
 }
 
+/**
+ * Parse NAL units of found picture and decode some basic information.
+ *
+ * @param s parser context.
+ * @param avctx codec context.
+ * @param buf buffer with field/frame data.
+ * @param buf_size size of the buffer.
+ */
+static inline int parse_nal_units(AVCodecParserContext *s,
+                                  AVCodecContext *avctx,
+                                  const uint8_t *buf, int buf_size)
+{
+    HEVCContext   *h  = &((HEVCParseContext *)s->priv_data)->h;
+    GetBitContext *gb = &h->HEVClc->gb;
+    SliceHeader   *sh = &h->sh;
+    const uint8_t *buf_end = buf + buf_size;
+    int state = -1, i;
+    HEVCNAL *nal;
+
+    /* set some sane default values */
+    s->pict_type         = AV_PICTURE_TYPE_I;
+    s->key_frame         = 0;
+    s->picture_structure = AV_PICTURE_STRUCTURE_UNKNOWN;
+
+    h->avctx = avctx;
+
+    if (!buf_size)
+        return 0;
+
+    if (h->nals_allocated < 1) {
+        HEVCNAL *tmp = av_realloc_array(h->nals, 1, sizeof(*tmp));
+        if (!tmp)
+            return AVERROR(ENOMEM);
+        h->nals = tmp;
+        memset(h->nals, 0, sizeof(*tmp));
+        h->nals_allocated = 1;
+    }
+
+    nal = &h->nals[0];
+
+    for (;;) {
+        int src_length, consumed;
+        buf = avpriv_find_start_code(buf, buf_end, &state);
+        if (--buf + 2 >= buf_end)
+            break;
+        src_length = buf_end - buf;
+
+        h->nal_unit_type = (*buf >> 1) & 0x3f;
+        h->temporal_id   = (*(buf + 1) & 0x07) - 1;
+        if (h->nal_unit_type <= NAL_CRA_NUT) {
+            // Do not walk the whole buffer just to decode slice segment header
+            if (src_length > 20)
+                src_length = 20;
+        }
+
+        consumed = ff_hevc_extract_rbsp(h, buf, src_length, nal);
+        if (consumed < 0)
+            return consumed;
+
+        init_get_bits8(gb, nal->data + 2, nal->size);
+        switch (h->nal_unit_type) {
+        case NAL_VPS:
+            ff_hevc_decode_nal_vps(h);
+            break;
+        case NAL_SPS:
+            ff_hevc_decode_nal_sps(h);
+            break;
+        case NAL_PPS:
+            ff_hevc_decode_nal_pps(h);
+            break;
+        case NAL_SEI_PREFIX:
+        case NAL_SEI_SUFFIX:
+            ff_hevc_decode_nal_sei(h);
+            break;
+        case NAL_TRAIL_N:
+        case NAL_TRAIL_R:
+        case NAL_TSA_N:
+        case NAL_TSA_R:
+        case NAL_STSA_N:
+        case NAL_STSA_R:
+        case NAL_RADL_N:
+        case NAL_RADL_R:
+        case NAL_RASL_N:
+        case NAL_RASL_R:
+        case NAL_BLA_W_LP:
+        case NAL_BLA_W_RADL:
+        case NAL_BLA_N_LP:
+        case NAL_IDR_W_RADL:
+        case NAL_IDR_N_LP:
+        case NAL_CRA_NUT:
+            sh->first_slice_in_pic_flag = get_bits1(gb);
+
+            if (h->nal_unit_type >= 16 && h->nal_unit_type <= 23) {
+                s->key_frame = 1;
+                sh->no_output_of_prior_pics_flag = get_bits1(gb);
+            }
+
+            sh->pps_id = get_ue_golomb(gb);
+            if (sh->pps_id >= MAX_PPS_COUNT || !h->pps_list[sh->pps_id]) {
+                av_log(h->avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", sh->pps_id);
+                return AVERROR_INVALIDDATA;
+            }
+            h->pps = (HEVCPPS*)h->pps_list[sh->pps_id]->data;
+
+            if (h->pps->sps_id >= MAX_SPS_COUNT || !h->sps_list[h->pps->sps_id]) {
+                av_log(h->avctx, AV_LOG_ERROR, "SPS id out of range: %d\n", h->pps->sps_id);
+                return AVERROR_INVALIDDATA;
+            }
+            if (h->sps != (HEVCSPS*)h->sps_list[h->pps->sps_id]->data) {
+                h->sps = (HEVCSPS*)h->sps_list[h->pps->sps_id]->data;
+                h->vps = h->vps_list[h->sps->vps_id];
+            }
+
+            if (!sh->first_slice_in_pic_flag) {
+                int slice_address_length;
+
+                if (h->pps->dependent_slice_segments_enabled_flag)
+                    sh->dependent_slice_segment_flag = get_bits1(gb);
+                else
+                    sh->dependent_slice_segment_flag = 0;
+
+                slice_address_length = av_ceil_log2_c(h->sps->ctb_width *
+                                                      h->sps->ctb_height);
+                sh->slice_segment_addr = get_bits(gb, slice_address_length);
+                if (sh->slice_segment_addr >= h->sps->ctb_width * h->sps->ctb_height) {
+                    av_log(h->avctx, AV_LOG_ERROR, "Invalid slice segment address: %u.\n",
+                           sh->slice_segment_addr);
+                    return AVERROR_INVALIDDATA;
+                }
+            } else
+                sh->dependent_slice_segment_flag = 0;
+
+            if (sh->dependent_slice_segment_flag)
+                break;
+
+            for (i = 0; i < h->pps->num_extra_slice_header_bits; i++)
+                skip_bits(gb, 1); // slice_reserved_undetermined_flag[]
+
+            sh->slice_type = get_ue_golomb(gb);
+            if (!(sh->slice_type == I_SLICE || sh->slice_type == P_SLICE ||
+                  sh->slice_type == B_SLICE)) {
+                av_log(h->avctx, AV_LOG_ERROR, "Unknown slice type: %d.\n",
+                       sh->slice_type);
+                return AVERROR_INVALIDDATA;
+            }
+            s->pict_type = sh->slice_type == B_SLICE ? AV_PICTURE_TYPE_B :
+                           sh->slice_type == P_SLICE ? AV_PICTURE_TYPE_P :
+                                                       AV_PICTURE_TYPE_I;
+
+            if (h->pps->output_flag_present_flag)
+                sh->pic_output_flag = get_bits1(gb);
+
+            if (h->sps->separate_colour_plane_flag)
+                sh->colour_plane_id = get_bits(gb, 2);
+
+            if (!IS_IDR(h)) {
+                sh->pic_order_cnt_lsb = get_bits(gb, h->sps->log2_max_poc_lsb);
+                s->output_picture_number = h->poc = ff_hevc_compute_poc(h, sh->pic_order_cnt_lsb);
+            } else
+                s->output_picture_number = h->poc = 0;
+
+            if (h->temporal_id == 0 &&
+                h->nal_unit_type != NAL_TRAIL_N &&
+                h->nal_unit_type != NAL_TSA_N &&
+                h->nal_unit_type != NAL_STSA_N &&
+                h->nal_unit_type != NAL_RADL_N &&
+                h->nal_unit_type != NAL_RASL_N &&
+                h->nal_unit_type != NAL_RADL_R &&
+                h->nal_unit_type != NAL_RASL_R)
+                h->pocTid0 = h->poc;
+
+            return 0; /* no need to evaluate the rest */
+        }
+        buf += consumed;
+    }
+    /* didn't find a picture! */
+    av_log(h->avctx, AV_LOG_ERROR, "missing picture in access unit\n");
+    return -1;
+}
+
 static int hevc_parse(AVCodecParserContext *s,
                       AVCodecContext *avctx,
                       const uint8_t **poutbuf, int *poutbuf_size,
                       const uint8_t *buf, int buf_size)
 {
     int next;
-    ParseContext *pc = s->priv_data;
+    ParseContext *pc = &((HEVCParseContext *)s->priv_data)->pc;
 
     if (s->flags & PARSER_FLAG_COMPLETE_FRAMES) {
         next = buf_size;
@@ -88,6 +273,8 @@ static int hevc_parse(AVCodecParserContext *s,
         }
     }
 
+    parse_nal_units(s, avctx, buf, buf_size);
+
     *poutbuf = buf;
     *poutbuf_size = buf_size;
     return next;
@@ -116,10 +303,43 @@ static int hevc_split(AVCodecContext *avctx, const uint8_t *buf, int buf_size)
     return 0;
 }
 
+static int hevc_init(AVCodecParserContext *s)
+{
+    HEVCContext  *h  = &((HEVCParseContext *)s->priv_data)->h;
+    h->HEVClc = av_mallocz(sizeof(HEVCLocalContext));
+    h->skipped_bytes_pos_size = INT_MAX;
+
+    return 0;
+}
+
+static void hevc_close(AVCodecParserContext *s)
+{
+    int i;
+    HEVCContext  *h  = &((HEVCParseContext *)s->priv_data)->h;
+    ParseContext *pc = &((HEVCParseContext *)s->priv_data)->pc;
+
+    av_freep(&h->skipped_bytes_pos);
+    av_freep(&h->HEVClc);
+    av_freep(&pc->buffer);
+
+    for (i = 0; i < FF_ARRAY_ELEMS(h->vps_list); i++)
+        av_freep(&h->vps_list[i]);
+    for (i = 0; i < FF_ARRAY_ELEMS(h->sps_list); i++)
+        av_buffer_unref(&h->sps_list[i]);
+    for (i = 0; i < FF_ARRAY_ELEMS(h->pps_list); i++)
+        av_buffer_unref(&h->pps_list[i]);
+
+    for (i = 0; i < h->nals_allocated; i++)
+        av_freep(&h->nals[i].rbsp_buffer);
+    av_freep(&h->nals);
+    h->nals_allocated = 0;
+}
+
 AVCodecParser ff_hevc_parser = {
     .codec_ids      = { AV_CODEC_ID_HEVC },
-    .priv_data_size = sizeof(ParseContext),
+    .priv_data_size = sizeof(HEVCParseContext),
+    .parser_init    = hevc_init,
     .parser_parse   = hevc_parse,
-    .parser_close   = ff_parse_close,
+    .parser_close   = hevc_close,
     .split          = hevc_split,
 };